Adam

Adam is a variant of stochastic gradient descent that gives every parameter its own step size. It keeps two exponential moving averages per parameter, one of the gradient and one of its elementwise square, and steps by the first divided by the square root of the second. Kingma and Ba (2015) named it for adaptive moment estimation. The Vision Transformer, for example, was pre-trained with Adam, whereas ResNets for ImageNet classification were trained with SGD and momentum.

Problem

SGD uses one learning rate for every weight, yet gradients differ greatly in size between layers, between biases and weights, and between weights attached to often-bright and to almost-always-dark pixels. Kingma and Ba note that weight sharing gives the layers of a convolutional network very different gradients, so SGD users often give convolution layers a smaller learning rate by hand. A rate small enough for the steepest direction is too small for the rest, the ill-conditioning described for gradient descent. Adam rescales each coordinate by an estimate of its own gradient magnitude: a cheap diagonal preconditioner built from the gradients backpropagation already provides.

Inputs and Outputs

• Input: a stochastic objective (a network’s mini-batch loss) and its gradients, initial parameters $\theta_0$, a step size $\alpha$, decay rates $\beta_1, \beta_2 \in [0, 1)$, and a small constant $\epsilon$.
• Output: the trained parameters. Between steps Adam carries two vectors the size of $\theta$.

Intuition

The first average, $m$, is a smoothed gradient that plays the role of momentum; the second, $v$, tracks how large each weight’s gradient typically is. When a weight’s gradient points the same way step after step, the ratio is near $\pm 1$ and the weight moves by about $\alpha$; when the gradient flips sign from batch to batch, $m$ averages toward zero while $v$ does not, and the step shrinks. Kingma and Ba note that this ratio falls near an optimum, a form of automatic annealing.

Algorithm

Algorithm 1 of Kingma and Ba (2015), with all operations elementwise:

1. Set $m_0 = 0$, $v_0 = 0$, $t = 0$.
2. Repeat: $t \leftarrow t + 1$; compute the mini-batch gradient $g_t = \nabla_\theta f_t(\theta_{t-1})$.
3. Update the first moment, $m_t = \beta_1 m_{t-1} + (1 - \beta_1)\, g_t$, and the second raw moment, $v_t = \beta_2 v_{t-1} + (1 - \beta_2)\, g_t^2$.
4. Correct the bias: $\hat m_t = m_t / (1 - \beta_1^t)$, $\hat v_t = v_t / (1 - \beta_2^t)$.
5. Step: $\theta_t = \theta_{t-1} - \alpha\, \hat m_t / (\sqrt{\hat v_t} + \epsilon)$.

Mathematical Formulation

Why bias correction is needed

Unrolling the recursion from $v_0 = 0$ gives $v_t = (1 - \beta_2)\sum_{i=1}^{t} \beta_2^{t-i} g_i^2$. If the second moment $\mathbb{E}[g^2]$ is stationary, taking expectations and summing the geometric series gives

$$\mathbb{E}[v_t] = \mathbb{E}[g^2]\,(1 - \beta_2) \sum_{i=1}^{t} \beta_2^{t-i} = \mathbb{E}[g^2]\,(1 - \beta_2^t),$$

as derived in Kingma and Ba’s Section 3; likewise $\mathbb{E}[m_t] = \mathbb{E}[g]\,(1 - \beta_1^t)$. Initializing at zero biases both averages toward zero, and dividing by $1 - \beta^t$ removes the bias. With $\beta_2 = 0.999$ the factor is $0.001$ at the first step and still only $0.63$ after 1,000 steps. By our arithmetic, without correction a constant gradient gives steps of $\alpha (1 - \beta_1^t)/\sqrt{1 - \beta_2^t}$: $3.16\alpha$ at first, about $6.6\alpha$ near step 12, and still $1.26\alpha$ at step 1,000. Kingma and Ba stress that the correction matters most for $\beta_2$ close to 1; in their variational-autoencoder experiment, training without it was unstable for such values, especially in the first epochs.

