Adam: momentum, scale and bias correction

Why we need Adam

A neural network has millions of weights, and their gradients live on very different scales: embeddings of rare tokens get rare large kicks, the last layers get constant small ones, normalization layers yet others. Plain SGD takes the step $\eta g$, and a single rate $\eta$ is at once too large for some weights and too small for others. On top of that, the mini-batch gradient is noisy, and the loss surface is full of elongated ravines and plateaus.

Adam (adaptive moment estimation, Kingma and Ba, 2014; ICLR 2015) tackles both problems with one formula:

• the first moment $m_t$ — an exponential average of gradients: inertia that damps noise and zigzags (like Momentum);
• the second moment $v_t$ — an exponential average of squared gradients: the typical gradient size of each weight, by which the step is divided (like RMSProp);
• bias correction — division by $1-\beta^t$, which removes the systematic underestimate of the moments in the first steps.

The result: every weight’s step is of order $\eta$ regardless of its gradient’s scale, with no manual tuning per layer. That is why Adam is the default optimizer for most tasks, and its variant AdamW trains transformers and large language models.

Intuition: a compass and a speedometer

Imagine driving off-road by a jittery compass. Turning the wheel at every twitch would be silly — it is wiser to follow the averaged direction. That is $m_t$. And the speed should depend on how strong the twitches usually are: on a smooth road you can drive confidently, on bumps carefully. That is $\sqrt{v_t}$: Adam’s step is the direction divided by the typical strength of the jolts.

The figure shows the gradient stream of one weight: mean $+0.5$ until step 100 and $-0.5$ after, plus noise. The orange line is the smoothed gradient $\hat m_t$, the teal lines are the corridor $\pm\sqrt{\hat v_t}$, and below is the resulting step in units of $\eta$.

The formula, bone by bone

$g_t$

the mini-batch gradient at step $t$ (steps are numbered from 1);

$m_t$

the first moment: an exponential moving average (EMA) of the gradient, $m_0=0$;

$v_t$

the second moment: an EMA of the element-wise squared gradient, $v_0=0$; every weight has its own $m$ and $v$;

$\beta_1,\ \beta_2$

decay rates, $0.9$ and $0.999$ by default: memory of about $1/(1-\beta)$ steps — 10 and 1000;

$1-\beta^t$

bias correction: compensates for the moments starting at zero;

$\eta$

the learning rate, $10^{-3}$ by default; here it sets the typical step length in the units of the weights;

$\epsilon$

a small constant ($10^{-8}$) against division by zero.

Let us run three steps by hand for the gradients $0.5$, $-0.2$, $0.4$ with the standard $\beta_1=0.9$, $\beta_2=0.999$. The last column shows what the step would be without bias correction.

Note the third step: the gradients changed sign, so $\hat m_3=0.23$ is noticeably smaller than the typical gradient size $\sqrt{\hat v_3}=0.39$, and Adam steps only $0.6\eta$. The more consistent the gradients, the closer the step is to $\eta$.

Bias correction: where 1 − βᵗ comes from

The moments start at zero, and the first steps “remember” that zero. Unroll the recursion for $m_t$:

This is a weighted sum of gradients, but the weights add up not to 1 but to $(1-\beta_1)(1+\beta_1+\dots+\beta_1^{t-1})=1-\beta_1^t$. The missing share $\beta_1^t$ went to the initial zero. Hence the correction.

Properties of the step

Scale invariance

Multiply all gradients of a weight by $c>0$. Then $m_t$ is multiplied by $c$, $v_t$ by $c^2$, $\sqrt{v_t}$ by $c$, and the ratio $\hat m_t/\sqrt{\hat v_t}$ does not change. Adam’s step is a homogeneous function of degree zero in the gradients (as long as $\sqrt{\hat v_t}\gg\epsilon$). Each weight’s effective rate $\eta/(\sqrt{\hat v_t}+\epsilon)$ adapts by itself: large for weights with small gradients, small for weights with large ones.

The step and the signal-to-noise ratio

Let a weight’s gradient be stationary with mean $\mu$ and standard deviation $\sigma$. Then $\hat m_t\approx\mu$, $\hat v_t\approx\mu^2+\sigma^2$ (the second moment is the squared mean plus the variance), and the mean step is

With a clean signal ($\sigma\ll|\mu|$) the step is $\approx\eta$; in strong noise ($\sigma\gg|\mu|$) it is proportional to the signal-to-noise ratio. This is automatic “braking”: near the minimum the mean gradient tends to zero while the noise remains — and Adam shortens its steps by itself.

Special cases

