The update rule is the whole game

Training reduces to one repeated question: given the current weights θ and the gradient g = ∇θ L of the loss, how far and in what direction should we move? The gradient points uphill, so we move against it. Everything else — momentum, adaptivity, weight decay — is a refinement of that single subtraction. The reason the refinements matter is that a raw gradient is a local, noisy, badly-scaled signal. It is local because it only describes the loss surface at one point; noisy because we estimate it from a minibatch, not the whole dataset; and badly scaled because different parameters — an embedding row that fires for one rare token versus a LayerNorm gain used on every token — receive gradients that differ by orders of magnitude.

A good optimizer is a signal-processing layer sitting between the raw gradient and the weight update. It averages away noise (momentum), it equalizes scale (RMSProp/Adam), and it keeps the effective step size in a predictable range so a single global learning rate η can govern a whole heterogeneous network. Hold that framing and every formula below reads as one more filter on the path from g to Δθ.

Advertisement

Plain SGD: subtract a scaled gradient

The baseline could not be simpler. At each step, for every parameter:

g_t = ∇_θ L(θ_t)      # minibatch gradient
θ_(t+1) = θ_t − η · g_t      # η = learning rate

One hyperparameter, one subtraction, no state carried between steps. On convex or well-conditioned problems this is provably sound and, with a decaying learning rate, converges. It is also memory-optimal: SGD stores nothing beyond the weights and the current gradient, which is why it never appears in the optimizer-memory tables that dominate large-model training.

The trouble is the single scalar η. It multiplies every parameter’s gradient equally, but the ‘right’ step for a large-gradient direction is small (or you overshoot and oscillate) while the right step for a tiny-gradient direction is large (or you crawl). With one η you cannot satisfy both, so you tune it to the loudest direction and let the quiet ones learn slowly. In a transformer, where gradient magnitudes vary enormously across embeddings, attention projections, and normalization parameters, that compromise is expensive — plain SGD converges slowly and is fragile to the learning-rate choice. Every later optimizer is an attempt to give each parameter its own effective step without hand-tuning a rate per tensor.

Advertisement

Momentum: an EMA of the gradient

The first fix targets noise and ravines. Instead of stepping on the raw gradient, accumulate a running average of gradients and step on that:

v_t = β · v_(t-1) + (1 − β) · g_t     # β ≈ 0.9
θ_(t+1) = θ_t − η · v_t

This v_t is an exponential moving average (EMA) of the gradient. Unrolling it, v_t = (1−β) Σ β^k g_(t−k): recent gradients dominate, older ones fade geometrically. With β = 0.9 the average has an effective memory of roughly 1/(1−β) = 10 steps. (Some texts write the recurrence as v_t = β v_(t-1) + g_t and fold the 1−β into η; the behavior is identical.)

Why it helps: across steps, the noisy components of the minibatch gradient point in random directions and partially cancel in the average, while the consistent downhill component reinforces. In a long, narrow valley — the geometry of most deep loss surfaces — raw SGD zig-zags across the walls; momentum builds velocity along the floor and damps the sideways bounce, like a heavy ball rolling past small bumps. Nesterov momentum adds a small ‘look-ahead’ refinement by evaluating the gradient at the point momentum is about to carry you. Momentum was the workhorse of the pre-transformer era and still trains many vision models, but it does not touch the scale problem — every parameter still shares one η.

RMSProp: an EMA of the squared gradient

The second fix targets scale directly. Track a running average of each parameter’s squared gradient and divide the step by its square root:

s_t = β · s_(t-1) + (1 − β) · g_t²   # elementwise square
θ_(t+1) = θ_t − η · g_t / (√s_t + ε)

Here s_t estimates the (uncentered) second moment — roughly the average magnitude-squared of recent gradients for that specific parameter. Dividing by √s_t normalizes the step: a parameter whose gradients are consistently large gets its step shrunk, while a parameter whose gradients are consistently tiny gets its step amplified. The result is that every parameter moves at a comparable effective rate regardless of its raw gradient scale — the per-parameter adaptivity that a single global η could never provide.

