ML

Optimizers: Adam, AdamW, LION & Beyond

The mathematics of neural network optimization: Adam's moment estimation, weight decay vs L2 regularization, LION's sign-based updates, Sophia's second-order information, learning rate scheduling, and optimizer selection.

Advanced

Prerequisites

Table of Contents

  1. Learning Objectives
  2. Notation
  3. Core Intuition
  4. SGD with Momentum
  5. Adam: Adaptive Moments
  6. AdamW: Decoupled Weight Decay
  7. LION: Sign-Based Optimization
  8. Sophia: Second-Order Information
  9. Learning Rate Scheduling
  10. Optimizer Selection Guide
  11. Common Pitfalls
  12. Summary
  13. Exercises

Learning Objectives

  1. Derive Adam from first principles (moment estimation + bias correction).
  2. Explain why AdamW differs from Adam + L2 and when it matters.
  3. Describe LION's memory-efficient sign updates.
  4. Analyze Sophia's diagonal Hessian approximation.
  5. Design learning rate schedules for different training phases.

Notation

  • mtm_t — first moment (gradient mean)
  • vtv_t — second moment (gradient variance)
  • β1,β2\beta_1, \beta_2 — decay rates for moments

Core Intuition

SGD takes the same step size in every direction — inefficient when gradients vary wildly across parameters (common in deep networks). Adam adapts the step size per parameter based on gradient history: parameters with consistently large gradients get smaller steps (already optimized); parameters with small, noisy gradients get larger steps (need more exploration). This makes training much faster and more stable.

Interactive: Learning Rate Schedules

Training StepLR

Schedules:

Insight: Warmup+Cosine is the most common modern schedule. Warmup prevents early instability, and cosine decay provides smooth annealing. Step decay was popular in CNNs but cosine dominates in transformer training.

SGD with Momentum

Vanilla SGD: θt+1=θt−ηgt\theta_{t+1} = \theta_t - \eta g_t.

With momentum:

mt=βmt−1+gt,θt+1=θt−ηmt.(1)m_t = \beta m_{t-1} + g_t, \quad \theta_{t+1} = \theta_t - \eta m_t. \tag{1}

Effect: Accelerates in consistent gradient directions; dampens oscillations in inconsistent directions.

Nesterov momentum: Look ahead before computing gradient:

gt=∇L(θt−ηβmt−1),mt=βmt−1+gt.(2)g_t = \nabla\mathcal{L}(\theta_t - \eta\beta m_{t-1}), \quad m_t = \beta m_{t-1} + g_t. \tag{2}

Adam: Adaptive Moments

Kingma & Ba (2015):

mt=β1mt−1+(1−β1)gt,vt=β2vt−1+(1−β2)gt2.(3)m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t, \quad v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2. \tag{3}

Bias correction:

m^t=mt1−β1t,v^t=vt1−β2t.(4)\hat{m}_t = \frac{m_t}{1-\beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1-\beta_2^t}. \tag{4}

Update:

θt+1=θt−ηm^tv^t+ϵ.(5)\theta_{t+1} = \theta_t - \eta\frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}. \tag{5}

Interpretation: Step size is η/v^t\eta / \sqrt{\hat{v}_t}. Parameters with large gradient variance get smaller steps (cautious). Parameters with small variance get larger steps (confident).

Standard hyperparameters: β1=0.9,β2=0.999,ϵ=10−8\beta_1=0.9, \beta_2=0.999, \epsilon=10^{-8}.

Memory: 2x model size (stores mm and vv).

AdamW: Decoupled Weight Decay

Loshchilov & Hutter (2019): L2 regularization in Adam is NOT the same as weight decay:

Adam + L2 (wrong):

gt=∇L+λθt,then apply Adam normalization.(6)g_t = \nabla\mathcal{L} + \lambda\theta_t, \quad \text{then apply Adam normalization.} \tag{6}

The λθt\lambda\theta_t term gets divided by v^t\sqrt{\hat{v}_t} — different parameters decay at different rates.

AdamW (correct):

θt+1=θt−ηm^tv^t+ϵ−ηλθt.(7)\theta_{t+1} = \theta_t - \eta\frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} - \eta\lambda\theta_t. \tag{7}

Weight decay is applied DIRECTLY to parameters, not through the adaptive mechanism.

Why it matters: AdamW gives consistent regularization across parameters. Empirically better for transformers (0.5-1% improvement over Adam + L2).

Standard λ\lambda: 0.01-0.1 for transformers.

LION: Sign-Based Optimization

Chen et al. (2023, Google Brain): Uses only the SIGN of the gradient:

