技術的概要：Pro-KLShampoo

問題定義

現代の大規模言語モデル（LLM）の前学習は、勾配の行列構造を利用するオプティマイザに大きく依存しています。ここで 2 つの明確なフロンティアが出現しました：

1. 明示的なクリフォネ因子分解前処理条件化：KL-Shampoo がその代表例であり、真の勾配共分散とクリフォネ構造の近似との間のカルバック・ライブラー（KL）ダイバージェンスを最小化することにより前処理条件行列を推定します。
2. 勾配モーメントの直交化：Muon がその典型例であり、スペクトルノルム下での最急降下を実現するためにモーメントを直交化します。

これらのアプローチは通常、孤立して開発されています。KL-Shampoo は KL 最小化を通じて強力な理論的基盤を提供しますが、完全なクリフォネ因子を維持する必要があるため、高い計算コストとメモリコストを伴います。一方、Muon などの直交化手法は効率的ですが、KL 推定の視点を持たず、異なる幾何学的枠組み（スペクトルノルム最急降下）の下で動作します。本論文は、KL-Shampoo の統計的効率性と直交化手法の計算効率性の両方を達成するために、これらのアプローチを統合する必要性に対処します。

手法：Pro-KLShampoo

著者はPro-KLShampoo（直交化によるホワイトニングを回復した射影 KL-Shampoo）を提案します。中核的な革新は、KL-Shampoo の前処理条件行列の構造に関する観察に由来します：その固有値スペクトルは「スパイク・アンド・フラット（spike-and-flat）」形状を示します。これは、いくつかの支配的な固有値（「スパイク」）に続き、ほぼ一様な尾部（「フラット」）が続く構造です。この構造は、層や学習段階を超えて経験的に観察され、ランク-$\rho$ の信号＋ノイズ勾配モデルの下では理論的に正確です。

1. 構造制限

Pro-KLShampoo は、KL-Shampoo のクリフォネ因子の一方（特に $m \le n$ の長方形層における右側因子 $R$）を、スパイク・アンド・フラット構造に整合するパラメトリック族に制限します：

• 射影部分空間（$r$ 次元）：前処理条件行列は、$U \in \text{St}(n, r)$ によって張られる追跡された部分空間上で完全なスペクトル構造を保持します。
• 補空間（$n-r$ 次元）：残りの方向は、単一の共有固有値 $\mu_\perp$ に縮退されます。

この制限により、右側因子の保存量が $O(n^2)$ から $O(nr)$ に、QR 分解のコストが $O(n^3)$ から $O(r^3)$ に削減されます。

2. 直交化による回復

スカラー尾部への単純な制限は、尾部が完全に一様でない場合に情報を失います。補空間における完全な KL-Shampoo 前処理条件行列の代数的形式を回復するために、Pro-KLShampoo は補勾配成分に対して直交化（極分解）を適用します。

更新則は以下の通りです：
$$\Delta W = -\alpha_{kl} L^{-1/2} G U S^{-1/2} U^\top - c_a \text{polar}\left(\mu_\perp^{-1/2} L^{-1/2} G P_\perp\right)$$
ここで：

• 第 1 項は、追跡された部分空間内における完全な KL-Shampoo 更新を表します。
• 第 2 項は、補空間に対して直交化を適用します。
• 著者は恒等式（式 12）を証明しており、この直交化が補空間における完全な KL-Shampoo の方向ごとのホワイトニングの代数的形式を回復することを示しています。これにより、Muon 風の直交化と KL 推定フレームワークが効果的に接続されます。

3. 部分空間追跡

最適な部分空間 $U$ は、行列 $\Phi_L = E[G^\top L^{-1} G]$ の固有空間として特定されます。スパイク・アンド・フラット仮説の下では、$\Phi_L$ のトップ-$r$ 固有空間が最適な選択となります。アルゴリズムは、$\Phi_L$ のRunning 推定値に対する QR 分解を介した 1 段階の冪乗法を用いて、この部分空間を追跡します。

主要な貢献

1. 構造観察とパラメトリック制限：KL-Shampoo の前処理条件行列における「スパイク・アンド・フラット」固有値構造（低ランク勾配モデル下で経験的に頑健かつ理論的に正確）を特定し、KL 目的関数をパラメトリック族に制限することでこれを活用します。
2. 直交化による代数的回復：補部分空間への直交化の適用が、完全な KL-Shampoo 前処理条件行列の代数的形式を回復することを示します。これにより、Muon 風の直交化と KL 推定の間に理論的架け橋が提供されます。
3. 収束保証：著者は、演算子ノルム滑らかさの下での理想化されたアルゴリズムに対して、$O(1/\sqrt{T})$ のレートでの非凸収束保証を提供します。
4. 経験的パフォーマンス：GPT-2（124M/350M）および LLaMA（134M/450M）モデルにおける広範な評価により、Pro-KLShampoo が検証損失、1 GPU あたりのピークメモリ、および一致した損失レベルに到達するまでのウォールクロック時間において、テストされたすべての部分空間ランク（$r \in \{32, 64, 128\}$）で KL-Shampoo を一貫して上回ることを示しています。

実験結果

著者は、4 つの前学習設定において、Pro-KLShampoo を AdamW、Muon、および KL-Shampoo と比較評価しました。

• 検証損失：Pro-KLShampoo は、すべてのモデルスケールにおいてテストされたすべてのランクで最低の検証損失を達成しました。
• メモリ効率：1 GPU あたりのピークメモリ使用量は KL-Shampoo よりも低く、Muon と同等でした。
• ウォールクロック時間：Pro-KLShampoo は KL-Shampoo よりも速く特定の検証損失レベルに到達しました。
  • GPT-2 においては、速度向上は主に 1 ステップあたりの時間の削減（$O(n^3)$ の QR 分解を $O(r^3)$ に置き換えること）によって引き起こされました。
  • LLaMA においては、速度向上は 1 ステップあたりの効率に加えて、より速い収束（より低い訓練損失曲線）によって引き起こされました。
• アブレーション研究：
  • 直交化を除去した場合（Smok-Hop 変種）、特に大規模スケール（LLaMA 450M）において性能が低下し、方向ごとのホワイトニングを回復するための直交化の必要性が確認されました。
  • 部分空間のみの変種も、補空間のみの変種も、完全な Pro-KLShampoo の性能には達しませんでした。これは、両方のコンポーネントが不可欠であることを示しています。

意義と主張

本論文は、Pro-KLShampoo が以前に孤立していた 2 つの最適化フロンティアを成功裡に統合したと主張しています。勾配共分散の内在的な「スパイク・アンド・フラット」構造を利用することで、KL-Shampoo の統計的効率性を達成しつつ、計算効率性を維持するために直交化を活用しています。

著者は、スカラー尾部更新の直交化による置換がヒューリスティックではなく、完全な KL-Shampoo 形式の代数的回復であると強調しています。これにより、この手法は KL ベースの前処理条件化の利点（ステップサイズの頑健性など）を保持しつつ、完全なクリフォネ因子の維持に関連する計算オーバーヘッドを大幅に削減できます。この研究は、低次元スペクトル構造が深層ネットワーク最適化において普遍的であり、LLM 前学習のためのより効率的なオプティマイザを設計するために効果的に活用できることを示唆しています。

論文は限界に関して控えめな姿勢を維持しており、ランク $r$ は経験的に設定されること、および混合重み $\alpha_{kl}$ の層ごとの適応的較正、およびニュートン・シュルツ反復とネステロフ・モーメントを含む実用的な変種の収束分析については今後の研究が必要であると指摘しています。