• $\beta_1=\beta_2=0$: $\hat m_t=g_t$, $\sqrt{\hat v_t}=|g_t|$, the step is $\eta\operatorname{sign}(g_t)$ — sign gradient descent (signSGD).
• $\beta_1=0$: RMSProp with bias correction.
• $\epsilon\gg\sqrt{\hat v_t}$: the step is $\approx(\eta/\epsilon)\,\hat m_t$ — SGD with momentum and rate $\eta/\epsilon$.
• Adam is a diagonal preconditioner: it scales each coordinate separately and cannot rotate the axes (see Figure 6 on the SGD page).

The mathematics behind the formula

Let us collect what hides behind the four lines.

• EMA as a weighted average. $\hat m_t=\sum_i w_i g_i$ with weights $w_i=(1-\beta_1)\beta_1^{t-i}/(1-\beta_1^t)$, which are positive and sum to 1. The weight halves roughly every $\ln2/(1-\beta)$ steps: 7 for $\beta_1=0.9$ and 693 for $\beta_2=0.999$.
• The effective window $1/(1-\beta)$: $m$ averages about the last 10 gradients, $v$ about the last 1000. The direction reacts fast, the scale slowly and therefore stably.
• The second moment is not the variance. $v$ estimates $\mathbb E[g^2]=\mu^2+\sigma^2$, the “uncentred” size of the gradient. That is why with a clean signal the step is $\approx\eta$ rather than infinite.
• The link to Newton’s method. It is sometimes said that $\sqrt v$ approximates the curvature. That is only partly true: $v$ estimates the diagonal of the empirical Fisher matrix (from mini-batch gradients), which for some losses is close to the Hessian near the minimum, but Adam takes its square root. It is safer to see Adam as step normalization rather than as an approximation to Newton.

An optimizer race in 3D

Let us compare five methods on two surfaces. The first is a plateau with a pit: $f=1-e^{-q(x,y)}$, $q=0.2(x-1)^2+1.2(y-0.5)^2$. Far from the pit $f\approx1$ and the gradient is tiny, while the pit itself is elongated. The second is the Rosenbrock banana.

AdamW: weight decay done right

Classic $L_2$ regularization adds $\tfrac\lambda2\|\theta\|^2$ to the loss, i.e. $\lambda\theta$ to the gradient. For SGD this is the same as multiplying the weights by $1-\eta\lambda$ at every step (weight decay). For Adam it is not: the term $\lambda\theta$ enters $m$ and $v$ and gets divided by $\sqrt{\hat v_t}$. Loshchilov and Hutter (2017; ICLR 2019) proposed to decouple decay from the gradient:

In AdamW the moments are computed from the loss gradient alone, and all weights decay at the same rate $\eta\lambda$ per step. In Adam with $L_2$ each weight’s decay is divided by its $\sqrt{\hat v_t}$: weights with large gradients are barely regularized, while weights with no loss gradient shrink at a constant speed of $\approx\eta$ per step, regardless of $\lambda$.

Adam in training: warmup and schedules

Even an adaptive method needs a rate schedule. The most common scheme for transformers is a linear warmup from zero to $\eta_{\max}$ over the first hundreds or thousands of steps, followed by a cosine decay to $\eta_{\min}$:

Why decay helps: from the section on properties, Adam’s step near the minimum is proportional to $\eta$ times the signal-to-noise ratio. The smaller $\eta$, the smaller the cloud in which the weights jitter around the minimum.

Why warmup is needed: in the first steps $\hat v_t$ is estimated from a few gradients and is very noisy, while the weights have just been initialized and the gradients are especially large. Liu et al. (RAdam, 2019) showed that the variance of the effective rate $\eta/\sqrt{\hat v_t}$ is large early in training and proposed to rectify it analytically; warmup solves the same problem more simply.

In code: from scratch and in PyTorch

A from-scratch implementation takes twenty lines. The decoupled flag switches weight decay between AdamW and $L_2$ inside the gradient.

pythonimport numpy as np

class Adam:
    """Adam; with weight_decay > 0 and decoupled=True it is AdamW. As in torch.optim."""
    def __init__(self, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0, decoupled=True):
        self.lr, (self.b1, self.b2), self.eps = lr, betas, eps
        self.wd, self.decoupled = weight_decay, decoupled
        self.m = self.v = None
        self.t = 0

    def step(self, theta, g):
        if self.m is None:
            self.m, self.v = np.zeros_like(theta), np.zeros_like(theta)
        if self.wd and not self.decoupled:
            g = g + self.wd * theta                        # L2 inside the gradient (Adam)
        self.t += 1
        self.m = self.b1 * self.m + (1 - self.b1) * g      # momentum
        self.v = self.b2 * self.v + (1 - self.b2) * g * g  # scale
        m_hat = self.m / (1 - self.b1 ** self.t)           # bias correction
        v_hat = self.v / (1 - self.b2 ** self.t)
        if self.wd and self.decoupled:
            theta = theta - self.lr * self.wd * theta      # AdamW: decay applied separately
        return theta - self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