mt=β1mt−1+(1−β1)gt,(8)m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t, \tag{8} θt+1=θt−η⋅sign(β2mt−1+(1−β2)gt)−ηλθt.(9)\theta_{t+1} = \theta_t - \eta \cdot \text{sign}(\beta_2 m_{t-1} + (1-\beta_2)g_t) - \eta\lambda\theta_t. \tag{9}

Key properties:

  • Memory efficient: Only stores mtm_t (1x model size vs 2x for Adam).
  • Uniform update magnitude: Every parameter changes by exactly ±η\pm\eta per step.
  • Better on large models: Matches or exceeds AdamW on ViT and LLMs.

When to use: Large models where memory is constrained; works especially well for vision transformers and LLMs above 7B.

Sophia: Second-Order Information

Liu et al. (2023): Use diagonal Hessian to adapt step size:

θt+1=θt−ηgtmax⁡(ht,γ),(10)\theta_{t+1} = \theta_t - \eta\frac{g_t}{\max(h_t, \gamma)}, \tag{10}

where hth_t is the diagonal Hessian estimate (curvature).

Hessian estimation (Hutchinson): ht=E[(g⊙g′)⋅sign]h_t = \mathbb{E}[(g \odot g') \cdot \text{sign}] computed cheaply via one extra backward pass per kk steps.

Advantage: In high-curvature directions (sharp minima), takes smaller steps. In flat directions, takes larger steps. More informed than Adam's variance-based adaptation.

2x speedup over AdamW in wall-clock time for reaching the same loss (fewer steps needed due to better step sizes).

Learning Rate Scheduling

Warmup + cosine decay (standard for transformers):

η(t)={ηmax⁡⋅t/Twt<Twηmin⁡+ηmax⁡−ηmin⁡2(1+cos⁡π(t−Tw)T−Tw)t≥Tw(11)\eta(t) = \begin{cases}\eta_{\max} \cdot t/T_w & t < T_w \\ \eta_{\min} + \frac{\eta_{\max}-\eta_{\min}}{2}(1+\cos\frac{\pi(t-T_w)}{T-T_w}) & t \geq T_w\end{cases} \tag{11}

WSD (Warmup-Stable-Decay): Warmup → constant → cooldown:

  • Warmup: 1-5% of training.
  • Stable: 80-90% at peak LR.
  • Decay: 10-15% cosine to ηmin⁡\eta_{\min}.
  • Advantage: Can resume training from the stable phase without schedule mismatch.

Typical values:

  • Peak LR: 3×10−43 \times 10^{-4} for pre-training; 10−510^{-5} for fine-tuning.
  • Warmup: 2000 steps (pre-training); 100 steps (fine-tuning).
  • ηmin⁡\eta_{\min}: 0.1 × ηmax⁡\eta_{\max}.

Optimizer Selection Guide

ScenarioRecommendedReason
LLM pre-trainingAdamW or LIONProven at scale
LLM fine-tuningAdamWStable, well-understood
Vision TransformerLION or AdamWLION saves memory
CNN (ResNet)SGD + momentumStill competitive
Memory-constrainedLION (1x state)33% memory savings
Maximum speedSophia2x fewer steps

Common Pitfalls

Pitfall 1. Using Adam without weight decay for transformers. Without weight decay, transformers overfit significantly. Always use AdamW with λ=0.01\lambda=0.01-0.10.1.

Pitfall 2. Skipping warmup. Without warmup, initial large gradients (random initialization) cause Adam's second moment to miscalibrate, leading to unstable early training.

Pitfall 3. Using the same learning rate for all model sizes. Optimal LR decreases with model size: η∝1/N\eta \propto 1/\sqrt{N} approximately. A 70B model needs lower LR than a 7B model.

Summary

  • Adam: Adaptive per-parameter LR via gradient moments; standard choice.
  • AdamW: Decoupled weight decay (correct regularization for transformers).
  • LION: Sign-based; 33% less memory; competitive quality.
  • Sophia: Diagonal Hessian; 2x fewer steps; more compute per step.
  • LR schedule: Warmup + cosine/WSD is the universal recipe.
  • Optimizer choice matters most at scale (pre-training). Less impact at fine-tuning.

Exercises

Exercise 1. Derive Adam's bias correction (equation 4) from the assumption that m0=v0=0m_0 = v_0 = 0.

Exercise 2. For a parameter with gradient history [0.1,0.1,0.1,10.0][0.1, 0.1, 0.1, 10.0]: compute Adam vs SGD updates and explain why Adam is more conservative on the outlier step.

Exercise 3. Compute the memory savings of LION vs AdamW for a 70B model in BF16.

Exercise 4. Design a WSD learning rate schedule for a 2T-token pre-training run. Specify warmup, stable, and decay durations and LR values.

Exercise 5. Explain why SGD still outperforms Adam for ResNets on ImageNet (hint: implicit regularization of large-batch SGD).

Next

Curriculum Learning & Data Scheduling