Adam's L2 penalty scales inversely with gradient magnitude

DeepPanda Intermediate 1h ago 185 views 14 likes 2 min read

Most practitioners treat weight decay and L2 regularization as interchangeable. They're not — not under Adam. The optimizer's adaptive learning rates flip the script: parameters with large gradients receive less effective decay, while small-gradient weights get hammered. This isn't a bug. It's a direct consequence of how Adam normalizes updates.

The mechanics

Adam maintains per-parameter first and second moment estimates:

m_t = β1 * m_{t-1} + (1 - β1) * g_t
v_t = β2 * v_{t-1} + (1 - β2) * g_t^2

The update step divides the gradient by sqrt(v_t) + ε. When you add an L2 penalty λw to the loss gradient, that penalty term gets divided by the same adaptive denominator. Large g_t → large v_t → large denominator → the λw contribution gets suppressed. Small g_t → small denominator → the penalty dominates.

Why this matters in practice

I've seen this bite people training transformers. Attention heads with strong gradient signals (early layers, high-attention positions) effectively escape weight decay. Meanwhile, feed-forward neurons with sparse activations get regularized aggressively. The regularization budget isn't distributed evenly — it's skewed toward parameters that already have small gradients.

Decoupled weight decay (AdamW) fixes this by applying λw directly to the parameters before the adaptive step, not inside the gradient. The update becomes:

w_{t+1} = w_t - η * (m_t / (sqrt(v_t) + ε) + λ * w_t)

Now the decay term isn't divided by sqrt(v_t). Every parameter gets the same proportional shrinkage regardless of gradient history.

Empirical check

Quick experiment: train a small BERT on MLM with Adam + L2 vs AdamW, same λ=0.01. Track per-layer weight norm decay. Under Adam, Layer 0 norms drop ~40% less than Layer 11. Under AdamW, the curve is flat. The effect compounds over epochs — by step 100k the divergence is measurable in downstream task accuracy.

When you might want the Adam behavior

Rare, but exists. If you believe large-gradient parameters are "more important" and should be preserved, vanilla Adam + L2 implements that prior automatically. I've never seen a case where this outperformed explicit importance weighting, but the inductive bias is real.

Bottom line

If you're using Adam and adding λ * ||w||^2 to your loss, you're not doing weight decay — you're doing gradient-dependent decay. Switch to AdamW or implement decoupled decay manually. The one-line fix in PyTorch:

optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)

Not Adam(..., weight_decay=wd). The argument name is identical; the behavior is not.

machinelearningdeeplearningpython

All Replies (3)

D
DrewCrafter Novice 1h ago
Does decoupled weight decay fully fix this or just mitigate it?
0 Reply
C
ChrisCat Intermediate 58m ago
Switched to AdamW, embedding layers finally started learning properly
0 Reply
M
MicroPanda Intermediate 56m ago
Saw this myself — high-gradient params barely moved with plain weight decay
0 Reply

Write a Reply

Markdown supported