Transformer 这个架构刚上手时最容易卡住的地方不是 Self-Attention 本身而是它周边那几个看起来不起眼、但缺了就完全跑不起来的组件。我第二周啃 Multi-Head、位置编码和 Add Norm 这三块的时候前前后后返工了三次才真正把逻辑理顺。这篇就把我踩过的坑、推导过的公式、以及实际写代码时那些文档里不会写的细节完整地摊开讲一遍。如果你已经看懂了单头注意力但一到多头就懵、一看到位置编码里那堆 sin/cos 就头大、或者对残差连接和 LayerNorm 到底加在哪、为什么加始终一知半解那这篇基本就是给你写的。我会从设计动机讲到公式推导再落到可以直接抄的 PyTorch 代码最后附上我调试时遇到的典型报错和排查思路。1. 为什么单头注意力不够用Multi-Head 的设计动机拆解1.1 从一个视角看句子到多个视角并行先说结论单头注意力并不是能力不足而是表达能力被压缩了。假设我们有一个 512 维的输入单头注意力做的事情是——把这 512 维整体投影成一组 Q、K、V然后算出一套注意力权重最后加权求和得到一个 512 维的输出。整个过程里句子中每个词对其它词的关注程度只有一套权重分布。问题就出在这一套上。语言里的关系是高度多义的。举个人话的例子小明把书放在了桌子上因为它很重——这里的它指代的是书还是桌子模型在判断指代关系时需要同时考虑重量这个语义线索和位置邻近这个结构线索。单头注意力只能给出一套混合了所有线索的权重等于把多个判断维度强行揉成一团。Multi-Head 的思路很直接既然一套不够那就并行搞 h 套。每一套也就是一个 head在自己的子空间里独立地学习一种关注模式。有的 head 可能专门盯语法依赖主谓宾有的 head 可能专门盯指代消解还有的 head 可能关注的是位置相邻的词。最后把这 h 套结果拼起来再过一个线性层融合。这里有个关键点很多人会误解多头不是把维度复制 h 份而是把维度切分 h 份。原始论文里 d_model 512h 8那么每个 head 的维度 d_k d_v 512 / 8 64。总计算量和单头在 512 维上做基本持平但表达能力的宽度完全不一样了。1.2 维度切分背后的计算账我把这笔账算清楚因为这是理解多头最核心的一步。假设输入序列长度 n 10d_model 512单头方案Q、K、V 各是 10×512 的矩阵注意力分数矩阵是 10×10计算量约 10×10×512。多头方案h8每个 head 的 Q、K、V 是 10×64注意力分数矩阵还是 10×108 个 head 加起来计算量约 8×10×10×64 10×10×512。看到没计算量几乎一样。但多头方案里那 8 个 10×10 的注意力矩阵是互相独立的每个都在自己的 64 维子空间里算相似度。这就是用同样的算力换更丰富的表达。注意这里说的计算量持平是理论值。实际工程里因为要 reshape、transpose加上 batch 维度显存占用和 kernel 调度会有差异但量级是一致的。1.3 一个容易踩的认知坑我一开始以为多头就是把单头跑 8 次然后平均。这个理解是错的会导致后面代码完全写歪。正确的理解是多头是一次投影、一次切分、并行计算、最后拼接融合。不是跑 8 次独立的注意力而是在一次矩阵运算里通过 reshape 把 h 个 head 塞进 batch 维度并行处理。这个认知差异直接决定了你代码是写成 for 循环慢且错还是写成张量操作快且对。2. Multi-Head 的实现细节与代码落地2.1 Q、K、V 的投影矩阵怎么设标准做法是用三个独立的线性层 W_q、W_k、W_v形状都是 (d_model, d_model)。注意是 d_model 到 d_model不是到 d_k。因为切分是在投影之后通过 reshape 完成的不是投影时就投到 64 维。import torch import torch.nn as nn import math class MultiHeadAttention(nn.Module): def __init__(self, d_model512, num_heads8, dropout0.1): super().__init__() assert d_model % num_heads 0, d_model 必须能被 num_heads 整除 self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads self.W_q nn.Linear(d_model, d_model) self.W_k nn.Linear(d_model, d_model) self.W_v nn.Linear(d_model, d_model) self.W_o nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) def forward(self, query, key, value, maskNone): batch_size query.size(0) # 1. 线性投影 Q self.W_q(query) # (B, n, d_model) K self.W_k(key) V self.W_v(value) # 2. 切分成多头: (B, n, d_model) - (B, h, n, d_k) Q Q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) K K.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) V V.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # 3. 缩放点积注意力 scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn torch.softmax(scores, dim-1) attn self.dropout(attn) # 4. 加权求和并合并多头 context torch.matmul(attn, V) # (B, h, n, d_k) context context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) # 5. 输出投影 return self.W_o(context)这段代码里有几个细节值得单独拎出来说。第一view和transpose的顺序不能反。我见过有人先 transpose 再 view结果维度全乱。正确顺序是先把 d_model 拆成 (h, d_k)再 transpose 把 h 换到前面。因为 view 是按内存连续顺序切的你得先保证切出来的 h 和 d_k 是相邻的语义。第二contiguous()不能省。transpose 之后张量在内存里不连续直接 view 会报错。加.contiguous()强制拷贝一份连续内存。这个坑我调了半小时才定位到。第三缩放因子是sqrt(d_k)不是sqrt(d_model)。因为点积是在 d_k 维上做的方差随 d_k 增长所以要除以 sqrt(d_k) 把方差拉回 1 附近否则 softmax 会进入饱和区梯度接近 0。2.2 为什么缩放因子必须是 sqrt(d_k)这个推导我建议每个人都自己推一遍不然永远记不牢。假设 Q 和 K 的每个分量都是均值 0、方差 1 的独立随机变量那么它们的点积 q·k Σ(q_i × k_i) 的均值是 0方差是 d_k因为 d_k 个独立项相加每项方差为 1。当 d_k 64 时点积的标准差就是 8。这么大的数值送进 softmax会导致最大的那个值接近 1其余接近 0分布极度尖锐梯度几乎消失。除以 sqrt(d_k) 8 之后方差重新变回 1softmax 工作在正常区间。实操心得如果你把缩放因子写错比如漏了训练初期 loss 会卡在一个较高的值下不去而且梯度范数异常小。这是排查缩放问题的一个信号。2.3 多头到底学到了什么一个可观察的现象训练完之后你可以把每个 head 的注意力矩阵可视化出来看。我在一个小型翻译任务上观察到的现象是浅层的 head 倾向于关注相邻位置类似 n-gram中层的 head 开始出现明显的语法依赖模式比如动词关注主语深层的 head 反而变得分散像是在做全局信息聚合。这个观察和很多论文里的结论一致也解释了为什么多头比单头强——它给了模型在不同层、不同子空间里学习不同粒度关系的机会。3. 位置编码让模型知道谁在前谁在后3.1 为什么 Self-Attention 天生无序这是 Transformer 最反直觉的一点。Self-Attention 的计算本质是加权求和而加权求和是置换等变的——你把输入序列里两个词的位置换一下输出也只是对应地换一下注意力权重的计算完全不关心顺序。换句话说如果不用位置编码我爱你和你爱我在模型眼里是同一个东西。RNN 靠时间步的递归天然带顺序CNN 靠卷积核的滑动窗口带局部顺序而 Self-Attention 是全局并行的没有任何内置的顺序概念。所以必须显式地把位置信息注入进去。3.2 正弦位置编码的公式与直觉原始论文用的是固定非学习的正弦编码$$PE_{(pos, 2i)} \sin(pos / 10000^{2i/d_{model}})$$ $$PE_{(pos, 2i1)} \cos(pos / 10000^{2i/d_{model}})$$其中 pos 是位置i 是维度索引。偶数维用 sin奇数维用 cos。这个设计有几个精妙之处我逐个拆第一每个维度对应一个不同频率的波。从 i0 到 id_model/2频率从高到低排列。低维i 小变化快能区分相邻位置高维i 大变化慢能编码长距离的位置差异。这就像用二进制编码数字——低位变化快高位变化慢组合起来能唯一表示任意位置。第二相对位置可以被线性表示。这是最关键的性质。对于任意固定的偏移 kPE(posk) 可以写成 PE(pos) 的线性函数通过一个旋转矩阵。这意味着模型可以通过线性变换学到相对位置的概念而不只是绝对位置。这对泛化到训练时没见过的序列长度很重要。第三不需要训练参数。固定编码不占参数、不增加训练负担而且理论上可以外推到任意长度虽然实际外推效果有限。class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000, dropout0.1): super().__init__() self.dropout nn.Dropout(dropout) 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) pe pe.unsqueeze(0) # (1, max_len, d_model) self.register_buffer(pe, pe) def forward(self, x): # x: (B, n, d_model) x x self.pe[:, :x.size(1), :] return self.dropout(x)3.3 为什么用exp(log)而不是直接写10000**(2i/d_model)这是个数值稳定性的小技巧。直接算10000 ** (2i/d_model)在 i 较大时会有精度问题而且除法在指数上容易溢出。改写成exp(2i/d_model * log(10000))之后指数部分变成一个线性项数值范围可控。这个写法在几乎所有开源实现里都能看到属于约定俗成但文档不写的细节。3.4 固定编码 vs 可学习编码怎么选现在很多模型比如 BERT 系列用的是可学习的位置嵌入而不是正弦编码。两者怎么选我的经验是维度正弦固定编码可学习位置嵌入参数量0max_len × d_model外推能力理论可外推实际有限完全无法外推到 max_len 之外训练稳定性高不依赖数据需要足够数据才能学好适用场景序列长度变化大、数据少固定长度、数据充足如果你的任务序列长度固定且数据量大可学习嵌入通常效果略好如果长度变化大或者数据有限正弦编码更稳。我自己的项目里做变长文本分类时用正弦编码做固定长度生成时用可学习嵌入。注意正弦编码加到 embedding 上时是相加不是拼接。相加会损失一部分信息因为 embedding 和 PE 混在一起了但保持了维度不变是标准做法。拼接虽然保留了信息但会让 d_model 翻倍后续所有计算量都上去不划算。4. Add Norm残差连接与层归一化的组合拳4.1 残差连接解决的是什么问题Transformer 里每个子层注意力、前馈网络外面都套了一个残差连接output Sublayer(x) x。这个设计来自 ResNet核心作用是缓解深层网络的梯度消失。原理很直白反向传播时梯度要穿过很多层。如果没有残差梯度是连乘的每层都乘一个小于 1 的因子几十层之后梯度就趋近于 0 了。有了残差梯度有一条高速公路可以直接回传——因为 output f(x) x 对 x 求导是 f(x) 1那个 1 保证了梯度至少能以 1 的系数传回去不会消失。在 Transformer 里残差还有第二个作用保护原始信息。注意力层可能会过度加工输入把一些细节抹掉。残差连接让原始 embedding 始终有一条通路直达深层模型可以选择性地使用加工后的信息还是原始信息。4.2 LayerNorm 为什么不用 BatchNorm这是面试高频题也是实际选型的关键。BatchNorm 是在 batch 维度上做归一化对每个特征统计整个 batch 的均值和方差LayerNorm 是在特征维度上做归一化对每个样本统计它自己所有特征的均值和方差。Transformer 不用 BatchNorm 的原因有三第一序列长度可变。BatchNorm 需要在一个 batch 内统计但 NLP 任务里不同样本长度不同padding 会污染统计量。LayerNorm 对每个样本独立计算不受 batch 组成影响。第二batch size 通常较小。大模型训练时 batch 往往只有几十甚至几个样本BatchNorm 的统计量噪声太大不稳定。LayerNorm 不依赖 batch size。第三推理时的一致性。BatchNorm 训练和推理的行为不一致推理用滑动平均的统计量LayerNorm 训练推理完全一致部署更简单。class AddNorm(nn.Module): def __init__(self, d_model, dropout0.1): super().__init__() self.norm nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, sublayer_output): # Post-LN: 先相加再归一化 return self.norm(x self.dropout(sublayer_output))4.3 Post-LN 还是 Pre-LN一个影响训练稳定性的选择原始论文用的是 Post-LN先残差相加再 LayerNorm但后来大家发现 Post-LN 在深层模型里训练很不稳定需要 warmup 才能收敛。于是有了 Pre-LN先 LayerNorm再进子层最后残差相加# Pre-LN 写法 def forward(self, x): x x self.attention(self.norm1(x)) x x self.ffn(self.norm2(x)) return x两者的区别我用一句话总结Post-LN 把归一化放在残差之后梯度回传路径上多了归一化的关卡Pre-LN 把归一化放在子层之前残差路径是干净的恒等映射梯度可以无阻碍回传。实测下来Pre-LN 在 12 层以上的模型里明显更稳几乎不需要 warmup 就能训练。现在主流的大模型基本都用 Pre-LN 或其变体。但要注意Pre-LN 的最终输出没有经过归一化所以通常在最后一层后面补一个 LayerNorm。对比项Post-LNPre-LN归一化位置残差相加之后子层输入之前训练稳定性差需 warmup好几乎不需 warmup梯度路径经过 LayerNorm干净恒等映射最终输出已归一化需额外补 LayerNorm原始论文采用未采用实操心得如果你在复现原始论文的模型但训练 loss 一直震荡先别怀疑学习率把 Post-LN 换成 Pre-LN 试试十有八九能稳住。这个改动成本极低收益极大。5. 完整组装与实操调试记录5.1 把三块拼成一个 Encoder Layer单个 Encoder Layer 的结构是多头注意力 → Add Norm → 前馈网络 → Add Norm。前馈网络就是两层线性变换夹一个激活函数中间维度通常是 d_model 的 4 倍。class FeedForward(nn.Module): def __init__(self, d_model512, d_ff2048, dropout0.1): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.linear2 nn.Linear(d_ff, d_model) self.dropout nn.Dropout(dropout) self.activation nn.ReLU() def forward(self, x): return self.linear2(self.dropout(self.activation(self.linear1(x)))) class EncoderLayer(nn.Module): def __init__(self, d_model512, num_heads8, d_ff2048, dropout0.1): super().__init__() self.attention MultiHeadAttention(d_model, num_heads, dropout) self.ffn FeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout1 nn.Dropout(dropout) self.dropout2 nn.Dropout(dropout) def forward(self, x, maskNone): # Pre-LN 结构 attn_out self.attention(self.norm1(x), self.norm1(x), self.norm1(x), mask) x x self.dropout1(attn_out) ffn_out self.ffn(self.norm2(x)) x x self.dropout2(ffn_out) return x注意这里 self-attention 的 Q、K、V 都来自同一个输入都是 norm1(x)这就是自注意力的含义。如果是 Decoder 里的交叉注意力Q 来自 decoderK、V 来自 encoder。5.2 我调试时遇到的三个典型报错报错一RuntimeError: view size is not compatible with input tensors size and stride这个就是我前面提到的 contiguous 问题。transpose 之后张量内存不连续view 失败。解决方法是 transpose 后加.contiguous()。这个报错信息其实挺明确的但第一次遇到容易懵。报错二AssertionError: d_model must be divisible by num_heads维度不整除。我一开始设 d_model512、num_heads12512/12 除不尽。要么改 num_heads 为 8 或 16要么改 d_model。这个约束是硬性的因为要均匀切分。报错三训练 loss 变成 NaN这个最头疼。我排查下来有两个原因一是缩放因子漏了或者写错导致 softmax 溢出二是 mask 用了 0 而不是 -1e9导致 padding 位置被算进了注意力。mask 的正确做法是把需要屏蔽的位置填一个极大的负数如 -1e9softmax 之后这些位置的概率就趋近于 0。5.3 常见问题速查表现象可能原因排查方向loss 卡住不降缩放因子错误检查是否除以 sqrt(d_k)loss 变 NaNmask 值不对 / 学习率过大检查 mask 是否用 -1e9降低 lr显存爆掉batch 太大 / 未用 mask 压缩减小 batch检查 padding 处理训练震荡Post-LN 无 warmup换 Pre-LN 或加 warmup注意力全均匀初始化太小 / 训练不足检查初始化增加训练步数外推长度失败用了可学习位置嵌入换正弦编码5.4 一个容易被忽略的初始化细节Transformer 的权重初始化对训练影响很大。原始论文用的是 Xavier 初始化但实践中很多人发现如果把注意力输出层的初始权重调小比如乘以 1/sqrt(2L)L 是层数训练会更稳。原理是残差连接会让方差逐层累积初始权重太大时深层激活值会爆炸。这个技巧在 GPT 系列里被明确采用属于知道的人都在用、但论文里不一定写的经验。6. 几个我反复验证过的实操建议关于 Multi-Head我的建议是先用小配置跑通再放大。d_model64、num_heads4 这样的小模型在 CPU 上几秒就能跑一个 batch调试起来极快。等逻辑完全正确了再放大到 512、8 头。我见过太多人一上来就上大配置结果一个报错调半天效率极低。关于位置编码max_len 要留足余量。我一开始设 max_len 等于训练集最大长度结果推理时来了个更长的样本直接越界。后来改成训练最大长度的 2 倍再没出过问题。正弦编码本身不占多少显存多留点没坏处。关于 Add Norm新项目直接上 Pre-LN。除非你要严格复现某篇老论文否则没有理由用 Post-LN。Pre-LN 训练稳定、不需要精心调 warmup省下来的时间够你多跑好几组实验。最后分享一个我验证模型是否正确的小技巧用一个极小的数据集比如 10 个句子让模型过拟合。如果模型能把这 10 个句子完美记住loss 降到接近 0说明前向传播和反向传播的逻辑都是对的。如果连过拟合都做不到那一定是实现有 bug别急着调超参先回去查代码。这个方法帮我定位过好几次隐藏的维度错误比看 loss 曲线猜原因高效得多。