Training Optimization(学習最適化)

article technology high #training-optimization#optimizer#adamw#lion#sophia#shampoo#soap#muon#second-order#sign-momentum#sgd#nesterov#adagrad#rmsprop#adafactor#came#zero#8bit-adam#sam#edge-of-stability#kfac#psgd#fp8-optimizer#llm-training
Created: 2026-05-26 Updated:

AdamW を 2024〜2026 デファクトとして、SGD/Nesterov・Adam 系統・Lion/Sophia・Shampoo/SOAP・ZeRO/8-bit・SAM/Edge of Stability・10 落とし穴・MUON/fp8 フロンティアを整理。

Training Optimization(学習最適化)

LLM 学習のオプティマイザは AdamW (Loshchilov 2017, arXiv:1711.05101) が 2024〜2026 デファクトで、GPT-3 (arXiv:2005.14165) / Llama-3 / Mistral / DeepSeek-R1 (arXiv:2501.12948) すべて採用。系譜は (1) SGD+momentum (Polyak 1964 / Nesterov 1983) の加速、(2) AdaGrad → RMSProp → Adam → AdamW (Duchi 2011 / Hinton 2012 / Kingma 2015, arXiv:1412.6980 / Loshchilov 2017) の適応学習率、(3) Lion (Chen 2023, arXiv:2302.06675) / Sophia (Liu 2023, arXiv:2305.14342) のメモリ効率・二次法フロンティアの三本柱。本記事はオプティマイザの内部数学・更新則・収束性に絞り、数値設定は tech-50、学習パラダイムは tech-49、並列訓練は tech-48 を参照。

モメンタム系列 — SGD / Polyak / Nesterov

SGDθ_{t+1} = θ_t − η·g_t (g_t はミニバッチ勾配)。L-smooth 凸で O(1/√T)、強凸で線形収束。条件数 κ=L/μ が大きいと狭軸振動と平坦軸の遅進行が共起し、これが momentum の起点。Polyak heavy-ball (1964) は v_t = β·v_{t-1} + η·g_tθ_{t+1} = θ_t − v_tκ → √κ の quadratic speedupNesterov accelerated gradient (1983) は lookahead 点 θ_t + β·v_{t-1} で勾配評価し O(1/k²) の一階法理論最適。Sutskever 2013 が 1-step 形に畳み PyTorch SGD(nesterov=True) 採用。LLM 事前学習では Adam 系に置換され実例稀だが、MUON (§フロンティア) で再評価中。

適応学習率系列 — AdaGrad → RMSProp → Adam → AdamW

埋め込み行列は登場行のみ更新される疎勾配、Attention/MLP は dense — 全 param に同一 η は最適でない。AdaGrad (Duchi 2011) は G_t = G_{t-1} + g_t² 累積で η/√(G_t+ε) を per-param 適用。疎勾配の rare 行は実効 LR を維持するが、G_t 単調増加で t→∞ で実効 LR→0 (“学習が死ぬ”)RMSProp (Hinton 2012、Coursera Lecture 6e — arXiv 不在の珍例) は EMA v_t = ρ·v_{t-1} + (1−ρ)·g_t² でこの死問題を回避。

Adam (Kingma 2015, arXiv:1412.6980) は RMSProp に 1 次モーメントと bias correction を加えた:

m_t = β1·m_{t-1} + (1−β1)·g_t
v_t = β2·v_{t-1} + (1−β2)·g_t²
m̂_t = m_t / (1 − β1^t),   v̂_t = v_t / (1 − β2^t)
θ_{t+1} = θ_t − lr · m̂_t / (√v̂_t + ε)

Bias correction の本質: m_0=0 起点で E[m_t]=(1−β1^t)·E[g_t]、t=1 で補正係数 10× と無視不能。built-in simulated warmup を与えるが LLM スケールでは別途明示 warmup 必須。メモリコスト: 7B params で m+v を fp32 維持すると 56 GB の optimizer state。

AdamW 更新則と decoupled weight decay の本質

Adam+L2AdamW の致命的差は重み減衰の適用位置。Adam+L2 は g_t_reg = g_t + λ·θ_{t-1} を先に加算し adaptive scale を通すため、実効減衰は λ·θ/√v̂頻繁更新パラメータほど弱く非一様。これが画像分類で SGD+momentum に汎化で負ける主因と Loshchilov & Hutter は同定。AdamW は減衰を adaptive scale に出す:

