SGD Is Dead, Long Live Adam

ニューラルネットの訓練は結局のところ「損失関数を最小化するパラメータを見つける」ゲーム。そのゲームのルールを決めるのがオプティマイザ。 SGDから始まって、Momentum、Adagrad、RMSprop、Adam、AdamW、そして最近のLionやSophiaまで、全部「前のやつの弱点を潰す」という進化の連鎖になっている。数式で追いかけるという記事。

1. 問題設定

パラメータ $\theta \in \mathbb{R}^d$ を持つモデルの損失関数 $\mathcal{L}(\theta)$ を最小化したい。

訓練データ全体の勾配は計算コストが高すぎるから、ミニバッチ $\mathcal{B} \subset \mathcal{D}$ からの確率的勾配を使う。

$g_t$ は真の勾配 $\nabla \mathcal{L}(\theta_t)$ の不偏推定量。ここから先、全てのオプティマイザはこの $g_t$ をどう加工して更新量を決めるかの話になる。

2. Vanilla SGD

更新則

一番シンプル。勾配の方向にそのまま進む。

$\eta$ は学習率。これだけ。だけどこれだけじゃダメ。

問題点は損失関数の等高線が細長い楕円（条件数が大きい）だと、勾配方向が最適解への最短経路と大きくずれてジグザグに振動する。学習率を大きくすると発散し、小さくすると収束が遅すぎる。全パラメータに同じ学習率を使うから、頻出パラメータと稀少パラメータで最適な学習率が違う場合に対処できない。

3. SGD with Momentum

これってSGDに慣性を追加したやつなんだよね。 まぁボールが丘を転がるイメージ。前の更新方向を覚えておいて、慣性で進み続ける。谷底を横切るジグザグを抑えて、一貫した方向に加速する。 ゴルフボールがバンカーの谷の一番深いところを発見するような感じ。

更新則

$\beta$ が慣性の強さ（典型的には $\beta = 0.9$）。$v_t$ は勾配の 指数移動平均 になっている。展開すると、

過去の勾配が指数的に減衰しながら蓄積されるから、一貫した方向の勾配は強化されて、振動方向の勾配は相殺される。

Nesterov Accelerated Gradient（NAG）

Nesterov (1983) はMomentumの「先読み」バージョンを提案した。まず慣性で仮の位置に進んでから、そこでの勾配を計算する。

「先に行ってみて、そこの景色を見てから修正する」という戦略。凸最適化での収束レートが $O(1/t^2)$ に改善される（Momentumは $O(1/t)$）。理論的にはきれいだけど、実務では普通のMomentumと大差ないことも多い。

4. Adagrad

ここからが面白い。SGD系はすべてのパラメータに同じ学習率を使う。だけどNLPの埋め込み層を考えてみると、"the" のベクトルは毎バッチで更新されるのに、"serendipity" のベクトルはめったに更新されない。稀少なパラメータにはもっと大きなステップを踏ませたい。

更新則 (Duchi et al., 2011)

勾配の二乗を累積して、パラメータごとに学習率を適応させる。

$\odot$ はアダマール積、$\epsilon \approx 10^{-8}$ はゼロ除算防止。除算もすべて要素ごと。

頻繁に大きな勾配をもらうパラメータは $G_t$ が大きくなって学習率が下がり、稀にしか勾配をもらわないパラメータは $G_t$ が小さいまま学習率が高く保たれる。スパースなデータに強い。

だけど致命的な欠陥がある。$G_t$ は単調増加。訓練が進むにつれて学習率は $\eta / \sqrt{G_t} \to 0$ に収束して、事実上学習が止まる。「Adagradの学習率消失問題」。深いネットワークでは使い物にならない。

凸問題ではリグレットバウンド $R(T) = O(\sqrt{T})$、収束レート $O(1/\sqrt{T})$ が証明されている。だけどDeep Learningの非凸問題ではこの保証は役に立たない。

5. RMSprop

Hintonの講義スライドから生まれた

Geoffrey Hinton が Coursera の講義で提案した（2012年、論文なし）。Adagrad の $G_t$ が単調増加する問題を、指数移動平均 で解決する。

$\gamma = 0.9$ が標準。Adagrad との違いは、$G_t$ が全履歴の 和 だったのに対して、$E[g^2]_t$ は 指数重み付き平均 になっていること。古い勾配情報は自動的に忘却されるから、学習率が際限なく縮むことがない。

数学的に言えば、Adagrad は勾配の $L^2$ ノルムで割っていたのに対して、RMSprop は RMS（Root Mean Square） で割っている。名前の由来。

6. AdaDelta

