Exclusive Self-Attention Beyond AdamW: Stability and Generalization under Muon Optimization

技術要約：AdamWを超えた排他的自己注意（Exclusive Self-Attention）

問題提起
Transformerにおける標準的な自己注意メカニズムでは、位置 $i$ におけるトークンの出力が、自身のバリューベクトル ($v_i$) を含むバリューベクトルの加重混合となることを許容している。この設計は表現力に富む一方で、残差経路（residual pathways）がすでに局所的なトークン情報を保持しているにもかかわらず、アテンションの出力が $v_i$ と強く整列したままになるという「表現上のショートカット」を生み出す可能性がある。この現象は「アテンション類似バイアス（attention similarity-bias）」と呼ばれ、Zhai [2] による排他的自己注意（Exclusive Self-Attention; XSA）によって対処された。XSAは、アテンション出力から自身のバリューベクトルに平行な成分を明示的に除去するものである。

しかし、オリジナルのXSA研究は、AdamWオプティマイザを用いたNanoGPTトレーニングスタック内でのみ、この手法を評価していた。著者らは、Muonのような代替オプティマイザとの互換性が未解決の問いであることを明示的に指摘している。これは、現代の小規模言語モデル研究（例：nanochat スタック）が、行列パラメータにはMuon、埋め込みにはAdamWを用いるといったハイブリッドな最適化戦略や、バリュー埋め込み、回転式位置エンコーディング（RoPE）といったアーキテクチャ的特徴をますます活用していることから、極めて重要なギャップとなっている。XSAの幾何学的制約が、これらの現代的なMuonベースのトレーニングレジームに転用された際、安定性または有益性を維持できるのかは不明である。

手法
本研究は、nanochat トレーニングスタックを用いた制御された再現および転用実験を行う。手法の詳細は以下の通りである：

1. 実装: パラメータ化された形式のXSAをnanochatに統合した。操作は $z_i = y_i - \alpha \cdot \text{proj}_{v_i}(y_i)$ と定義される。ここで、$y_i$ は標準的なアテンション出力、$v_i$ はトークンのバリューベクトルであり、$\alpha \in [0, 1]$ が投影強度を制御する。$\alpha=1$ のとき、自己整列成分は完全に除去される。
2. 正当性の検証: 実装は、$\alpha=0$ においてXSAパスが標準的なアテンションと等価であることを数値的に検証した。短時間のトレーニング走行により、ベースラインと $\alpha=0$ のパスが一致した初期損失挙動を示すことを確認した。
3. 実験設定: 同一のトレーニング条件（データ順序、オプティマイザファミリー、学習率スケジュールの同一性、およびコンテキスト長）の下で、2つのモデル深度を評価した：
   • d12: 12レイヤー、隠れ次元 768、12ヘッド。
   • d24: 24レイヤー、隠れ次元 1024、16ヘッド。
   • オプティマイザ: 行列パラメータはMuonで最適化され、その他のパラメータはAdamWを使用する。
   • データ: ClimbMix-400B。
4. 指標: 本研究では、トレーニング損失、検証時のbits-per-byte (bpb)、ダウンストリームの集計スコア（Base CORE および ChatCORE）、およびアテンションの診断（出力と自己バリュー間のコサイン整列、対角挙動）を追跡する。

主な貢献

• Muonスタックへの再現: オリジナルの研究で残されたオプティマイザの互換性に関する問いに直接答える形で、XSAをMuonベースのトレーニングスタックに初めて統合した。
• パラメータ化された実装: 標準的なアテンション、完全なXSA、および強度パラメータ $\alpha$ による部分的な投影を可能にする柔軟な実装、ならびにレイヤー選択制御を実現した。
• 正当性とインストルメンテーション: $\alpha=0$ の極限値の検証、および初期のトレーニング比較を歪めていたロギングのアーティファクト（ステップインデックスおよび平滑化の慣習）の修正を行った。
• 2段階の深度評価: 包括的な指標セットを用いて、d12およびd24の深度における比較分析を行った。
• メカニズムおよびシステム分析: アテンションの自己整列に関する診断、および実行時のオーバーヘッドの評価を実施した。

結果
実証結果は、XSAはMuonベースのスタックにおいて安定しているものの、モデルの深さや特定の指標に応じてダウンストリームの利益が混在することを示している：

• 最適化の安定性: XSAおよびベースラインモデルのトレーニング損失の軌跡は、両方の深度において密接に一致している。発散、NaN、または不安定性の兆候は見られず、XSAがMuonの更新幾何学を乱さないことが確認された。
• 検証パフォーマンス (Bits-per-Byte): XSAは、保持された検証用bpbを一貫して改善した：
  • d12: 0.8484 から 0.8457 へ減少（相対改善率 0.32%）。
  • d24: 0.7193 から 0.7171 へ減少（相対改善率 0.31%）。
• ダウンストリーム指標 (CORE):
  • d24: XSAは Base CORE (0.2593 から 0.2606、+0.50%) と ChatCORE (0.3638 から 0.3682、+1.21%) の両方を改善した。
  • d12: bpbは改善したが、Base CORE は 0.1602 から 0.1508 へ減少した (-5.87%)。
• メカニズム診断: アテンションの診断では、XSAの下で自己整列および対角挙動の減少が示されており、これは文脈の混合を促進するという意図されたメカニズムと一致している。
• 実行時オーバーヘッド: オーバーヘッドは控えめであると特徴付けられる。初期の推定では約5%とされていたが、繰り返しのスループット測定により、その差は短時間の実行におけるシステムレベルの変動に近いことが示され、禁止的なペナルティはないことが示唆された。

意義と主張
本論文は、排他的自己注意（XSA）が、Muonベースのトレーニングスタックにとって互換性のあるアーキテクチャ的帰納バイアスであることを結論付けているが、その有用性は普遍的ではない。主要な意義は、XSAの表現的制約が、トレーニングを不安定化させることなく、現代的な最適化レジームにクリーンに転送できることを示した点にある。

しかし、著者は、ダウンストリームの利得がスケールおよび指標に依存することを強調している。両方の深度における検証bpbの改善は表現上の利益を示唆しているが、d12におけるBase COREの低下は、自己整列を除去することの利益が、モデルの深さ、局所情報を保持するための残差/バリュー経路の能力、あるいは特定の評価タスクに依存している可能性を示している。本研究は、XSAを標準的なアテンションに対する普遍的に優れた代替物として主張するのではなく、特定の文脈（本研究においては特にd24）において利益をもたらす、安定した低攪乱な制約として提示している。

本研究は、バイト単位の再現というよりも、ロバスト性チェックとしての役割を果たしており、現代の言語モデル研究におけるアーキテクチャの変更、オプティマイザの幾何学、および評価指標の間の複雑な相互作用を浮き彫りにしている。