💻 computer science

Multi-Head Attention Residuals

本論文は、異なる特徴部分空間が複数のルーティングヘッドを介して深さの履歴に独立して注意を向けることを可能にし、それによって部分空間間の不一致を解消し、標準的なTransformerと比較して様々なモデルスケールにおいて検証損失およびダウンストリームタスクの性能を大幅に向上させる、パラメータフリーのアーキテクチャ変更であるMulti-Head Attention Residuals (MHAR) を導入する。

原著者： Cheng Luo, Zefan Cai, Junjie Hu

公開日 2026-07-31

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Cheng Luo, Zefan Cai, Junjie Hu

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

巨大で多層構造の図書館を建設しているところを想像してみてください。そこでは、新しいフロアが増えるたびに、新しい知識の部屋が追加されていきます。人工知能の世界では、これらの「図書館」は「トランスフォーマー」と呼ばれ、私たちが今日使っている多くのスマートなツールの背後にある脳となっています。物語や数学の問題を理解するためには、図書館は10階目で作業をしている間も、1階で何が起きたかを覚えておく必要があります。従来、これらの図書館は、すべてのフロアを繋ぐ単一の狭い廊下を持っていました。10階の部屋が過去を振り返ろうとする際、その部屋はすぐ下の階の部屋を覗き見ることしかできません。その部屋が、1階からそこまでのすべてを完璧に要約してくれていることを信じるしかないのです。それは、映画の全編を思い出すために、目の前に座っている人に「何が起きたか」を尋ねるようなものです。もしその人が詳細を見逃していたら、あなたもそれを見逃すことになります。

これを解決するために、研究者たちは最近、「アテンション・レジデュアル（Attention Residuals）」と呼ばれる新しいアイデアを試みました。単に下の階を見るのではなく、図書館のすべての部屋に、どの以前のフロアにも見ることができる魔法の望遠鏡を与えたのです。しかし、一つ問題がありました。その望遠サーはレンズを一つしか持っていなかったのです。部屋の中のあらゆる詳細――あらゆる色、音、数字――は、どのフロアに焦点を当てるかを決めるために、その同じ一つのレンズを通さなければなりませんでした。それは、全く異なる好みの友人たち（ある人はアクション映画、ある人はドキュメンタリー、またある人はアニメが好き）が、全員で一つのチャンネルに同意しなければならないようなものでした。彼らは妥協を強いられ、その妥協によって、多くの場合、誰も本当に必要としていたものを見ることができなくなりました。

この論文は、巧妙なアップグレードである「マルチヘッド・アテンション・レジデュアル（Multi-Head Attention Residuals: MHAR）」を紹介しています。研究者たちは、映画館が異なるジャンルを同時に上映するために複数のスクリーンを持っているように、図書館の望遠鏡も複数のレンズを持つべきだと気づきました。彼らは、その単一の共有レンズを、いくつかの小さく独立したレンズへと分割したのです。これにより、「アクション映画」の部分はアクション満載のフロアを見ることができ、一方で「ドキュメンタリー」の部分は歴史のフロアを見ることができ、互いに争ったり妥協したりすることなく、それぞれの役割を果たすことができるようになりました。

チームはこのアイデアをテストするために、1億パラメータの小規模なモデルから、最大10億パラメータの巨大なモデルまで、さまざまなサイズの図書館を構築しました。その結果、小さな図書館では、従来の単一レンズ・システムでも十分に機能することがわかりました。しかし、図書館がより大きく複雑になるにつれて、単一のレンズはボトルネックとなり、標準的なシステムよりもAIの性能を悪化させてしまいました。しかし、新しいマルチレンズ・システムは、図書館が成長するにつれてどんどん良くなっていきました。最大のスケールにおいて、新しい手法は、標準的なモデルと比較して、AIの混乱（検証損失として測定）を約0.063減少させましたが、旧式の単一レンズ法は、実際には0.140ほど性能を低下させてしまいました。

研究者たちはさらに、このアップグレードが、より大きな望遠鏡を作ったり、図書館にさらなる重みを加えたりする必要はなく、既存のガラスをより賢く配置する方法であることも示しました。彼らは、余計な「見回し」によって図書館が停滞しないよう、特別な超高速ツールさえも構築し、速度を元の設計とほぼ同等の速さに保ちました。最後に、彼らはこのテクニックが、すでに構築済みの巨大な80億パラメータの図書館においても有効であることを証明し、ゼロから作り直すことなく、数学の問題を解いたり難しい科学の質問に答えたりする能力を向上させました。核心となる発見は、AIモデルがより広く複雑になるにつれて、彼らは何を記憶すべきかについて合意できなくなるということであり、複数の独立した「振り返る方法」を与えることが、その潜在能力を解き放つ鍵であるということです。