Effective step and scale invariance

Ignoring $\epsilon$, the step is $\Delta_t = \alpha\, \hat m_t / \sqrt{\hat v_t}$. Kingma and Ba argue that $|\Delta_t|$ can approach its larger bound, $\alpha(1 - \beta_1)/\sqrt{1 - \beta_2}$, only when a gradient was zero at every earlier step, and is typically about $\alpha$ or less, so $\alpha$ sets the scale of steps in parameter space, like a trust region. Multiplying every gradient by $c$ multiplies $\hat m_t$ by $c$ and $\hat v_t$ by $c^2$, which cancel: the update is invariant to rescaling the gradient, and, coordinate by coordinate, to a diagonal rescaling.

Precursors

AdaGrad (Duchi, Hazan & Singer, 2011) divides by the root of the sum of past squared gradients, $\theta_{t+1} = \theta_t - \eta\, g_t / \sqrt{\sum_{i \le t} g_i^2}$ in its diagonal form; this suits sparse gradients but keeps shrinking the steps. RMSProp, which Kingma and Ba attribute to a 2012 Coursera lecture by Tieleman and Hinton, replaces the sum by an exponential moving average (decay 0.9 in the lecture). Kingma and Ba designed Adam to combine the strengths of both, adding a moving average of the gradient (instead of momentum applied to the rescaled gradient) and the bias correction RMSProp lacks.

Parameters

• Step size $\alpha$, default $0.001$. Because it sets the step in parameter units, Kingma and Ba argue that its order of magnitude can often be chosen in advance; it still needs tuning and a schedule.
• $\beta_1$, default $0.9$: the gradient is averaged over about $1/(1 - \beta_1) = 10$ steps.
• $\beta_2$, default $0.999$: the squared gradient is averaged over about 1,000 steps.
• $\epsilon$, default $10^{-8}$, prevents division by zero; Reddi, Kale and Kumar (2018) remark that its choice can be critical.

These defaults are Kingma and Ba’s recommendation for the problems they tested, and PyTorch’s.

Complexity

A step costs $O(d)$ for $d$ parameters. Adam stores two vectors, $m$ and $v$, where SGD with momentum stores one, so its state is twice the model’s size.

Implementation

The example trains a two-layer network on scikit-learn’s 8×8 digit images: 64 pixels, 32 ReLU units, 10 classes, cross-entropy loss, batches of 32 for 20 epochs. Its hand-written gradients agree with PyTorch’s autograd to rounding error. The script measures the first step with and without bias correction, sweeps the learning rate of SGD with momentum 0.9 and of Adam on raw pixel values (0 to 16) and on standardized pixels, and multiplies every gradient by 1,000.

import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

SEED = 2
np.seterr(all="ignore")                      # diverging runs overflow; train() stops them

digits = load_digits()                        # 1797 grayscale 8x8 images, pixel values 0..16
X_raw, X_raw_te, y_tr, y_te = train_test_split(
    digits.data, digits.target, test_size=0.25, random_state=0, stratify=digits.target)
mean, std = X_raw.mean(axis=0), X_raw.std(axis=0)
std[std == 0] = 1.0                           # four border pixels are 0 in every training image
X_std, X_std_te = (X_raw - mean) / std, (X_raw_te - mean) / std
N, D, H, K = 1347, 64, 32, 10                 # images, pixels, hidden units, classes
Y_tr = np.eye(K)[y_tr]

def init(seed):
    r = np.random.default_rng(seed)
    return [r.normal(0, np.sqrt(2 / D), (D, H)), np.zeros(H),
            r.normal(0, np.sqrt(2 / H), (H, K)), np.zeros(K)]

