技術要約：Zeta – 行列最適化のための座標適応型プリコンディショニングによる二重ホワイトニング

問題提起

大規模ニューラルネットワークの学習は、要素ごとの適応を超えて、重みパラメータの構造的特性を利用する行列認識型オプティマイザへと、ますますシフトしています。その代表例が Muon です。Muonは、ニュートン・シュルツ反復（Newton–Schulz iteration）を用いたスペクトルホワイトニングによって、重みの更新に対して直交制約を課し、特異値を均等化することで、特定の方向が更新を支配することを防ぎます。

しかし、本論文は、Muonや同様のスペクトル手法における、極めて重要でありながら過小評価されている脆弱性を指摘しています。それは、ニュートン・シュルツ反復は入力のコンディショニング（条件付け）に対して非常に敏感であるという点です。深層ネットワークの学習における生のモメンタム行列は、しばしば激しい座標ごとのスケール不均一性を示すことがあります。これらの行列では、少数の座標が不釣り合いに大きなマグニチュードを持ち、それが特異スペクトルを歪ませてしまいます。このような無条件の行列がニュートン・シュルツのステップに入ると、直交化プロセスが劣化し、勾配の幾何学を忠実に反映しなくなる、歪んだ更新方向が生じてしまいます。著者らは、スペクトルホワイトニングには、スペクトルホワイトニングが信頼性を持って機能するために、スケール正規化された入力が必要であると主張しています。

手法：Zeta オプティマイザ

この構造的なギャップに対処するため、著者らは、座標ホワイトニングとスペクトルホワイトニングを厳密に順序付けられた非チューナブルなパイプラインとして構成した、二重ホワイトニング・オプティマイザである Zeta を提案します。核心となる洞察は、スペクトルホワイトニングは統計的な安定化の上に構築されなければならないということです。つまり、座標ホワイトニングが、スペクトルホワイトニングが確実に機能するために必要な統計的な等方性を確立するのです。

Zetaのパイプラインは次のように動作します：

1. 座標ホワイトニング（プリコンディショニング）：
   スペクトル段階に入る前に、Zetaはモメンタム行列に対して座標ごとの正規化を適用します。これは、Adamと同様に、一次モーメント（$M_t$）と二次モーメント（$V_t$）の指数移動平均を保持します。その後、モメンタム行列は、実行中の二次モーメントによって要素ごとに正規化されます：
   $$\tilde{G}_t = \frac{M_t}{\sqrt{V_t} + \epsilon}$$
   このステップは、行列内のスケールの格差を均質化し、行列を統計的に等方的な状態へと変換します。理論的には、これは近似的な対角プリコンディショナとして機能し、スペクトルを歪ませる原因となるスケールの不均衡に対処します。

2. スペクトルホワイトニング（直交化）：
   この均質化された行列 $\tilde{G}_t$ が、次にニュートン・シュルツ反復に投入されます。入力から極端な要素ごとのスケール格差が取り除かれているため、その特異スペクトルはより良好に条件付けられています。反復は忠実な直交近似へと収束し、特異値を均等化して、バランスの取れた更新方向を生み出します。

3. 更新の構築：
   得られた直交行列 $U_t$ は、スペクトルの方向を変えることなく、標準的なオプティマイザの設定（例：AdamWブランチ）と更新のマグニチュードを適合させるために、RMS正規化された係数を用いてスケーリングされます。

4. デュアルパス実行：
   行列値パラメータ（例：アテンションおよびフィードフォワード投影）に対しては、完全な二重パイプラインが適用されます。非行列パラメータ（例：バイアス、LayerNormのスケール）に対しては、スペクトルステップがスキップされ、AdamWスタイルの更新にフォールバックされます。

理論的貢献

本論文は、設計を裏付ける3つの主要な理論的保証を提供しています：

• 近似ホワイトニング： 座標正規化ステップは、ヒューリスティックなスケーリングを超え、理想的なホワイトニング変換の効率的な対角近似であることを証明しています（定理1）。
• 直交化誤差の低減： 座標ホワイトニングが、ニュートン・シュルツ反復への入力の条件数を厳密に改善することを証明しています。その結果、生のモメンタムに対して操作する場合と比較して、スペクトル段階の直交化誤差が厳密に減少します（定理2）。
• スケールの保存： 最終的なRMSスケーリングは、更新のスペクトル方向を維持しながら、ネットワークの他の部分とグローバルなマグニチュードを同期させます（定理3）。

実証結果

Zetaは、言語モデリング（0.6Bから8Bパラメータ）、混合エキスパート（MoE）アーキテクチャ、およびビジョンタスクにおいて評価され、一貫して強力なベースライン（AdamW、Muon、AdaMuon）と同等またはそれを上回る性能を示しました。

• 収束速度： プリトレーニングタスクにおいて、ZetaはMuonおよびAdamWよりも速い収束を示しました。例えば、1.7Bパラメータのモデルにおいて、Zetaは同じ損失レベルに到達するために、AdamWに対して1.67倍の高速化（相対的）を達成し、Muon（1.52倍）およびAdaMuon（1.61倍）を上回りました。8Bモデルでは、Muonの1.19倍に対し、Zetaは1.25倍の高速化を達成しました。
• 効率性： Zetaのイテレーションあたりの時間コストは、AdamWと同等（例：16 GPUを用いた8Bモデルで、24.46秒に対し25.47秒）であり、二重ホワイトニングのプロセスが大幅な実行時間のペナルティを課さないことを示しています。
• MoEパフォーマンス： 疎なエキスパート活性化が大幅な勾配の不均一性を生み出すMoEアーキテクチャにおいて、Zetaは明確な優位性を維持し、競合モデルよりも少ない変動でより低い損失へと収束しました。
• ダウンストリームタスク： 12の言語理解ベンチマーク（MMLU、GSM8kなど）において、Zetaで訓練されたモデルは、1.7Bおよび8Bの両方のスケールで最高の平均精度を達成しました。同様に、CIFAR-100ビジョンのタスクにおいて、ZetaはViT-TinyおよびViT-Baseの両方で最高のTop-1検証精度を達成しました。
• ハイパーパラメータの堅牢性： モメンタム係数（$\beta_1, \beta_2$）のグリッドサーチにより、Zetaがハイパーパラメータの選択に対して堅牢であることが示されました。テストされた範囲内での最終的な訓練損失の変動は1.2%未満でした。

意義と主張

本論文は、既存の行列認識型オプティマイザにおける根本的なミスマッチ、すなわち、制御されていない座標ごとのスケールを持つ信号に対して、強力な幾何学的操作（スペクトルホワイトニング）を適用しているという問題を解決すると主張しています。カイ二乗一様性テストを通じて要素内のスケール不均衡の蔓延を検証し、座標ホワイトニングがこれを効果的に修正することを実証することで、著者らは、直交化の前にスケール不均衡を解消することが、安定したスペクトル手法の前提条件であることを確立しました。

Zetaの意義は、以前は競合していた2つの最適化パラダイム（座標ごとの適応スケーリングとスペクトル直交化）を、単一の、論理的に不可欠なパイプラインへと統合したことにあります。操作の順序は、チューナブルなハイパーパラメータではなく、数学的な依存関係です。すなわち、入力が適切に条件付けられることを保証するために、スペクトルホワイトニングの前に座標ホワイトニングが行われなければなりません。このアプローチは、行列レベルの最適化による幾何学的な利点を損なうことなく、より速い収束とより良い汎化を実現します。