Adam optimiser step

Background

Adam (Adaptive Moment Estimation) is the default optimiser for training deep networks and transformers. It keeps a per-parameter running estimate of the gradient's first moment (the mean, $m$) and second moment (the uncentered variance, $v$), then scales each parameter's step by those estimates — so parameters with consistently large or noisy gradients get smaller, better-conditioned updates. It is the optimiser behind most LLM pre-training; AdamW, its weight-decay variant, differs by a single term. Knowing it cold is table-stakes for an ML interview.

Problem statement

Implement adam_step(params, grads, m, v, t, lr, beta1, beta2, eps) that applies one Adam update, in place, to every parameter tensor. For each parameter $p$ with gradient $g$ and moment buffers $m, v$:

$$\begin{aligned} m &\leftarrow \beta_1\, m + (1-\beta_1)\, g \\ v &\leftarrow \beta_2\, v + (1-\beta_2)\, g^2 \\ \hat{m} &= \frac{m}{1-\beta_1^{\,t}}, \qquad \hat{v} = \frac{v}{1-\beta_2^{\,t}} \\ p &\leftarrow p - \mathrm{lr}\cdot \frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon} \end{aligned}$$

Defaults from the original paper: lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8.

Input

• params — list[np.ndarray], the model parameters. Mutated in place.
• grads — list[np.ndarray], same shapes as params; the gradient $g$ for each.
• m — list[np.ndarray], first-moment buffers (start at 0). Mutated in place.
• v — list[np.ndarray], second-moment buffers (start at 0). Mutated in place.
• t — int, $t \ge 1$, the current step (1-indexed). Drives bias correction.
• lr, beta1, beta2, eps — scalars: learning rate, first/second-moment decay rates, and the stabiliser $\epsilon$.

Output

Returns None. Updates params, m, and v in place, so the caller's references see the new values after the call.

Examples

Example 1 — first step from zero state

Input:  params=[1.0], grads=[0.1], m=[0.0], v=[0.0],
        t=1, lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8
Output: m=[0.01], v=[1e-5], params≈[0.999]

Explanation: $m = 0.1\cdot0.1 = 0.01$, so $\hat{m} = 0.01/(1-0.9) = 0.1$. $v = 0.001\cdot0.01 = 10^{-5}$, so $\hat{v} = 10^{-5}/(1-0.999) = 0.01$. Then $p \leftarrow p - 10^{-3}\cdot 0.1/\sqrt{0.01} = p - 10^{-3}$, giving $p \approx 0.999$.

Example 2 — bias correction at $t=1$

Input:  params=[10.0], grads=[2.0], m=[0.0], v=[0.0],
        t=1, lr=1.0, beta1=0.9, beta2=0.999, eps=1e-12
Output: params≈[9.0]

Explanation: at $t=1$ the $(1-\beta^t)$ denominator exactly cancels the $(1-\beta)$ numerator, so $\hat{m}=g=2$ and $\hat{v}=g^2=4$. The step is $\mathrm{lr}\cdot 2/\sqrt{4} = 1.0$, so $p \approx 9.0$. The first step is $\approx \mathrm{lr}\cdot\operatorname{sign}(g)$ regardless of $\beta$ — this is what keeps Adam from stalling at step 1.

Constraints

• Within each tuple, params[i], grads[i], m[i], v[i] share the same shape.
• $t \ge 1$ and is the true step count — bias correction divides by $1-\beta^{\,t}$.
• Update in place (*=, +=, -=). Rebinding m/v discards the optimiser state across steps.
• $\epsilon$ is added after the square root: $\sqrt{\hat{v}} + \epsilon$.
• Tests compare with tolerance atol ≈ 1e-6.

Notes

• Bias correction. m and v are initialised to 0, biasing them toward 0 early on; dividing by $1-\beta^{\,t}$ corrects this and vanishes as $t$ grows.
• State persistence. m and v are optimiser state that lives across steps — the in-place requirement is what lets a training loop reuse them.