Adam Optimizer › Limitations and Practical Considerations

Adam is a powerful and broadly applicable optimizer, but it has several limitations that practitioners encounter regularly.

Adam stores two additional tensors per parameter: $m_t$ and $v_t$. For large models, this triples the memory required for optimizer state compared to SGD (which stores one tensor per parameter for momentum, or zero for vanilla SGD). A 7 billion parameter model requires roughly 7B floats for the parameters themselves. Adam requires another 14B floats (two moments), for a total of 21B floats of optimizer state. At float32, that is 84 GB, which is prohibitive on a single GPU.

Remedies include:

• 8-bit Adam: The bitsandbytes library implements Adam with 8-bit quantized optimizer states, reducing memory by roughly 4x with minimal accuracy impact.
• ZeRO: Distributed optimizer state sharding across GPUs, so each GPU holds only a fraction of the moments.
• Adafactor: An adaptive optimizer with $O(1)$ optimizer state that factors the second moment matrix, used for very large models.

As noted in the comparison section, Adam can find sharper minima than SGD on some image classification benchmarks, leading to worse test accuracy despite better training accuracy. This problem is less prevalent in NLP, where the datasets are large enough that sharpness does not correlate strongly with generalization. If you observe a training-test gap that seems larger than expected, two practical remedies are AdamW (which decouples weight decay and tends to find broader minima) and data augmentation (which artificially widens the training distribution, making sharp minima less likely to perfectly fit the training set).

The theoretical convergence guarantee for Adam relies on assumptions (convexity, or specific non-convex conditions) that are not always met. AMSGrad modifies the second moment estimate to use the maximum of all past estimates, guaranteeing convergence in the online convex setting. However, empirical comparisons typically show Adam converging as well as or better than AMSGrad on practical tasks, so AMSGrad has seen limited adoption. The theoretical non-convergence of Adam in worst-case settings does not appear to manifest in the loss landscapes of practical neural networks, which have much more benign structure than the adversarial examples used in the theoretical analysis.

A subtle but important issue is that when you add weight_decay to PyTorch's torch.optim.Adam, you are not getting true weight decay. You are getting L2 regularization, which adds $\lambda \theta$ to the gradient before the update. With Adam's adaptive scaling, this L2 penalty is also divided by $\sqrt{\hat{v}_t}$, which means the actual regularization strength varies by parameter and is weaker for frequently updated parameters (which often need regularization most). This interaction defeats much of the purpose of weight decay as a regularizer. This is why AdamW, introduced by Loshchilov and Hutter in 2017, applies weight decay directly to the parameters after the Adam step rather than through the gradient. AdamW is the subject of the next chapter, and it is the recommended optimizer for essentially all modern neural language modeling.

A subtle practical issue is that Adam interacts with learning rate schedules in a non-obvious way. When the learning rate schedule reduces $\alpha$ by a large factor late in training (as cosine annealing does), the moment estimates $m_t$ and $v_t$ were built up under the higher learning rate regime. The second moment $v_t$ reflects the gradient magnitudes that caused the Adam updates when $\alpha$ was large. As $\alpha$ decreases, the adaptive learning rates per parameter shrink proportionally, but the relationship between $v_t$ and the current training dynamics may no longer be well-calibrated. Some practitioners reset the Adam optimizer state when making large learning rate jumps (such as at the beginning of a second training stage) to avoid carrying over stale moment estimates. This is sometimes called "warm restarting" the optimizer, and it can produce a brief spike in loss followed by improved convergence in the new learning rate regime.