The ε (typically 1e-8) is a floor that stops division by zero and caps the step for parameters with near-zero gradient history. RMSProp is essentially the ‘denominator’ of Adam. Notice it only rescales — it uses the raw gradient g_t in the numerator, so it fixes scale but not noise. Momentum fixed noise but not scale. Each solves half the problem; the obvious move is to combine them.

Adam: momentum in the numerator, RMSProp in the denominator

Adam (‘adaptive moment estimation’) keeps both running averages — the mean of the gradient and the mean of the squared gradient — and uses the first as the direction and the second as the per-parameter scale:

m_t = β1 · m_(t-1) + (1 − β1) · g_t     # 1st moment (mean of g)
v_t = β2 · v_(t-1) + (1 − β2) · g_t²    # 2nd moment (mean of g²)

The numerator uses the momentum-style average m_t (smoothed direction), and the denominator uses the RMSProp-style average v_t (per-parameter scale). Loosely, the update direction is m_t / √v_t, which behaves like a signal-to-noise ratio: when a parameter’s gradients are consistent, m_t is large relative to √v_t and the step is confident; when they are erratic, the numerator partly cancels while the denominator stays large, and the step shrinks automatically.

The two decay rates play different roles. β1 = 0.9 gives the direction a ~10-step memory — responsive. β2 = 0.999 gives the scale a ~1000-step memory — deliberately slow, because the magnitude estimate should be stable and not lurch with a single noisy batch. That asymmetry is the reason Adam feels robust: a fast, forgiving direction over a slow, steady scale. But both averages start at zero, and that initialization introduces a bias we have to correct before the update is trustworthy.

The bias-correction problem: why m_hat and v_hat exist

Both m_t and v_t are initialized to zero. Because an EMA blends the new value with the previous accumulator, early estimates are dragged toward that zero start — they are biased toward zero, most severely on the first few steps. Concretely, after one step m_1 = (1−β1) g_1 = 0.1 g_1: the first-moment estimate is ten times smaller than the gradient it is supposed to average. For v_1 = (1−β2) g_1² = 0.001 g_1² the shortfall is a thousand-fold.

Taking expectations, one can show E[m_t] = (1 − β1^t) E[g_t] when the gradient distribution is roughly stationary, and likewise for v_t with β2. So dividing by that exact factor removes the bias:

m_hat_t = m_t / (1 − β1^t)
v_hat_t = v_t / (1 − β2^t)

At t = 1, m_hat = m_1 / 0.1 = g_1 and v_hat = v_1 / 0.001 = g_1² — the corrected estimates equal the actual quantities they estimate, exactly as they should. As t grows, β^t → 0, the correction factor → 1, and it quietly switches itself off. Without this step the early updates are badly miscalibrated — and, as the worked example shows next, in a surprising direction.

The full Adam update, assembled

Putting the four pieces together, one Adam step at iteration t is:

g_t     = ∇_θ L(θ_t)              # 1. gradient
m_t     = β1 · m_(t-1) + (1 − β1) · g_t     # 2. 1st moment (EMA of g)
v_t     = β2 · v_(t-1) + (1 − β2) · g_t²    # 3. 2nd moment (EMA of g²)
m_hat_t = m_t / (1 − β1^t)              # 4. bias-correct
v_hat_t = v_t / (1 − β2^t)              # 5. bias-correct
θ_(t+1) = θ_t − η · m_hat_t / (√v_hat_t + ε)   # 6. update

# Defaults: η=1e-3, β1=0.9, β2=0.999, ε=1e-8

