Convergence Rate Analysis of the AdamW-Style Shampoo: Unifying One-sided and Two-Sided Preconditioning

以下は、論文「Convergence Rate Analysis of the AdamW-Style Shampoo: Unifying One-sided and Two-Sided Preconditioning（AdamW 風 Shampoo の収束率解析：片側と両側前処理の統合）」の詳細な技術的要約です。

1. 問題定義

本論文は、AlgoPerf 神経網訓練コンペティションで優勝した最先端アルゴリズムであるAdamW 風 Shampooオプティマイザの理論的収束解析に取り組んでいます。Shampoo（SOAP や Muon などのその変種を含む）は、フル行列前処理器を利用することで、対角前処理法（Adam など）よりも優れた経験的パフォーマンスを示してきましたが、非凸最適化に対するその理論的収束保証は欠如していたか、限定的でした。

具体的には、既存の文献には以下の 3 つの主要なギャップが存在しました：

1. 非凸性：ほとんどの先行解析は凸設定に制限されていました。
2. 前処理の範囲：以前の非凸解析（例：Xie ら、2026）は片側前処理に限定されており、実用的な実装で用いられるより複雑かつ効果的な両側前処理を無視していました。
3. 実用的コンポーネント：既存の理論は、現代の大規模言語モデル（LLM）訓練に不可欠なモメンタムやデカップルド重み減衰（AdamW 風）といった重要な実用的特徴をしばしば省略していました。

本論文は、非凸フレームワークの下で片側および両側前処理の両方を統合し、一般的な AdamW 風 Shampoo アルゴリズムに対する厳密な収束率を確立することを目的としています。

2. 手法

著者らは、AdamW 風 Shampoo オプティマイザを表すアルゴリズム 1を解析します。更新則は以下の要素を含みます：

• 指数移動平均（EMA）：勾配の 1 次モーメント（$M_k$）および 2 次モーメント（$L_k, R_k$）を追跡します。
• 両側前処理：更新ステップは $X_{k+1} = (1-\lambda\eta)X_k - \eta L_{k,\epsilon}^{-\frac{1}{2p}} M_k R_{k,\epsilon}^{-\frac{1}{2q}}$ であり、ここで $L_k$ と $R_k$ は勾配の累積外積（$E[GG^T]$ および $E[G^TG]$ の近似）です。
• 統合された指数：パラメータ $p$ と $q$（$1/p + 1/q = 1$）により、アルゴリズムは両側（$p, q < \infty$）と片側（$p=1, q=\infty$、またはその逆）の前処理間を切り替えることができます。

主要な技術的革新：
非凸設定における両側行列前処理の複雑さを処理するために、著者らは以下の新しい証明技法を開発しました：

1. Schatten-p Hölder 不等式：Adam 解析で用いられる標準的なベクトルベースの Hölder 不等式の代わりに、著者らは勾配の核ノルム（$\|\cdot\|_*$）を評価するためにSchatten-p Hölder 不等式を利用します。これは、前処理器が行列パラメータの特異値に作用するため、極めて重要です。
2. 行列 Cauchy-Schwarz 不等式：更新ステップ（$L^{-1/2p} M R^{-1/2q}$）のスペクトルノルムを評価するために、左右の前処理器間の相互作用を処理する行列固有の Cauchy-Schwarz 不等式を採用します。
3. スペクトルフロア解析：前処理器（$\hat{\epsilon}$）の下限に関する条件（条件 4）を導入します。彼らは、実際には前処理器の有効なスペクトルフロアが数値的正則化 $\epsilon$（例：$10^{-12}$）よりもはるかに大きいと主張し、理論的 bound を経験的観察と整合させます。

3. 主要な貢献

• 統合された収束率：本論文は、片側および両側前処理の両方を網羅する、非凸最適化における AdamW 風 Shampoo の最初の収束率を確立しました。
• 最適率の導出：勾配の平均期待核ノルムが以下の率で収束することを証明しました：
  $$\frac{1}{K} \sum_{k=1}^K \mathbb{E}[\|\nabla f(X_k)\|_*] \leq O\left( \frac{\sqrt{m+n} \cdot C}{K^{1/4}} \right)$$
  ここで、$m, n$ は行列の次元、$K$ は反復回数、$C$ は最適 SGD 率の定数に一致します。
• 実用的設定の理論的正当化：数値的正則化子 $\epsilon$ を極めて小さな値（例：$10^{-12}$）に設定しても収束を妨げない理由を説明します。理論は、有効な前処理器フロアが数値的 $\epsilon$ ではなく、勾配ノイズのスケールによって決定されることを示しています。
• ノルムの関係性：著者らは、導出された核ノルム率が、$\|\nabla f(X)\|_* \approx \Theta(\sqrt{\min(m,n)}) \|\nabla f(X)\|_F$ という理想的な条件下での SGD の最適フロベニウスノルム率（$O(K^{-1/4})$）と類似していることを示しました。

4. 結果

• 理論的 bound：
  • 最悪ケースの bound は $\sqrt{m+n}$ に依存し、一部の対角法やフル行列の下限で見られる $\sqrt{mn}$（総次元 $d$）の依存性よりも著しくtight（厳密）です。
  • 収束率は、$\sqrt{m+n}$ の因子を除けば、非凸確率的最適化の既知の下限（$O(K^{-1/4})$）と一致します。
• 経験的検証（GPT-2 事前学習）：
  • 勾配ノルム比：GPT-2 事前学習の実験では、比 $\|\nabla f(X)\|_* / \|\nabla f(X)\|_F$ が理論的上限 $\sqrt{\min(m,n)}$ に一貫して近い値を維持することが示されました。これは、実際には核ノルム率がフロベニウスノルム率と実質的に同等であるという仮定を検証するものです。
  • スペクトルフロア：前処理器（$L_k, R_k$）の最小固有値は、数値的 $\epsilon$（$10^{-12}$）よりも桁違いに大きく、ノイズ依存スケール $\hat{\sigma}^2/(m+n)$ と同程度であることが観測されました。これは、「非現実的」なより大きなスペクトルフロアという仮定が、現実世界の LLM 訓練において成り立つことを確認しました。

5. 意義

• 理論と実践の架け橋：この研究は、AlgoPerf コンペティションで優勝した特定の Shampoo 変種に対する最初の厳密な理論的基盤を提供し、なぜそれが実際にはこれほど効果的に機能するのかを説明します。
• 対角法に対する進歩：行列構造パラメータに対して、非対角（フル行列）前処理が対角法（Adam など）よりも優れていることを理論的に検証し、Shampoo によって捉えられるパラメータ間の相関がより良い収束方向をもたらすことを示しました。
• ハイパーパラメータへの指針：この解析は、ハイパーパラメータの選択（例：モメンタム $\theta$ と 2 次モーメント減衰 $\beta$ の関係、および $\epsilon$ の役割）に対する理論的正当性を提供し、大規模モデルのオプティマイザ調整における実践者の指針となります。
• 将来のオプティマイザの基盤：片側および両側前処理を統合し、モメンタム/重み減衰を組み込むことで、このフレームワークは深層学習における高度な 2 次近似オプティマイザの解析のための新たな基準を設定します。

要約すると、本論文は Shampoo の理論的解析を成功裏に統合し、現代の深層学習モデルに固有の複雑な行列構造を捉えながら、非凸問題に対して最適の収束率を達成することを証明しました。