Foundation reference

Adam & Adaptive Gradient Methods

Adam updates each parameter using running averages of its gradient and squared gradient. The published GPT-3 and Llama training recipes use Adam or AdamW (Brown et al., 2020; Touvron et al., 2023).

A newer notebook covers this concept. This older page keeps its summary, math and sources.

Concept 3 of 100OptimizationStudy phase 3
Key equation
θt+1=θt−αm^tv^t+ε\theta_{t+1} = \theta_t - \alpha \frac{\hat m_t}{\sqrt{\hat v_t} + \varepsilon}
Reading map and next steps

Study guide

Where to start

Concept 3 of 100AdamOptimization · Study phase 3: Optimization & generalization
Question

How does dividing each coordinate by √v̂ₜ change which solution an overparameterized network reaches, and why does Adam sometimes generalize worse than SGD (Wilson et al., 2017)?

How to study this pagePredict · Check · Generalize · Continue
PredictSay what you expect.

Before you move a control in the demo, say what you expect to change.

CheckCheck against the math and the papers.

Compare what the demo shows with the key equation and with the papers cited below.

GeneralizeName what stays the same.

Identify the part of the idea that survives a change of notation or model, so you can recognize it in another paper or system.

ContinueContinue in the newer notebook.

The newer notebook covers the same concept and has an interactive demo.

Why It Matters for Modern Models

  • Adam updates each parameter using running averages of its gradient and squared gradient. The published GPT-3 and Llama training recipes use Adam or AdamW (Brown et al., 2020; Touvron et al., 2023).
  • Not every large model uses them; PaLM, for example, was trained with Adafactor (Chowdhery et al., 2023).
  • Adam-family optimizers are also the usual choice for fine-tuning, RLHF and diffusion-model training, where gradient scales differ widely from one parameter to another.

Questions to Think About

  • How do bias correction and exponential averaging behave when the objective keeps changing, as in curriculum training or RLHF?

Interactive Demos

Core Math

If you prefer intuition first, try the demos above, then come back for the math.

Key Equation
θt+1=θt−αm^tv^t+ε\theta_{t+1} = \theta_t - \alpha \frac{\hat m_t}{\sqrt{\hat v_t} + \varepsilon}

For gradient gt=∇θLt(θt)g_t = \nabla_\theta L_t(\theta_t):

mt=β1mt−1+(1−β1)gtvt=β2vt−1+(1−β2)gt2m^t=mt/(1−β1t),v^t=vt/(1−β2t)θt+1=θt−αm^tv^t+ε\begin{aligned} m_t &= \beta_1 m_{t-1} + (1-\beta_1) g_t \\ v_t &= \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \\ \hat m_t &= m_t / (1-\beta_1^t),\quad \hat v_t = v_t / (1-\beta_2^t) \\ \theta_{t+1} &= \theta_t - \alpha \frac{\hat m_t}{\sqrt{\hat v_t} + \varepsilon} \end{aligned}

Here mtm_t is an exponential moving average of the gradient and vtv_t an exponential moving average of its elementwise square (an uncentred second-moment estimate); dividing by 1−β1t1-\beta_1^t and 1−β2t1-\beta_2^t corrects the bias from initializing both at zero.

Reddi et al. (2018) gave a simple convex problem on which Adam fails to converge to the optimum, and proposed the AMSGrad variant.

Key Papers

Adam: A Method for Stochastic Optimization

Kingma & Ba2015ICLR
Read paper →

On the Convergence of Adam and Beyond

Reddi et al.2018ICLR
Read paper →

Connections

Where to go next

A newer notebook covers Adam, with an interactive demo.

Continue in the newer notebook
Other connections (4)

Which limitations does it address, or leave to a later idea?

  • Stochastic MLE: a limitation of Maximum Likelihood that Adam addresses. Open Maximum Likelihood

    Training by maximum likelihood means minimizing an average cross-entropy over a large dataset, in practice with noisy minibatch gradients. Adam was designed for such stochastic objectives (Kingma & Ba, 2015): momentum smooths the noisy gradients, and per-parameter step sizes help with sparse gradients, such as those of rare-word embeddings.

Where does it break down?

  • Edge of Stability: where Adam breaks down, taken up in Sharpness. Open Sharpness

    With full-batch gradient descent the sharpness, the largest eigenvalue of the Hessian, rises until it hovers near 2/η, where η is the step size; the loss keeps falling, though not monotonically, even though classical analysis says steps this large should diverge (Cohen et al., 2021). Adam behaves the same way for its preconditioned sharpness, which settles near 38/η when β₁ = 0.9 (Cohen et al., 2022).

Where else does the same idea appear?

  • Adaptive step sizing: the same technique appears in Efficiency. Open Efficiency

    Adam's per-parameter step sizes let one learning rate work across parameters whose gradients differ greatly in scale. Parameter-efficient fine-tuning methods such as LoRA and QLoRA are usually trained with Adam-type optimizers (Hu et al., 2022; Dettmers et al., 2023).

  • Two learning rates: the same technique appears in GANs. Open GANs

    GANs are commonly trained with Adam, and many setups give the generator and the discriminator different learning rates. Heusel et al. (2017) call this the two time-scale update rule and prove that, under mild assumptions, it converges to a local Nash equilibrium, a result that carries over to Adam.