Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers

技術要約：Muonで訓練されたTransformerにおける表現–読み出しインターフェースでの、Grokking後の崩壊について

問題提起
本論文は、モジュロ演算（具体的には加算および減算）で訓練されたTransformerモデルにおいて、「grokking（グロッキング）」——汎化性能が訓練精度飽和からるかに後に現れる現象——の後に発生する、特定の不安定性を調査している。先行研究では、Muon（直交化オプティマイザ）がAdamWよりもgrokkingを加速させることが確立されているが、本研究は、Muonで訓練されたモデルが、grokking後に汎化性能の再帰的かつ急激な崩壊に頻繁に陥ることを報告している。決定的なことに、この失敗はMuonに限ったことではなく、AdamWモデルも不安定性を示す。本研究は、これら両方のオプティマイザが5つのシードに対して厳密に安定しているわけではないことを示している。選択されたAdamWのベースラインは、5つのシードのうち4つで閾値未満の評価（最小27.59%に到達）を記録した一方で、Muonは5つのシードすべてでそれらを記録した。核心となる問題は、訓練セットが解決された状態であるにもかかわらず、これらのモデルが汎化能力を失うことである。また、標準的なメカニズム的進捗指標（スペクトル支持およびパワー分布）は、回路が正常であることを示し続け、崩壊を検知できない。

手法
著者らは、隠れ層の重み行列をMuonによって更新し、トークン/位置埋め込みおよび出力アンエンベディング（読み出し）をAdamWによって更新するという、分割オプティマイザ訓練レジームを採用している。本研究では、デコーダーのみのTransformer（深さ1、2、4）を用い、モジュロ演算タスク（$a+b \pmod p$ および $a-b \pmod p$、ここで $p=113$、および $p=97$ のバリアント）を利用している。

主な手法構成要素は以下の通りである：

1. マッチド・ブランチング（Matched Branching）: 因果関係を特定するために、ビット単位で同一のパラメータ状態から軌跡を分岐させる。一方のブランチは標準的な最適化を継続し、他のブランチは特定のパラメータグループ（隠れ行列、埋め込み、または読み出し）を凍結して必要性をテストする。
2. フーリエ分解: 操作グリッド上の2次元フーリエモードを用いて、モデルの内部残差状態を分解する。著者らは、特定の「ファミリー」のモード（例：加算ファミリー $(k, k)$、減算ファミリー $(k, -k)$）を定義し、フィルタリング介入を用いて、これらのファミリーのタスク性能に対する十分性と必要性を測定する。
3. 介入スイート: 学習されたアルゴリズムの正確な性質と、不整合の源泉を特定するために、位相スクランブリング、周波数再配置、およびクロス読み出し置換などの介入を適用する。
4. アブレーション: スペクトルの分散と不安定性の役割をテストするために、Muonの正規直交化パスのアブレーションを行う。

主要な知見と結果

• 崩壊の性質: この失敗は「隠れ層–読み出しの不整合（hidden–readout misalignment）」である。訓練損失がゼロに近づくと勾配は消失するが、2つのオプティマイザは異なる反応を示す。Muonの更新ステップサイズは勾配の大きさにほとんど依存しない（弾力性 $\approx -0.03$）が、AdamWのステップサイズは比例し続ける（弾力性 $\approx +1.5$）。これにより、隠れ表現と読み出しは、パラメータあたり8.0倍の速度で乖離していく。訓練セットは依然として解決されているが、読み出しが、共に訓練されたはずの表現をデコードできなくなっている。
• マスキング vs 回路の失敗: 本論文は2つの失敗モードを区別している。回路の失敗においては、フーリエファミリー自体がタスクの計算を停止する。回路のマスキングにおいては、タスクに整合したファミリー（例：加算ファミリー）は単独ではタスクを完璧に解き続ける（精度100%）が、その寄与がフルモデルにおける敵対的な剰余（adversarial remainder）によって圧倒される。著者らは、タスクに整合した成分のみをリスケールすることでフルモデルの精度が完全に回復することを示し、回路は健全であったが「票を奪われた（outvoted）」ことを証明した。
• オプティマイザの違い: Muonはgrokkingに有意に早く到達するが（中央値 13,011 ステップ vs AdamWの 30,486 ステップ）、すべてのテスト構成において普遍的に不安定である。AdamWも5つのシードにおいて不安定であるが、崩壊の深刻度（閾値未満の評価回数で測定）はMuonよりも2桁低い。単一シードのスウィープにおいて、AdamW構成は学習率に応じて深刻度が変化する条件付きの不安定性を示したが、5つのシードにわたる再現実験により、選択された「安定な」ベースラインがシード間で厳密に安定しているわけではないことが確認された。Muonの不安定性は、その更新ルールが表現を広いスペクトルに拡散させること（326個の実効共役ペア vs AdamWの約5個）に起因しており、これがタスクに整合したファミリーの総パワーに対するシェアを小さくしている。
• ファミリーのタスク選択: オプティマイザではなく、タスクがフーリエファミリーを選択する。加算で訓練されたモデルは $(k, k)$ モードを使用し、減算で訓練されたモデルは $(k, -k)$ モードを使用する。ファミリー内の基底は正確であるが共有はされず、同じシードから訓練された2つのモデルは、異なる周波数と位相を使用するように分岐し、互いの表現をデコードできず、読み出しがほぼ直交するようになる。
• 深さと整合性: 深いモデルでは、フーリエコードは最後から2番目のブロックのMLPによって書き込まれるが、アンエンベディングにとって「判読可能」になるのは最終ブロックである。不整合は、具体的には最終ブロックの残差とアンエンベディングの間のインターフェースで発生する。
• 防止策: 回路が形成された後に、隠れ行列または非隠れパラメータ（埋め込みと読み出し）のいずれかを凍結することで、崩壊を完全に防ぐことができる。「Stable Muon」では、grokking後に埋め込みと読み出しを凍結することで、5つのシードおよび451,400ステップのポストgrokking期間にわたって厳密な安定性を達成したが、標準的なMuonは繰り返し失敗する。

意義と主張
本論文は、分割オプティマイザ訓練およびポストgrokkingのダイナミクスにおける根本的な限界を特定していると主張している。残差ストリームには特権的な基底は存在せず、損失関数が平坦化すると、損失は表現と読み出しの間の特定の整合性を維持するための圧力を提供しない。著者らは以下の通り主張している：

1. 標準的な指標は不十分である: スペクトル指標（支持およびパワー分布）は、崩壊の間も不変であり得る。なぜなら、失敗の本質はアルゴリズム自体の喪失ではなく、基底の変化（不整合）だからである。
2. Grokkingは可逆的なマスキングである: 汎化から崩壊への遷移は、記憶から汎化への遷移の逆プロセスである。どちらも、タスクに整合したファミリーの表現パワーのシェアが、敵対的な剰余に対して上昇または低下することを伴う。
3. インターフェースが失敗の所在である: 不安定性は、隠れ表現や読み出し単独の特性ではなく、それらの結合的な識別（joint identification）の特性である。片側を固定すればシステムは安定する。
4. Muonの速さは代償を伴う: Muonは回路の発見を加速するが、その更新ルールは、標準的な訓練条件下では読み出しとの不整合に対してより脆弱な表現を作り出す。そのため、介入（凍結）が必要となる。

本研究は新しい応用を提案するものではなく、メカニスティック・インタープリタビリティにおける特定の失敗モードに関するメカニズム的説明を提供するものである。これは、モデルが正しい回路を保持していても、表現のドリフトによってそれを利用できなくなる可能性があることを強調している。