AdamW Math Deep Dive

Two ways to shrink a weight

There are two distinct mechanisms people call “weight decay,” and conflating them is the entire source of the confusion.

The first is L2 regularization: you change the objective. Add a penalty on the parameter norm and let the gradient of that penalty flow through whatever optimizer you happen to be using.

L’(θ) = L(θ) + (λ’/2) · ||θ||^2
g’   = ∇L’(θ) = g + λ’θ        where g = ∇L(θ)

The second is weight decay proper: you leave the objective alone and multiplicatively shrink the parameter as a separate step in the update rule.

θ ← (1 − λ) · θ   — then apply whatever the optimizer says

One changes the loss surface; the other changes the update rule. Whether they coincide depends entirely on what sits between the gradient and the parameter.

Why they coincide under plain SGD

For vanilla SGD the two mechanisms really are interchangeable, and the one-line proof is worth doing because it shows exactly which assumption AdamW later violates. Take a gradient-descent step on the L2-augmented objective:

θ ← θ − α(g + λ’θ)
  = θ − αg − αλ’θ
  = (1 − αλ’)θ − αg

Compare that with the decoupled form θ ← (1 − λ)θ − αg. They are the same update whenever λ = αλ’. So under SGD, L2 and weight decay are the same algorithm under a reparameterization of the coefficient — and even here the mapping already couples decay to the learning rate: halve α and you have silently halved your effective decay. The step that made this work was pulling α out as a common scalar factor, which is only legal because SGD multiplies the entire gradient by the same number.

The preconditioner breaks the equivalence

Adam does not multiply the gradient by a scalar. It multiplies it, elementwise, by a per-parameter preconditioner 1/(√v̂ + ε) built from the running second moment. Redo the algebra with the L2 term folded into the gradient (set β_1 = 0 to keep the notation clean):

θ ← θ − α · (g + λ’θ) / (√v̂ + ε)
  = θ − αg/(√v̂ + ε) − [αλ’/(√v̂ + ε)] · θ

The decay coefficient is no longer the constant αλ’. It is αλ’/(√v̂_i + ε) — a different number for every parameter, inversely proportional to that parameter’s recent gradient magnitude. There is no reparameterization of λ’ that recovers uniform decay, because no single scalar can cancel a vector.

The direction of the distortion is the damning part: weights with large, noisy gradients get less shrinkage, quiet weights get more. A second, subtler contamination: λ’θ sits inside g’, so it also feeds v, inflating the denominator and damping the real gradient signal along with the decay.

The decoupled update, written out

AdamW’s fix (Loshchilov and Hutter, 2019) is to take the decay out of the gradient entirely, so it never touches the moment estimates and never passes through the preconditioner:

g_t     = ∇L(θ_{t-1})                  # no λθ term
m_t     = β_1 m_{t-1} + (1 − β_1) g_t
v_t     = β_2 v_{t-1} + (1 − β_2) g_t^2
m̂_t     = m_t / (1 − β_1^t)
v̂_t     = v_t / (1 − β_2^t)
θ_t     = θ_{t-1} − α_t · m̂_t/(√v̂_t + ε)   − α_t · λ · θ_{t-1}
                ↳ adaptive step        ↳ decoupled decay

Shapes are trivial — m, v, g, and θ all share the parameter’s shape, and every operation is elementwise — but the placement of the decay term is everything. It sits outside the division, so its coefficient is the same scalar α_tλ for every weight in the group.

PyTorch writes it in the other order — shrink first, θ ← θ(1 − α_tλ), then subtract the adaptive step — but the two are algebraically identical, because the step vector is computed from g_t and does not depend on θ. Expanding gives θ − α_tλθ − α_t s either way.

A worked numeric example: two weights, two fates

Take α = 1e-3, decay coefficient 0.1, and two weights whose second moments differ by two orders of magnitude — an ordinary spread inside a transformer.

Under Adam with L2, weight B is decayed 100× harder than weight A, purely because its gradients are smaller. Hold those second moments fixed for 1000 steps and compound the multiplicative factors:

Adam+L2, A:  (1 − 0.001)^1000  ≈ 0.368     → shrunk to ~37%
Adam+L2, B:  (1 − 0.1)^1000    ≈ 1.7e-46   → annihilated
AdamW,   A:  (1 − 1e-4)^1000   ≈ 0.905     → shrunk to ~90%
AdamW,   B:  (1 − 1e-4)^1000   ≈ 0.905     → identical

Real second moments move, so nothing is literally annihilated — but the spread is real and persistent. AdamW applies one regularizer; Adam with L2 applies a different one to every tensor in your model.

The equilibrium weight scale

A cleaner way to see what decoupling buys is to ask where a weight comes to rest. Set the expected AdamW update to zero:

α · m̂/(√v̂ + ε)  =  −α · λ · θ*
⇒  θ*  =  −(1/λ) · m̂/(√v̂ + ε)

The learning rate cancels. And because m̂/(√v̂ + ε) is a normalized quantity — a signal-to-noise ratio bounded around ±1 for a consistent gradient — AdamW imposes an approximate ceiling |θ*| ≤ 1/λ that is the same for every decayed weight. With λ = 0.1, that is a scale of order 10.

Run the same argument for Adam with L2 and you get θ* = −g/λ’: the resting magnitude tracks the raw gradient, so high-gradient weights settle large and quiet ones settle near zero. AdamW enforces a uniform scale budget; L2-inside-Adam enforces a gradient-proportional one. A heuristic fixed-point argument, not a theorem — but it captures the practical difference.

Hyperparameter coupling: learning rate, decay, and the schedule

The second thing decoupling buys is a more separable search space. With L2 inside Adam, α and λ’ interact through the preconditioner, and a grid search over the pair produces a diagonal ridge: change one and you must chase the other. Loshchilov and Hutter’s central empirical result is that AdamW’s (α, λ) surface is far more axis-aligned, so you can tune the two roughly independently — α as a speed knob, λ as a shape knob.

Two caveats that trip people up. First, decoupled does not mean learning-rate-independent: PyTorch’s per-step decay is lr × weight_decay, so it still scales with α. The original paper writes the update as θ − η_t(α · m̂/(√v̂+ε) + λθ), where the schedule multiplier η_t scales decay but the base learning rate does not — so a λ quoted from the paper is not a weight_decay you can paste into PyTorch. Second, because α_t follows your schedule, a cosine decay to zero also decays your regularization pressure to zero at the end of training.

Nor is λ portable across run shapes: compounding (1 − αλ) over T steps gives roughly exp(−αλT), so halving the batch size at a fixed token budget doubles the step count and doubles your regularization. The paper’s normalized coefficient λ = λ_norm · √(b / (B·T)) exists to absorb exactly that.