Adam and AdamW

What it is and when you reach for it

Adam (adaptive moment estimation) is a first-order stochastic optimizer that keeps, per parameter, an exponential moving average of the gradient (the first moment) and of the squared gradient (the second moment), and takes steps proportional to the ratio of the two. The ratio makes the update roughly unit-scale regardless of how large or small the raw gradients are, which is why one learning rate works across layers whose gradient magnitudes differ by orders of magnitude. AdamW is the variant you should almost always use: it applies weight decay directly to the parameters instead of folding an L2 penalty into the gradient, which under Adam's adaptive denominator turns out to be a genuinely different algorithm. Reach for AdamW whenever you are training a deep network with heterogeneous gradient scales: transformers, diffusion models, and most modern architectures use it by default. Plain SGD with momentum still competes on convolutional vision models and is cheaper in memory, but for attention-based models AdamW is not just the default, it is close to load-bearing: transformer training with vanilla SGD is markedly less stable and slower to converge at any practical learning rate.

The math, from Adagrad to Adam

Adagrad: per-parameter learning rates

SGD moves every parameter with the same step size η: θ ← θ − ηg. But gradient magnitudes are wildly non-uniform across a deep network, and a single η must be small enough for the largest-gradient parameter, starving all the others. Adagrad (Duchi, Hazan, Singer, 2011) fixed this by accumulating the sum of squared gradients per parameter, G_t = Σ_i≤t g_i², and dividing: θ ← θ − η g_t / (√G_t + ε). Parameters that have seen large gradients get small steps, rarely-updated parameters get large ones, which is exactly what sparse features want. The flaw is that G_t only grows: the effective learning rate decays monotonically toward zero whether or not the loss landscape asks for it, and training stalls.

RMSProp: forgetting the past

RMSProp (Hinton, in a lecture slide, never formally published) replaces the growing sum with an exponential moving average: v_t = ρv_t-1 + (1 − ρ)g_t². Now the denominator tracks the recent root-mean-square gradient rather than the whole history, so the effective step size can recover when gradients shrink. What RMSProp lacks is momentum on the numerator and any account of the EMA's startup transient.

Adam: two moments

Adam (Kingma and Ba, 2014) keeps EMAs of both moments:

m_t = beta1 * m_{t-1} + (1 - beta1) * g_t        first moment  (mean of g)
v_t = beta2 * v_{t-1} + (1 - beta2) * g_t^2      second moment (mean of g^2)

m_hat = m_t / (1 - beta1^t)                       bias correction
v_hat = v_t / (1 - beta2^t)

theta = theta - lr * m_hat / (sqrt(v_hat) + eps)

The numerator is momentum: averaging gradients over roughly 1/(1 − β₁) ≈ 10 recent steps smooths minibatch noise. The denominator is RMSProp: dividing by the recent RMS gradient normalizes the step per coordinate. If the gradient of one coordinate is consistently around some magnitude c, then m̂ ≈ c and √v̂ ≈ c, so the update is about η in size no matter what c is. Adam's step size is controlled by the signal-to-noise ratio of the gradient, not by its magnitude, and is soft-bounded by the learning rate itself. That invariance to gradient scale is what makes one learning rate portable across layers, and largely across model widths.

Bias correction, derived

Both EMAs start at zero, so early estimates are biased toward zero. The correction term is not a heuristic, it falls straight out of unrolling the recursion. With m₀ = 0:

m_t = (1 − β₁) Σ_i=1..t β₁^t−i g_i.

Take expectations and suppose the gradient is roughly stationary over the window, E[g_i] ≈ E[g]. The geometric sum Σ_i=1..t β₁^t−i = (1 − β₁^t) / (1 − β₁), so

E[m_t] ≈ E[g] · (1 − β₁^t).

The EMA underestimates the true mean by exactly the factor 1 − β₁^t: this is where β^t comes from. Dividing by it makes the estimator unbiased at every t, and the factor decays to 1, so correction only matters early. The same algebra with β₂ gives v̂. The reason you cannot skip it: v is corrected by 1 − β₂^t, which at t = 1 with β₂ = 0.999 is 0.001, a thousandfold underestimate sitting under a square root in the denominator.

A worked step

