Adam and AdamW: adaptive optimisation and weight decay
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.
---
SGD: the starting point
Before getting into Adam, we should take a moment to think about what an optimiser does. Vanilla Gradient Descent moves each parameter θ\\thetaθ in the direction of steepest descent at rate η\\etaη:
θt+1=θt−η gt\\theta\_{t+1} = \\theta\_t - \\eta \\, g\_tθt+1\=θt−ηgt
Where the true gradient is defined over the whole dataset as gt=∇θL(θt)g\t = \\nabla\\\theta \\mathcal{L}(\\theta\_t)gt\=∇θL(θt). 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.
practice
EasyQ129
Implement SGD with momentum.
Momentum partially addresses this, where instead of using the raw gradient to determine the steps, we maintain an exponential moving average of past gradients:
mt=βmt−1+(1−β)gtm\t = \\beta m\{t-1} + (1-\\beta) g\_tmt\=βmt−1+(1−β)gt
θt+1=θt−η mt\\theta\_{t+1} = \\theta\_t - \\eta \\, m\_tθt+1\=θt−ηmt
The momentum term mtm\_tmt 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.
◎ thinkWhy does momentum cancel in oscillating directions?
▸ show answer
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.
Adam: adaptive learning rates
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 mtm\_tmt. As before the exponential moving average of the gradient:
mt=β1mt−1+(1−β1)gtm\t = \\beta\_1 m\{t-1} + (1 - \\beta\_1) g\_tmt\=β1mt−1+(1−β1)gt
Second moment vtv\_tvt. This is the exponential moving average of the squared gradient:
vt=β2vt−1+(1−β2)gt2v\t = \\beta\_2 v\{t-1} + (1 - \\beta\_2) g\_t^2vt\=β2vt−1+(1−β2)gt2
Typical β1=0.9,β2=0.999\\beta\_1 = 0.9, \\beta\_2 = 0.999β1\=0.9,β2\=0.999. 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 vtv\_tvt, and a parameter with small or sparse gradients will have a small vtv\_tvt.
The update step normalises the momentum term by the square root of this estimate:
θt+1=θt−η⋅mtvt+ε\\theta\_{t+1} = \\theta\_t - \\eta \\cdot \\frac{m\_t}{\\sqrt{v\_t} + \\varepsilon}θt+1\=θt−η⋅vt+εmt
The denominator vt\\sqrt{v\_t}vt 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 vt\\sqrt{v\_t}vt and a smaller effective step, while small recent gradients give a small vt\\sqrt{v\_t}vt 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.
Optimisation trajectories on an anisotropic landscape
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 vt\\sqrt{v\_t}vt, giving small steps in the steep direction and large steps in the shallow one, so it heads almost straight in.
Bias correction
There is a small problem with the way the moment estimates are defined here. Both mtm\_tmt and vtv\_tvt are initialised at zero, then at step 1:
m1=(1−β1)g1m\_1 = (1 - \\beta\_1) g\_1m1\=(1−β1)g1
With β1=0.9\\beta\_1 = 0.9β1\=0.9 this means m1=0.1⋅g1m\_1 = 0.1 \\cdot g\_1m1\=0.1⋅g1. The true first moment of the gradient is g1g\_1g1, but we are estimating it as 0.1⋅g10.1 \\cdot g\_10.1⋅g1, 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 mtm\_tmt with the assumption that gradients come from a stationary distribution:
E\[mt\]=(1−β1t) E\[g\]\\mathbb{E}\[m\_t\] = (1 - \\beta\_1^t) \\, \\mathbb{E}\[g\]E\[mt\]\=(1−β1t)E\[g\]
∂ deriveWhere does the (1−β1t)(1 - \\beta\_1^t)(1−β1t) factor comes from?
▸ show derivation
So mtm\_tmt underestimates the true first moment by a factor of (1−β1t)(1 - \\beta\_1^t)(1−β1t). Dividing by this factor gets back the unbiased estimate:
m^t=mt1−β1t\\hat{m}\_t = \\frac{m\_t}{1 - \\beta\_1^t}m^t\=1−β1tmt
practice
MediumQ118
Implement a full Adam optimizer step.
The same logic applies to vtv\_tvt to leave us with:
v^t=vt1−β2t\\hat{v}\_t = \\frac{v\_t}{1 - \\beta\_2^t}v^t\=1−β2tvt
This is an important correction for the first moment, but the second moment really shows the scale of the issue. At t=1t=1t\=1 with β2=0.999\\beta\_2 = 0.999β2\=0.999: the correction factor is 1−0.999=0.0011 - 0.999 = 0.0011−0.999\=0.001, 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 (1−β2t)(1 - \\beta\_2^t)(1−β2t) approaches 1.
First and second moment bias correction over time
The corrected Adam update is:
θt+1=θt−η⋅m^tv^t+ε\\theta\_{t+1} = \\theta\_t - \\eta \\cdot \\frac{\\hat{m}\_t}{\\sqrt{\\hat{v}\_t} + \\varepsilon}θt+1\=θt−η⋅v^t+εm^t
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, vNo surprises here, the interesting part comes when we start trying to add regularisation.
Weight decay: what is it supposed to do?
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:
Lreg(θ)=L(θ)+λ2∥θ∥2\\mathcal{L}\_{\\text{reg}}(\\theta) = \\mathcal{L}(\\theta) + \\frac{\\lambda}{2} \\|\\theta\\|^2Lreg(θ)\=L(θ)+2λ∥θ∥2
The gradient of the penalty term is λθ\\lambda \\thetaλθ, so adding it to the global gradient and running SGD gives:
θt+1=θt−η(gt+λθt)=(1−ηλ)θt−ηgt\\theta\_{t+1} = \\theta\_t - \\eta(g\_t + \\lambda \\theta\_t) = (1 - \\eta\\lambda) \\theta\_t - \\eta g\_tθt+1\=θt−η(gt+λθt)\=(1−ηλ)θt−ηgt
If we look at the factor (1−ηλ)(1 - \\eta\\lambda)(1−ηλ), 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 ηλ\\eta\\lambdaηλ each step, regardless of gradient history and regardless of which parameter it is.
note
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 λ\\lambdaλ is just a global hyperparameter that needs tuning, but generally works well at around 0.01.
Naive Adam + L2
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 g\~t=gt+λθt−1\\tilde{g}\t = g\_t + \\lambda \\theta\{t-1}g\~t\=gt+λθt−1, the moment estimates become:
mt=β1mt−1+(1−β1)(gt+λθt−1)m\t = \\beta\_1 m\{t-1} + (1 - \\beta\1)(g\_t + \\lambda \\theta\{t-1})mt\=β1mt−1+(1−β1)(gt+λθt−1)
vt=β2vt−1+(1−β2)(gt+λθt−1)2v\t = \\beta\_2 v\{t-1} + (1 - \\beta\2)(g\_t + \\lambda \\theta\{t-1})^2vt\=β2vt−1+(1−β2)(gt+λθt−1)2
With the entire update as before:
θt=θt−1−η⋅m^tv^t+ε\\theta\t = \\theta\{t-1} - \\eta \\cdot \\frac{\\hat{m}\_t}{\\sqrt{\\hat{v}\_t} + \\varepsilon}θt\=θt−1−η⋅v^t+εm^t
What we have to do now is to isolate the contribution of λθ\\lambda \\thetaλθ to the update. In steady state, the exponential moving average of a constant input converges to that input: the λθt−1\\lambda \\theta\{t-1}λθt−1 term contributes approximately λθt−1\\lambda \\theta\{t-1}λθt−1 to m^t\\hat{m}\_tm^t. So the weight-decay contribution to the parameter update is:
Δθtdecay≈−η⋅λθt−1v^t+ε\\Delta \\theta\t^{\\text{decay}} \\approx -\\eta \\cdot \\frac{\\lambda \\theta\{t-1}}{\\sqrt{\\hat{v}\_t} + \\varepsilon}Δθtdecay≈−η⋅v^t+ελθt−1
Compare this to what we had previously for SGD (we'll derive it properly for AdamW in a moment):
Δθtdecay, SGD=−ηλθt−1\\Delta \\theta\t^{\\text{decay, SGD}} = -\\eta \\lambda \\theta\{t-1}Δθtdecay, SGD\=−ηλθt−1
The difference is the 1v^t+ε\\frac{1}{\\sqrt{\\hat{v}\_t} + \\varepsilon}v^t+ε1 factor. In Adam+L2, the weight-decay contribution passes through the same adaptive scaling as the gradient, meaning the effective decay rate is:
effective decay (Adam+L2)≈ηλv^t+ε\\text{effective decay (Adam+L2)} \\approx \\frac{\\eta \\lambda}{\\sqrt{\\hat{v}\_t} + \\varepsilon}effective decay (Adam+L2)≈v^t+εηλ
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, v^t\\sqrt{\\hat{v}\_t}v^t is large, so the effective decay is small; for a parameter with small recent gradients, v^t\\sqrt{\\hat{v}\_t}v^t is small, and the effective decay is large.
Effective weight decay per step: AdamW vs Adam+L2
practice
MediumQ325
Implement the coupled Adam+L2 step and see exactly where λθ enters the computation.
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: g\~t=gt+λθt\\tilde{g}\_t = g\_t + \\lambda \\theta\_tg\~t\=gt+λθt also goes into vtv\_tvt via the squared term, inflating v^t\\hat{v}\_tv^t for all parameters. This further corrupts the adaptive scaling.
How AdamW fixes it
practice
MediumQ36
Implement a full AdamW update step.
The fix (proposed by Loshchilov and Hutter) was to decouple weight decay from the gradient entirely. Instead of folding λθ\\lambda \\thetaλθ into the gradient before the moment update, it should be applied directly to the parameters after the update:
θt+1=(1−ηλ)θt−η⋅m^tv^t+ε\\theta\_{t+1} = (1 - \\eta\\lambda) \\theta\_t - \\eta \\cdot \\frac{\\hat{m}\_t}{\\sqrt{\\hat{v}\_t} + \\varepsilon}θt+1\=(1−ηλ)θt−η⋅v^t+εm^t
This returns the update step to the original Adam variant, and the weight-decay component falls out as:
Δθtdecay=−ηλθt\\Delta \\theta\_t^{\\text{decay}} = -\\eta\\lambda \\theta\_tΔθtdecay\=−ηλθt
This is constant and matches exactly what SGD + weight decay gives, and the adaptive part is back to only affecting the gradient step.
Where λθ enters the computation for each optimiser
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, vThe 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.
· · ·
What happens in practice
A few defaults that work well in practice: a learning rate of η=10−3\\eta = 10^{-3}η\=10−3 (the standard Adam starting point), tuned down to 10−410^{-4}10−4 for fine-tuning; weight decay λ=0.01\\lambda = 0.01λ\=0.01 (the PyTorch AdamW default); and β1=0.9\\beta\_1 = 0.9β1\=0.9, β2=0.999\\beta\_2 = 0.999β2\=0.999, ε=10−8\\varepsilon = 10^{-8}ε\=10−8, which work for the vast majority of cases.
AdamW decouples η\\etaη and λ\\lambdaλ, so you can tune each independently without them interfering. In Adam+L2 they are entangled through v^t\\hat{v}\_tv^t, 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.
problems in this essay · 4
Q129SGD with MomentumEasyQ118Adam Optimizer StepMediumQ325Adam + L2 Step (Coupled)MediumQ36AdamW StepMedium
feedback