An Adam Optimizer Step
Apply one Adam update with worked numbers, tracking the bias-corrected first and second moment estimates.
Problem
Adam is an adaptive optimizer that maintains running averages of the gradient (first moment) and its square (second moment), then scales each parameter's step by these. It combines momentum with per-parameter learning rates, which makes it robust across problems with little tuning.
Update rule
With gradient g at step t: m = beta1 m + (1-beta1) g, v = beta2 v + (1-beta2) g^2. Bias-correct m_hat = m/(1-beta1^t) and v_hat = v/(1-beta2^t), then step by -lr m_hat / (sqrt(v_hat) + eps). The bias correction matters most in early steps when the running averages are still near zero.
import numpy as np
beta1,beta2,eps,lr=0.9,0.999,1e-8,0.1
m=v=0.0; theta=1.0
for t in range(1,4):
g=2*theta # gradient of theta^2, minimum at 0
m=beta1*m+(1-beta1)*g
v=beta2*v+(1-beta2)*g*g
mh=m/(1-beta1**t); vh=v/(1-beta2**t)
theta-=lr*mh/(np.sqrt(vh)+eps)
print('t',t,'theta',round(theta,4))
Result
Minimizing theta^2 from theta=1, each Adam step moves theta toward zero. The bias correction inflates the early moment estimates so the first steps are not artificially tiny. Because v_hat normalizes the step by the gradient scale, Adam takes similar-sized steps regardless of how large the raw gradients are, which is why it works out of the box on very different loss surfaces.
- Bias correction is essential for the first few steps; without it Adam starts far too slowly.
- The per-parameter scaling adapts to differing gradient magnitudes across weights automatically.
- Kronos trains its neural surrogates and control policies primarily with Adam and its variants.