Adam: A Method for Stochastic Optimization

Adam is the optimizer most deep networks are trained with, and its update rule is short enough to implement in a dozen lines. This walkthrough builds it from plain SGD and watches it optimise a simple function.

The problem

Plain stochastic gradient descent takes the same step size for every parameter: $\theta \leftarrow \theta - \alpha \, g$, where $g$ is the gradient. Two difficulties follow:

• One rate does not fit all parameters. Some directions are steep and need small steps; others are flat and could move faster. A single $\alpha$ compromises.
• Mini-batch gradients are noisy. Each step sees a different batch, so raw gradients jitter, and the path zig-zags.

We would like each parameter to get its own effective step size, adapted from the gradient history, with the noise smoothed out.

The idea

Keep two running averages of the gradient as training proceeds:

• the first moment $m$, a smoothed gradient (momentum), which damps noise;
• the second moment $v$, a smoothed gradient squared, a per-parameter scale of how large gradients have been.

Divide the smoothed gradient by the square root of the second moment. A parameter with consistently large gradients gets a smaller effective step; a quiet parameter gets a larger one. Each direction adapts on its own.

The update rule

At step $t$, with gradient $g_t$ and decay rates $\beta_1, \beta_2$:

$$m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t \qquad v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2$$

Because $m$ and $v$ start at zero, early estimates are biased toward zero, so they are corrected:

$$\hat{m}_t = \frac{m_t}{1 - \beta_1^{\,t}} \qquad \hat{v}_t = \frac{v_t}{1 - \beta_2^{\,t}}$$

Then the parameter update:

$$\theta_t = \theta_{t-1} - \alpha \, \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$$

The paper's defaults are $\beta_1 = 0.9$, $\beta_2 = 0.999$, and $\epsilon = 10^{-8}$ to avoid dividing by zero.

A minimal implementation

Minimise $f(\theta) = \theta_0^2 + 10\,\theta_1^2$, a bowl that is ten times steeper along one axis than the other, which is exactly where a single learning rate struggles:

import numpy as np

def grad(theta):
    # gradient of theta0^2 + 10*theta1^2
    return np.array([2 * theta[0], 20 * theta[1]])

theta = np.array([5.0, 5.0])
m = np.zeros(2)
v = np.zeros(2)
alpha, b1, b2, eps = 0.1, 0.9, 0.999, 1e-8

for t in range(1, 201):
    g = grad(theta)
    m = b1 * m + (1 - b1) * g
    v = b2 * v + (1 - b2) * g**2
    m_hat = m / (1 - b1**t)          # bias correction
    v_hat = v / (1 - b2**t)
    theta = theta - alpha * m_hat / (np.sqrt(v_hat) + eps)

print(np.round(theta, 4))            # both components near 0

What to observe

• Both components approach zero together, even though the two axes have very different curvature, the per-parameter scaling handles the steep axis and the flat one at once.
• Drop the bias correction (use m and v directly): the first several steps are noticeably smaller, because the zero-initialised averages start shrunk toward zero. The correction is what fixes the early steps.
• Feed the same function to plain SGD with one rate large enough for the flat axis, and it oscillates or diverges along the steep one.

Limitations

• Adam often trains fast but can generalise slightly worse than well-tuned SGD with momentum on some vision tasks, the choice is empirical.
• Naive L2 weight decay interacts badly with the adaptive scaling; the common fix is AdamW, which decouples weight decay from the gradient update.
• It is still first-order and local: it finds a nearby minimum, not a global one, and remains sensitive to the base learning rate $\alpha$.

Further reading

• Adam: A Method for Stochastic Optimization (arXiv:1412.6980)
• Decoupled Weight Decay Regularization, AdamW (arXiv:1711.05101)