θ_t = θ_{t-1} − lr · m̂_t/(√v̂_t + ε) − lr · λ · θ_{t-1}

全 param に一様 λ を保証し adaptive 動学と正則化を独立制御。LLM 標準は β1=0.9, β2=0.95, ε=1e-8, λ=0.1 (詳細は tech-50)。埋め込み・LayerNorm を WD 除外するのが GPT-3 系譜の標準実装。

Sign-based 系列 — Lion / Sophia の memory-efficient 設計

2023〜2026 で AdamW の m+v 二重状態を単一状態化する系統が台頭。Lion (EvoLved Sign Momentum) (Chen 2023, arXiv:2302.06675) は進化的プログラム探索で発見された人手非設計オプティマイザ:

update = sign(β1·m_{t-1} + (1−β1)·g_t)
m_t    = β2·m_{t-1} + (1−β2)·g_t
θ_t    = θ_{t-1} − lr · (update + λ·θ_{t-1})

(1) m のみで AdamW の 50% メモリ、(2) sign() で更新規模は常に ±lr、(3) AdamW 比 3〜10× 低い LR が必要。ViT で +1〜2% top-1、LLM では同等以上の perplexity を 50% メモリで報告するが、production 採用は限定的 — 勾配蓄積・mixed precision との相性と μP 転移の clean さで AdamW に劣る (>7B での default 化はまだ)。

Sophia (Liu 2023, arXiv:2305.14342) は Hutchinson 推定で対角 Hessian 近似する scalable 二次法。ĥ_t = g_t ⊙ g_t を k=10 step ごと EMA、θ_t = θ_{t-1} − lr · m_t / max(γ·h_t, ε)max(γ·h_t, ε) クリップが 0 近傍 Hessian 推定の発散防止の要。GPT-2 124M で AdamW 比 2× token 効率を主張するが 7B+ での独立再現は限定的 (mid-2024 時点)。

二次法・前処理法 — Shampoo / SOAP / K-FAC

二次法は強凸で O(κ) → O(log κ) 改善するが Hessian の O(n²) 記憶・O(n³) 逆行列が n=7B で非現実 — これが Kronecker 分解の動機。Shampoo (Gupta 2018, arXiv:1802.09568) は行列 W∈ℝ^{m×n} で:

L_t = L_{t-1} + G_t·G_t^T,   R_t = R_{t-1} + G_t^T·G_t
W_t = W_{t-1} − η · L_t^{-1/4} · G_t · R_t^{-1/4}

行列層 Hessian の L ⊗ R 近似で記憶 O(m²+n²)、逆行列ルートは T_inv=100〜1000 step ごと固有値分解で償却。Distributed Shampoo (Anil 2020, arXiv:2002.09018) は T5 で内部採用。>1B で支配的にならない理由: (a) L/R 追加メモリ、(b) 固有値分解 +10〜30% per-step、(c) preconditioner all-reduce、(d) 良 schedule + AdamW で長期は利得相殺。SOAP (Vyas 2024, arXiv:2409.11321) は Shampoo の固有基底内で Adam を走らせる統合方式で NanoGPT speedrun の小規模 SOTA を獲得。K-FAC (Martens 2015, arXiv:1503.05671) は Fisher 行列の Kronecker 近似で実装複雑、LLM 採用稀。

メモリ効率変種 — ZeRO / 8-bit Adam / Adafactor / CAME

ZeRO (Rajbhandari 2019, arXiv:1910.02054) は DDP の全 GPU で θ+g+(m,v) を冗長保持する問題をシャーディングで解決。

Stageシャード対象メモリ削減 (N_d GPUs)
1optimizer states (m, v)~4×
2optimizer + gradients~8×
3optimizer + gradients + params~N_d×

