让Transformer每层能更智能地读取历史信息,提升模型性能。
Multi-Head Attention Residuals

- 将注意力残差拆分为多头,每头独立选择历史层
- 在100M到1B参数模型上均降低验证损失,最大提升0.149
- 头数设为8效果最佳,且可中途转换不丢失精度
Transformers通过单一加性残差流在深度间传递信息:每个子层仅读取最新状态。注意力残差放宽此限制,使每层可通过学习的softmax访问历史。但该查询在整个宽度上共享,导致所有特征子空间必须通过同一分布读取历史。当子空间对读取哪层存在分歧时,这一强制妥协的代价上升,且随模型宽度增长。本文提出多头注意力残差(MHAR):将路由查询重塑为H个子空间独立的头,每个头拥有自己的深度历史softmax。读取变为分块对角形式,重塑操作零参数、开销极小,且当H=1时精确还原原注意力残差。在经去重、质量过滤、侧重科学与代码的Nemotron基语料上从头训练,MHAR在100M、350M和1B模型上分别将验证损失降低-0.061、-0.149和-0.140,优于四种方法。增益随规模递增。头数是真实设计维度而非随意调参:验证损失随H呈U形变化,各尺度最优值为H=4或8,采用H=8于大模型;超过此值(如H=16)会部分抵消收益。直接探测训练后查询确认子空间分歧是根本驱动。融合Triton路由核使训练吞吐量从0.2-0.5倍提升至0.55-0.88倍,峰值内存接近基线。使用保持身份的增量残差转换支持8B模型中段训练,带来GSM8K+3.2和GPQA+3.1提升。
原文摘要 · Abstract (English)
Transformers propagate information across depth through a single additive residual stream: every sublayer reads only the most recent state. Attention residuals relax this by letting each sublayer attend, through a learned softmax. However, that read uses a single query shared across the entire width, so every feature subspace must read the depth history through one distribution. The cost of this forced compromise grows with how much the subspaces disagree about which layers to read, and disagreement grows with model width. We introduce Multi-Head Attention Residuals (MHAR): the routing query is reshaped into H per-subspace heads, each with its own softmax over the depth history. The read becomes block-diagonal, the reshape adds zero parameters and negligible compute, and H = 1 recovers attention residuals exactly. Trained from scratch on a deduplicated Nemotron-based anneal corpus that is quality-filtered and STEM- and code-heavy, MHAR improves validation loss over a standard Transformer at 100M, 350M, and 1B (-0.061, -0.149, and -0.140). It achieves the best result among four methods in every setting, with the gain increasing from 100M to the larger scales. The head count is a real design axis rather than a free knob: validation loss is U-shaped with respect to H, with a flat optimum at H = 4 or H = 8 across scales. We adopt H = 8 for large-scale models; over-splitting beyond this point (H = 16) consistently gives back part of the gain. A direct probe of the trained queries confirms that learned subspace disagreement is the underlying driver. Fused Triton routing kernels increase attention-residual training throughput from 0.2-0.5x to 0.55-0.88x of the baseline while maintaining near-baseline peak memory. An identity-preserving conversion using delta attention residuals supports 8B mid-training, yielding improvements of +3.2 on GSM8K and +3.1 on GPQA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。