One parameter, constant gradient g = 0.5, defaults β₁ = 0.9, β₂ = 0.999, at t = 1: m₁ = 0.1 × 0.5 = 0.05 and v₁ = 0.001 × 0.25 = 0.00025. Without correction the ratio is 0.05/√0.00025 = 0.05/0.0158 ≈ 3.16, so the very first step is 3.16 times the learning rate: the second moment is far more underestimated than the first, and the square root amplifies the mismatch. With correction, m̂ = 0.05/0.1 = 0.5 and v̂ = 0.00025/0.001 = 0.25, giving 0.5/√0.25 = 1.0 exactly: a step of precisely one learning rate. At t = 2 the corrected estimates are again exactly 0.5 and 0.25. For a constant gradient the corrected ratio is identically 1, which is the cleanest way to see that the correction recovers the intended semantics from step one instead of after a thousand steps of warm-up.

Where epsilon goes

The paper writes the update as m̂ / (√v̂ + ε), with ε outside the square root, and that is what PyTorch implements. The other placement, m̂ / √(v̂ + ε), looks interchangeable and is not: when v̂ is tiny the first form bounds the step by m̂/ε while the second bounds it by m̂/√ε, a factor of 10⁴ apart at ε = 10^-8. Optax exposes both knobs separately (eps outside, eps_root inside, default 0), and TensorFlow's historical RMSProp put it inside, which is one reason naive hyperparameter ports across frameworks silently change the optimizer. ε is also not just overflow armor: for coordinates whose gradients are smaller than about ε, the denominator saturates and Adam degrades gracefully toward plain momentum SGD. Some training recipes deliberately raise ε to 10^-6 or 10^-5 to damp the adaptivity, and in low-precision training a larger ε guards against v underflowing to zero. Treat it as a real hyperparameter with a placement convention, not as a numerical footnote.

AdamW: why L2 in the gradient is not weight decay

Classical weight decay multiplies weights toward zero each step: θ ← (1 − ηλ)θ − ηg. Classical L2 regularization adds λθ to the gradient of the loss. Under plain SGD these are the same operation, since θ − η(g + λθ) = (1 − ηλ)θ − ηg, and that identity is why the two names were used interchangeably for decades. Loshchilov and Hutter (2017) pointed out that the identity dies the moment the update is preconditioned. Feed g + λθ into Adam and the decay term travels through both EMAs and gets divided by √v̂:

L2 inside the gradient:   theta -= lr * (m_hat(g + lambda*theta)) / (sqrt(v_hat) + eps)
                                       ^ decay is rescaled per-coordinate by 1/sqrt(v_hat)

Decoupled weight decay:   theta -= lr * m_hat(g) / (sqrt(v_hat) + eps)  +  lr * lambda * theta
                                       ^ decay is applied at full strength, uniformly

The consequence is backwards from what regularization intends: parameters with historically large gradients have large √v̂, so their decay term is shrunk the most, while rarely-updated parameters are decayed hardest. Worse, since the decay contribution to the update is normalized by the same RMS machinery as the gradient, its effective strength stops scaling with λ in any clean way, and the optimal λ becomes entangled with the learning rate. AdamW moves the decay out of the gradient entirely, multiplying θ by (1 − ηλ) as a separate step, so the adaptive machinery only ever sees the loss gradient. The paper's experiments show the decoupled version generalizes better and, usefully in practice, makes learning rate and weight decay close to independently tunable. This is why torch.optim.Adam(weight_decay=...) and torch.optim.AdamW(weight_decay=...) are different optimizers, not one optimizer with a renamed argument: the former adds λθ to the gradient, the latter decays the parameter directly.

Implementation, twice

The PyTorch version mirrors torch.optim.AdamW's exact update order: decay the parameter first, update the moments, then take the bias-corrected step, with the corrections folded into scalars so the stored m and v remain the raw EMAs. The JAX version follows optax's convention: a pair of pure functions, init(params) producing the state and update(grads, state, params) producing updates to be added to the parameters. Optax applies decoupled decay by adding λθ to the scaled update rather than pre-multiplying the parameter; expanding both shows they compute the same (1 − ηλ)θ − η·adam_step, differing only in floating point rounding order.

import math
import torch