Zeiler (2012) はRMSprop を独立に再発明して、さらに一歩進めた。学習率 $\eta$ 自体を不要にする。

分子に「過去の更新量のRMS」を使うことで、学習率のハイパーパラメータが消える。単位解析（dimensional analysis）的にも正しくなる（Adagrad/RMSpropでは勾配の単位と更新量の単位が合わない）。$\rho = 0.95$ が標準。

面白いアイデアだけど、実務ではAdamに完全に喰われた。次はadamの話。

7. Adam

AdamはKingma & Ba (2014) が提案した、現時点でもっとも広く使われているオプティマイザ。Momentum（一次モーメント）と RMSprop（二次モーメント）を組み合わせた。

更新則

バイアス補正

$m_0 = v_0 = 0$ で初期化するから、訓練初期は $m_t, v_t$ がゼロに偏る。展開すると、

だから補正する。

$t$ が小さいとき $1 - \beta^t$ は $1$ より大幅に小さい（$\beta_1 = 0.9$ で $t = 1$ なら $0.1$）から、補正項が10倍のブーストをかける。$t \to \infty$ で $1 - \beta^t \to 1$ になって補正は消える。

パラメータ更新

デフォルトは $\beta_1 = 0.9$、$\beta_2 = 0.999$、$\eta = 10^{-3}$、$\epsilon = 10^{-8}$。

直感的理解

$\hat{m}_t$ は「どの方向に進むか」（Momentum）、$\sqrt{\hat{v}_t}$ は「どのくらい慎重に進むか」（適応的学習率）。勾配の分散が大きいパラメータは慎重に、安定しているパラメータは大胆に更新する。これはパラメータごとの 信頼区間 を暗黙的に計算しているとも解釈できる。

別の見方をすると、Adamは損失ランドスケープの局所的な曲率に適応している。曲率が大きい方向（急な谷）では歩幅を縮め、曲率が小さい方向（なだらかな斜面）では歩幅を広げる。ニュートン法の対角近似とも解釈できる。Adamがほぼ全てのタスクで「そこそこ動く」のはこれが理由。

8. Adamの問題とAdamW

Weight Decayの罠

L2正則化の損失は $\mathcal{L}_{\text{reg}} = \mathcal{L} + \frac{\lambda}{2}\|\theta\|^2$。SGDではL2正則化とweight decayは等価。

だけどAdamでは $\nabla \mathcal{L} + \lambda\theta$ が一次/二次モーメントの推定に混入して、正則化項が適応的スケーリングの影響を受けてしまう。大きな勾配を持つパラメータの正則化が弱まり、小さな勾配のパラメータの正則化が強くなる。これはおかしい。

Decoupled Weight Decay (Loshchilov & Hutter, 2017)

AdamWはweight decayをAdam の更新ステップから 分離 する。

weight decay項 $(1 - \eta\lambda)\theta_t$ がモーメント推定の外側にある。これによりweight decayがパラメータの大きさにのみ比例し、適応的スケーリングの影響を受けない。

Transformer系モデル（BERT、GPT、ViT）の訓練はほぼ全てAdamWを使っている。「Adam使ってます」と言う人の大半は実際にはAdamWを使っている。PyTorchの torch.optim.AdamW がそのまま標準。デファクト。

9. ここまでの整理

数式で比較するとこうなる。

10. LAMB / LARS：大バッチ訓練

バッチサイズを $8192$ 以上にすると、従来のオプティマイザは収束が劣化する。バッチが大きいと勾配の分散が小さくなって汎化性能が落ちる（Sharp Minimaに収束しやすい）問題もあるけど、もっと直接的な問題は レイヤー間で更新量のスケールがバラバラ になること。

LARS (You et al., 2017)

Layer-wise Adaptive Rate Scaling。各レイヤー $l$ の更新量をパラメータノルムと勾配ノルムの比（trust ratio）でスケーリングする。

更新量がパラメータの大きさに対して「信頼できる」範囲に収まるように調整する。

LAMB (You et al., 2019)

LARS のアイデアをAdamWに適用したもの。BERTのバッチサイズ $32768$ での訓練を76分に短縮した。

$\|\theta_t\| / \|r_t\|$ がtrust ratio。レイヤーごとの更新量をパラメータノルムに比例させて安定化する。

11. Lion：遺伝的アルゴリズムが見つけたオプティマイザ

Chen et al. (2023, Google Brain) はオプティマイザの更新則自体をプログラムとして表現して、遺伝的アルゴリズム（AutoML-Zero）で探索した。見つかったのがLion（EvoLved Sign Momentum）。人間は設計していないんだよね。面白くない？

更新則

