AdamW against L2 in the gradient
Medium15 min10 hidden tests
A team fine-tuning a vision model wants to know why switching from Adam with L2 regularization to AdamW changed their results. Write decay_two_ways(w, grad, m, v, t, lr, weight_decay, ...) that takes one step both ways from the same weights and Adam state. L2 in the gradient adds weight_decay * w to the gradient and then runs a normal Adam step. AdamW runs the Adam step on the plain gradient and separately shrinks each weight by lr * weight_decay * w. In the second example, AdamW shrinks both weights by the same 0.1%, while L2 barely touches the weight whose past gradients were large.
Change weight 2’s past gradients
L2: Adam on grad + wd·w. AdamW: Adam on grad, then w −= lr·wd·wAdamW shrinks both weights by lr × wd × w = 0.01 × 0.10 × 1 = 0.001. L2 adds wd × w to the gradient instead, and Adam then divides it by √v̂: weight 1 shrinks 0.001461, weight 2 0.000015, about 95 times less, because its past gradients were larger.
AdamW 0.001 for both. L2: 0.001461 and 0.000015.
Examples
First step from zero state
decay_two_ways([1.0, -2.0], [0.5, 0.5], [0.0, 0.0], [0.0, 0.0], 1, 0.1, 0.1)
returns {'l2': [0.9, -2.1], 'adamw': [0.89, -2.08]}
No gradient: only the decay acts, and Adam's scaling bends the L2 version
decay_two_ways([1.0, 1.0], [0.0, 0.0], [0.0, 0.0], [0.0001, 1.0], 10, 0.01, 0.1)
returns {'l2': [0.998539, 0.999985], 'adamw': [0.999, 0.999]}
Requirements
w,grad,mandvare 1-D lists or arrays of the same length.tcounts steps from 1.weight_decayis 0 or more.- The Adam step is the standard one with bias correction:
m' = beta1 * m + (1 - beta1) * g,v' = beta2 * v + (1 - beta2) * g ** 2, steplr * m_hat / (sqrt(v_hat) + eps), withm_hat = m' / (1 - beta1 ** t)andv_hat = v' / (1 - beta2 ** t). - L2: use
g = grad + weight_decay * win that step. AdamW: useg = grad, then also subtractlr * weight_decay * w, withwthe weights before this step. - Return
{"l2": array, "adamw": array}: the new weights from each method, both float arrays of the input's length.
The idea behind it
Weight decay pulls weights toward 0 to fight overfitting. With plain gradient descent, adding an L2 penalty 0.5 * wd * ||w||^2 to the loss and shrinking the weights each step are the same thing. With Adam they differ. If the decay wd * w joins the gradient, Adam divides it by sqrt(v_hat) like any other gradient. A weight with large past gradients then gets almost no decay, and one with tiny gradients gets a lot. Loshchilov and Hutter's AdamW, from 2017, decouples the two: Adam handles the loss gradient, and the decay is applied straight to the weights, w -= lr * wd * w. Every weight shrinks by the same share. AdamW is the default optimizer for training transformers, usually with a weight decay near 0.1. In PyTorch, Adam(weight_decay=...) is the L2 version and AdamW is the decoupled one.