In plain SGD, L2 regularisation and weight decay are the same thing. Adding a penalty to the loss, or shrinking the weights directly after each step, will update the weights in identical ways. As a result, the two names get used interchangeably, but when we apply the same trick to Adam the equivalence breaks down.
The adaptive step rescales each parameter's update by its own moving gradient, but it scales the L2 penalty as well. This means how much a parameter decays depends on how large its recent gradients have been. AdamW fixes this by applying the decay to the weights directly, outside the adaptive step, so it stays the same fixed fraction for every parameter.
Below we will derive this in detail and show why the two diverge.
Before getting into Adam, we should take a moment to think about what an optimiser does. Vanilla Gradient Descent moves each parameter in the direction of steepest descent at rate :
Where the true gradient is defined over the whole dataset as . In practice we use stochastic gradient descent (SGD), which lets us approximate the full-dataset loss by processing one minibatch at a time. This introduces noise, but lets us take many steps cheaply. Incidentally the noise is often a helpful regularising force.
The core problem here with SGD is that a single learning rate for all parameters is generally inefficient. Loss landscapes are rarely smooth and easily navigable. Some directions are steep and need small steps to remain stable, while others are shallow and need large steps to avoid getting stuck. Unfortunately SGD treats them all the same.
Momentum partially addresses this, where instead of using the raw gradient to determine the steps, we maintain an exponential moving average of past gradients:
The momentum term is the same shape as the model parameters, so it doubles the amount of memory required. It accumulates gradients across consistent directions and cancels in oscillating ones.
This gives faster progress in the shallow directions and less noise in the steep ones.
So this allows us to modulate the update depending on how consistent the gradient updates are, but still doesn’t give us control over the learning rate of each parameter.
The core insight behind Adam is that each parameter can have its own effective learning rate, which is calibrated to the magnitude of recent gradients.
There are cases we’ve not really considered thus far. What about parameters that are receiving large and consistent gradients? Unlike those on a plateau, these updates can quickly overshoot, especially as momentum builds, so they can afford to have their updates shrunk a bit. Equally parameters with small or irregular gradients need larger relative steps to make progress. In Adam these are accounted for with two moment estimates rather than just the one.
First moment . As before the exponential moving average of the gradient:
Second moment . This is the exponential moving average of the squared gradient:
Typical . This tracks the variance (or more precisely, the uncentered second moment) of recent gradients per parameter. Like the first moment, this is also an additional value stored for each trainable parameter. A parameter with large, consistent gradients will have a large , and a parameter with small or sparse gradients will have a small .
The update step normalises the momentum term by the square root of this estimate:
The denominator is approximately the root-mean-square of recent gradients for that parameter. Dividing by it has the effect of making the step size inversely proportional to the magnitude of the gradient: large recent gradients give a large and a smaller effective step, while small recent gradients give a small and a larger one.
The trajectory diagram below shows SGD and Adam on a simple quadratic landscape where one direction is much steeper than the other.
SGD has one learning rate for both directions, so it is capped by the steep one: any larger and the steep direction would diverge. That leaves it zig-zagging across the steep axis while it crawls along the shallow one toward the minimum. Adam divides each coordinate by its own , giving small steps in the steep direction and large steps in the shallow one, so it heads almost straight in.
There is a small problem with the way the moment estimates are defined here. Both and are initialised at zero, then at step 1:
With this means . The true first moment of the gradient is , but we are estimating it as , i.e. ten times too small. This is an initialisation bias, and in order to return the estimate to the unbiased form, we need to apply a correction. If we take the expected value of with the assumption that gradients come from a stationary distribution:
So underestimates the true first moment by a factor of . Dividing by this factor gets back the unbiased estimate:
The same logic applies to to leave us with:
This is an important correction for the first moment, but the second moment really shows the scale of the issue. At with : the correction factor is , meaning we divide by 0.001, scaling up by a factor of 1000. This is particularly important in making sure the early steps of training are accurately modelling the moments. Over time the correction starts to wash out, but it still takes many steps before approaches 1.
The corrected Adam update is:
We can write the step for the Adam optimiser out very concisely in code:
1def adam_step(params, grads, m, v, t, lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8):
2 m = beta1 * m + (1 - beta1) * grads
3 v = beta2 * v + (1 - beta2) * grads**2
4 m_hat = m / (1 - beta1**t)
5 v_hat = v / (1 - beta2**t)
6 return params - lr * m_hat / (np.sqrt(v_hat) + eps), m, v
No surprises here, the interesting part comes when we start trying to add regularisation.
Before asking whether weight decay works correctly in Adam, we need to be clear about what weight decay is. Generally speaking, wanting parameters that minimise the loss while staying small is a sensible heuristic. This means parameters span a smaller dynamic range, they are less likely to result in very large activations, and the network tends to generalise better.
The standard way to encode this is to add an L2 penalty to the objective:
The gradient of the penalty term is , so adding it to the global gradient and running SGD gives:
If we look at the factor , this means that every step the weights are multiplied by a number slightly less than 1. This is very literally the weights decaying: over time they shrink by a fixed fraction each step, regardless of gradient history and regardless of which parameter it is.
For SGD, L2 regularisation and weight decay are exactly the same thing. The equivalence breaks for adaptive optimisers.
Each parameter decays in the same way, by an amount that is proportional to the size of the parameter. Naturally, large outlying weights receive a stronger penalty than smaller central weights. The regularisation strength is just a global hyperparameter that needs tuning, but generally works well at around 0.01.
Understandably, the obvious thing to do with Adam is to apply the same trick, add the regularisation term to the loss, and allow the weight decay to be folded into the parameter update step. The problem is not initially obvious from looking at the equations. It only surfaces when you work out the effective weight decay for a parameter. With , the moment estimates become:
With the entire update as before:
What we have to do now is to isolate the contribution of to the update. In steady state, the exponential moving average of a constant input converges to that input: the term contributes approximately to . So the weight-decay contribution to the parameter update is:
Compare this to what we had previously for SGD (we'll derive it properly for AdamW in a moment):
The difference is the factor. In Adam+L2, the weight-decay contribution passes through the same adaptive scaling as the gradient, meaning the effective decay rate is:
This varies per parameter, and the same adaptive properties that helped with the gradients have the opposite effect for the weight decay: for a parameter with large recent gradients, is large, so the effective decay is small; for a parameter with small recent gradients, is small, and the effective decay is large.
This is backwards from what you want. The parameters that are most active in the network, and therefore probably most important to regularise strongly, receive the least regularisation. The parameters that are rarely updated end up getting squeezed hardest.
The argument also has a second-order component: also goes into via the squared term, inflating for all parameters. This further corrupts the adaptive scaling.
The fix (proposed by Loshchilov and Hutter) was to decouple weight decay from the gradient entirely. Instead of folding into the gradient before the moment update, it should be applied directly to the parameters after the update:
This returns the update step to the original Adam variant, and the weight-decay component falls out as:
This is constant and matches exactly what SGD + weight decay gives, and the adaptive part is back to only affecting the gradient step.
In code, the difference is minimal:
1def adamw_step(params, grads, m, v, t, lr=1e-3, beta1=0.9, beta2=0.999,
2 eps=1e-8, weight_decay=0.01):
3 m = beta1 * m + (1 - beta1) * grads
4 v = beta2 * v + (1 - beta2) * grads**2
5 m_hat = m / (1 - beta1**t)
6 v_hat = v / (1 - beta2**t)
7 return params - lr * m_hat / (np.sqrt(v_hat) + eps) - lr * weight_decay * params, m, v
The weight decay is applied to params
, not routed through the gradient. Practically speaking, this is a tiny change to the code, but without working through the derivation, bugs like these can hide in plain sight for years.
A few defaults that work well in practice: a learning rate of (the standard Adam starting point), tuned down to for fine-tuning; weight decay (the PyTorch AdamW default); and , , , which work for the vast majority of cases.
AdamW decouples and , so you can tune each independently without them interfering. In Adam+L2 they are entangled through , so adjusting the learning rate also shifts the effective regularisation. One of the practical benefits of AdamW is simply that the hyperparameters mean what you think they mean.