Optimizers: SGD, momentum, Adam

The idea

Every optimizer in common use is gradient descent with refinements. The ladder, one rung at a time:

SGD. $\theta \leftarrow \theta - \eta\, g_t$ , with $g_t$ the gradient on a random minibatch. Unbiased but noisy, and one global $\eta$ has to suit every parameter in the network at once.

Watch the learning rate itself decide the outcome on a single parameter — too small crawls, too large overshoots and oscillates, too large still and it diverges:

η = 0.30

Each step scales distance-to-minimum by (1 − 2η) ≈ 0.40 — at η = 0.30 that is converging smoothly.

+ Momentum. Keep a velocity — an exponentially weighted average of recent gradients — and step along it:

$v_t = \beta v_{t-1} + g_t, \qquad \theta \leftarrow \theta - \eta\, v_t$

With $\beta = 0.9$ , a direction the gradients agree on compounds toward a terminal velocity of $g/(1-\beta) = 10g$ , while components that alternate sign cancel. Effect: faster along consistent directions, damped oscillation across steep narrow ravines, minibatch noise averaged out.

The ravine language is literal — here is a basin with exactly that kind of ridge running across it, and a ball spiraling down into it:

Loading the interactive 3D loss-surface visualization…

+ Per-parameter scale (RMSProp idea). Track a running average of the squared gradient and divide by its root, so each parameter gets a step size inversely proportional to its typical gradient magnitude. Parameters with rare, small gradients (embedding rows, early layers) stop being starved by a global $\eta$ chosen for the loudest layer.

Adam = both + bias correction.

$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$

\begin{aligned} \hat{m}_t &= \frac{m_t}{1-\beta_1^t}, \\ \hat{v}_t &= \frac{v_t}{1-\beta_2^t}, \\ \theta &\leftarrow \theta - \eta\, \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}. \end{aligned}

Defaults: $\beta_1 = 0.9$ , $\beta_2 = 0.999$ , $\epsilon = 10^{-8}$ . The $\hat{m}/\sqrt{\hat{v}}$ ratio is roughly gradient-scale-free, which is why Adam works out of the box across wildly different architectures.

AdamW fixes one subtlety: L2 regularization added into the gradient gets divided by $\sqrt{\hat{v}}$ like everything else, so exactly the weights with large gradient variance feel less decay. AdamW applies decay directly to the weights, $\theta \leftarrow (1 - \eta\lambda)\theta$ , outside the adaptive machinery — the standard for transformers.

Worked example

Bias correction, with numbers. Single parameter, constant gradient $g = 1$ , defaults above. At $t = 1$ :

$m_1 = 0.1 \cdot 1 = 0.1, \qquad v_1 = 0.001 \cdot 1 = 0.001$

Uncorrected, the step would be $\eta\, m_1/\sqrt{v_1} = \eta \cdot 0.1/0.0316 \approx 3.16\,\eta$ — a 3x overshoot on step one, purely an artifact of the zero-initialized averages (and it gets worse for smaller $\beta_1$ /larger $\beta_2$ gaps). With correction:

\begin{aligned} \hat{m}_1 &= \frac{0.1}{1 - 0.9} = 1, \\ \hat{v}_1 &= \frac{0.001}{1 - 0.999} = 1, \\ \text{step} &= \eta \cdot \frac{1}{\sqrt{1} + \epsilon} \approx \eta. \end{aligned}

Exactly the sensible thing: first step of size $\eta$ in the gradient direction.

Scale invariance, same setup. Repeat with $g = 0.001$ everywhere: $\hat{m}_1 = 0.001$ , $\hat{v}_1 = 10^{-6}$ , step $= \eta \cdot 0.001/0.001 = \eta$ . Identical. Adam takes $\eta$ -sized steps whether the raw gradients are of order $1$ or $10^{-3}$ — with SGD, those two scenarios would need learning rates 1000x apart. That is the practical content of "adaptive".

Cheat sheet

Optimizer	Update	Adds	State/param
SGD	$\theta \leftarrow \theta - \eta g$	—	0
SGD + momentum	$v \leftarrow \beta v + g$ ; $\theta \leftarrow \theta - \eta v$	direction smoothing	1
RMSProp	$\theta \leftarrow \theta - \eta g / (\sqrt{s} + \epsilon)$ , $s$ EMA of $g^2$	per-param scale	1
Adam	momentum + RMS scale + bias correction	both	2
AdamW	Adam + decoupled decay $(1 - \eta\lambda)\theta$	correct weight decay	2

Defaults: $\beta = 0.9$ (momentum); Adam $\beta_1 = 0.9$ , $\beta_2 = 0.999$ , $\epsilon = 10^{-8}$ ; AdamW $\lambda \approx 0.01$ - $0.1$ .
LR starting points: SGD+momentum (vision) $0.1$ with batch 256 + decay; Adam generic $10^{-3}$ ; transformer pretraining $\sim 10^{-4}$ - $6\times10^{-4}$ with warmup; fine-tuning $10^{-5}$ - $5\times10^{-5}$ .
Terminal velocity: constant gradient → momentum amplifies steps by $1/(1-\beta)$ ; remember 10x at $\beta = 0.9$ .
Schedules: warmup (linear, first 1-5% of steps) then cosine or linear decay is the modern default recipe.

Symptoms → fixes

Loss diverges / NaNs in first steps → LR too high or no warmup; add warmup, clip gradient norm to ~1.
Loss plateaus early, gradients healthy → LR too low or decayed too soon.
Adam trains, SGD stalls → heterogeneous gradient scales; keep Adam or fix normalization.
Overfitting with Adam + L2 "not helping" → decay is being adaptively rescaled; switch to AdamW.
Extra GPU memory 2-3x parameter count during training → optimizer state ( $m$ , $v$ ) — expected, not a leak.