AdamW › Why Adam Breaks the Equivalence

Adam is an adaptive optimizer: it scales each parameter's gradient by a factor that depends on the history of squared gradients for that parameter. This is Adam's greatest strength, allowing it to effectively handle parameters with very different gradient scales. But this same adaptive scaling is precisely what breaks the equivalence between L2 regularization and weight decay.

Recall from the previous chapter that Adam maintains two moment estimates and applies them to scale the gradient:

$$\theta_{t+1} = \theta_t - \frac{\alpha}{\sqrt{\hat{v}_t} + \epsilon} \hat{m}_t$$

where $\hat{m}_t$ is the bias-corrected first moment (a smoothed version of the current gradient) and $\hat{v}_t$ is the bias-corrected second moment (a smoothed version of the squared gradient, capturing gradient magnitude history).

The key property is that $\hat{v}_t$ differs across parameters. A parameter that has historically received large gradients has a large $\hat{v}_t$, which shrinks its effective learning rate. A parameter that has received small gradients has a small $\hat{v}_t$, which gives it a larger effective learning rate. This per-parameter adaptive scaling is what lets Adam handle heterogeneous gradient distributions.

Now, what happens when you apply L2 regularization with Adam? The combined gradient passed to Adam includes the regularization term:

$$g_t^{\text{reg}} = g_t + \lambda \theta_t$$

where $g_t = \nabla_{\theta} L(\theta_t)$ is the original gradient. Adam then uses this regularized gradient to update both moment estimates and the parameters:

$$m_t = \beta_1 m_{t-1} + (1 - \beta_1)(g_t + \lambda \theta_t)$$ $$v_t = \beta_2 v_{t-1} + (1 - \beta_2)(g_t + \lambda \theta_t)^2$$

The adaptive scaling term $\hat{v}_t$ now accumulates the regularization signal $\lambda \theta_t$ along with the loss gradient. The effective learning rate applied to the regularization penalty is:

$$\frac{\alpha}{\sqrt{\hat{v}_t} + \epsilon}$$

This effective rate differs from parameter to parameter depending on the history of squared gradients. Parameters that have received large gradients from the loss will have large $\hat{v}_t$ values, which reduces their effective learning rate, including the rate at which regularization shrinks them. Parameters with small historical gradients have small $\hat{v}_t$ values, giving them a larger effective learning rate, and consequently stronger effective regularization.

The result is that L2 regularization does not apply uniformly across parameters when used with Adam. The effective regularization strength for parameter $i$ is roughly:

$$\lambda_{\text{eff},i} \approx \frac{\lambda}{\sqrt{\hat{v}_{t,i}} + \epsilon}$$

This effective strength is high for rarely-updated parameters (small $\hat{v}_{t,i}$) and low for frequently-updated parameters (large $\hat{v}_{t,i}$). In a language model, word embeddings for common tokens receive dense, large gradient updates, giving them large $\hat{v}$ values and therefore weak effective regularization. Embeddings for rare tokens receive sparse, small updates, giving them small $\hat{v}$ values and therefore strong effective regularization.

This is the opposite of what good regularization should do. High-frequency embeddings, which are updated on almost every training step and are most at risk of overfitting to the training data distribution, receive the weakest regularization. Low-frequency embeddings, which are rarely updated and represent underspecified representations, receive the strongest regularization. The regularization strength has become an artifact of gradient history rather than a deliberate design choice.

This is the observation Loshchilov and Hutter formalized in their 2017 paper "Decoupled Weight Decay Regularization." Their analysis had two parts: a theoretical derivation showing why L2 regularization and weight decay are not equivalent for adaptive optimizers, and empirical experiments on image classification and language modeling tasks showing that the decoupled version outperforms the coupled version. They found that AdamW outperforms Adam+L2 on a variety of benchmarks when both are carefully tuned, and that the improvement is more pronounced as regularization becomes more important, specifically in regimes with smaller datasets relative to model size, or with more aggressive regularization strengths.

An important nuance in their analysis: the optimal hyperparameters for Adam+L2 and AdamW are not interchangeable. Because Adam+L2 applies non-uniform effective regularization, a weight_decay value that seems to work in Adam is not numerically equivalent to the same value in AdamW. The Loshchilov-Hutter paper showed that the optimal weight decay for AdamW tends to be larger in absolute terms, because the effective regularization per parameter is more predictable and you can apply stronger regularization without the unintended side effect of over-regularizing rarely-updated parameters.