7B × 64 GPU の ZeRO-3 で per-GPU optimizer state は 56 GB → 875 MB。ZeRO-Infinity (arXiv:2104.07857) は NVMe オフロード。8-bit Adam (Dettmers 2022, arXiv:2110.02861) は (m,v) を block_size=2048 ごとのブロック単位量子化で 4 byte → 1 byte、bitsandbytes 実装、~5% per-step オーバヘッド × 4× メモリ削減。SFT/DPO 用途が主、事前学習は fp32 state 標準。Adafactor (Shazeer 2018, arXiv:1804.04235) は行列 2 次モーメントを V_t ≈ r_t · c_t^T / (1^T·c_t) のランク 1 分解で V×d → V+d 削減、T5/PaLM 採用、momentum 欠如で AdamW より遅め。CAME (Luo 2023, arXiv:2307.02047) は Adafactor の因子分解誤差を confidence matrix で追跡し self-calibration を加える post-training 向け。

収束理論と損失景観 — implicit bias / SAM / Edge of Stability

Implicit bias of SGD (Keskar 2017, arXiv:1609.04836): 小 batch SGD は flat minima へ導く implicit bias を持ち分布シフトにロバスト。Adam は 1/√v の per-param scaling で sharp minima 傾向だが AdamW+WD+warmup+cosine でかなり相殺される。SAM (Foret 2020, arXiv:2010.01412) は min_θ max_{||ε||≤ρ} L(θ+ε) を minimax 定式し ε* ≈ ρ·g/||g|| で近似、1 step につき 2 fwd+2 bwd (倍コスト)。perturb と update は同一バッチ厳守 — 異バッチで sharpness 推定にノイズが入る。LLM 事前学習効果は小さいが SFT/finetune で gain 大。

Edge of Stability (EOS) (Cohen 2021, arXiv:2103.00065): NN 訓練中に λ_max が漸進増大し 2/η 付近で安定振動。古典理論は λ_max > 2/η で発散と予測するが catapult dynamics で訓練成功。LLM loss spike の正体は概ねこれPaLM (arXiv:2204.02311) の rollback (100 step 巻き戻し+バッチスキップ+続行) はこの理解の標準対応。

落とし穴 — 10 つの頻発失敗モード

#落とし穴症状対策
1Adam bias-correction 過大 (β2=0.999)50〜200 step で spikeβ2=0.95 で 1/(1−β2) を 1000→20 へ
2Checkpoint で optimizer state 未保存resume 直後の loss 異常(m, v, step_t) 完全保存・1-step delta で検証
3Weight tying × decoupled WD で embedding 過剰収縮rare token zero-shot 低下embedding は weight_decay=0 (GPT-3/Llama 系譜)
4fp16 optimizer state × bf16 master のミスマッチ数千 step 後 NaN(m, v) を fp32 維持、Megatron-LM 標準
5Lion の LR スウィープを Adam スケールで実施発散もしくは収束遅延AdamW LR から 3〜10× 下げて再スウィープ
6SAM の perturb と update で異バッチ汎化改善消失、振動同一ミニバッチを両 fwd-bwd で使用
7Shampoo の preconditioner 更新間隔過長phase 転換後に劣化T_inv を phase 境界で短縮、condition number 監視
8LoRA + AdamW で frozen 重みに WD 漏れ凍結 backbone が縮小trainable と frozen を別 param group へ (weight_decay=0)
9Lion + grad accumulation の sign 集約誤差accumulation 有無で curve 不一致micro-batch sign の majority vote 集約
10Adafactor + 外部 LR schedule の二重 scaling最終性能低下relative_step=False で絶対 LR schedule に統一

2026 フロンティア

MoE オプティマイザ: DeepSeek-V3 (arXiv:2412.19437) の Top-K routing で expert 勾配が疎、v_t 蓄積が N/K× 遅く実効 LR が expert 不均一 — expert-aware gradient clipping が探索中。fp8 optimizer state: (m, v)E4M3 で持ち Transformer Engine の amax history window=16 で動的 scaling、MS-AMP O3 (arXiv:2310.18313) は全 fp8 化。MUON (NanoGPT speedrun 2024〜2025): Nesterov SGD の更新行列を Newton-Schulz 反復で orthogonalize し等方更新化、Shampoo preconditioner を gradient post-processing として軽量化。PSGD-Kron は Newton-Schulz オンライン推定の近縁。Online μP + auto-LR: 訓練中に LR を online 推定、loss spike auto-recovery (grad_norm/EMA > 10 で batch skip + 100 step rollback + LR×0.9) は Megatron-Core / DeepSpeed で標準装備化中。

See also

Local graph