デフォルトは $\beta_1 = 0.9$、$\beta_2 = 0.99$。

更新量は $\text{sign}(\cdot)$ で すべてのパラメータが同じ絶対値 $\eta$ で更新される。Adamのように $\sqrt{\hat{v}_t}$ で割るステップがないから、メモリが約50%削減（二次モーメント $v_t$ を保持しない）。

$\text{sign}$ 演算は暗黙的な正則化として機能する。大きな勾配も小さな勾配も同じ $\pm 1$ になるから、外れ値的に大きな勾配の影響が抑えられる。

実用上のポイントとして、LionはAdamWの学習率の1/3〜1/10に設定する。weight decayは逆に3〜10倍に。大バッチ（1024以上）で特に効果的。小バッチだとAdamWの方が安定することが多い。

12. Sophia：二次情報を安く使う

ニュートン法の更新則は $\theta_{t+1} = \theta_t - H^{-1} g_t$（$H$ はヘシアン）。曲率情報を使うから収束が速い。だけどヘシアン $H \in \mathbb{R}^{d \times d}$ の逆行列はパラメータ数 $d$ が数十億のLLMでは計算不能。

Sophia の戦略 (Liu et al., 2023)

対角ヘシアンの推定値 $\hat{h}_t$ をプリコンディショナーとして使い、クリッピングで安定化する。

$h_t$ はヘシアンの対角成分の推定値。$\text{clip}(x, \gamma) = \max(-\gamma, \min(\gamma, x))$ は要素ごとのクリッピング。

ヘシアンの推定方法

フルヘシアンは $O(d^2)$ で無理だから、2つの軽量推定法を提案している。

Hutchinson推定量: ランダムベクトル $u \sim \mathcal{N}(0, I)$ に対して、

$Hu$ はヘシアン-ベクトル積で、自動微分で $O(d)$ で計算できる。

Gauss-Newton-Bartlett推定量（自己回帰LM向け）: 損失が交差エントロピーのとき、

ミニバッチの個々の勾配の二乗平均。実質的にFisher情報行列の対角近似。

Adamでは $m_t / \sqrt{v_t}$ でスケーリングするから、$v_t$ が非常に小さいと更新量が爆発しうる。Sophiaのクリッピングは最悪ケースの更新量を $\gamma$ に制限する。非凸問題でヘシアンが急変しても安全。

GPT-2の事前学習でAdamWの 2倍の速度（同じperplexityに到達するステップ数が半分）。ヘシアン推定は10ステップに1回でよく、オーバーヘッドは約5%。

13. Muon：勾配を直交化する

Muon（Multiplied by Orthogonalized Nesterov momentum、あるいは単に Momentum + Orthogonalization）は、勾配のモーメンタムを直交化してからパラメータを更新する。たぶん今のところ最新のやつ。

コアアイデア

重み行列 $W \in \mathbb{R}^{m \times n}$ に対する勾配 $G_t$ のモーメンタム $M_t$ を計算して、$M_t$ を 極分解（polar decomposition） で直交化する。

$(M_t^\top M_t)^{-1/2}$ はグラム行列の逆平方根。$\tilde{M}_t$ は $M_t$ の特異値を全て $1$ に揃えた行列。つまり更新方向の「向き」だけ残して「大きさ」を均一化する。

Newton-Schulz反復

$(M^\top M)^{-1/2}$ を直接計算するのは $O(n^3)$ で高い。Muonは Newton-Schulz反復 で近似する。

5回程度の反復で十分な精度が得られる。各反復は行列積2回（$O(mn \cdot \min(m,n))$）。

実装上は係数を最適化した quintic variant が使われる。

$a = 3.4445, b = -4.7750, c = 2.0315$。3回の反復で相対誤差 $10^{-7}$ 以下。

なぜ直交化が効くのか

勾配のスペクトル（特異値）は通常、数個の大きな成分に支配されている。Adamは各要素を独立にスケーリングするから、このスペクトルの偏りに対処できるけど、要素間の相関は無視する。Muonは行列全体のスペクトルを均一化するから、支配的な方向だけでなく全方向に均等に更新が行き渡る。これは スペクトルノルム制約 の下での最適化に相当する。

注意点として、Muonは行列パラメータ（Linear層、Attention層）に適用する。埋め込み層やLayerNormのような1次元パラメータにはAdamWを使う。ハイブリッド構成が標準。

14. 全体比較

15. Pythonで書くとこうなる

import numpy as np

