The Adam Optimizer

Adam keeps two running averages per parameter and uses them to give every parameter its own step size. It is the default optimizer in most deep-learning code. After this lesson your program applies Adam, including its bias correction, to a single parameter.

Two problems with a single global step

• Noise. Mini-batch gradients jitter from step to step, so the path zig-zags instead of heading for the minimum.
• Scale. Some parameters see gradients of size 100 and others of size 0.001. One learning rate is far too big for the first and far too small for the second.

Two moving averages

For each parameter Adam tracks the gradient g's running mean and running mean of squares:

m  <-  b1 * m + (1 - b1) * g        direction: a smoothed gradient (momentum)
v  <-  b2 * v + (1 - b2) * g * g    magnitude: how large gradients have been lately

The usual values are b1 = 0.9 and b2 = 0.999. m plays the role of momentum: noise cancels and a consistent direction survives. v remembers the typical gradient size; dividing by its square root turns "gradient" into "gradient in units of its own typical size".

Bias correction

Both averages start at 0, so early on they are biased toward 0: after one step m holds only a tenth of the gradient. Adam divides by how much of the average has been filled in so far, using the step count t (starting at 1 on the first step):

m_hat = m / (1 - b1^t)
v_hat = v / (1 - b2^t)

As t grows, b1^t and b2^t shrink to 0 and the correction fades out. On the very first step it exactly undoes the shrinkage: m_hat is g and v_hat is g*g.

The update

theta  <-  theta - lr * m_hat / ( sqrt(v_hat) + eps )

eps (about 1e-8) only prevents dividing by zero. Because the step is the smoothed gradient divided by the typical gradient size, its magnitude is close to lr regardless of how big or small the raw gradients are. That is why one learning rate such as 0.001 works across so many models.

A worked trace

theta = 2, lr = 0.1, b1 = 0.9, b2 = 0.999, gradients 0.5, then -0.5:

t = 1  g = 0.5    m = 0.05      v = 0.00025    m_hat = 0.5       v_hat = 0.25
       step = 0.1 * 0.5 / (0.5 + eps) = 0.1000          theta = 1.9000
t = 2  g = -0.5   m = -0.005    v = 0.00049975 m_hat = -0.02632  v_hat = 0.25
       step = 0.1 * (-0.02632) / 0.5 = -0.00526         theta = 1.9053

The first step moves the full lr. When the gradient flips sign, the average m nearly cancels and the step shrinks to about 5% of the first: Adam refuses to chase noise.

Common mistakes

• Bias-correcting with the step count before incrementing it. At t = 0 the denominator 1 - b1^0 is 0.
• Squaring after averaging ((mean g)^2) instead of averaging the square.
• Putting eps inside the square root, or forgetting it, which breaks on g = 0.
• Not resetting m, v and t when a new run starts.
• Using the wrong t for v_hat: both corrections use the same t.

Your exercise

The input is a script of INIT, LR, BETAS, EPS, STEP and PRINT commands. The starter keeps the state in a four-slot list holding theta, m, v and t, and supplies a helper for integer powers. Implement adam_step, which updates that state in place for one gradient.