Multi-Head Attention Residuals：Transformer の残差経路を多頭化する新提案

導入

Transformer の残差接続は、しばしば単なる安定化の仕組みとして語られます。しかし実際には、層をまたいで情報を運ぶ主要な通路でもあります。標準的な構造では、各サブレイヤーは基本的に直前の状態を読み取り、その結果を残差ストリームへ加えます。論文「Multi-Head Attention Residuals」は、この深さ方向の情報伝達が、モデル幅の拡大に対して十分に柔軟なのかを問い直しています。

核心ポイント

• 単一の読み取り分布が制約になる：attention residuals は、各サブレイヤーが learned softmax によって過去の深さ履歴を参照できるようにします。ただし、その routing query はモデル幅全体で共有されます。つまり、異なる特徴サブスペースも同じ層選択分布を使わなければなりません。
• 幅が広いほど不一致が増える可能性：モデルが広くなると、内部の特徴サブスペースはより専門化します。あるサブスペースは浅い層の情報を、別のサブスペースは深い層の情報を必要とするかもしれません。1 つの softmax では、その違いを折衷するしかありません。
• MHAR は残差ルーティングを多頭化する：提案手法では、routing query を H 個のサブスペース head に reshape します。各 head は深さ履歴に対して独自の softmax を持ち、隠れ次元の異なるブロックが異なる層を読めるようになります。
• 追加コストは小さい：要約によれば、この reshape はパラメータを増やさず、計算量の増加もごく小さいものです。また H=1 の場合は従来の attention residuals と完全に一致します。
• 複数スケールで改善を報告：重複除去と品質フィルタリングを行った Nemotron ベースの anneal corpus、特に STEM とコードに重いデータでゼロから学習した結果、標準 Transformer に対して 100M、350M、1B でそれぞれ -0.061、-0.149、-0.140 の検証損失改善が示されています。比較された 4 手法の中でも、各設定で最良だったとされています。
• head 数には最適域がある：H を増やせば常に良くなるわけではありません。検証損失は U 字型を示し、H=4 または H=8 付近が安定した最適域です。H=16 では一部の利得が失われると報告されています。

意義と影響

MHAR の面白さは、Transformer の大きな骨格を変えずに、残差ストリームの読み取り方だけを細かくする点にあります。token 方向の attention では、複数 head が異なるパターンを見ることはすでに標準です。MHAR は同じ考え方を深さ方向へ拡張し、特徴ブロックごとに必要な層履歴を選べるようにします。

一方で、今回の結果は特定のコーパス、規模、実装条件に基づいています。より多様なデータ構成やアーキテクチャで同じ傾向が確認されるかは、今後の検証が必要です。実装面でも、論文要約では fused Triton routing kernels によって attention-residual の学習スループットをベースラインの 0.55-0.88x まで改善し、ピークメモリをほぼベースラインに保ったとされています。効率的な kernel が実用上の鍵になるでしょう。

総じて、MHAR は「残差は 1 本の受動的な通路」という前提を見直す低侵襲な提案です。幅の広い LLM でサブスペース間の層選択の違いが重要になるなら、残差ルーティングの多頭化は有望な設計軸になり得ます。

出典：Hugging Face Daily Papers