A useful property falls out of this form. When a parameter’s gradient is steady, m_hat ≈ √v_hat in magnitude, so the ratio m_hat / √v_hat ≈ ±1 and the step size is bounded by roughly η — independent of the gradient’s raw magnitude. This self-normalization is why Adam’s learning rate transfers across layers and even across models: η sets an approximate trust region on how far any parameter moves per step, and the adaptive denominator handles the per-parameter scaling underneath. It is also why the same η = 1e-3-ish setting is a sane starting point for an enormous range of architectures, something no fixed-scale SGD rate can claim.

A worked single step, by hand

Take one parameter with value θ = 0.500 and gradient g_1 = 0.20 at the first step, with the default hyperparameters and m_0 = v_0 = 0.

m_1     = 0.9·0 + 0.1·0.20      = 0.020
v_1     = 0.999·0 + 0.001·0.20²  = 0.00004
m_hat   = 0.020   / (1 − 0.9)   = 0.020 / 0.1   = 0.200
v_hat   = 0.00004 / (1 − 0.999) = 0.00004/0.001 = 0.040
√v_hat  = 0.200
step    = η · m_hat/(√v_hat+ε) = 1e-3 · 0.200/0.200 = 1e-3
θ_new   = 0.500 − 0.001 = 0.499

The corrected step is exactly η = 0.001: on the first iteration m_hat/√v_hat = g/|g| = 1, confirming the trust-region intuition. Now watch what bias correction bought us. Using the uncorrected moments, the step would have been 1e-3 · 0.020/√0.00004 = 1e-3 · 0.020/0.006325 = 3.16e-3 — more than three times too large. The asymmetry is the culprit: v sits under a square root, so its thousand-fold early bias becomes a ~31.6× deflation of the denominator, dwarfing the ten-fold bias in the numerator. Bias correction cancels both cleanly and hands back a calibrated η-sized step. Skip it and the opening steps of training lurch, exactly when the weights are most fragile.

AdamW: decoupling weight decay

Regularization usually means weight decay — gently shrinking weights toward zero each step to discourage overly large values. The classic way is L2 regularization: add (λ/2)·θ² to the loss, which adds λ·θ to the gradient. That works fine for SGD, but inside Adam it misbehaves. The decay term flows through the adaptive denominator, so it gets divided by √v_hat along with everything else — parameters with large gradient history are decayed less, and parameters with tiny gradients are decayed more. The amount of regularization a weight receives ends up coupled to its gradient magnitude, which is not what anyone intends.

AdamW fixes this by decoupling the decay from the gradient-based update — apply it as a separate, direct shrink:

θ_(t+1) = θ_t − η · m_hat/(√v_hat + ε)  −  η · λ · θ_t

Now every weight decays by the same fraction ηλ per step, untouched by v_hat. The separation restores weight decay to its intended, uniform effect and, empirically, generalizes better. AdamW — not vanilla Adam — is the de-facto optimizer for training essentially every modern transformer, with λ around 0.1 a common choice.

Why adaptive methods dominate transformer training

Transformers are a near-perfect adversary for a single global learning rate. A single model mixes token embeddings (sparse, huge gradients on the few tokens in a batch), attention query/key/value projections, feed-forward matrices, and LayerNorm gains and biases — and the natural gradient scale of these tensors differs by orders of magnitude. Worse, the scales drift during training. Plain SGD would demand a per-tensor learning-rate schedule tuned by hand; Adam derives the equivalent automatically from each parameter’s own v_hat.

There is also a sparsity argument. A token embedding row only receives a gradient when its token appears; over many steps its EMA v_hat stays small, so when a gradient does arrive the step is amplified — the parameter keeps learning despite rare updates, where SGD would starve it. Add the empirical record: the attention-is-all-you-need line of models, GPT-style pretraining, BERT, and virtually every large language model since were trained with Adam or AdamW, because SGD either fails to converge at these scales or needs impractically careful tuning and warmup. The robustness is the point: Adam is forgiving of learning-rate choice and architecture changes, which is priceless when a single pretraining run costs a fortune and cannot be casually re-tuned. That reliability, more than any single-run speed edge, is why the field standardized on it.