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

Problem

AdamW differs from "Adam plus L2 regularization" in one specific way: L2 regularization adds weight_decay * param into the gradient, so it gets divided by Adam's adaptive per-parameter learning rate like every other gradient term does. AdamW instead subtracts lr * weight_decay * param from the parameter directly, decoupled from the gradient-based update entirely (Loshchilov & Hutter, 2017).

It also corrects for the fact that the first and second moment estimates (running averages of the gradient and its square) start at zero and are biased toward zero for the first few steps — without that correction, early updates are artificially small.

Implement a single adamw_step: given a parameter, its gradient, the current moment estimates, and the 1-indexed step count, return the updated parameter and moments after one AdamW update.

Source: src/39_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]: ...

>>> 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: param = param - lr * weight_decay * param. This is the one line that makes it AdamW rather than Adam-with-L2.
2. Update the biased first moment: exp_avg = beta1 * exp_avg + (1 - beta1) * grad.
3. Update the biased second moment: exp_avg_sq = beta2 * exp_avg_sq + (1 - beta2) * grad**2.
4. Compute bias-corrected estimates: divide exp_avg by 1 - beta1**step and exp_avg_sq by 1 - beta2**step.
5. Apply the update: param = param - lr * m_hat / (v_hat.sqrt() + eps).
6. Return the updated (param, exp_avg, exp_avg_sq) — note exp_avg/exp_avg_sq are returned before bias correction; the correction is a read-time adjustment, not a permanent change to the stored moments.

The key insight is that bias correction matters most exactly when it's easy to forget it matters: at step=1, exp_avg equals (1 - beta1) * grad, a small fraction of the true gradient, so without dividing by 1 - beta1**step the very first update would be far too small relative to later steps.

Reference solution

import torch


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]:
    beta1, beta2 = betas

    # decoupled weight decay: applied directly to the parameter, never
    # folded into `grad`
    param = param - lr * weight_decay * param

    # biased first/second moment estimates
    exp_avg = beta1 * exp_avg + (1 - beta1) * grad
    exp_avg_sq = beta2 * exp_avg_sq + (1 - beta2) * grad.pow(2)

    # bias correction: moments start at 0 and are biased toward 0 early on
    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

• Decoupled decay applied to param directly — the single structural difference from Adam+L2; get this line wrong (fold it into grad instead) and the optimizer silently becomes plain Adam with L2.
• beta1**step / beta2**step — Python's ** on a plain float step count, computed fresh each call since step changes every iteration.
• exp_avg_sq.sqrt() — the per-parameter adaptive learning-rate denominator; combined with eps added after the square root (unlike RMSNorm's eps-inside convention) to avoid dividing by exactly zero.
• torch.optim.AdamW — PyTorch's own implementation, used here as an independent oracle by running it over the same gradient sequence and comparing final parameters.

How to Recognize This Pattern