アテンション情報でMoEのルーティングを強化する「AAR」——数学推論で+3.37ポイント改善

MoEのルーターが抱える「文脈不足」問題

Mixture-of-Experts（MoE）型言語モデルでは、各トークンをどの専門家（エキスパート）に割り振るかを決める「ルーター」が中心的な役割を担う。しかし従来のルーターは、トークンの隠れ状態（hidden state）のみを判断材料とするため、モデルが実際にどのような文脈を処理しているかという情報を十分に活用できていないという課題があった。

AARの提案：アテンション重みをルーターに接続する

この課題に対し、arxivに投稿された論文ではAttention-Aware Routing（AAR）という新しい手法が提案されている。AARは、スライディングウィンドウ内のアテンション重みから時系列的・スペクトル的な特徴量を抽出し、それをルーターの入力に追加するというアプローチをとる。これらの特徴量は隠れ状態とは独立（disentangled）しており、モデルの文脈的な状態を要約したものとして機能するとされる。

重要な点として、ベースとなるトランスフォーマーのパラメータは完全に凍結したまま、ルーティングに関するパラメータのみを学習する設計となっている。これにより、ルーティングの変化だけを唯一の変数として切り離して評価できる実験設定が実現されている。

GSM8Kで+3.37ポイントの改善を報告

OLMoEを用いた評価では、ルーティングのみをファインチューニング（SFT）したベースラインと比較して、数学的推論ベンチマークのGSM8Kにおいて+3.37ポイントの改善が得られたと報告されている。

ルーティングとアテンションが形成する「結合回路」

論文ではパフォーマンス向上にとどまらず、ルーティングとアテンションの間に結合回路（coupled circuit）が存在することも示されている。具体的には、層$l$でのルーティングの変化が残差ストリームを通じて伝播し、層$l+1$におけるアテンションシンク（attention sink）を増幅させるという現象が観測されたとされる。これはアテンション機構そのものを直接更新することなく、アテンションの形状を変化させることを意味しており、モデル内部の情報フローの相互依存性を示す知見として注目される。

誤答を短くし、正答の長さは維持

AARにはもう一つ興味深い副次効果が報告されている。誤った回答については生成が長く発散する傾向が抑制され、不正解の回答が短くなる一方で、正解の回答の長さは変わらないという非対称的な挙動が確認されたという。これはルーティングの改善が、モデルの「確信度」に応じた生成制御にも寄与している可能性を示唆している。

層の深さへの強い依存性：検索と推論のトレードオフ

AARは深さ（depth）に対して強い感度を持つことも明らかにされた。全層に一律に適用した場合、事実の検索（factual retrieval）タスクの性能が低下する一方、数学的推論における改善はネットワークの深い層に導入した場合に持続するとされる。

この感度は、検索能力と推論能力が層の深さにわたってトレードオフの関係にあることを示しており、論文では層を選択的に適用する「layer-selective AAR」が、各層においてアテンションが保持するルーティング関連情報を調べる制御されたプローブとして機能すると述べられている。

まとめ

AARは、MoEモデルのルーターにアテンション由来の文脈情報を加えることで、ベースモデルを変更せずにルーティングの質を向上させる手法として提案されている。数学推論での性能改善に加え、ルーティングとアテンションの結合メカニズムや層深度に応じた能力分布という新たな知見をもたらしており、MoEアーキテクチャの解釈可能性研究においても重要な一歩となる可能性がある。詳細はarXiv（arXiv:2609.20974）で公開されている。

---

出典: Attention-Aware Routing: Coupling Routing and Attention in MoEs