def loss_grad(p, X, Y):
    """Two-layer ReLU network, softmax cross-entropy averaged over the images in X."""
    W1, b1, W2, b2 = p
    a = X @ W1 + b1
    h = np.maximum(a, 0)
    s = h @ W2 + b2
    s -= s.max(axis=1, keepdims=True)
    P = np.exp(s)
    P /= P.sum(axis=1, keepdims=True)
    G = (P - Y) / len(X)                      # d loss / d scores
    dA = (G @ W2.T) * (a > 0)
    return -np.mean(np.log(P[Y > 0])), [X.T @ dA, dA.sum(0), h.T @ G, G.sum(0)]

def accuracy(p, X, y):
    W1, b1, W2, b2 = p
    return np.mean((np.maximum(X @ W1 + b1, 0) @ W2 + b2).argmax(axis=1) == y)

class Adam:
    """Algorithm 1 of Kingma & Ba (2015); wd > 0 adds L2 (wd * theta to the gradient)
    or, with decoupled=True, AdamW's decay theta <- theta - lr * wd * theta (PyTorch form)."""
    def __init__(self, params, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8, wd=0.0,
                 decoupled=False, bias_correction=True):
        self.p, self.lr, self.b1, self.b2, self.eps = params, lr, b1, b2, eps
        self.wd, self.decoupled, self.bc = wd, decoupled, bias_correction
        self.m = [np.zeros_like(q) for q in params]
        self.v = [np.zeros_like(q) for q in params]
        self.t = 0

    def step(self, grads):
        self.t += 1
        for q, g, m, v in zip(self.p, grads, self.m, self.v):
            if self.wd and self.decoupled:
                q -= self.lr * self.wd * q            # decoupled weight decay (AdamW)
            elif self.wd:
                g = g + self.wd * q                   # L2 penalty enters the moments
            m[:] = self.b1 * m + (1 - self.b1) * g    # first moment (mean)
            v[:] = self.b2 * v + (1 - self.b2) * g * g  # second raw moment
            m_hat = m / (1 - self.b1 ** self.t) if self.bc else m
            v_hat = v / (1 - self.b2 ** self.t) if self.bc else v
            q -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

class SGDMomentum:
    def __init__(self, params, lr, mu=0.9):
        self.p, self.lr, self.mu = params, lr, mu
        self.buf = [np.zeros_like(q) for q in params]

    def step(self, grads):
        for q, g, b in zip(self.p, grads, self.buf):
            b[:] = self.mu * b + g
            q -= self.lr * b

def train(X, make_opt, seed=SEED, epochs=20, B=32, grad_scale=1.0):
    """Return the parameters and the training loss after each epoch (inf once diverged)."""
    rng, p = np.random.default_rng(seed), init(seed)
    opt = make_opt(p)
    losses = []
    for epoch in range(epochs):
        order = rng.permutation(N)
        for i in range(0, N, B):
            idx = order[i:i + B]
            _, g = loss_grad(p, X[idx], Y_tr[idx])
            opt.step([grad_scale * gi for gi in g])
            if not np.isfinite(p[0]).all():
                return p, losses + [np.inf]
        losses.append(loss_grad(p, X, Y_tr)[0])
    return p, losses

def summary(losses):
    """Final loss, or a label for a run that blew up. A chaotic run's final value depends
    on floating-point rounding, so it can differ between NumPy builds; we do not print it."""
    if not np.isfinite(losses[-1]):
        return "diverged"
    if losses[-1] > 10 * min(losses):
        return "unstable"
    return f"{losses[-1]:.2g}"

# 1. Bias correction: the size of the very first step, |delta theta| / lr.
p = init(SEED)
_, g = loss_grad(p, X_std[:32], Y_tr[:32])
for bc in [True, False]:
    q = [x.copy() for x in p]
    Adam(q, bias_correction=bc).step(g)
    steps = np.concatenate([np.abs(a - b).ravel() for a, b in zip(q, p)]) / 1e-3
    print(f"first step / lr, bias correction {str(bc):5s}: median {np.median(steps[steps > 0]):.3f}")

