多頭注意力殘差

摘要

Transformer透過單一的加性殘差流在深度維度上傳播資訊：每個子層僅讀取最近的狀態。注意力殘差透過讓每個子層經由學習到的softmax進行注意力運算，放寬了此限制。然而，該讀取操作使用一個跨整個寬度共享的單一查詢，因此每個特徵子空間都必須透過同一個分佈來讀取深度歷史。此強迫性取捨的代價會隨著子空間在應讀取哪些層面上的分歧而增長，而分歧又隨模型寬度增加而加劇。我們提出多頭注意力殘差（Multi-Head Attention Residuals, MHAR）：將路由查詢重塑為H個按子空間劃分的頭，每個頭對深度歷史各有其softmax。讀取操作因此成為區塊對角（block-diagonal），重塑不增加任何參數且計算量可忽略，而H = 1時則精確還原為注意力殘差。在經過去重、品質過濾且以STEM和程式碼為主的Nemotron退火語料庫上從頭訓練，MHAR在100M、350M和1B規模下相較於標準Transformer均改善了驗證損失（分別為-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倍，同時維持接近基線的峰值記憶體。利用δ注意力殘差進行保恆等變換（identity-preserving conversion）支援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.