3.5 — Adam: momentum + RMSProp + bias correction

1. Why this session exists

Adam (Kingma & Ba, 2014)[1] is the optimiser you will see in 95% of deep-learning code written since 2015 — every GPT, every BERT, every diffusion model. It is also the most assembled optimiser in this chapter: a clean stack of the three ideas we just built. Once you see it this way, the paper's equations stop looking like magic and become inevitable.

The recipe:

1. Maintain an EMA $m_t$ of the gradient. (Momentum, 3.3.)
2. Maintain an EMA $v_t$ of the squared gradient. (RMSProp, 3.4.)
3. Bias-correct both. (3.2.)
4. Take a step in $-\alpha \cdot \hat m_t / (\sqrt{\hat v_t} + \varepsilon)$.

That's the entire algorithm. The rest of this session is one derivation, one careful look at the bias correction, and a handful of practical notes.

2. The Adam update, derived

At step $t$, compute the (mini-batch) gradient $g_t = \nabla L(\theta_t)$. Then:

(a) First-moment EMA — a velocity, mean of recent gradients:

(b) Second-moment EMA — element-wise mean of squared gradients:

where $g_t^{\odot 2}$ denotes the element-wise square.

(c) Bias correction (from 3.2): both EMAs are initialised at zero, so they are biased toward zero in early steps. Correct:

(d) Update:

All operations after the gradient are element-wise. $\varepsilon \sim 10^{-8}$ for numerical stability of the division.

That's Adam, equation (1) of the paper. Read it: each parameter coordinate gets its own direction (sign of the momentum) and its own step size (inversely proportional to its recent gradient magnitude).

3. What each piece does

Re-cast each component in terms of what it solves:

• The $m_t$ term smooths the gradient over time, providing variance reduction (cf. 3.1's batch-noise argument) and acceleration in consistent directions (cf. 3.3's valley argument).
• The $\sqrt{\hat v_t}$ term rescales per coordinate, equalising the magnitude of steps across parameters with very different gradient statistics (cf. 3.4).
• The bias correction $1/(1 - \beta_i^t)$ fixes the cold-start bias of zero-initialised EMAs (cf. 3.2). Without it, the first few hundred steps of training would be artificially small.

The combined unit analysis is clean: $\hat m_t$ has the units of $\theta$ per unit time; $\sqrt{\hat v_t}$ has the units of $|\theta|$ per unit time; their ratio is dimensionless; $\alpha$ supplies the units. The effective per-coordinate learning rate is

bounded above by $\alpha / \varepsilon$ (which is huge — $10^5$ for $\alpha = 10^{-3}, \varepsilon = 10^{-8}$) but in practice always much smaller because $\hat v_t^{(i)}$ dominates $\varepsilon$ after the first few steps.

A more useful bound: the signal-to-noise form. Note

with equality iff the gradient is constant in sign. So at each coordinate, Adam's step is at most $\alpha$ in absolute value (modulo the $\varepsilon$ correction). The per-parameter LR is bounded by $\alpha$, which is one of the reasons Adam is so easy to tune — you almost never have to worry about it blowing up a coordinate.

4. The bias correction in practice

To see the correction at work, take $\beta_1 = 0.9$, $\beta_2 = 0.999$, a constant gradient $g = 1$, and start from $m_0 = v_0 = 0$.

At $t = 1$:

So $\hat m_1 / \sqrt{\hat v_1} = 1$, and the first Adam step is exactly $-\alpha \cdot 1$. Without bias correction, the first step would be $-\alpha \cdot 0.1 / \sqrt{0.001} = -\alpha \cdot 3.16$ — overshooting in the other direction by a factor of 3, because $v_1$ is underestimating the variance more than $m_1$ is underestimating the mean.

The correction is what makes early steps well-scaled.

5. Compute and memory

Per parameter, Adam stores two extra scalars ($m$ and $v$) and performs $\sim 10$ float operations per step (two EMA updates, two corrections, one division, one $\sqrt{}$, one step). Memory cost: 3× the model parameters ($\theta, m, v$). For a 7B-parameter model in float32, that is $\sim 84$ GB of optimiser state alone, plus the model and gradients. This is why "Adam optimiser state" is the dominant memory cost in LLM training and why much of the recent work (Shazeer's Adafactor, bitsandbytes 8-bit Adam, ZeRO sharding) is about compressing or sharding $(m, v)$.

For convolutional and small-scale models the memory is rarely an issue. For LLMs it is a first-class concern.

6. Default hyperparameters and convergence intuition

The Adam paper's recommended defaults — used unchanged by nearly every major model — are

For LLMs, $\alpha$ is usually scaled down to $10^{-4}$ or $3 \times 10^{-4}$, and $\beta_2$ is sometimes pushed to $0.95$ (more reactive second moment, helps with noisy LM gradients).

Adam does not have a convergence guarantee for non-convex losses. Reddi et al. (2018, "On the Convergence of Adam") even showed it can diverge in some convex problems where AdaGrad converges. The fix in that paper (AMSGrad) keeps $\hat v_t$ monotonically non-decreasing. It is correct theory and almost no-one uses it in practice — vanilla Adam works empirically.

7. Pseudocode (PyTorch-equivalent)

# state per parameter: m, v
for t, batch in enumerate(loader, start=1):
    g = compute_gradient(theta, batch)

    m = beta1 * m + (1 - beta1) * g
    v = beta2 * v + (1 - beta2) * (g * g)

    m_hat = m / (1 - beta1 ** t)
    v_hat = v / (1 - beta2 ** t)

    theta = theta - alpha * m_hat / (np.sqrt(v_hat) + eps)

Ten lines. Drop in $\beta_1 = 0.9, \beta_2 = 0.999, \alpha = 10^{-3}, \varepsilon = 10^{-8}$ and you have the most-used optimiser of the last decade.

8. Takeaways

• Adam = momentum ($m_t$) + RMSProp ($v_t$) + bias correction.
• Update: $\theta \leftarrow \theta - \alpha\, \hat m_t / (\sqrt{\hat v_t} + \varepsilon)$, element-wise.
• Per-coordinate step is bounded above by $\alpha$, which makes the LR robust to tune.
• Memory cost: 3× model parameters. The main bottleneck in LLM training optimiser state.
• Defaults that work: $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\alpha = 10^{-3}$ (or $10^{-4}$ for LLMs), $\varepsilon = 10^{-8}$.
• Adam has no general non-convex convergence guarantee; we use it because it works.