# 2. Learning-rate sweep, 20 epochs, raw 0..16 pixels versus standardized pixels.
lrs = [1e-4, 3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1]
print("\nfinal training loss after 20 epochs (2.3 = ln 10, chance level)")
print(f"{'':27s}" + "".join(f"{lr:>9g}" for lr in lrs))
for name, X in [("raw pixels", X_raw), ("standardized", X_std)]:
    for opt_name, make in [("SGD + momentum", lambda lr: lambda p: SGDMomentum(p, lr)),
                           ("Adam", lambda lr: lambda p: Adam(p, lr))]:
        cells = [summary(train(X, make(lr))[1]) for lr in lrs]
        print(f"{name:13s} {opt_name:14s}" + "".join(f"{c:>9s}" for c in cells))

# 3. Scale invariance: multiply every gradient by 1000.
print()
for opt_name, make in [("SGD + momentum, lr 0.1", lambda p: SGDMomentum(p, 0.1)),
                       ("Adam, lr 0.01", lambda p: Adam(p, 0.01))]:
    out = []
    for c in [1.0, 1000.0]:
        p, losses = train(X_std, make, grad_scale=c)
        L = losses[-1]
        out.append(f"x{c:g}: loss {L:.3g}, test acc {accuracy(p, X_std_te, y_te) if L < np.inf else 0:.3f}")
    print(f"{opt_name:23s} " + "   ".join(out))

Output:

first step / lr, bias correction True : median 1.000
first step / lr, bias correction False: median 3.162

final training loss after 20 epochs (2.3 = ln 10, chance level)
                              0.0001   0.0003    0.001    0.003     0.01     0.03      0.1
raw pixels    SGD + momentum      0.5     0.21    0.085     0.03    0.049      1.7      2.3
raw pixels    Adam                1.6     0.41    0.087    0.016   0.0079        1        2
standardized  SGD + momentum      1.6     0.78     0.29     0.11    0.026   0.0051  0.00064
standardized  Adam                1.2     0.37    0.082    0.011    0.001  0.00015 unstable

SGD + momentum, lr 0.1  x1: loss 0.000638, test acc 0.967   x1000: loss inf, test acc 0.000
Adam, lr 0.01           x1: loss 0.00103, test acc 0.964   x1000: loss 0.00103, test acc 0.964

With correction, every weight with a nonzero gradient first moves by between $0.998\alpha$ and $\alpha$; without it, by up to $0.1/\sqrt{0.001} = 3.16$ times more.

For SGD, the useful learning rates move with the pixel scale: the best rate is 0.1 on standardized pixels, but raw pixels make the first-layer gradients larger, so 0.03 already fails and the best rate, 0.003, ends about 50 times higher. Adam’s best rate is 0.01 on raw and 0.03 on standardized pixels, and on raw pixels it beats SGD (0.0079 against 0.03; in nine of ten seeds). Standardizing still helps both; our explanation is that Adam rescales coordinates but cannot undo the correlations between weights that uncentered pixels create. The best rates held within one grid step over ten seeds. Entries of 1 or more are failed runs. Adam at 0.1 on standardized pixels oscillates chaotically, so its final loss depends on rounding and differs between NumPy builds; the script prints “unstable” for any run that ends ten times above its best epoch. With every gradient multiplied by 1,000, SGD diverges while Adam is unchanged to the printed digits (over ten seeds, at most 0.02% apart, because $\epsilon$ is not rescaled).

The second part feeds identical gradients to the NumPy class and to torch.optim.Adam and torch.optim.AdamW (run with PyTorch 2.14.1, continuing the script above):

import torch

