20. AdamW From Scratch (Bias Correction + Decoupled Weight Decay)

Problem

AdamW differs from "Adam plus L2 regularization" in one specific way that matters in practice: instead of folding weight decay into the gradient — where it would get divided by the adaptive per-parameter denominator like every other gradient term — AdamW applies weight decay directly to the parameter, decoupled from the gradient-based update entirely. It also corrects for the fact that the first and second moment estimates start at zero and are biased toward zero for the first several steps; without that correction, early updates would be artificially small.

Implement a single AdamW parameter-update step from its update equations, given the current moment estimates and the step count.

Source: src/20_adamw_from_scratch.py

def adamw_step(
    param: torch.Tensor,
    grad: torch.Tensor,
    exp_avg: torch.Tensor,
    exp_avg_sq: torch.Tensor,
    step: int,
    lr: float = 1e-3,
    betas: tuple[float, float] = (0.9, 0.999),
    eps: float = 1e-8,
    weight_decay: float = 0.01,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    ...

Examples:
>>> p = torch.tensor([1.0, 2.0])
>>> g = torch.tensor([0.1, -0.2])
>>> m = torch.zeros(2)
>>> v = torch.zeros(2)
>>> new_p, new_m, new_v = adamw_step(p, g, m, v, step=1)
>>> new_p.shape
torch.Size([2])

Step-by-Step Approach

1. Apply decoupled weight decay first, directly to the parameter, before touching the gradient at all: param = param - lr * weight_decay * param. This is the entire "W" in AdamW.
2. Update the first-moment (mean of gradient) running average: exp_avg = beta1 * exp_avg + (1 - beta1) * grad.
3. Update the second-moment (mean of squared gradient) running average: exp_avg_sq = beta2 * exp_avg_sq + (1 - beta2) * grad**2.
4. Compute the bias-correction terms from the step count: 1 - beta1**step and 1 - beta2**step. These are small early on (correcting hard) and approach 1 as training progresses (correcting less).
5. Divide each moment estimate by its bias-correction term to get m_hat and v_hat — this is what undoes the zero-initialization bias.
6. Apply the final update: param = param - lr * m_hat / (sqrt(v_hat) + eps). Note eps is added after the square root here, unlike LayerNorm/RMSNorm where it goes inside — a different convention for a different reason (it exists purely to prevent division by zero when a gradient has been exactly zero so far).
7. Validate against torch.optim.AdamW itself: run both your manual loop and the real optimizer over the same sequence of fixed gradients and confirm the final parameters match to float precision.

The key insight is the order of operations: weight decay is applied to the raw parameter before the adaptive gradient step, not mixed into the gradient beforehand — that separation is precisely what "decoupled" means and what distinguishes AdamW from Adam-with-L2.

Reference solution

def adamw_step(param, grad, exp_avg, exp_avg_sq, step, lr=1e-3,
               betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01):
    beta1, beta2 = betas

    # DECOUPLED weight decay: shrink the parameter directly, *before* the
    # gradient-based update -- this is the "W" in AdamW. Plain Adam+L2 would
    # instead add `weight_decay * param` into the gradient, which then gets
    # divided by the adaptive per-parameter denominator below, so large-
    # second-moment parameters get decayed less. Decoupling removes that
    # unwanted interaction.
    param = param - lr * weight_decay * param

    # exponential moving averages of the gradient (1st moment) and the
    # squared gradient (2nd moment)
    exp_avg = beta1 * exp_avg + (1 - beta1) * grad
    exp_avg_sq = beta2 * exp_avg_sq + (1 - beta2) * grad.pow(2)

    # BIAS CORRECTION: exp_avg/exp_avg_sq start at 0 and are biased toward 0
    # for early steps -- dividing by (1 - beta^step) undoes that bias. As
    # step -> infinity, both correction terms -> 1 and this becomes a no-op.
    bias_correction1 = 1 - beta1 ** step
    bias_correction2 = 1 - beta2 ** step
    m_hat = exp_avg / bias_correction1
    v_hat = exp_avg_sq / bias_correction2

    param = param - lr * m_hat / (v_hat.sqrt() + eps)
    return param, exp_avg, exp_avg_sq

Key Functions & Tricks

• param - lr * weight_decay * param — decoupled decay, applied to the parameter directly instead of injected into the gradient.
• beta1 ** step / beta2 ** step — the exponential bias term that bias correction divides out; requires the caller to pass a correct 1-indexed step count.
• v_hat.sqrt() + eps — eps added after the square root here (contrast with LayerNorm/RMSNorm's eps-inside-the-sqrt convention).
• torch.nn.Parameter + torch.optim.AdamW + manual .grad assignment — the pattern used to build an independent oracle: drive the real optimizer with the exact same fixed gradients and diff the result against the manual loop.
• opt.zero_grad() between steps — needed in the oracle to prevent gradient accumulation across the fixed-gradient steps.

How to Recognize This Pattern