有时候你真想回到那个天真的年代——觉得一个注意力头就够用了。我至今记得那个深夜，线上服务突然OOM，日志里满屏的 CUDA out of memory，就因为我把头数从 8 砍到了 1，企图节省显存。结果呢？模型直接退化成了词袋。那感觉就像你为了减肥把大脑皮层切掉了四分之三——人还活着，但只会流口水。
所以，多头注意力。这名字听起来就像某种金融衍生品，对吧？但拆开了看，它其实是一种相当暴力的美学：把输入劈成好几份，各自去关注不同的子空间，最后再拼接起来。
单头的惨痛教训：一次线上的OOM事故
先泼盆冷水。单头注意力不是不能用，但局限大到令人发指。想象你在一个吵闹的鸡尾酒会上，只能听见一个频段的声音——你可能会抓住某个关键词，但完全丢失了语气、方向、甚至是谁在说话。Transformer 里的自注意力机制就是那个耳朵，而单头意味着你只有一个耳蜗。 信息瓶颈就卡在这里：context vector 的维度必须等于模型维度 d_model，所以当你强制压缩成一个头时，每个位置的表示被迫塞进一个狭窄的空间，所有类型的关联性（语法、语义、位置）全搅在一起。
那次事故后我跑了 ablation study。用 d_model=512 的模型，单头配置下，WMT14 英德翻译的 BLEU 直接从 27.3 掉到 21.8，loss 降不下去，哪怕训练步数加倍。最可笑的是，显存确实降了 30%，但 GPU 利用率反而更高了——因为反向传播时梯度方差巨大，收敛慢得像在用拨号上网加载高清图片。
拆解多头：矩阵乘法里藏着什么魔法？
别被论文里优雅的公式骗了。所谓的多头，代码层面就是 把 Q、K、V 线性变换到更低的维度，并行算 scaled dot-product attention，然后 concat 起来再过一次线性层。 比如你设 h=8 个头，d_k=d_v=64，那么每个头处理的其实是原向量的一个 64 维切片——这八份切片在初始化时就随机分好了，训练过程中逐渐各自收敛到不同的投影子空间。
我写这段代码时最容易犯的错是 transpose 与 reshape 的顺序。有一次我竟然把 batch 维和 head 维搞混了，结果每个样本都在关注其他样本的序列……模型学到的 pattern 诡异极了，输出总是以“根据《圣经》记载”开头，无论输入是啥。后来发现是因为梯度在错误的维度上被平均，相当于强制共享了不该共享的上下文。 正确的实现必须用 view() 或 reshape() 拆分出 head 维，然后 transpose(1,2) 将 (batch, seq, head, dim) 变成 (batch, head, seq, dim)，这样 PyTorch 的 batch matmul 才能自然并行。
至于那个著名的“为什么要缩放点积？”——答案是 稳定 softmax 梯度。当 d_k 较大时，点积的方差会线性增长，导致 softmax 输出趋近 one-hot，梯度接近零。除以 sqrt(d_k) 本质上是做了一次方差归一化。我第一次读论文时觉得这简直是多余的细节，直到自己训了一个头维 128 的模型，验证集的困惑度每隔几百步就突然跳到 NaN……加入 scaling 后，训练曲线平滑得像德芙广告。
压测说话：8头 vs 单头，吞吐量差了十倍？
我们在 A100 上做过一次严格对比。模型：Transformer-base，6 层，批大小 32，序列长度 128。单头配置（d_k=512）训练吞吐 1200 tokens/sec，多头（h=8, d_k=64）吞吐 10800 tokens/sec。接近 9 倍的差距！而且这还没算模型质量。为什么？因为 GPU 的 Tensor Core 最爱规则的矩阵尺寸，8 个 64×512 的矩阵乘法比一个 512×512 的矩阵乘法更容易填满计算单元、减少 warp divergence。另外，反向传播时，多头的梯度是各头独立计算的，可以更充分地利用内存带宽。
还有一个隐形红利：多头天然“集成”了不同的注意力模式。可视化权重时你会发现，有些头专门负责相邻位置的语法一致（如冠词-名词），有些头捕捉长距离的指代，还有一两个头似乎在关注无用信息——我们叫它“摆烂头”。但奇妙的是，剪掉这些摆烂头反而会掉点，因为它们或许充当了某种正则化作用。这比单头那种把所有模式糅杂在一起的“浆糊”强太多了。
血泪填坑指南：三个差点让我删库跑路的陷阱
陷阱一：头数不是越大越好，甚至不是 2 的幂。 很多人默认 8、16，但如果你不对 d_k 仔细斟酌，会栽大跟头。我们试过 h=32，d_k=16（总模型维 512），结果每个头的表示维度太小，对 key 的线性投影几乎沦为降维损失。模型在简单的语言建模任务上就过拟合，测试集 ppl 反而上升。后来改 h=12，d_k≈43，不是整数？直接用 d_k = (d_model + h - 1) // h 做零填充，再 mask 掉填充位置——效果竟然比标准配置还好。所以，请根据数据量和任务复杂度用网格搜索找 (h, d_k) 的最佳平衡点，别迷信 8 或 16。
陷阱二：训练初期注意力分布崩溃。 这是大 batch 训练时的常见死法。如果不加任何偏置或 dropout，前几步更新后，softmax 后的权重会极端化，导致某些位置获得几乎全部注意力，梯度骤降。解决方案很简单：在 scaled scores 矩阵上加上一个下三角的 mask（对于 autoregressive）还不够，需要给一个显式的注意力 dropout。我们用的 rate 0.1，在 Q*K^T 之后立即进行，不是按元素随机置零，而是整行丢弃一部分 attention weights，迫使模型不能依赖单一路径。这个 trick 救了我们的机器翻译模型，bleu 从 18 飙回 28。
陷阱三：多头融合时的线性层初始化不当。 那个拼接后的线性层 W_o 要是用默认的 Xavier 初始化，你会发现验证 loss 在初期震荡剧烈。这是因为多头输出的 concat 向量的方差是单个头输出的 h 倍，如果 W_o 不够“保守”，会放大噪声。我的做法：用 nn.init.xavier_uniform_ 之后再乘以 0.1 的缩放因子。 这个魔改让 100 步内的 loss 下降更稳定。别信源码里的默认参数，那都是给理想数据准备的。
写到这里，我突然意识到，当初恨透的多头注意力，如今反而成了工具箱里最趁手的兵器。它粗糙、暴力、难以调参，但一旦驯服，那种从数据中自动解耦出多重语义的工程美感，真的会让人上瘾。所以，下次遇到说“单头够用”的同事，请把这篇文章甩给他——顺便看看他会不会掉进我掉过的坑。
