AdamW from Scratch

The problem

AdamW (Adam with decoupled weight decay) is the optimizer behind practically every transformer in this roadmap. You will implement one step of it, exactly as PyTorch's torch.optim.AdamW does.

The idea, explained in detail on the theory page: every weight keeps two running averages, of its gradient ($m$, the direction) and of its squared gradient ($v$, the typical size). The step is their ratio, so every weight moves by about the learning rate, whatever the size of its gradient. Before the step, every weight is pulled a little towards zero (weight decay).

One step, in five lines

All operations are elementwise: every weight is handled on its own, with its own $m$ and $v$.

1. Weight decay, a small pull of the parameters towards zero:

2. Running average of the gradient:

3. Running average of the squared gradient:

4. Bias correction, because $m$ and $v$ start at zero and are too small during the first steps:

5. The step, applied to the parameters from line 1:

where:

• $\theta$ (theta) is param and $g$ is grad;
• $m$ and $v$ are the two running averages, all zeros before the first step;
• $\eta$ (eta) is lr, $\lambda$ (lambda) is weight_decay, $\epsilon$ (epsilon) is eps, and $(\beta_1, \beta_2)$ are betas;
• $t$ is the step number: $1, 2, 3, \dots$, and $\beta_1^t$ means $\beta_1$ to the power $t$.

Example. Two parameters, param = [1.0, -2.0], with gradients of very different size, grad = [0.5, -50.0]. First step ($t = 1$), with lr = 0.1, weight_decay = 0.01 and the default betas.

1. Weight decay multiplies by $1 - 0.1 \cdot 0.01 = 0.999$: [0.999, -1.998].
2. $m = 0.1 \cdot g$ = [0.05, -5.0].
3. $v = 0.001 \cdot g^2$ = [0.00025, 2.5].
4. Bias correction divides $m$ by $1 - 0.9 = 0.1$ and $v$ by $1 - 0.999 = 0.001$: $\hat m$ = [0.5, -50.0], $\hat v$ = [0.25, 2500.0].
5. $\hat m / \sqrt{\hat v}$ = [0.5 / 0.5, -50 / 50] = [1, -1], times lr: the step is [0.1, -0.1]. New param = [0.899, -1.898].

Both parameters moved by exactly lr, although the second gradient is 100 times larger. That is Adam's point: the step size does not depend on the scale of the gradient.

What to implement

adamw_step(param, grad, m, v, t, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01): return (new_param, new_m, new_v). The defaults are PyTorch's. Before the first step m and v are zeros and t = 1.

Requirements

• Decoupled weight decay: it shrinks the parameters directly (line 1). It is not added to the gradient, so it never enters $m$ or $v$.
• Bias correction uses the step number $t$, which starts at 1. With $t = 0$, $1 - \beta^0 = 0$ and the correction divides by zero.
• Do not modify the input arrays.

Hint

The example above is a good self-check: on the very first step, $\hat m = g$ and $\hat v = g^2$, so $\hat m / \sqrt{\hat v}$ is $+1$ or $-1$ for every parameter.