Muonはgrokkingを加速した後も破綻する：9つの設定すべてで汎化崩壊が発生

Muonはモーメンタム更新をまず直交化するoptimizerで、近年はTransformerの隠れ重み行列によく用いられ、埋め込み層と出力ヘッドはAdamWに任せる構成が採用されている。新たな研究では、正規化層を持たないdecoder-only Transformerを使って`(a+b) mod 113`をテストし、この役割分担によってテスト精度95%というgrokkingの閾値をより速く超えられることを確認した。しかし、その後は9つのMuonハイパーパラメータ設定すべてで閾値を下回り、一部の実行では精度が繰り返し1%未満まで低下してから回復した。AdamWも完全に安定していたわけではないが、崩壊はより高い学習率に強く依存していた。

著者らは、原因を表現—読み出しインターフェースに特定した。最終残差表現`h`と出力行列`W_U`は、その積`W_U h`を通じてのみlogitsを決定するため、両者は関数を変えずに共同で基底変換できる。別々の更新ダイナミクスを適用すると、もともと整合していた基底が徐々にずれる可能性がある。モデルが訓練セットを解いた後、勾配はすでに約`10^-6`まで低下していたが、Muonパラメータ群におけるパラメータ当たりの移動速度は、なおAdamW群の8倍だった。まったく同一の状態から分岐させて隠れ行列を凍結した場合も、埋め込みと出力ヘッドを凍結した場合も、破綻を回避できた。後者を固定した5回の実行では、grokking後のステップを累計451,400回観測しても、再び閾値を下回ることはなかった。

Fourier介入により、演算回路そのものの損傷と「マスキング」も区別された。一部のcheckpointでは、タスクに整合した周波数成分だけなら精度は依然として100%だったが、残りの表現に圧倒され、モデル全体の精度は45.85%にとどまった。その成分を単純に増幅すると、精度は99.9%まで回復した。これは、周波数サポートや表現パワーだけを追跡していると、読み出し時の競合を見落とす可能性があることを示している。

工学的に注意すべき点は、この研究が小規模なモジュラー演算Transformerのみを対象としており、大規模言語モデルでも同様の崩壊が起きると直接推論することはできない点だ。公開repositoryには訓練コード、介入分析、191件のCSV、20件のHTMLサマリーが含まれているが、checkpointはコミットされておらず、完全な再現には再訓練が必要となる。次のステップとして、より現実的なpre-trainingおよびfine-tuningのスケールで、パラメータ群間の相対的な更新量を調べるとともに、凍結や共有基底制約が、固定的なoptimizer routingよりも堅牢かどうかを評価すべきだ。