class AdamW(torch.optim.Optimizer):
    """From-scratch AdamW matching torch.optim.AdamW's update order:
    decay first, then moments, then the bias-corrected step.
    decoupled=False turns it into torch.optim.Adam's L2 coupling,
    which is the two-line difference the whole AdamW paper is about.
    """

    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999),
                 eps=1e-8, weight_decay=1e-2, decoupled=True):
        super().__init__(params, dict(lr=lr, betas=betas, eps=eps,
                                      weight_decay=weight_decay,
                                      decoupled=decoupled))

    @torch.no_grad()
    def step(self):
        for group in self.param_groups:
            lr, (b1, b2) = group["lr"], group["betas"]
            eps, wd = group["eps"], group["weight_decay"]
            for p in group["params"]:
                if p.grad is None:
                    continue
                g = p.grad
                state = self.state[p]
                if not state:
                    state["step"] = 0
                    state["m"] = torch.zeros_like(p)
                    state["v"] = torch.zeros_like(p)
                m, v = state["m"], state["v"]
                state["step"] += 1
                t = state["step"]

                if wd != 0.0:
                    if group["decoupled"]:
                        p.mul_(1 - lr * wd)        # AdamW: shrink the weight itself
                    else:
                        g = g.add(p, alpha=wd)     # Adam+L2: decay rides the gradient

                # Moment EMAs. lerp_(g, 1-b1) is b1*m + (1-b1)*g in one
                # fused op, exactly as in PyTorch's implementation.
                m.lerp_(g, 1 - b1)
                v.mul_(b2).addcmul_(g, g, value=1 - b2)

                # Bias correction folded into scalars: the stored m, v stay
                # uncorrected, so state memory holds the plain EMAs.
                step_size = lr / (1 - b1 ** t)
                denom = (v.sqrt() / math.sqrt(1 - b2 ** t)).add_(eps)
                p.addcdiv_(m, denom, value=-step_size)

import jax
import jax.numpy as jnp

def adamw(lr, b1=0.9, b2=0.999, eps=1e-8, weight_decay=1e-2):
    """optax-style transformation: init(params) -> state,
    update(grads, state, params) -> (updates, state).
    Apply with params + updates (optax.apply_updates convention).
    """

    def init(params):
        return dict(mu=jax.tree.map(jnp.zeros_like, params),
                    nu=jax.tree.map(jnp.zeros_like, params),
                    count=jnp.zeros([], jnp.int32))

    def update(grads, state, params):
        count = state["count"] + 1
        mu = jax.tree.map(lambda m, g: b1 * m + (1 - b1) * g,
                          state["mu"], grads)
        nu = jax.tree.map(lambda v, g: b2 * v + (1 - b2) * g * g,
                          state["nu"], grads)
        # Bias correction as scalars; mu/nu stay uncorrected in the state.
        c1 = 1 - b1 ** count.astype(jnp.float32)
        c2 = 1 - b2 ** count.astype(jnp.float32)

        def one(m, v, p):
            adam_step = (m / c1) / (jnp.sqrt(v / c2) + eps)
            # Decoupled decay added to the update, as optax.adamw does;
            # algebraically identical to PyTorch's p *= 1 - lr*wd.
            return -lr * (adam_step + weight_decay * p)

        updates = jax.tree.map(one, mu, nu, params)
        return updates, dict(mu=mu, nu=nu, count=count)

    return init, update

Both versions are step-for-step checkable against the production optimizers on a fixed seed, which is the single best debugging tool an optimizer implementation has: any divergence at step k with identical inputs is a wrong line, not noise.

import copy
import torch
import torch.nn as nn

torch.manual_seed(0)
net_ref = nn.Sequential(nn.Linear(64, 128), nn.Tanh(), nn.Linear(128, 10))
net_mine = copy.deepcopy(net_ref)

opt_ref = torch.optim.AdamW(net_ref.parameters(), lr=1e-3,
                            betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01)
opt_mine = AdamW(net_mine.parameters(), lr=1e-3,
                 betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01)

for step in range(100):
    x = torch.randn(32, 64)               # same batch feeds both copies
    for net, opt in ((net_ref, opt_ref), (net_mine, opt_mine)):
        opt.zero_grad(set_to_none=True)
        net(x).pow(2).mean().backward()
        opt.step()

for a, b in zip(net_ref.parameters(), net_mine.parameters()):
    # Same op order means this is typically bitwise-equal on CPU;
    # the tolerance covers fused/foreach paths that reorder arithmetic.
    assert torch.allclose(a, b, rtol=0, atol=1e-7)
print("matches torch.optim.AdamW for 100 steps")

import jax
import jax.numpy as jnp
import optax

key = jax.random.PRNGKey(0)
params = {"w": jax.random.normal(key, (64, 10)) * 0.1,
          "b": jnp.zeros(10)}

def loss_fn(p, x):
    return jnp.mean((x @ p["w"] + p["b"]) ** 2)

ref = optax.adamw(1e-3, b1=0.9, b2=0.999, eps=1e-8, weight_decay=0.01)
init, update = adamw(1e-3, weight_decay=0.01)

p_ref, p_mine = params, params
s_ref, s_mine = ref.init(params), init(params)

