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.

Which tensors should actually be decayed

Decoupling makes the decay uniform across the parameters you apply it to, which makes the choice of which parameters more consequential, not less. The standard GPT-style recipe splits the model into two parameter groups: everything with two or more dimensions gets weight_decay = 0.1; everything one-dimensional gets 0.0.

The rationale is that decay is a prior on the scale of a linear map. Shrinking a matrix in W_Q, W_K, or an MLP projection genuinely constrains the function class. Shrinking a bias just translates the preactivation distribution, and shrinking a LayerNorm or RMSNorm gain toward zero actively fights the normalizer — those gains exist to set the scale, and pulling them to zero fights the normalizer and can destabilize training. Embeddings are the contested case: they are 2-D and so decayed by default, but their gradients are extremely sparse — a rare token gets ten gradient updates and thousands of decay steps — so its vector bleeds toward zero. Some recipes exclude them.

AdamW on a CPU SLM budget

The decay itself is free. It is one fused multiply-add per parameter, riding along in a loop that already computes a square root and a division — on a CPU the cost disappears into the memory traffic. What is not free is Adam’s state.

per parameter (fp32):  θ 4B + g 4B + m 4B + v 4B = 16 B
125M-param SLM:        125e6 × 16 B ≈ 2.0 GB
  of which optimizer state (m, v):  1.0 GB

Three quarters of that is not the model — it is bookkeeping, and it is exactly what Adafactor and 8-bit optimizers attack. CPU training is also memory-bandwidth-bound: the step streams four full-size arrays, so fusing the decay into that same pass matters more than its FLOP count suggests. Under mixed precision, apply the decay to the fp32 master weights, not the bf16 copy — a relative shrink of 1e-4 sits below bf16’s roughly three decimal digits of resolution and would round away to nothing.

Pitfalls that silently change your regularizer

Every item here fails quietly — no exception, no warning, just a slightly worse model.

• Passing weight_decay to torch.optim.Adam. That is the L2-in-the-gradient path, not AdamW. Same argument name, different algorithm, no error. Use torch.optim.AdamW.
• Double regularizing. An explicit L2 term in the loss plus a nonzero weight_decay runs both mechanisms at once; only one is uniform.
• Clipping asymmetry. Clipping acts on g before the optimizer, so under AdamW the decay is never clipped — under Adam with L2 it was.
• Epsilon placement. √v̂ + ε and √(v̂ + ε) are different algorithms, and a larger ε softens the preconditioner back toward SGD.
• Decaying frozen parameters. In LoRA or partial fine-tuning a parameter with no gradient is still shrunk every step if it sits in a decayed group.