Adam Optimizer › Adam Variants

Adam has inspired a family of related optimizers, each addressing a specific limitation. Understanding the variants helps you choose the right optimizer for your use case.

AMSGrad (Reddi, Kale, Kumar, 2018) was proposed to fix the theoretical non-convergence problem in Adam. Recall that Adam's second moment estimate is an exponential moving average, which can "forget" past large gradients. AMSGrad uses the maximum of all past second moment estimates instead:

$$\hat{v}_t^{\max} = \max(\hat{v}_{t-1}^{\max}, \hat{v}_t)$$ $$\theta_{t+1} = \theta_t - \frac{\alpha}{\sqrt{\hat{v}_t^{\max}} + \epsilon} \hat{m}_t$$

By using $\hat{v}_t^{\max}$ in the denominator instead of $\hat{v}_t$, AMSGrad ensures that the effective learning rate never increases. Once a large gradient has been seen, the corresponding dimension's effective learning rate stays compressed. This monotonically non-increasing property is what enables the convergence guarantee.

In practice, AMSGrad often converges more slowly than Adam because the non-increasing effective learning rate can make it too conservative in the latter part of training. Most practitioners find that Adam with a well-tuned schedule outperforms AMSGrad, which is why AMSGrad has seen limited adoption despite its theoretical appeal. You can enable it in PyTorch with torch.optim.Adam(..., amsgrad=True).

Nadam (Dozat, 2016) combines Adam with Nesterov momentum. Recall that Nesterov momentum modifies the standard momentum update by computing the gradient at the "lookahead" position rather than the current position. Nesterov momentum typically converges slightly faster than standard momentum because the gradient is evaluated at a more future-looking position, meaning the update step already accounts for where you will be after applying the current velocity.

Nadam incorporates this improvement into Adam by substituting the Nesterov first moment estimate into the update rule. The difference from Adam is subtle but can matter in settings where the gradient changes direction quickly. PyTorch provides Nadam as torch.optim.NAdam. For most applications, the improvement over Adam is marginal, but it is worth knowing about as the theoretically cleaner variant.

AdamW (Loshchilov and Hutter, 2017) is arguably the most important Adam variant and deserves a full chapter of its own (which follows this one). The key insight is that Adam's weight_decay parameter does not implement true weight decay: it adds $\lambda \theta$ to the gradient before the Adam update, which means the regularization term is divided by $\sqrt{\hat{v}_t}$ along with the gradient. This weakens the regularization for frequently updated parameters and strengthens it for rarely updated ones, which is the opposite of what you want.

AdamW implements decoupled weight decay: it applies the weight decay step directly to the parameters after the Adam update, independently of the gradient:

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

This simple decoupling produces substantially better generalization in language modeling and has become the standard optimizer for training transformers. All modern large language models use AdamW rather than vanilla Adam.

Adafactor (Shazeer and Stern, 2018) addresses the memory overhead problem. For a weight matrix of shape $[r, c]$, Adam stores two additional matrices of the same shape ($m_t$ and $v_t$), tripling the memory for that parameter. Adafactor factors the second moment matrix into two rank-one vectors of shape $[r]$ and $[c]$, reducing the optimizer state from $O(rc)$ to $O(r + c)$.

For large embedding tables or weight matrices, this can be a massive memory saving. A 50,000-word vocabulary with 1024-dimensional embeddings requires a matrix of shape $[50000, 1024]$. Adam stores two copies of this, totaling $2 \times 50000 \times 1024 = 102$ million floats for the optimizer state alone. Adafactor reduces this to $(50000 + 1024) \times 2 \approx 102,000$ floats. The tradeoff is that the factored approximation is less accurate than the true per-element second moment, which can slow convergence and require more careful hyperparameter tuning. Adafactor is used in training T5 and some other very large models where memory is the binding constraint.