SGD vs Adam vs AdamW Optimizer Math

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 Δθ.

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.

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.

The hidden cost: optimizer state is 2x your parameters

Adam’s power has a price paid in memory. For every parameter you must store two extra full-precision numbers — the first moment m and the second moment v — that persist across steps. That is the ‘2x params’ of optimizer state: a model with P parameters carries 2P optimizer values on top of the weights themselves. Momentum SGD keeps only one buffer (1P); plain SGD keeps none.

In a realistic mixed-precision setup the full accounting is starker. The standard recipe keeps, per parameter: a 16-bit weight (2 bytes) and 16-bit gradient (2), plus an fp32 master copy of the weight (4) and fp32 Adam moments m and v (4 + 4). That totals 16 bytes per parameter, of which 12 — the fp32 master weight plus the two moments — are the optimizer’s footprint, dwarfing the 4 bytes of ‘live’ half-precision weight and gradient. A 7-billion-parameter model therefore needs roughly 7e9 × 16 = 112 GB just for weights, gradients, and optimizer state — before a single activation is stored. The moments you added for stability have become the dominant memory line item.

The bridge to ZeRO and optimizer offload

That 16-bytes-per-parameter reality is precisely what large-scale training systems are built to dismantle, and it is why the memory arithmetic matters as much as the update rule. If the optimizer state is the biggest tenant in memory, the obvious question is: does every GPU really need a full copy of it? ZeRO (the Zero Redundancy Optimizer) answers no. In standard data-parallel training every worker redundantly holds the entire model, gradients, and optimizer state; ZeRO partitions them across workers, so with N data-parallel ranks each holds only 1/N of the optimizer states (stage 1), then also 1/N of the gradients (stage 2) and parameters (stage 3), gathering shards on demand. Because the Adam moments are the fattest slice, sharding them first (stage 1) delivers the biggest immediate win.

Offload attacks the same target on a single node: park m and v — and even the fp32 master weights — in CPU RAM or NVMe, and run the Adam update on the CPU, keeping only the active computation on the accelerator. It trades bandwidth for capacity, letting a modest machine train a model whose optimizer state would never fit in accelerator memory. Both techniques exist because Adam chose stability over frugality — and for CPU-bound and small-lab training, understanding where those two moment buffers live is often the difference between a run that fits and one that does not.

Choosing an optimizer and avoiding the traps

Practical guidance falls out of the math. For transformer and language-model training, reach for AdamW first — defaults of β1 = 0.9, β2 = 0.999, ε = 1e-8, weight decay λ ≈ 0.1, and a learning rate in the 1e-4 to 1e-3 range with warmup — are a strong starting point that rarely needs heroics. Consider momentum SGD only when memory is tight and you can afford to tune, or for problem classes (some vision setups) where it is known to generalize slightly better.

The common traps are worth naming. Forgetting bias correction poisons the opening steps, as the worked example showed. Using L2 regularization in Adam and expecting AdamW behavior silently couples decay to gradient scale — pick the decoupled form on purpose. Setting ε too large quietly turns Adam back toward SGD by flooring the denominator. Skipping learning-rate warmup lets the early, high-variance v_hat produce erratic steps before the second-moment estimate stabilizes. And forgetting the memory model — budgeting for weights but not for the 2P of moment state — is how a run that ‘should fit’ dies with an out-of-memory error. The optimizer is not a detail you set and forget; it is the engine, and its moments are both why training is stable and why it is heavy.