Training Optimization(学習最適化)
AdamW を 2024〜2026 デファクトとして、SGD/Nesterov・Adam 系統・Lion/Sophia・Shampoo/SOAP・ZeRO/8-bit・SAM/Edge of Stability・10 落とし穴・MUON/fp8 フロンティアを整理。
article technology ja 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 speedup。Nesterov 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+L2 と AdamW の致命的差は重み減衰の適用位置。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) |
|---|---|---|
| 1 | optimizer states (m, v) | ~4× |
| 2 | optimizer + gradients | ~8× |
| 3 | optimizer + 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 つの頻発失敗モード
| # | 落とし穴 | 症状 | 対策 |
|---|---|---|---|
| 1 | Adam bias-correction 過大 (β2=0.999) | 50〜200 step で spike | β2=0.95 で 1/(1−β2) を 1000→20 へ |
| 2 | Checkpoint で optimizer state 未保存 | resume 直後の loss 異常 | (m, v, step_t) 完全保存・1-step delta で検証 |
| 3 | Weight tying × decoupled WD で embedding 過剰収縮 | rare token zero-shot 低下 | embedding は weight_decay=0 (GPT-3/Llama 系譜) |
| 4 | fp16 optimizer state × bf16 master のミスマッチ | 数千 step 後 NaN | (m, v) を fp32 維持、Megatron-LM 標準 |
| 5 | Lion の LR スウィープを Adam スケールで実施 | 発散もしくは収束遅延 | AdamW LR から 3〜10× 下げて再スウィープ |
| 6 | SAM の perturb と update で異バッチ | 汎化改善消失、振動 | 同一ミニバッチを両 fwd-bwd で使用 |
| 7 | Shampoo の preconditioner 更新間隔過長 | phase 転換後に劣化 | T_inv を phase 境界で短縮、condition number 監視 |
| 8 | LoRA + AdamW で frozen 重みに WD 漏れ | 凍結 backbone が縮小 | trainable と frozen を別 param group へ (weight_decay=0) |
| 9 | Lion + grad accumulation の sign 集約誤差 | accumulation 有無で curve 不一致 | micro-batch sign の majority vote 集約 |
| 10 | Adafactor + 外部 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
- tech-50 Training Hyperparameters — 数値設定 (LR / β / batch / clip / μP / fp8)
- tech-49 Training Paradigm — pretraining + SFT/DPO/RLVR
- tech-48 Distributed Training — DP/TP/PP/EP・ZeRO・FSDP
- Source mindmap canvas
Backlinks
10 backlinks
- related Scaling and Operations(スケーリング則と LLMOps)
- related Reinforcement Learning(強化学習)
- related Containerization and Orchestration(コンテナ化とオーケストレーション)
- related Neural Networks and Deep Learning(ニューラルネットワークとディープラーニング)
- related Machine Learning(機械学習)
- related Kernel and Memory Optimization(カーネル・メモリ最適化)
- related Model Serialization Format(モデルシリアライゼーションフォーマット)
- related Model Runtime and Conversion(モデルランタイムと変換)
- related Model Compression Techniques(モデル圧縮)
- related Model Compression and Optimization(モデル圧縮と最適化)