def compare(torch_opt, numpy_opt_kwargs, steps=300, B=32):
    """Feed identical mini-batch gradients to torch.optim and to the NumPy Adam class;
    return the largest parameter difference seen and the final NumPy parameters."""
    p_np = init(SEED)
    p_t = [torch.tensor(q.copy(), requires_grad=True) for q in p_np]   # float64
    opt_np = Adam(p_np, **numpy_opt_kwargs)
    opt_t = torch_opt(p_t)
    rng, worst = np.random.default_rng(SEED), 0.0
    for t in range(steps):
        idx = rng.choice(N, B, replace=False)
        _, g = loss_grad(p_np, X_std[idx], Y_tr[idx])
        for q, gi in zip(p_t, g):
            q.grad = torch.tensor(gi)
        opt_t.step()
        opt_np.step(g)
        worst = max(worst, max(np.abs(q.detach().numpy() - r).max() for q, r in zip(p_t, p_np)))
    return worst, p_np

print(f"torch {torch.__version__}")
lam = 0.01
cases = [("Adam",                    lambda p: torch.optim.Adam(p, lr=1e-3), {}),
         ("Adam, weight_decay=0.01", lambda p: torch.optim.Adam(p, lr=1e-3, weight_decay=lam), {"wd": lam}),
         ("AdamW, weight_decay=0.01", lambda p: torch.optim.AdamW(p, lr=1e-3, weight_decay=lam),
          {"wd": lam, "decoupled": True}),
         ("AdamW, weight_decay=1",    lambda p: torch.optim.AdamW(p, lr=1e-3, weight_decay=1.0),
          {"wd": 1.0, "decoupled": True})]
final = {}
for name, make_t, kw in cases:
    worst, final[name] = compare(make_t, kw)
    print(f"{name:25s} max |torch - numpy| over 300 steps = {worst:.1e}")

# Where the two kinds of weight decay act: first-layer weights of the pixels that are
# 0 in every training image (zero loss gradient) versus all other pixels.
dead = X_raw.std(axis=0) == 0
W0 = init(SEED)[0]
print(f"\n{dead.sum()} always-zero pixels; mean |W1| over their rows / over the other rows")
print(f"{'initial':25s} {np.abs(W0[dead]).mean():.3f} / {np.abs(W0[~dead]).mean():.3f}")
for name, _, _ in cases:
    W1 = final[name][0]
    print(f"{name:25s} {np.abs(W1[dead]).mean():.3f} / {np.abs(W1[~dead]).mean():.3f}")

Output:

torch 2.14.1
Adam                      max |torch - numpy| over 300 steps = 1.1e-16
Adam, weight_decay=0.01   max |torch - numpy| over 300 steps = 2.2e-16
AdamW, weight_decay=0.01  max |torch - numpy| over 300 steps = 1.2e-14
AdamW, weight_decay=1     max |torch - numpy| over 300 steps = 4.4e-16

4 always-zero pixels; mean |W1| over their rows / over the other rows
initial                   0.147 / 0.140
Adam                      0.147 / 0.153
Adam, weight_decay=0.01   0.022 / 0.132
AdamW, weight_decay=0.01  0.147 / 0.153
AdamW, weight_decay=1     0.109 / 0.121

The NumPy class matches PyTorch’s Adam, Adam with weight_decay, and AdamW to rounding error, so PyTorch’s two weight-decay options are the two rules compared under AdamW below. The four always-zero pixels show the difference: their first-layer weights get no loss gradient and, without decay, never move. With an L2 term their only gradient is $\lambda\theta$, which Adam normalizes like any other, so even a tiny $\lambda$ moves them by up to about $\alpha$ per step: they fall from 0.147 to 0.022 in 300 steps, while informative pixels’ weights, whose loss gradients dominate $\lambda\theta$, barely feel the penalty. AdamW shrinks every weight by the same factor $1 - \alpha\lambda$ per step: negligibly for $\lambda = 0.01$, and by $(1 - 10^{-3})^{300} \approx 0.74$ for $\lambda = 1$, from 0.147 to 0.109.

Properties and Behavior