技術要約：マルチヘッド・アテンション・レジデュアル（MHAR）

1. 問題提起

標準的なTransformerは、単一の加算的な残差ストリームを通じて深さ方向に情報を伝播させます。ここでは、各サブレイヤーは直前の状態のみを読み取ります（ $h_\ell = h_{\ell-1} + f_\ell(h_{\ell-1})$ ）。これにより、初期の層で計算された有用な特徴量は累積和の中に保持されますが、それらを個別にアドレス（指定）することはできなくなります。

アテンション・レジデュアル（Attention Residuals）（Kimi, 2025）は、すべての過去のサブレイヤー出力をメモリとして扱うことでこれを緩和し、各サブレイヤーが全深度履歴から学習されたソフトマックス重み付きの組み合わせを検索できるようにします。しかし、このメカニズムは構造的なボトルネックを抱えています。それは、単一の共有されたルーティング・クエリ（ $q \in \mathbb{R}^d$ ）を使用して深度履歴をスコアリングするという点です。その結果、 $d$ 個のすべての特徴量座標は、単一の分布 $\alpha$ を通じて、同じ割合で同じ層から読み取ることを強制されます。

論文では、これが**「強制された妥協（forced compromise）」**を生み出すと主張しています。モデルの幅（width）が増大するにつれ、異なる特徴部分空間（subspace）は、どの歴史的レイヤーが最も重要であるかについて、互いに異なる嗜好を発達させる可能性が高いと考えられます。単一のクエリは、これら異なる嗜好を同時に満たすことはできず、異なる深度への嗜好を単一の読み取りへと崩壊させてしまいます。この妥協のコストはモデルの幅とともに増大し、大規模化において単一クエリによるルーティングを有害にする可能性があります。

2. 手法：マルチヘッド・アテンション・レジデュアル（MHAR）

著者らは、パラメータや計算量を大幅に増やすことなく、この「強制された妥協」を排除する、アテンション・レジデュアルの厳密な一般化である**マルチヘッド・アテンション・レジデュアル（MHAR）**を提案します。

コア・メカニズム

単一のクエリ $q$ の代わりに、MHARはルーティング・クエリを $H$ 個の「部分空間ごとのヘッド」へと再構成し、各ヘッドが深度履歴に対して独自の独立したソフトマックスを持ちます。

再構成（Reshaping）: 単一のクエリベクトル $q \in \mathbb{R}^d$ を、 $(H, d/H)$ の形状を持つテンソルへと再構成します。
独立したルーティング: 各ヘッド $h$ について、モデルはソースベクトルの該当するスライスに基づいた、個別の深度分布 $\alpha^{(h)}$ を計算します。
ブロック対角の読み取り: 取得される混合物はブロック対角となります。すなわち、各部分空間が独立して深度履歴を読み取ります。
復元: $H=1$ と設定することで、元のシングルヘッド・アテンション・レジデュアルと正確に一致します。

実装の詳細

パラメータ効率: この変換はリシェイプ（再構成）であり、シングルヘッドのベースラインと比較して追加のパラメータはゼロです。
計算効率: オーバーヘッドは無視できる程度であり、深度軸に対する $H$ 個の小さなソフトマックス計算（長さ $\le 2L+1$ ）のみを含みます。
システム最適化: 深度ルーティングがメモリ帯域制限（ $O(L^2 B T d)$ のトラフィックを必要とする）を受けることを認識し、著者らは融合されたTritonカーネルを提供しています。これらのカーネルは、オンラインソフトマックスと融合されたRMSNormを実行し、キーやスタックされたソースの中間生成を排除します。これにより、ルーティング時間は torch.compile より2倍、Eager実装の参照実装より6.4倍高速化され、トレーニングのスループットをベースラインの $0.55\text{--}0.88\times$ まで引き上げます。

ハイパーパラメータの設定

ヘッド数 $H$ は、モデルのKVヘッドの数と等しく設定されます。著者らは、これがモデルの消費粒度とルーティングの粒度を一致させる、ハイパーパラメータ調整を必要としない準最適なデフォルト設定であることを発見しました。

