技術要約: CacheMuon

問題提起

Muonは、ディープラーニングにおける行列値の隠れ層パラメータを対象とした、幾何学的な性質を考慮したオプティマイザである。座標ごとのオプティマイザ（Adamなど）とは異なり、Muonは標準的なモーメンタム更新を、モーメンタム行列の極分解因子から導出された直交化された方向へと置き換える。Muonは大規模言語モデルにおいて強力な経験的性能とスケーラビリティを示しているが、その実用的な効率性は、直交化ステップによってボトルネックとなっている。

現在の実装では、高コストな特異値分解（SVD）を避けるために、ニュートン・シュルツ反復（具体的にはグラム・ニュートン・シュルツ法）を用いて極分解因子を近似している。しかし、この近似であっても、毎回の最適化ステップで複数の行列乗算を必要とする。その結果、Muonの計算コストの大部分は、互いに密接に関連する行列直交化問題を繰り返し解くことに費やされている。近年の研究では、単一の直交化呼び出しの効率化（例：PolarExpress、Gram Newton–Schulz）が行われているが、これらの手法は、最適化プロセスに内在する時間的な相関を無視して、毎ステップで直交化変換をゼロから再計算している。

手法: CacheMuon

著者らは、訓練中のモーメンタム行列の滑らかな進化を利用するテンポラル（時間的）なプリコンディショニング手法であるCacheMuonを提案する。核となる仮説は、もしモーメンタム行列が緩やかに変化するのであれば、前のステップで計算された直交化変換は、現在のステップにおいても情報量として有用であり続けるというものである。

コアメカニズム

CacheMuonは、以下のキャッシュメカニズムを導入することで、Muonの直交化ルーチンを修正する：

1. フレッシュ・ソルブ（更新ステップ）： 特定の「更新（refresh）」ステップ（トリガーによって決定される）において、手法は完全なFresh Gram Newton–Schulz (FreshGNS) ソルバーを実行する。これにより、近似的な極分解因子 $\tilde{D}_t$ と、極めて重要なことに、キャッシュされた左変換行列 $Q_t$ が生成される。
2. キャッシュの再利用（中間ステップ）： 更新ステップではないステップでは、完全なソルバーを実行する代わりに、キャッシュされた変換 $Q_{a_t}$ （ここで $a_t$ は最新の更新ステップのインデックス）を再利用して、極分解因子を近似する： $\hat{D}_t = Q_{a_t} X_t$ （ここで $X_t$ は正規化されたモーメンタム行列である）。
3. 残差ベースのトリガー： キャッシュされた変換がいつ正確性を失うかを判断するために、著者らは軽量な残差ベースの再利用トリガーを導入している。キャッシュを再利用する前に、アルゴリズムは候補となる更新の直交性残差を計算する：
   $$\text{err}_t = \frac{\| \hat{D}^{\text{cand}}_t (\hat{D}^{\text{cand}}_t)^\top - I \|_F}{\|I\|_F}$$
   もし $\text{err}_t \leq \gamma$ （規定された閾値）であれば、キャッシュされた変換は受理される。そうでなければ、キャッシュを更新するためにフレッシュ・ソルブがトリガーされる。

理論的分析

本論文では、CacheMuonを不正確な（inexact）Muon更新として分析している。キャッシュされた方向 $\hat{D}_t$ と正確なMuonの方向 $D_t$ の間の偏差は、加法的な誤差項 $\delta_t$ によって抑えられる：
$$\| \hat{D}_t - D_t \|_{\text{sp}} \leq \delta_t = L_Q \tau_t + \varepsilon_{\text{ref}}$$
ここで、

• $\varepsilon_{\text{ref}}$ は、フレッシュ・ソルバーの誤差（有限のニュートン・シュルツ反復によるもの）である。
• $\tau_t$ は、現在のものに対するキャッシュされたグラム行列の「古さ（staleness）」である。
• $L_Q$ は、ソルバーによって誘導される多項式写像から導出されるリプシッツ定数である。

標準的な仮定（滑らかさ、有界な分散、有界な古さ）の下で、著者らは非凸最適化に対する確率的収束保証を証明している。収束レートは、集約された不正確性パラメータ $\bar{\delta} = L_Q \bar{\tau} + \varepsilon_{\text{ref}}$ に依存する。$\bar{\delta} < 1$ である限り、手法は収束し、そのレートは $O((1-\bar{\delta})^{-1} N^{-1/4})$ とスケーリングする。

主な貢献

1. CacheMuon アルゴリズム： 訓練中の時間的な滑らかさを利用して、直交化変換をキャッシュすることで、極分解因子の計算コストを償却するMuonのバリアント。
2. 収束理論： CacheMuonを不正確なオプティマイザとして扱う理論的枠組みを提供し、フレッシュ・ソルバーの誤差とキャッシュの古さによって制御される確率的収束保証を与える。
3. 残差ベースのトリガー： 更新の忠実度のプロキシとして機能する、キャッシュの更新時期を動的に決定する、実用的かつ低コストなヒューリスティック。
4. 実証的検証： 言語モデル（GPT-2）およびビジョン（ResNet-18）のワークロードにおける、制御可能な品質–効率のフロンティアの提示。

実験結果

著者らは、GPT-2 Large/Small (OpenWebText) および ResNet-18 (CIFAR-10) においてCacheMuonを評価した。

• 品質–効率のトレードオフ： 残差閾値 $\gamma$ を変化させることで、著者らは滑らかなフロンティアを実証している。
  • 保守的な再利用 ($\gamma=2$)： フレッシュなMuon（GramMuonおよびPolarExpMuon）の性能を密接に追跡しながら、GPT-2 Largeにおいて直交化FLOPsを（GramMuonに対して）13.6%、（PolarExpMuionに対して）30.9% 削減する。
  • 積極的な再利用 ($\gamma=15$)： 検証損失のわずかな低下（例：GPT-2 Largeにおいて3.096 vs 3.037）のみで、直交化FLOPsを最大65.0% 削減（GramMuonに対して）する。
• 診断的洞察： 実験により、残差トリガーが、キャッシュされた更新とフレッシュな更新の間の整合性と効果的に相関していることが示された。訓練の初期段階では残差が高く、頻繁な更新を強制する。訓練が安定するにつれて残差は低下し、より高いキャッシュヒット率と大きな効率化を可能にする。
• スケーラビリティ： 本手法は、異なるモデルスケール（GPT-2 Small vs Large）やモダリティ（Vision vs Language）にわたって競争力のある性能を維持している。

意義と主張

本論文は、CacheMuonが、基礎となる幾何学的更新ルールを変更することなく、Muonの計算オーバーヘッドを削減するための実用的なアプローチを提供することを主張している。直交化を純粋に空間的な問題としてではなく、時間的な問題として扱うことで、冗長な計算を削減する。

著者らは、本手法が直交化FLOPsを大幅に削減する一方で、現在のハイレベルな実装における実際のウォールクロック・スピードアップは、システムレベルのオーバーヘッド（例：残差チェックのためのホスト・デバイス間同期）により控えめであることを強調している。彼らは、より最適化された低レベルの実装（例：プローブをデバイスカーネルに融合させること）によって、これらの算術的な節約のより大きな割合を実際の実行時の利得に変換できると考えている。

最終的に、本研究は、時間的な再利用が幾何学的なオプティマイザを最適化するための有望な方向性であり、計算効率と訓練品質のバランスを取るための調整可能なノブを提供することを示唆している。