技術的サマリー：楽観的双対平均法が現代のオプティマイザを統合する

問題定義

深層学習の最適化は、主に 2 つの軸に沿って進化してきた。幾何学と適応（モデル構造を反映する暗黙的なノルムや前処理行列を用いた更新の設計、例：Adam、Muon）と、構成要素の組み合わせ（安定性のための勾配フィードバック、モメンタム、スケジューリングの組み合わせ、例：Lion、スケジューリングフリー手法）である。多くの最先端オプティマイザ（Muon、Lion、AdEMAMix、NAdam など）は異なって見えるが、それらの関係を説明し、ハイパーパラメータ調整に対する原理的な指針を提供する統合された理論的枠組みは欠けている。

具体的かつ持続的な課題は、重み減衰の調整である。複数エポックのトレーニングにおいて、重み減衰は反復点のノルムを制約する正則化器として機能する。しかし、過学習が懸念されにくい単一エポックのトレーニングであっても、重み減衰はしばしば有益である。最近のスケーリング則は、重み減衰をモデル次元 $d$ に比例して $1/d$ に設定することを提案しているが、モデル次元 $d$ とトレーニング期間 $n$ への依存関係を分離することは困難である。さらに、最適な重み減衰の値はオプティマイザによって大きく異なる（例：Lion は AdamW よりもはるかに大きな減衰を必要とする）ため、オプティマイザごとの高コストな調整が必要となる。

手法：SODA

著者らは、勾配の双対処理と反復点の原始処理を統合する、楽観的双対平均法（ODA）の一般化であるSODA（Stochastic Optimistic Dual Averaging）を導入する。

理論的枠組み

SODA は、標準的な双対平均更新を一般化する：
$$x_{k+1} \in \arg \min_{x \in X} \gamma_k \langle d_k, x \rangle + h_k(x)$$
ここで、原始外挿列 $(y_k)$ と楽観的双対変数 $(\bar{m}_{k+1})$ を導入する。アルゴリズムは以下のものを維持する：

1. 双対処理：モメンタムに基づく勾配平均 $m_{k+1}$ と、現在の勾配 $\nabla f(y_k, \xi_k)$ を取り入れた楽観的平均 $\bar{m}_{k+1}$。
2. 原始処理：パラメータ $\lambda_k$ と $\bar{\lambda}_k$ によって制御される反復点更新 $x_{k+1}$ と外挿点 $y_{k+1}$。

更新則は以下の通りである：
$$\begin{aligned} m_{k+1} &= (1 - \alpha_k)m_k + \alpha_k \nabla f(y_k, \xi_k) \\ \bar{m}_{k+1} &= (1 - \bar{\alpha}_k)m_{k+1} + \bar{\alpha}_k \nabla f(y_k, \xi_k) \quad (\text{楽観性}) \\ z_{k+1} &\in \partial h_k^*(-\gamma_k \bar{m}_{k+1}) \\ x_{k+1} &= (1 - \lambda_k)x_k + \lambda_k z_{k+1} \\ y_{k+1} &= (1 - \bar{\lambda}_k)x_{k+1} + \bar{\lambda}_k z_{k+1} \quad (\text{原始外挿}) \end{aligned}$$
ここで、$h_k$ は幾何学を定義する正則化器（ミラーマップ）であり（例：Lion には $\ell_\infty$、Muon にはスペクトル幾何学）、である。

既存オプティマイザの統合

本論文は、特定のパラメータと幾何学を選択することで、SODA がいくつかの現代のオプティマイザを特殊ケースとして回復することを示している：

• NAdam / Lion / Muon：これらは、「近代化」レジーム（勾配評価点に対する原始外挿がない $\bar{\lambda}_k = 0$ および楽観性 $\bar{\alpha}_k > 0$）の下で回復される。
  • Lion：$\ell_\infty$ 幾何学（$h_k = K^*$）と特定のモメンタムパラメータを使用する。
  • Muon：スペクトル幾何学とネステロフ・モメンタム（$\bar{\alpha}_k = \alpha_k$）を使用する。
  • NAdam：適応的スケーリングを伴う平滑化された $\ell_\infty$ 幾何学を使用する。
• 論文の表 1は、これらの手法を、勾配平均（$\alpha_k$）、楽観性（$\bar{\alpha}_k$）、幾何学（$h_k$）の対応する選択に明示的にマッピングしている。

SODA ラッパー

