マルチヘッドアテンション残差

要旨

Transformerは、単一の加算的残差ストリームを通じて情報を深さ方向に伝播する。すなわち、各サブレイヤーは最新の状態のみを読み取る。アテンション残差は、各サブレイヤーが学習されたソフトマックスを通じてすべての過去の状態にアテンションできるようにすることで、この制約を緩和する。しかし、その読み取りは幅全体で共有される単一のクエリを使用するため、すべての特徴部分空間が単一の分布を通じて深さ方向の履歴を読み取らなければならない。この強制された妥協のコストは、各部分空間がどの層を読み取るかについてどの程度不一致であるかに応じて増大し、その不一致はモデル幅とともに拡大する。本稿では、Multi-Head Attention Residuals（MHAR）を導入する。ルーティングクエリを、深さ方向の履歴に対してそれぞれ独自のソフトマックスを持つH個の部分空間別ヘッドに再形成する。読み取り操作はブロック対角的になり、この再形成はパラメータを一切追加せず、計算コストも無視できる程度である。H=1の場合、アテンション残差と正確に一致する。重複除去され、品質フィルタリングされ、STEMおよびコードに重点を置いたNemotronベースのアニールコーパスでゼロから学習したMHARは、100M、350M、1Bの各規模で標準Transformerに対する検証損失を改善する（それぞれ-0.061、-0.149、-0.140）。これは、すべての設定で4つの手法の中で最良の結果を達成し、その利得は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.