AdamW

AdamW is Adam with weight decay applied directly to the weights, outside Adam's per-parameter scaling (Loshchilov & Hutter, 2019). Adam (Adaptive Moment Estimation; Kingma & Ba, 2015) keeps running averages of the gradient and of its square, called the first and second moments, and sets each parameter's update by dividing the bias-corrected first moment by the sum of the square root of the bias-corrected second moment and a small constant. We write the update for a parameter, then compare direct weight shrinkage with adding a penalty to the loss. You will need gradients and exponentially weighted averages.

Kingma & Ba designed Adam for problems with noisy or sparse gradients. Adam and AdamW are now the usual choice for training transformers.

The Adam Algorithm

Adam keeps moving averages of both the gradient and its squared values:

\begin{aligned} m_t &= \beta_1 m_{t-1} + (1-\beta_1) g_t \\ v_t &= \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \end{aligned}

where $g_t = \nabla_w L_{B_t}(w_t)$ is the mini-batch gradient, the square is taken entry by entry, $t = 1, 2, \ldots$ counts updates, and $m_0 = v_0 = 0$ . Because both averages start at zero, they are biased toward zero early in training, so Adam applies a bias correction:

\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1 - \beta_2^t}

The update rule is:

w_{t+1} = w_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}

Each parameter's step is divided by $\sqrt{\hat{v}_t} + \epsilon$ , where $\hat{v}_t$ estimates that parameter's mean squared gradient. Kingma & Ba's default settings are $\beta_1 = 0.9$ , $\beta_2 = 0.999$ , $\epsilon = 10^{-8}$ . Large language models are often trained with $\beta_2 = 0.95$ instead; LLaMA is one example (Touvron et al., 2023).

AdamW: Decoupled Weight Decay

Weight decay shrinks every weight a little at each step. For plain SGD this is the same as adding an $L_2$ penalty to the loss, with the coefficient rescaled by the learning rate. For Adam it is not. An $L_2$ penalty adds $\lambda w_t$ to the gradient, and Adam then divides that term by $\sqrt{\hat{v}_t}$ like the rest, so weights with large gradients are regularized less than intended (Loshchilov & Hutter, 2019). AdamW applies the decay to the weights directly, outside the adaptive scaling:

w_{t+1} = w_t - \eta \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda w_t \right)

where $\lambda$ is the weight-decay coefficient. This is the form PyTorch's AdamW uses. In Loshchilov & Hutter's notation the decay term is multiplied by a schedule multiplier $\eta_t$ rather than by the learning rate.

Loshchilov & Hutter showed that decoupling makes the best weight-decay value largely independent of the learning rate, and that it substantially improved Adam's test error on image classification, where Adam had usually trailed SGD with momentum. AdamW is widely used to train large language models; LLaMA, for example, used AdamW with weight decay 0.1 (Touvron et al., 2023).

Pros and Cons

Pros

Fast early progress on the loss with default settings
Handles sparse gradients and parameters whose gradients differ widely in scale
A common default in practice, especially for transformers

Cons

On some tasks, well-tuned SGD with momentum generalizes better (Wilson et al., 2017)
Scales each entry on its own, ignoring the structure of weight matrices
Stores two extra values per parameter ( $m$ and $v$ ), against one for SGD with momentum
Transformers trained with it usually need a learning-rate warmup (see below)

When to Use AdamW

AdamW is a strong default when:

You want reasonable results without extensive tuning
You are training transformers, including language models
Gradients are noisy or sparse

Training a transformer with Adam or AdamW usually starts with a short learning-rate warmup; LLaMA used 2,000 warmup steps (Touvron et al., 2023). Liu et al. (2020) trace the need for warmup to the high variance of Adam's per-parameter step size in the first updates, when $\hat{v}_t$ rests on only a few gradients. Xiong et al. (2020) show that the original post-norm Transformer layout makes warmup especially important.

One limitation, that each entry of a weight matrix is scaled on its own, motivates matrix-aware optimizers such as Muon, which treat the update of a whole weight matrix as one object.

Choosing Between SGD and Adam

There's no universally best optimizer:

Aspect	SGD + Momentum	AdamW
Tuning effort	Usually higher; sensitive to the learning rate	Usually lower; default settings often work
Generalization	Better on some tasks when well tuned (Wilson et al., 2017)	The usual choice for transformers
Early progress	Slower	Faster
Optimizer memory	One extra value per parameter (velocity)	Two extra values per parameter ( $m$ and $v$ )

Sources

Kingma, D. P., & Ba, J. (2015). Adam: A Method for Stochastic Optimization. ICLR. arxiv.org/abs/1412.6980
Loshchilov, I., & Hutter, F. (2019). Decoupled Weight Decay Regularization. ICLR. arxiv.org/abs/1711.05101
Liu, L., et al. (2020). On the Variance of the Adaptive Learning Rate and Beyond. ICLR. arxiv.org/abs/1908.03265
Xiong, R., et al. (2020). On Layer Normalization in the Transformer Architecture. ICML. arxiv.org/abs/2002.04745
Wilson, A. C., et al. (2017). The Marginal Value of Adaptive Gradient Methods in Machine Learning. NeurIPS. arxiv.org/abs/1705.08292
Touvron, H., et al. (2023). LLaMA: Open and Efficient Foundation Language Models. arXiv. arxiv.org/abs/2302.13971
Goodfellow, I., Bengio, Y., & Courville, A. (2016). Deep Learning, chapter 8: Optimization for Training Deep Models. MIT Press. deeplearningbook.org