SGD and Optimization Variants

Why Does This Exist?

Batch gradient descent needs a pass over every training example for a single parameter update. With $N = 10{,}000{,}000$ images, one step costs ten million forward passes and ten million backward passes. Training would grind to a halt.

Stochastic gradient descent replaces the true gradient with an estimate from a small random mini-batch of $B$ samples, typically 32 to 512. Throughput rises by factors of thousands. The price is noise: each estimate wobbles around the truth, and on stretched loss surfaces the wobble has structure. Where curvature is steep along one axis and shallow along another, plain SGD ricochets between the steep walls while barely advancing along the valley floor.

The variants exist to fix those two complaints. Momentum averages out the oscillation. RMSprop rescales each coordinate by its history. Adam does both, with a correction for the early steps. For the object they all consume, see gradients; for the loop around them, see gradient descent variants.

The one idea you need first

A gradient is one slope per parameter, pointing the steepest uphill way, so training steps against it. A mini-batch gradient is that object estimated on $B$ random examples: unbiased, meaning right on average, with variance that shrinks as $1/B$. Everything below is machinery for living with that variance.

Think of It Like This

How It Actually Works

Given data $\{(\mathbf{x}_i, y_i)\}_{i=1}^N$ and parameters $\mathbf{\theta}$, the true loss is the average $\mathcal{L}(\mathbf{\theta}) = \frac{1}{N} \sum_i \ell_i(\mathbf{\theta})$. At step $t$ a random mini-batch $\mathcal{B}_t$ of size $B$ gives the estimate $\mathbf{g}_t = \frac{1}{B} \sum_{i \in \mathcal{B}_t} \nabla \ell_i(\mathbf{\theta}_t)$, and the update is $\mathbf{\theta}_{t+1} = \mathbf{\theta}_t - \eta \mathbf{g}_t$. Because batches are uniform samples, $\mathbb{E}[\mathbf{g}_t] = \nabla \mathcal{L}(\mathbf{\theta}_t)$, with covariance scaling as $1/B$.

Momentum

Polyak momentum keeps a velocity $\mathbf{v}_t = \beta \mathbf{v}_{t-1} + \mathbf{g}_t$ with friction $\beta$ in $[0.9, 0.99)$, and steps $\mathbf{\theta}_{t+1} = \mathbf{\theta}_t - \eta \mathbf{v}_t$. Unrolled, the velocity is an exponential moving average $\sum_{\tau} \beta^{t-\tau} \mathbf{g}_\tau$. Gradients that flip sign cancel toward zero; gradients that persist reinforce, multiplying the effective step by $1/(1 - \beta)$, which is $10$ times at $\beta = 0.9$.

RMSprop

Different parameters need different step sizes. RMSprop tracks squared gradients $\mathbf{s}_t = \beta_2 \mathbf{s}_{t-1} + (1 - \beta_2) \mathbf{g}_t^2$ element-wise, then steps $\mathbf{\theta}_{t+1} = \mathbf{\theta}_t - \eta\,\mathbf{g}_t / (\sqrt{\mathbf{s}_t} + \epsilon)$ with $\epsilon \approx 10^{-8}$. Loud coordinates get turned down, quiet ones get turned up.

Adam

Adam (adaptive moment estimation) keeps both moments, $\mathbf{m}_t = \beta_1 \mathbf{m}_{t-1} + (1 - \beta_1) \mathbf{g}_t$ and $\mathbf{v}_t = \beta_2 \mathbf{v}_{t-1} + (1 - \beta_2) \mathbf{g}_t^2$, then corrects the zero start-up bias with $\hat{\mathbf{m}}_t = \mathbf{m}_t/(1 - \beta_1^t)$ and $\hat{\mathbf{v}}_t = \mathbf{v}_t/(1 - \beta_2^t)$, stepping $\mathbf{\theta}_{t+1} = \mathbf{\theta}_t - \eta\,\hat{\mathbf{m}}_t/(\sqrt{\hat{\mathbf{v}}_t} + \epsilon)$. Defaults are $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$. Second-order context for why curvature matters lives in the Jacobian and Hessian.

Worked example

One Adam step by hand, $t = 1$, from $\mathbf{\theta}_0 = [1.0, 2.0]$ with $\eta = 0.1$ and gradient $\mathbf{g}_1 = [0.2, 4.0]$, twenty times steeper along the second coordinate.

First moment: $\mathbf{m}_1 = 0.9 \times \mathbf{0} + 0.1 \times [0.2, 4.0] = [0.02, 0.40]$. Squared gradient: $\mathbf{g}_1^2 = [0.04, 16.00]$. Second moment: $\mathbf{v}_1 = 0.999 \times \mathbf{0} + 0.001 \times [0.04, 16.00] = [0.00004, 0.01600]$.

Bias correction at $t = 1$ divides out the start-up deficit exactly: $\hat{\mathbf{m}}_1 = [0.02, 0.40]/0.1 = [0.2, 4.0]$, $\hat{\mathbf{v}}_1 = [0.00004, 0.01600]/0.001 = [0.04, 16.00]$. Then $\sqrt{\hat{\mathbf{v}}_1} = [0.2, 4.0]$, the ratios are $[0.2/0.2,\ 4.0/4.0] = [1.0, 1.0]$, and $\mathbf{\theta}_1 = [1.0, 2.0] - 0.1 \times [1.0, 1.0] = [0.9, 1.9]$.

The ravine vanished on step one: a $20$-to-$1$ gradient imbalance became equal $0.1$ moves per coordinate. That equalisation is the whole point of the second moment.

Code

One Adam step as a small function, matching the worked example exactly. NumPy only.

1import numpy as np2
3
4def adam_step(theta: np.ndarray, grad: np.ndarray, m: np.ndarray, v: np.ndarray, t: int,5              lr: float = 0.1, b1: float = 0.9, b2: float = 0.999, eps: float = 1e-8) -> tuple:6    """One Adam update with bias correction; start m and v at zero."""7    m = b1 * m + (1 - b1) * grad8    v = b2 * v + (1 - b2) * grad ** 29    m_hat = m / (1 - b1 ** t)10    v_hat = v / (1 - b2 ** t)11    return theta - lr * m_hat / (np.sqrt(v_hat) + eps), m, v12
13
14theta = np.array([1.0, 2.0])15m = np.zeros(2)16v = np.zeros(2)17theta, m, v = adam_step(theta, np.array([0.2, 4.0]), m, v, 1)18print(np.round(theta, 4))  # -> [0.9 1.9]

Both coordinates move exactly $0.1$ despite the $20$-to-$1$ gradient ratio, which is the equalisation the worked example derived by hand.

Watch Out For

The Quick Version

• Mini-batch SGD estimates the true gradient on $B$ random examples: unbiased, with variance scaling as $1/B$.
• Momentum averages past gradients with decay $\beta$, cancelling oscillation and multiplying steady progress by $1/(1 - \beta)$.
• RMSprop divides each coordinate by its root-mean-square history, equalising loud and quiet parameters.
• Adam keeps both moments with exact bias correction, and is the default starting choice for deep networks.
• In the worked step, a gradient of $[0.2, 4.0]$ becomes equal moves of $0.1$ per coordinate.
• Never skip bias correction or warm-up, and use AdamW so weight decay is not rescaled.

What to Read Next

• Gradient Descent Variants
• Jacobian and Hessian