重要な実用的な貢献として、新しいハイパーパラメータを導入することなく、任意のベースオプティマイザ（例：Adam、Lion、Muon）をラップできるSODA ラッパー（アルゴリズム 1）がある。

• 設定：$\bar{\lambda}_k = 0$（「近代化」レジームに一致）、一様な反復点平均 $\lambda_k = 1/(k+2)$、ステップサイズのスケーリング $\gamma_k = (k+2)\eta_k$ を設定する。
• 重み減衰の排除：ラッパーは原点ではなく初期化 $z_0$ に更新を固定する。その結果生じる更新則は、実質的に $\mu_k = 1/(k+2)$ の重み減衰スケジューリングを実装する。
• 解釈：著者らは、この枠組みにおいて、重み減衰は単なる正則化器ではなく、原始平均の一種であると主張する。理論的に裏付けられたこの $1/k$ 減衰により、個別の重み減衰ハイパーパラメータを調整する必要がなくなる。

主要な貢献

1. 楽観性と双対化による統合：SODA は、Muon、Lion、NAdam などを双対平均の楽観的インスタンスとして回復する単一のテンプレートを提供し、それらの構造的関係を明確にする。
2. 理論的保証：この枠組みは、$L$-滑らかな凸関数に対する収束解析に基づいて理論的に裏付けられたハイパーパラメータ選択（特に $1/k$ 減衰スケジューリング）を導出する。重み減衰を原始平均のメカニズムとして再定義する。
3. パラメータフリーのラッパー：提案された SODA ラッパーは、ベースオプティマイザの重み減衰を調整する必要性を排除しつつ、性能を維持または向上させる。
4. 実証的検証：広範な実験により、モデルサイズやトレーニング期間にわたって一貫した改善が示された。

実験結果

実験は、FineWeb100データセットを使用してnanoGPTで行われ、SODA ラップされたオプティマイザを調整済みのベースラインと比較した。

• 期間横断的転送：Muon の最適な重み減衰は、トレーニング期間の増加とともに減少することが判明し、静的なスケーリング則（$1/d$ など）と矛盾した。SODA ラッパーの動的な $1/(k+2)$ 減衰は、Muon に最良の調整済み重み減衰を与えた場合でも、Muon を一貫して上回った（図 1）。
• 一般的な適用性：ラッパーは、追加の調整なしに、異なるトレーニング期間（1$\times$ および 4$\times$ Chinchilla ステップ）においてAdam、Muon、Scionの性能を向上させた（図 2）。
• スケーリング：モデルサイズを 64M から 1B パラメータにスケーリングする際（Chinchilla スケーリングに従う）、SODA（特に楽観的変種SODA$^\dagger$）はベースオプティマイザ（Scion/uScion）を一貫して上回り、モデルサイズが増大するにつれて性能の差が広がった（図 4）。
• 簡素化：SODA$^\dagger$ 設定は、すべての層で Adam を Lion に置き換え、単一のモメンタムハイパーパラメータを使用することで Muon ベースラインを簡素化したが、競争力のある、あるいは優れた性能を達成した。

意義と主張

本論文は、SODA がハイパーパラメータ調整を必要とせずに、異なるトレーニング期間やモデルスケール間で最適化戦略を転送するための原理的なメカニズムを提供すると主張している。

• 重み減衰の再定義：著者らは視点の転換を提唱する。重み減衰は過学習に対する正則化器としてだけでなく、最適化軌道を初期化に固定する原始平均の一種として理解されるべきである。
• 実用的影響：現代のオプティマイザを SODA 枠組みの下で統合することで、著者らは重み減衰の調整という面倒で高価なプロセスを排除する「プラグアンドプレイ」ラッパーを提供する。これは、モデルサイズの増大に伴い手動調整が prohibitively expensive（許容できないほど高価）になるにつれて、特に重要である。
• 謙虚さ：著者らは限界を認め、理論的解析は現在凸設定に限定されていること（ただし深層学習にとって経験的に有益である）と、単一エポックトレーニングに焦点を当てていることを指摘している。過学習の懸念から、複数エポックのトレーニングでは、正則化強度のためのハイパーパラメータが依然として必要になる可能性があると示唆している。

要約すると、SODA は楽観的双対平均のレンズを通じて現代のオプティマイザの設計空間を統合し、スケーリングにわたるトレーニングの安定性と性能を向上させる、理論的に裏付けられたパラメータフリーな手法を提供する。