Adam Optimizer › Common Failures and Fixes

Adam is robust but not failure-proof. These are the most common issues and their solutions.

Symptom: Loss is NaN or increases rapidly in the first few hundred steps.

Cause: Gradients are very large at initialization, causing step sizes that blow up the parameters. With poorly initialized weights, the first gradient can be orders of magnitude larger than what Adam stabilizes to after warm-up.

Fix: Use learning rate warm-up. Start with $\alpha$ at 1-10% of the target value and linearly increase it over the first $T_{\text{warmup}}$ steps (typically 1,000-10,000 steps). This gives Adam time to accumulate reliable moment estimates before taking large steps. Warm-up is essentially standard practice for transformer training.

Symptom: Training loss stops decreasing after a few thousand steps and stays flat.

Cause: Either the learning rate is too low, the moment estimates have accumulated high curvature noise, or the model is stuck in a flat region.

Fix: Check the learning rate first. Then consider learning rate decay: cosine annealing or step decay after a plateau. A learning rate restart (cyclic warm-up) can also help escape flat regions.

Symptom: Training with gradient accumulation over $K$ steps gives different results than training with batch size $K$ times larger.

Cause: When you accumulate gradients over $K$ steps and then take one Adam step, you are mixing per-step gradients differently than if you averaged them into one batch gradient. The moment estimates see $K$ small updates instead of one averaged update, leading to subtly different dynamics.

Fix: When using gradient accumulation for large effective batch sizes, scale $\alpha$ by $\sqrt{K}$ (not $K$) and keep the moment estimates consistent by only calling the optimizer step once per accumulation cycle.

Symptom: Adam achieves lower training loss than SGD but higher test error.

Cause: The adaptive learning rate may cause Adam to focus on dimensions with low variance, which can correspond to spurious correlations in the training data rather than generalizable features.

Fix: Switch to SGD with momentum for tasks where this generalization gap matters (primarily image classification with standard architectures). For NLP tasks, this problem is much less pronounced.