3. 主な貢献

MHARアーキテクチャ: ルーティング・クエリを独立したヘッドに分割することで、単一クエリによる深度ルーティングの「強制された妥協」を取り除く、パラメータフリーの修正手法。
効率的なカーネル: 深度ルーティングのメモリオーバーヘッドとレイテンシを大幅に削減し、実用的なスケールでの深度ルーティングを可能にする、融合されたTritonカーネル。
恒等性を保持する変換: 「デルタ・アテンション・レジデュアル」（既存のストリームにルーティング項を加える手法）を用いて、事前学習済みモデルにMHARを移植する方法。これにより、損失のスパイクなしに中間学習が可能になります。
メカニスティックな洞察: 学習済みクエリの直接的なプロービングにより、部分空間間の不一致がモデルの幅とともに増大すること、そしてそれがシングルヘッド・ルーティングの劣化の主要な要因であることを確認しました。

4. 実験結果

事前学習（スクラッチから）

FineWeb-Edu上で3つのスケール（100M, 350M, 1B）を用いて学習：

検証損失（Validation Loss）: MHARは、標準的なTransformerおよびシングルヘッド・アテンション・レジデュアルに対して、一貫して検証損失を改善しました。
- 100M: ベースラインに対し $-0.049$ nats。
- 350M: ベースラインに対し $-0.080$ nats。
- 1B: ベースラインに対し $-0.063$ nats。
スケーリング挙動: シングルヘッド・ルーティングに対するMHARの優位性は、スケールとともに拡大します。シングルヘッド・ルーティングは100Mでは改善（$-0.039 $）しますが、350M（$ +0.055 $）および1B（$ +0.140$）では有害な状態へと劣化します。
計算量等価性: MHARは、ベースラインに対して $1.3\text{--}1.5\times$ の計算量等価な利得を達成しています。

中間学習（8B）

能力指向のコーパスでアニーリングされた8Bモデル（Marin-8B）を用い、スケールを合わせたコントロール群と比較：

恒等性の保持: 変換は最適化のショックを引き起こしません。損失曲線は最初の100ステップにおいてベースラインと一致しています。
ダウンストリームの利得: MHARは、推論タスクにおいて統計的に有意な改善をもたらしました：
- GSM8K: $+3.2$ ポイント（ペアードMcNemar $p=0.004$ ）。
- GPQA: $+3.1$ ポイント（ペアードMcNemar $p=0.038$ ）。
- MMLU、MATH、およびコードタスクにおける利得は、この特定のセットアップでは統計的に有意ではありませんでした。

アブレーション研究

ヘッド数: 最適なヘッド数 $H$ は、KVヘッドの数に従います。未学習のレジームでは、より多くのヘッドが常に有利ですが、収束したレジームでは $H=KV$ が曲線の「膝（knee）」となります。
幅 vs 深度: 学習済みクエリのプロービングにより、「幅の不一致（width disagreement）」（部分空間間でのルーティング嗜好の相違）がモデルの幅とともに単調に増大することが明らかになり、「強制された妥協」の仮説が検証されました。
ヘッドの整列: ルーティング・ヘッドをアテンションのKVグループに整列させても、任意の部分空間ルーティングと比較して測定可能なメリットはありませんでした。これは、恩恵が特定の構造への整列ではなく、ヘッドの多重性に由来することを示唆しています。

5. 意義と主張

本論文は、マルチヘッド・アテンション・レジデュアルが、Transformerが深度履歴にアクセスする方法における根本的なアーキテクチャ上の限界を解決することを主張しています。異なる特徴部分空間が独立して異なるレイヤーに注意を向けることを可能にすることで、MHARはシングルクエリ・ルーティングが失敗する規模においても効果的にスケールします。

著者らは以下を強調しています：

改善は、追加のパラメータや計算量によるものではなく、むしろ構造的なボトルネックの除去によるものです。
融合カーネルによって深度ルーティングのメモリ制限が緩和されているため、大規模化においても実用的です。
この手法は汎用性が高く、スクラッチからの学習、および恒等性を保持した中間学習の両方で機能します。
マルチヘッド・ルーティングの必要性は、モデルの幅の直接的な帰結です。モデルが大きくなるにつれ、「単一クエリによる強制された妥協」は負債となり、MHARは深度を意識したアーキテクチャをスケールさせるための重要な要素となります。