# The three steps from the table: g = 0.5, −0.2, 0.4 (lr = 1, so the step is in units of η)
opt, th = Adam(lr=1.0), np.zeros(1)
for g in [0.5, -0.2, 0.4]:
    new = opt.step(th, np.array([g]))
    print(f"t={opt.t}: m̂={opt.m[0] / (1 - 0.9 ** opt.t):.4f}  step/η={th[0] - new[0]:.4f}")
    th = new
# t=1: m̂=0.5000  step/η=1.0000
# t=2: m̂=0.1316  step/η=0.3456
# t=3: m̂=0.2306  step/η=0.5955

# Scale invariance: gradients ×1000 give the same steps
rng = np.random.default_rng(0)
a, b, x, y = Adam(lr=0.01), Adam(lr=0.01), np.zeros(1), np.zeros(1)
for t in range(100):
    g = rng.normal(1.0, 0.5, size=1)
    x, y = a.step(x, g), b.step(y, 1000 * g)
print(x, y, abs(x - y))   # [-0.95081564] [-0.95081565] [8.64890637e-09] — the difference comes from ε alone

The output matches the table from “The formula, bone by bone”, and torch.optim.Adam moves the weight by exactly $\eta$ on the same first step. A typical training loop with AdamW, separate decay groups, warmup and cosine:

pythonimport math
import torch
import torch.nn.functional as F

decay = [p for p in model.parameters() if p.ndim >= 2]     # weight matrices
no_decay = [p for p in model.parameters() if p.ndim < 2]   # biases and normalization layers
opt = torch.optim.AdamW([{"params": decay, "weight_decay": 0.1},
                         {"params": no_decay, "weight_decay": 0.0}],
                        lr=3e-4, betas=(0.9, 0.95), eps=1e-8)

total, warmup = 10_000, 500
def lr_factor(step):                       # lr multiplier: warmup, then cosine down to 10%
    if step < warmup:
        return (step + 1) / warmup
    p = (step - warmup) / (total - warmup)
    return 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * p))
sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_factor)

for step, (x, y) in zip(range(total), loader):
    loss = F.mse_loss(model(x), y)
    opt.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # gradient clipping
    opt.step()
    sched.step()                          # the schedule steps after every update

Hyperparameters and practice

• Start with the defaults. $\eta=10^{-3}$ (or $3\cdot10^{-4}$ for transformers), $\beta=(0.9,\,0.999)$, AdamW with $\lambda=0.01$. Tune $\eta$ first, then $\lambda$ and the schedule.
• Gradient clipping by norm ($1.0$) is almost free and saves you from rare outliers.
• Save the optimizer state (opt.state_dict()) together with the weights: when resuming with zero $m$ and $v$ the first steps will be “cold” again.
• When SGD is better. On computer-vision tasks a well-tuned SGD with momentum often generalizes better (Wilson et al., 2017). For transformers, recommender systems with embeddings and problems with sparse gradients Adam wins.

Common mistakes

1. Adam(weight_decay=…) instead of AdamW. That is $L_2$ inside the gradient: the decay is divided by $\sqrt{\hat v_t}$ and acts unevenly (Figure 6).
2. A rate taken from SGD. $\eta=0.1$ is fine for SGD with momentum, but for Adam it means a step of about $0.1$ in weight units at every step — training will almost surely diverge.
3. Decay on all parameters. Biases and LayerNorm/BatchNorm weights are usually excluded from weight decay: shrinking them towards zero hurts rather than regularizes.
4. A poor $\epsilon$. A too small one rounds to zero in float16; a too large one ($\gtrsim\sqrt{\hat v}$) silently switches adaptation off — the step becomes $(\eta/\epsilon)\hat m$.
5. No warmup on a transformer. A large $\eta$ from the first steps, with a noisy estimate $\hat v_t$, often makes the loss explode early in training.
6. The schedule on the wrong step. scheduler.step() once per epoch where the schedule is defined in steps (or vice versa) stretches or squeezes warmup hundreds of times.

Links to other formulas

• SGD, Momentum and RMSProp: Adam’s numerator is EMA momentum, its denominator is $\sqrt v$ from RMSProp; that page also covers the stability condition and the condition number.
• Cross-entropy and softmax: the loss AdamW minimizes in language models; its gradient with respect to the logits, $s-y$, is the $g_t$ of the last layer.
• MSE and Huber: the scale of the target sets the scale of the gradient — Adam is insensitive to it, SGD is not.
• GELU and SiLU: the activations of transformers trained with AdamW.

Summary

Check yourself