for step in range(100):
    x = jax.random.normal(jax.random.fold_in(key, step), (32, 64))
    g_ref = jax.grad(loss_fn)(p_ref, x)
    u_ref, s_ref = ref.update(g_ref, s_ref, p_ref)
    p_ref = optax.apply_updates(p_ref, u_ref)

    g_mine = jax.grad(loss_fn)(p_mine, x)
    u_mine, s_mine = update(g_mine, s_mine, p_mine)
    p_mine = optax.apply_updates(p_mine, u_mine)

worst = jax.tree.map(lambda a, b: jnp.abs(a - b).max(), p_ref, p_mine)
print(worst)   # every entry at float32 rounding level, ~1e-7 or below

Using it on a real shape of problem

Real training runs never hand every parameter the same decay. The near-universal transformer recipe puts weight matrices in a decayed group and biases, LayerNorm gains, and usually embeddings in a zero-decay group, then wraps the learning rate in warmup plus cosine decay. On a 10M-parameter transformer language model with lr 3×10^-4, warmup over the first few hundred steps, and weight decay 0.1, expect the loss to fall fast through the warmup, then descend smoothly; a loss that spikes and recovers repeatedly usually means β₂ is too slow for your gradient noise (see the next section) or the warmup is too short. Exact numbers are model- and machine-dependent, but the shape of the curve is not.

import torch

# Two param groups: ndim >= 2 tensors (matmul weights, embeddings) decay,
# 1-D tensors (biases, norm gains) do not. This is the nanoGPT/minGPT split.
decay, no_decay = [], []
for name, p in model.named_parameters():
    if not p.requires_grad:
        continue
    (decay if p.ndim >= 2 else no_decay).append(p)

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, fused=True)  # LLM-style beta2

sched = torch.optim.lr_scheduler.SequentialLR(
    opt,
    [torch.optim.lr_scheduler.LinearLR(opt, 0.01, 1.0, total_iters=500),
     torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=9500)],
    milestones=[500])

for step, (x, y) in enumerate(loader):
    loss = model(x, y)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step(); sched.step()
    opt.zero_grad(set_to_none=True)

import jax
import optax

# Same split, optax-style: a mask pytree routes decay to ndim >= 2 leaves.
decay_mask = jax.tree.map(lambda p: p.ndim >= 2, params)

sched = optax.warmup_cosine_decay_schedule(
    init_value=0.0, peak_value=3e-4,
    warmup_steps=500, decay_steps=10_000, end_value=3e-5)

opt = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adamw(sched, b1=0.9, b2=0.95, eps=1e-8,
                weight_decay=0.1, mask=decay_mask))
opt_state = opt.init(params)

@jax.jit
def train_step(params, opt_state, batch):
    loss, grads = jax.value_and_grad(loss_fn)(params, batch)
    updates, opt_state = opt.update(grads, opt_state, params)
    return optax.apply_updates(params, updates), opt_state, loss

Applications

The default for transformers, and the hyperparameters that ship

AdamW trains GPT-style language models, BERT-style encoders, ViTs, diffusion models, and nearly everything with attention in it. The hyperparameters are remarkably stable across the field: lr around 10^-4 to 6×10^-4 for pretraining with warmup and cosine or linear decay, β₁ = 0.9, weight decay 0.1 on matrices, gradient clipping at norm 1.0. The interesting deviation is β₂. The library default is 0.999, an averaging window of roughly a thousand steps, but GPT-3, LLaMA, and most large-model reports since use β₂ = 0.95, a window of about twenty steps. The reason is stability under heavy-tailed gradient noise: with a thousand-step memory, a sudden gradient spike is divided by a stale, small √v̂ and produces a huge parameter jump, the classic loss-spike signature of large-batch LLM training. A twenty-step memory lets the denominator react to the spike within a few steps and absorb it. The cost is noisier adaptation on small models, where 0.999 remains fine. When you see a training-instability postmortem that ends with "we lowered beta2", this is the mechanism.

The memory bill, and what people do about it

