Transformer 火到今天这个程度我见过不少朋友一上来就抄个nn.TransformerEncoderLayer就开始训模型结果 loss 不降、注意力图全糊、显存爆得莫名其妙。说实话模型结构理解不透调参就是盲人摸象。这篇就干一件事——把 Transformer 的结构从骨架到每一根神经末梢全部拆开讲一遍包括每个模块为什么这么设计、张量形状怎么变、代码怎么写、训练时哪些坑自己踩过。不管你是刚看完 Attention Is All You Need 想动手复现还是已经用过 BERT、GPT、Swin Transformer、CLIP 的 text encoder但一直没搞明白底层的这篇都值得从头看一遍。1. Transformer 整体架构总览与设计哲学1.1 从 RNN 的痛点倒推 Transformer 的设计动机在 Transformer 出现之前处理序列数据的主流方案是 RNN、LSTM 和 GRU。它们有一个共同的硬伤必须按时间步串行计算。第 t 步的输出依赖第 t-1 步的隐状态这意味着在 GPU 上做不了并行。一条 512 长度的序列就要老老实实跑 512 次前向。训练数据规模一上去这个瓶颈直接卡死。另一方面LSTM 虽然靠门控机制缓解了长距离依赖问题但序列一长信息还是会衰减。从第 1 个词传到第 512 个词中间经过几百次非线性变换原始信息已经被稀释得差不多了。Attention 机制其实早在机器翻译领域就被用来做对齐但真正把它当作整个网络的核心骨架是 2017 年那篇论文干的事。Transformer 的核心洞察就是序列内部的依赖关系根本不需要靠递归来建模。任意两个位置之间的关联度可以用点积直接算出来一次矩阵乘法搞定所有位置对。这样一来序列长度维度完全可并行训练效率提升了几个数量级。注意并行化是 Transformer 能在大规模语料上训练的前提条件但并行化本身也带来了显存消耗随序列长度平方增长的代价这个后面会细讲。1.2 编码器-解码器骨架的宏观布局原始 Transformer 是一个标准的 encoder-decoder 结构服务于机器翻译任务。编码器由 N6 个相同的层堆叠而成每一层包含两个子层多头自注意力和前馈网络。解码器同样 N6 层但每层多了一个交叉注意力子层用来从编码器输出中提取信息。每个子层外围都套了残差连接和层归一化。原始论文用的是 Post-LN 结构也就是LayerNorm(x Sublayer(x))。这个细节很关键因为后来大量实践发现 Post-LN 在深层网络中训练不稳定需要 warmup 才能收敛而 Pre-LNx Sublayer(LayerNorm(x))就不需要那么激进的 warmup。现在主流的实现比如 GPT 系列、LLaMA用的都是 Pre-LN。所有子层和嵌入层的输出维度都是 d_model512这个统一维度是整个架构能堆叠的基础。前馈网络内部把维度扩张到 d_ff2048做完非线性变换再压回 512。注意力头数 h8每个头的维度 d_kd_vd_model/h64。这里有个容易忽略的细节多头注意力并没有增加总计算量。把 512 维切成 8 份各 64 维和直接用 512 维做单头注意力参数量和浮点运算量几乎一样。多头的价值在于让模型在不同的表示子空间里并行捕捉不同类型的关联模式。1.3 超参数配置与版本差异溯源原始论文的配置是 base 版本d_model512、h8、d_ff2048、N6、dropout0.1。big 版本把 d_model 提到 1024、h16、d_ff4096、N6。参数量的差距主要体现在 d_model 上因为注意力层的参数量是 4×d_model²前馈层是 2×d_model×d_ff。后来各个变体在结构上做了不同程度的修改。BERT 只用了编码器把 N 加到 12 或 24GPT 只用了解码器去掉了交叉注意力子层T5 回归完整的 encoder-decoder但简化了位置编码方案。Swin Transformer 则把注意力窗口限制在局部用移位窗口来扩大感受野解决了视觉任务中序列长度过大的问题。TabNet 走的是另一条路用稀疏注意力做表格数据本质上是把特征选择集成进了网络结构里。理解这些变体的前提是把原始结构吃透。下面几节我会逐个模块拆开讲。2. 输入表示词嵌入与位置编码的协同设计2.1 词嵌入层的维度选择与权重共享输入序列首先通过嵌入层映射成 d_model 维的稠密向量。词表大小通常在三万到十万之间嵌入矩阵的形状是[vocab_size, d_model]。以 d_model512、词表五万计算嵌入层参数量约 2560 万占了整个 base 模型约 6500 万参数的相当大一坨。原始论文里还做了一个权重共享的操作把解码器的输出投影层和嵌入层的权重矩阵绑定即输出层用嵌入矩阵的转置来做 logits 计算。这样做有两个好处一是减少参数量二是让输入和输出空间的语义表示保持一致。这个技巧在后来的语言模型里几乎成了标配GPT-2、T5 都沿用了。嵌入之后通常会乘一个sqrt(d_model)的缩放因子。这个操作在实际代码里很容易被漏掉或者写错位置。它的目的是让嵌入向量的方差和位置编码的方差保持在同一量级避免加法之后位置信息被淹没。class TokenEmbedding(nn.Module): def __init__(self, vocab_size, d_model): super().__init__() self.embed nn.Embedding(vocab_size, d_model) self.d_model d_model def forward(self, x): # 缩放因子不可省略 return self.embed(x) * math.sqrt(self.d_model)2.2 正弦位置编码的数学推导与实现细节自注意力本身是置换不变的你打乱输入顺序输出的集合是一样的。所以必须显式注入位置信息。原始论文选了一个不含参数的正弦方案$$PE_{(pos, 2i)} \sin(pos / 10000^{2i/d_{model}})$$ $$PE_{(pos, 2i1)} \cos(pos / 10000^{2i/d_{model}})$$其中 pos 是位置索引i 是维度索引。偶数维用 sin奇数维用 cos。不同的维度对应不同的频率从 2π 到 10000·2π 形成一组几何级数的波长。这样设计的好处是任意固定偏移 k 的位置编码可以表示成当前位置编码的线性变换模型理论上能学会相对位置的表示。我第一次手写这段代码的时候被div_term的计算绕进去了。关键是把10000^{2i/d_model}转化成指数形式再向量化def positional_encoding(max_len, d_model): pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) div_term torch.exp( torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model) ) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) return pe.unsqueeze(0) # [1, max_len, d_model]提示div_term只有 d_model/2 个元素广播到 position 上之后正好填满偶数和奇数位置各一半。写成arange(0, d_model, 2)而不是arange(0, d_model//2)这是个常见笔误。实际部署时位置编码通常注册成 buffer 而不是 Parameter因为它是固定的、不需要梯度更新。用register_buffer注册的优点是模型保存和加载时会自动带上转移到 GPU 时也会自动同步。2.3 为什么是加法而不是拼接一个自然的问题是为什么位置编码要和词嵌入相加而不是拼接成一个更长的向量拼接的方案在理论上也说得通d_model 变成原来两倍让模型自己去学融合方式。从信息论角度看加法是一种信息混叠操作更省维度。而论文作者在实验中发现加法方案的效果和拼接方案基本持平。既然效果一样加法参数量更少、计算更省自然就成了首选。不过这里有个隐含前提嵌入向量和位置编码的量级要匹配。如果嵌入向量的每个元素都是标准差为 1 的随机数而位置编码的值域在 [-1, 1] 之间加法之后位置信号占比极小模型很难学到。所以前面提到的sqrt(d_model)缩放不是可选项而是必须的。另外后来的 RoPE旋转位置编码和 ALiBi线性偏置注意力走的是完全不同的路线——直接把位置信息注入到注意力计算里而不是加在输入上。这些方案在长文本场景下表现更好但那是另一个专题了这里先聚焦原始正弦方案。