class Adam:
    def __init__(self, params, lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8):
        self.lr, self.beta1, self.beta2, self.eps = lr, beta1, beta2, eps
        self.m = [np.zeros_like(p) for p in params]
        self.v = [np.zeros_like(p) for p in params]
        self.t = 0

    def step(self, params, grads):
        self.t += 1
        for i, (p, g) in enumerate(zip(params, grads)):
            self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * g
            self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * g**2

            m_hat = self.m[i] / (1 - self.beta1**self.t)
            v_hat = self.v[i] / (1 - self.beta2**self.t)

            p -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

class AdamW:
    def __init__(self, params, lr=1e-3, beta1=0.9, beta2=0.999,
                 eps=1e-8, weight_decay=0.01):
        self.lr, self.beta1, self.beta2 = lr, beta1, beta2
        self.eps, self.wd = eps, weight_decay
        self.m = [np.zeros_like(p) for p in params]
        self.v = [np.zeros_like(p) for p in params]
        self.t = 0

    def step(self, params, grads):
        self.t += 1
        for i, (p, g) in enumerate(zip(params, grads)):
            # Decoupled weight decay（ここがAdamとの違い）
            p *= (1 - self.lr * self.wd)

            self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * g
            self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * g**2

            m_hat = self.m[i] / (1 - self.beta1**self.t)
            v_hat = self.v[i] / (1 - self.beta2**self.t)

            p -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

class Lion:
    def __init__(self, params, lr=1e-4, beta1=0.9, beta2=0.99,
                 weight_decay=0.1):
        self.lr, self.beta1, self.beta2 = lr, beta1, beta2
        self.wd = weight_decay
        self.m = [np.zeros_like(p) for p in params]

    def step(self, params, grads):
        for i, (p, g) in enumerate(zip(params, grads)):
            # sign of interpolation（メモリは1状態だけ）
            update = np.sign(self.beta1 * self.m[i] + (1 - self.beta1) * g)

            p *= (1 - self.lr * self.wd)
            p -= self.lr * update

            # モーメンタム更新（次ステップ用）
            self.m[i] = self.beta2 * self.m[i] + (1 - self.beta2) * g

class SophiaG:
    """Sophia with Gauss-Newton-Bartlett estimator"""
    def __init__(self, params, lr=1e-4, beta1=0.965, beta2=0.99,
                 gamma=2.0, eps=1e-12, weight_decay=0.1):
        self.lr, self.beta1, self.beta2 = lr, beta1, beta2
        self.gamma, self.eps, self.wd = gamma, eps, weight_decay
        self.m = [np.zeros_like(p) for p in params]
        self.h = [np.ones_like(p) for p in params]  # ヘシアン推定

    def step(self, params, grads, hessians=None):
        for i, (p, g) in enumerate(zip(params, grads)):
            self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * g

            if hessians is not None:
                self.h[i] = self.beta2 * self.h[i] + (1 - self.beta2) * hessians[i]

            ratio = self.m[i] / np.maximum(self.h[i], self.eps)
            update = np.clip(ratio, -self.gamma, self.gamma)

            p *= (1 - self.lr * self.wd)
            p -= self.lr * update

16. 学習率スケジュール

オプティマイザだけじゃ話は半分なんだよね。学習率スケジュールが性能を左右することも多い。

Warmup: 最初の数ステップで学習率を $0$ から目標値まで線形に増加。Adamの二次モーメント $v_t$ が安定するまでの時間を稼ぐ。

Cosine Decay: Warmup後、コサイン関数に従って学習率を徐々に下げる。

GPT系の訓練ではほぼこれ。RAdam はWarmupの自動化を試みたけど、結局手動Warmup + Cosineが勝つことが多い。

参考文献

• Kingma, D.P. & Ba, J. (2014). "Adam: A Method for Stochastic Optimization." ICLR 2015
• Loshchilov, I. & Hutter, F. (2017). "Decoupled Weight Decay Regularization." ICLR 2019
• Duchi, J., Hazan, E. & Singer, Y. (2011). "Adaptive Subgradient Methods for Online Learning." JMLR, 12, 2121-2159
• Zeiler, M. (2012). "ADADELTA: An Adaptive Learning Rate Method." arXiv:1212.5701
• Hinton, G. (2012). "Neural Networks for Machine Learning." Coursera Lecture 6e（RMSpropの提案）
• You, Y. et al. (2017). "Large Batch Training of Convolutional Networks." arXiv:1708.03888（LARS）
• You, Y. et al. (2019). "Large Batch Optimization for Deep Learning." ICLR 2020（LAMB）
• Chen, X. et al. (2023). "Symbolic Discovery of Optimization Algorithms." NeurIPS 2023（Lion）
• Liu, H. et al. (2023). "Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training." ICLR 2024
• Muon Optimizer (2024). Newton-Schulz iteration for gradient orthogonalization