Multi-Head Attention Residuals：让 Transformer 的残差通路也拥有多头路由

导语

Transformer 的残差连接看似只是“把上一层结果加回来”，但它实际上承担着跨层传递信息的主干作用。标准结构中，每个子层主要读取最新的隐藏状态；如果较早层的信息已经被后续变换稀释，模型只能通过这一条不断更新的残差流继续向前传播。论文《Multi-Head Attention Residuals》关注的正是这个被长期默认的设计：深度方向上的信息路由是否也应该像 token 维度的注意力一样，更细粒度、更可选择？

核心要点

• 问题来自“单一路由分布”：已有 attention residuals 让子层可以通过 learned softmax 回看不同深度的历史状态，但它使用一个跨整个宽度共享的 query。换言之，不同特征子空间即使想读取不同层，也必须共用同一套深度权重。
• 宽模型中妥协成本更高：论文认为，模型越宽，内部特征子空间的功能分化越明显，它们对“应该读取哪一层”的偏好也更容易不一致。单一 softmax 会把这些偏好压成折中方案，从而限制表达能力。
• MHAR 的做法很直接：Multi-Head Attention Residuals 将路由 query reshape 成 H 个面向子空间的 head，每个 head 都有自己的 softmax，用来在深度历史上选择信息。读取结构因此变成块对角形式。
• 几乎不增加额外负担：根据摘要，MHAR 的 reshape 不增加参数，计算开销也很小；当 H=1 时，它会退化为原来的 attention residuals，因此可被看作后者的自然泛化。
• 实验显示稳定收益：作者在去重、质量过滤、偏 STEM 和代码的 Nemotron-based anneal corpus 上从头训练模型。相对标准 Transformer，MHAR 在 100M、350M、1B 规模上分别带来 -0.061、-0.149、-0.140 的验证损失改进，并在四种方法比较中各设置下表现最好。
• 头数不是越多越好：验证损失随 H 呈 U 形，H=4 或 H=8 附近较优；H=16 会回吐部分收益。作者在大规模模型中采用 H=8。

意义与影响

这项工作有意思之处在于，它没有重新设计注意力主干，也没有引入复杂外部记忆，而是把“多头”思想移到了残差深度路由上。对于大语言模型而言，深层网络内部不同特征可能承担语法、事实、推理、代码模式等不同功能，让它们用各自的深度读取策略，符合模型宽度扩展后的直觉。

不过，素材给出的结果仍主要来自特定训练语料和若干规模设置。MHAR 是否能在更广泛的数据配比、架构变体和训练预算下保持收益，还需要更多复现。其工程侧也并非完全免费：论文提到通过 fused Triton routing kernels 将 attention-residual 训练吞吐提升到基线的 0.55-0.88x，并保持接近基线的峰值显存，这说明高效实现仍是落地关键。

总体看，MHAR 提供了一个低参数、低侵入的 Transformer 改造方向：让残差流不再只是单条被动通道，而成为可按子空间分头选择的深度信息路由机制。对于追求训练效率和模型质量平衡的 LLM 架构研究，这是一条值得继续验证的路线。

来源：Hugging Face Daily Papers