【大模型】Transformer 完整数据流伪代码：第二期* 单头 vs 多头注意力，到底差在哪？

上一篇我们用极简单头注意力讲透了 Transformer 编码器-解码器完整数据流。但细心的读者会发现原始论文用的是多头注意力那单头和多头到底是什么关系把多头的头数设为 1是不是就等于原生单头这一篇我们从代码改造、数学推导、特性对比、选型建议四个维度把这个问题彻底讲透。一、从单头到多头代码改造的 4 个核心变化点我们以上一篇的「原生极简单头注意力」为基准看标准多头做了哪些增量修改。基准原生极简单头上一篇入门版defself_attention(x,maskNone):极简原生单头默认无输出融合层Qlinear_Q(x)# [B, N, d_model]Klinear_K(x)Vlinear_V(x)returnscaled_dot_product_attention(Q,K,V,mask)升级标准多头注意力defmulti_head_attention(x,maskNone,num_heads8):batch_size,seq_len,d_modelx.shape# 1. 整体QKV投影和单头完全一致维度仍为d_modelQlinear_Q(x)Klinear_K(x)Vlinear_V(x)# 【新增1】维度拆分输出通道切分给多个头d_kd_model//num_heads QQ.reshape(batch_size,seq_len,num_heads,d_k).transpose(1,2)KK.reshape(batch_size,seq_len,num_heads,d_k).transpose(1,2)VV.reshape(batch_size,seq_len,num_heads,d_k).transpose(1,2)# 注意mask 需要扩展维度以适配多头# 原始 mask shape: [seq_len, seq_len] → 广播为 [batch, num_heads, seq_len, seq_len]# 【变化2】单路计算 → 多路并行独立计算attn_outputscaled_dot_product_attention(Q,K,V,mask)# 【新增3】维度还原多头结果拼接回d_modelattn_outputattn_output.transpose(1,2).reshape(batch_size,seq_len,d_model)# 【新增4】输出融合线性层 Wooutputlinear_out(attn_output)returnoutput改造变化总结视图变换新增reshape transpose仅重组张量形状不引入新参数目的是实现多子空间并行计算计算路径从单路注意力变为多头并行注意力各自独立计算权重结果拼接注意力计算完成后将多个头的输出拼接回完整维度融合层新增linear_outWo 矩阵负责跨头信息整合——这是「标准多头」和「极简单头」最核心的参数差异注意底层缩放点积注意力的计算逻辑完全没变单头、多头都调用同一个核心函数。二、数学推导num_heads1 的多头 ≠ 原生极简单头这是最容易踩的认知误区我们从参数量和效果本质两方面证明。1. 参数量对比注意以下对比基于工业界默认实现。原生单头注意力通常省略输出融合层 Wo而标准多头注意力始终包含 Wo。若给原生单头也加上 Wo则两者参数量相等。统一设定dmodel512d_{model}512dmodel​512原生极简单头无 Wo仅包含 Wq、Wk、Wv 三个投影矩阵参数量 3 × d_model × d_model 3 × 512 × 512 786,432多头num_heads1含 Wo包含 Wq、Wk、Wv Wo 四个投影矩阵参数量 3 × d_model × d_model d_model × d_model 4 × 512 × 512 1,048,576 结论num_heads1的多头比原生极简单头多了整整一组输出投影矩阵参数量多出 1/3数学上完全不等价。2. 为什么多头效果更好不是因为维度拆分本身而是多独立子空间解耦。单头只有一组权重所有语义关联语法、指代、长距离依赖等被迫挤压在同一个特征空间学习互相干扰多头相当于多个独立的「专家小组」每个头在自己的子空间里专注捕捉一类关联最后通过 Wo 融合结论表达能力显著更强。原始 Transformer 论文的消融实验也验证了这一点单头相比 8 头标准设置翻译任务 BLEU 值下降约 0.7-1.0多头确实带来明确的效果收益。三、优缺点对比表维度原生单头注意力多头注意力8头标准表达能力单一特征空间仅能学习一组注意力模式上限低多独立子空间可同时捕捉多种语义依赖上限高参数量少仅3组投影矩阵默认无WoQKV 参数量与带 Wo 的单头持平整体略高于极简单头计算复杂度低单路矩阵运算理论 FLOPs 与单头完全相同仅增加张量视图变换reshape/transpose的访存开销可忽略不计训练稳定性高收敛平稳相对更敏感头数过多易出现训练波动可解释性中等单一权重图直观但混合多种语义模式解读时易混淆较强不同头可自发 specialize 到不同功能句法/指代/语义便于针对性分析但需逐一检查各头工程实现极简代码量少边界处理简单维度变换多掩码、KV Cache 等边界处理更复杂四、选型指南什么时候用单头什么时候必须多头优先选择单头的场景轻量小模型百万级参数量以内多头收益微弱单头性价比更高简单任务文本分类、情感分析等浅层语义任务单头足够胜任原型快速验证快速验证算法思路简化代码优先跑通全流程极端算力受限边缘设备、低算力推理场景追求极致速度与显存占用必须使用多头的场景生成类任务文本生成、翻译、对话等复杂序列任务需捕捉多种语义依赖中大型模型千万参数以上单头表达能力会成为明显瓶颈长文本任务长文本任务中多头注意力更有能力同时捕捉局部语法关系和长距离语义依赖单头容易顾此失彼标准工程实现对齐工业界主流保证模型迁移与复现性实用建议无特殊理由时默认选择 8/12/32 头与 d_model 匹配保证 d_k 为整数的标准多头这是业界充分验证的最优解。原始论文也证实头数过多或过少都会降低性能d_k64 是经验平衡点。附录多头版完整 Transformer 伪代码参考按需查看核心改动为将单头替换为多头注意力其余整体流程与第一篇完全一致defencoder_single_layer(x):attn_outputmulti_head_attention(x,maskNone)# 多头自注意力xadd_norm(x,attn_output)ffn_outputfeed_forward(x)xadd_norm(x,ffn_output)returnxdefdecoder_single_layer(x,encoder_output,causal_mask):masked_attnmulti_head_attention(x,maskcausal_mask)# 多头掩码自注意力xadd_norm(x,masked_attn)cross_attnmulti_head_attention(x,k_and_vencoder_output)# 多头交叉注意力xadd_norm(x,cross_attn)ffn_outputfeed_forward(x)xadd_norm(x,ffn_output)returnx下一篇预告搞懂了多头注意力的底层逻辑下一篇我们将讲解推理阶段最重要的工程优化KV Cache仅在推理阶段生效训练阶段不需要看看大模型是怎么通过缓存历史 K、V 向量把推理速度提升数倍的。