Adam's price is state: two fp32 tensors the size of the model, 8 bytes per parameter, on top of the parameters and gradients. In the standard mixed-precision recipe (half-precision weights and grads, fp32 master weights, fp32 m and v) the optimizer-adjacent state comes to 16 bytes per parameter, the arithmetic at the heart of the ZeRO paper: a 7B-parameter model carries about 56 GB of optimizer state before a single activation is stored. That bill motivates three lines of engineering. Sharding: FSDP and ZeRO partition m, v, and master weights across data-parallel ranks, which is exactly what torchtitan does for Llama-style pretraining. Quantized state: bitsandbytes 8-bit Adam stores m and v as 8-bit values with block-wise scaling constants ( Dettmers et al., 2021), cutting state memory roughly 4x with matching accuracy on the benchmarks in the paper. And kernel fusion: the update touches every parameter with a handful of flops, so it is bandwidth-bound and dominated by kernel-launch overhead when run tensor-by-tensor; foreach batches tensors per kernel and fused=True in torch.optim.AdamW performs the whole update in a single multi-tensor kernel. Because the moments live in fp32 while compute runs in bf16 or fp16, the optimizer step is also where mixed-precision training pays its precision debts; the companion page on mixed precision and loss scaling covers why the master copy must stay fp32 and why fp16 gradients must be unscaled before this update runs.

Beyond Adam: second-order-lite research

Adam's preconditioner is diagonal: each coordinate is scaled independently. Current research reintroduces cross-coordinate structure without full second-order cost. Shampoo preconditions each weight matrix with Kronecker-factored statistics (one small matrix per dimension), and a distributed implementation won the external tuning track of the 2024 AlgoPerf optimizer benchmark. Muon replaces the second moment entirely for 2-D weights: it orthogonalizes the momentum matrix with a few Newton-Schulz iterations, so every direction in the update has comparable magnitude, while embeddings, norms, and scalars stay on AdamW. It set the nanoGPT speedrun records and has since been scaled to multi-billion-parameter pretraining runs. Both are best understood as answers to the same question Adam answers diagonally: how should the raw gradient be reshaped before it becomes a step. AdamW remains the production default, but it is no longer obviously the end of the story.

Against the real libraries

torch.optim adds, over the reference implementation above: the foreach and fused execution paths (multi-tensor and single-kernel updates), a capturable mode that keeps the step count on-device so the optimizer can be captured into CUDA graphs, amsgrad, differentiable optimizer steps for meta-learning, and battle-tested state-dict save/load including cross-device restore. The scheduler ecosystem hangs off it. optax adds composability as the core design: adamw is literally chain(scale_by_adam, add_decayed_weights, scale_by_learning_rate), and the same combinators give you clipping, schedules, gradient accumulation (MultiSteps), and per-subtree masking without touching optimizer internals; everything is a pure function, so the whole update jits and shards. bitsandbytes adds the 8-bit state representation with block-wise quantization, stable-embedding handling, and paged optimizer variants that spill state to CPU memory under pressure, the difference between fitting and not fitting a fine-tune on a single consumer GPU.

The from-scratch version is genuinely enough when you are doing optimizer research (you need to modify the update anyway), when the model is small enough that state memory and launch overhead are irrelevant, and as executable documentation. Verification against the libraries is the check shown above: identical parameters, identical data, fixed seed, run both optimizers side by side and assert per-parameter agreement to float32 rounding after 100 steps. Since the update order here mirrors the library exactly, agreement on CPU is typically bitwise; any looser tolerance you find yourself needing is a discrepancy worth explaining before trusting the implementation.

Traps and misconceptions

"weight_decay in torch.optim.Adam is weight decay." It is L2-in-the-gradient, the exact coupling the AdamW paper shows to be a different and worse regularizer under adaptive scaling. If you mean weight decay, use AdamW. The two differ precisely when √v̂ varies across parameters, which is always.

"Decay everything uniformly." Decaying LayerNorm gains and biases toward zero fights their function (a norm gain of zero kills the channel), and most recipes exempt them along with, usually, embeddings. The two-param-group split by ndim is the standard fix; forgetting it typically costs a little accuracy quietly rather than failing loudly.

"Bias correction is startup bookkeeping you can drop." Dropping it inflates the first step to (1−β₁)/√(1−β₂) ≈ 3.16x the learning rate at the defaults, as the worked example shows, exactly when the model is at its most fragile initialization. Implementations that skip it compensate with longer warmup, which is solving a problem you created.

"Epsilon is a numerical detail, any small value works." Placement (inside vs outside the square root) changes the small-gradient behavior by orders of magnitude, frameworks disagree on the convention, and several published models tune ε deliberately. Porting a recipe across frameworks without checking the epsilon convention is a classic source of silent divergence.

"Adam is scale-invariant, so gradient scaling never matters." The invariance is asymptotic and per-coordinate. A constant scale factor does cancel in m̂/√v̂ eventually, but ε breaks the invariance for small gradients, and in fp16 training the scale factor must still be removed before the update because otherwise ε and weight decay act on the wrong scale, and because scaled gradients can overflow the fp16 range on the way to the optimizer. Details in the mixed precision page.