多头注意力残差

摘要

Transformer通过单一的加性残差流沿深度方向传播信息：每个子层只读取最近的状态。注意力残差让每个子层通过一个学习到的softmax进行注意力操作，从而缓解了这一限制。然而，这种读取使用一个在整个宽度上共享的单一查询，因此每个特征子空间都必须通过同一个分布来读取深度历史。这种被迫妥协的代价随子空间在读取哪些层上的分歧程度而增长，且分歧随模型宽度增大。我们引入了多头注意力残差（MHAR）：将路由查询重塑为H个按子空间划分的头，每个头对深度历史拥有各自的softmax。这种读取变为分块对角形式；重塑不增加任何参数，计算开销可忽略不计；H=1时精确恢复注意力残差。在基于Nemotron的去重退火语料库（经过质量过滤，且以STEM和代码为主）上从头训练后，MHAR在100M、350M和1B规模上的验证损失改进分别为-0.061、-0.149和-0.140。四种方法中，它在所有设置下均取得最佳结果，且增益从100M到更大规模不断增加。头数是一个真实的设计维度，而非一个可自由调节的旋钮：验证损失关于H呈U形，在不同规模上都在H=4或H=8处具有平坦的最优点。我们为大模型采用H=8；超过这一点的过度拆分（H=16）始终会损失一部分增益。对训练后查询的直接探测证实，学习到的子空间分歧是根本驱动因素。融合的Triton路由内核将注意力残差训练吞吐量从基线的0.2–0.5倍提高到0.55–0.88倍，同时保持接近基线的峰值显存。一种使用Δ注意力残差的保持恒等映射的转换支持8B模型的中途训练，在GSM8K上带来+3.2、在GPQA上带来+3.1的提升。

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.