• Convergence. Kingma and Ba gave an $O(\sqrt{T})$ regret bound for online convex optimization, with a step size decaying as $\alpha/\sqrt{t}$ and a decaying $\beta_1$. Reddi, Kale and Kumar (2018) showed the proof to be wrong: it assumes that $\sqrt{v_{t+1}}/\alpha_{t+1} - \sqrt{v_t}/\alpha_t$ is positive semidefinite, so that effective learning rates never increase, which holds for SGD and AdaGrad but not for exponential moving averages. In their counterexample on $[-1, 1]$, the gradient is $C > 2$ once every three steps and $-1$ otherwise; with $\beta_1 = 0$ and $\beta_2 = 1/(1 + C^2)$, Adam scales the rare large gradient down by a factor of almost $C$, so it cannot outweigh the two wrong-way steps, and Adam converges to $x = +1$, the worst point. For any constant $\beta_1 < \sqrt{\beta_2}$, which the defaults satisfy, they constructed convex problems on which Adam fails to converge.
• Generalization. Wilson et al. (2017) built a linearly separable problem on which SGD reaches zero test error while AdaGrad, RMSProp and Adam approach one half, and reported that with equal tuning SGD and momentum generalized better on every task they tried. On CIFAR-10 with a VGG-style network, test error (mean ± standard deviation over five runs) was 7.65 ± 0.14% for SGD, 9.60 ± 0.19% for RMSProp, the best adaptive method, and 9.78 ± 0.25% for tuned Adam; the adaptive methods progressed faster early, then plateaued, and needed as much tuning. Loshchilov and Hutter (2019) identified Adam’s handling of weight decay as a major factor in this gap.
• Practice in vision. He et al. (2016) trained ResNets on ImageNet with SGD, momentum 0.9 and weight decay $10^{-4}$. Dosovitskiy et al. (2021) pre-trained the Vision Transformer with Adam ($\beta_1 = 0.9$, $\beta_2 = 0.999$, weight decay 0.1) and fine-tuned with SGD and momentum; noting that ResNets are typically trained with SGD, they found Adam slightly better for pre-training ResNets in their setting too (their Appendix D.1).

Limitations

• Convergence and generalization: no guarantee with constant decay rates, and test accuracy can lag behind SGD with momentum on image classification.
• Not tuning-free: the step size needs tuning and a schedule, and $\epsilon$ can matter.
• Memory: two state vectors per parameter.

Variants

• AdamW (Loshchilov & Hutter, 2019) decouples weight decay from the gradient. For SGD, $\theta - \eta(g + \lambda'\theta) = (1 - \eta\lambda')\theta - \eta g$, so an L2 penalty $\lambda'$ equals a weight decay $\eta\lambda'$. For adaptive methods it does not: the L2 gradient is divided by $\sqrt{\hat v}$, so weights with historically large gradients are regularized less. The paper applies the decay outside the adaptive step, $\theta_t = \theta_{t-1} - \eta_t\big(\alpha\,\hat m_t/(\sqrt{\hat v_t} + \epsilon) + \lambda\theta_{t-1}\big)$ with a schedule multiplier $\eta_t$; PyTorch’s AdamW also multiplies the decay by the learning rate. The authors reported that decoupling made the best weight decay less dependent on the learning rate and gave a 15% relative improvement in Adam’s test error on CIFAR-10 and ImageNet 32×32, letting it compete with SGD and momentum on image classification. torch.optim.Adam with weight_decay implements the L2 form (unless decoupled_weight_decay=True).
• AMSGrad (Reddi et al., 2018) normalizes by the running maximum $\hat v_t = \max(\hat v_{t-1}, v_t)$, so effective learning rates never increase; they proved a regret bound and reported similar or better results on their problems. Their algorithm has no bias correction; PyTorch’s amsgrad=True takes the maximum of $v$, then corrects.
• NAdam (torch.optim.NAdam), which PyTorch’s documentation attributes to Dozat’s Incorporating Nesterov Momentum into Adam, uses Nesterov momentum in the first moment.
• AdaMax (Kingma & Ba, 2015, Section 7.1) uses an infinity norm in the denominator.