Adam: trace direction and coordinate scale

Summary

Adam saves a smoothed gradient and a smoothed squared gradient for each parameter. It divides directional history by a coordinate-scale estimate. Bias correction compensates for zero-initialized buffers, not gradient noise in general.

Task: start ${\theta }_0=[1,1]$, $m_0=v_0=[0,0]$, ${\beta }_1={\beta }_2=0.5$, and learning rate $\eta =0.1$. Apply prescribed gradients $[2,4]$, then $[2,0]$. Use $ϵ=0$ only for this nonzero-denominator fixture; production implementations use positive epsilon. Does the second coordinate stop when its current gradient becomes zero?

The mechanism

All operations below are coordinatewise; the step counter $t$ counts optimizer updates, starting at one. Momentum saves a decaying sum; Adam’s first moment includes the additional factor $(1-{\beta }_1)$. Its second moment averages $g_t^2$, not squared deviations from a mean.

At step one, divide each buffer by $1-0.5=0.5$. Both update magnitudes are $0.1$, giving ${\theta }_1=[0.9,0.9]$. At step two, divide by $1-0.5^2=0.75$. The second displacement is

Thus ${\theta }_2\approx [0.8,0.842264973]$. The vanished current gradient does not erase saved directional history.

The complete state update is:

m_t = beta1 * m_{t-1} + (1 - beta1) * g_t
v_t = beta2 * v_{t-1} + (1 - beta2) * g_t^2

# bias correction (for early steps when m, v are biased toward 0)
m_hat = m_t / (1 - beta1^t)
v_hat = v_t / (1 - beta2^t)

theta_t = theta_{t-1} - lr * m_hat / (sqrt(v_hat) + eps)

The Adam paper uses this outside-the-square-root epsilon convention. Actual APIs may expose additional choices; inspect the specified implementation.

Why those correction factors? For constant gradient mean $\mu$ and zero initial state, $\mathbb{E}[m_t]=(1-{\beta }_1^t)\mu$. Likewise, constant raw second moment $\nu$ gives $\mathbb{E}[v_t]=(1-{\beta }_2^t)\nu$. The denominators remove those initialization factors. Under changing distributions, corrected buffers still represent weighted history, not unbiased estimates of every current quantity. The ratio of corrected moments is not thereby an unbiased update.

Remove first-moment memory

New task: set ${\beta }_1=0$, retaining the same gradients and ${\beta }_2=0.5$. Now ${\hat{m}}_2=[2,0]$, while ${\hat{v}}_2=[4,16/3]$ is unchanged. Check: ${\theta }_2=[0.8,0.9]$; the second coordinate no longer moves on step two. Use when: testing whether saved directional history causes movement after a gradient disappears.

Enter epsilon’s scale

New task: use first gradient $[10^{-9},1]$ and $ϵ=10^{-8}$ with the same learning rate. The first-step coordinate ratio is $g/(|g|+ϵ)$, not merely its sign. Check: ${\theta }_1\approx [0.990909091,0.900000001]$, rather than $[0.9,0.9]$. Use when: gradient magnitudes approach epsilon; scale cancellation is then only an approximation.

Failure and transfer

Clipping the current gradient does not erase $m_t$ or directly bound the complete Adam displacement. Even a wholly zero current gradient can leave a nonzero moment-driven step. Normalization also does not guarantee decreasing loss on each update.

Reconstruct: why is $v_t$ not the gradient variance?

What an interviewer expects you to say

Reconstruct buffers, step count, correction factors, and the denominator before discussing optimizer preferences. Two FP32 moment arrays require $8N$ bytes: $560$ decimal GB for $70$ billion parameters. That excludes weights, gradients, master copies, activations, and communication buffers.

theta_t = theta_{t-1} - lr * m_hat / (sqrt(v_hat) + eps) - lr * weight_decay * theta_{t-1}