Adam at the Edge of Stability: Adaptive Feedback, Provable Oscillation, and Gradient Reversal
Organizations: School of Engineering and Applied Sciences, Harvard University.
Abstract
The edge-of-stability (EoS) phenomenon of full-batch Adam has been widely observed, yet its underlying dynamical mechanism remains poorly understood. In this paper, we identify Adam's second-moment adaptation as a negative-feedback mechanism that drives the dynamics toward the stability boundary. We characterize this mechanism through the active curvature, namely, the preconditioned curvature along the preconditioned gradient direction, and establish rigorous characterizations in progressively richer settings: rank-one quadratics with momentum, diagonal quadratics, on which the active curvature separates from the sharpness, and general objectives. Importantly, the mechanism predicts gradient reversal of full-batch Adam near the edge: consecutive gradients repeatedly point in nearly opposite directions, as we observe across fully connected networks, ResNets, ViTs, LSTMs, GPT-2 medium, and Adam-family optimizers. Consistent with this picture, averaging iterates suppresses these fast oscillations and produces smoother and lower loss curves. Together, these results provide an important first step towards fully understanding the dynamical behavior of Adam's EoS through active curvature and gradient reversal.
Figures & tables
Appendix figures & tables45 assets
Supplementary material from the paper’s appendix.
Appendix
| loss | train acc | ||||
|---|---|---|---|---|---|
| MSE | 0 | 3.35 | 2.018 | 0.973 | |
| MSE | 0.3 | 2.25 | 2.013 | 0.992 | |
| MSE | 0.5 | 2.17 | 2.011 | 0.996 | |
| MSE | 0.9 | 2.01 | 1.965 | 0.996 | |
| MSE | 0.9 | 2.00 | 1.885 | 1.000 | |
| MSE | 0.9 | 2.01 | 1.970 | 1.000 |
| architecture | params | loss | |||
|---|---|---|---|---|---|
| ResNet-BN | 1.22M | 0 | 2.89 | 2.023 | 0.00406 |
| ResNet-BN | 1.22M | 0.9 | 2.02 | 1.957 | 0.000164 |
| ViT | 413k | 0 | 2.39 | 2.021 | 0.000777 |
| ViT | 413k | 0.9 | 1.99 | 1.955 | 0.000377 |
| LSTM | 6.70M | 0 | 2.46 | 2.102 | 0.139 |
| LSTM | 6.70M | 0.9 | 1.95 | 1.466 | 1.11 |
| variant | threshold | locked | |||
|---|---|---|---|---|---|
| Adam (uncorrected) | 2.01 | 1.95 | 100% | ||
| Adam, bias-corrected | 2.02 | 1.94 | 99% | ||
| AdamW | 2.00 | 1.92 | 94% | ||
| Adam + | 1.87 | 1.76 | 43% | ||
| Padam ( ) | 2.01 | 1.97 | 100% | ||
| Nadam | 2.09 | 2.00 | 100% |
| full-set loss | minibatch loss | |||||
|---|---|---|---|---|---|---|
| model | ||||||
| GPT-2 | 0 | full | 2.14 | 2.76 | 2.14 | 2.76 |
| GPT-2 | 0 | 1/2 | 2.11 | 15.09 | 2.59 | 19.18 |
| GPT-2 | 0 | 1/4 | 3.58 | 21.12 | 3.36 | 10.80 |
| GPT-2 | 0 | 1/8 | 4.22 | 29.69 | 6.36 | 58.58 |
| GPT-2 | 0.9 | full | 1.82 | 1.91 | 1.82 | 1.91 |
| setting | reversal steps | loss | ||||||
| seed 2, | 0 | 0.999 | 0.001 | -0.993 | +0.990 | -0.993 | 100% | 0.156 |
| 0.3 | 0.999 | 0.001 | -0.992 | +0.982 | -0.985 | 100% | 0.102 | |
| 0.7 | 0.999 | 0.001 | -0.982 | +0.938 | -0.941 | 100% | 0.0742 | |
| seed 1 | 0.9 | 0.999 | 0.001 | -0.982 | +0.943 | -0.402 | 99% | 0.0611 |
| 0.95 | 0.999 | 0.001 | -0.971 | +0.908 | +0.480 | 98% | 0.0608 | |
| , | 0 | 0.95 | 0.001 | -0.963 | +0.904 | -0.963 | 100% | 0.062 |
| architecture | reversal steps | ||||
|---|---|---|---|---|---|
| ResNet-BN | 0 | -0.998 | +0.994 | -0.998 | 100% |
| ResNet-BN | 0.9 | -0.966 | +0.872 | -0.944 | 100% |
| ViT | 0 | -0.996 | +0.990 | -0.996 | 100% |
| ViT | 0.9 | -0.975 | +0.905 | -0.966 | 100% |
| LSTM | 0 | -0.986 | +0.991 | -0.986 | 79% |
| LSTM | 0.9 | -0.498 | -0.286 | -0.027 | 50% |
| steps | reversal | ||||
|---|---|---|---|---|---|
| – | |||||
| – | |||||
| – | |||||
| – | |||||
| – | |||||
| – |
| or | |||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|
| variant | reversal | reversal | |||||||||
| Adam (uncorrected) | -0.987 | +0.947 | -0.909 | 99% | 1.97 | 2.01 | -0.998 | +0.995 | 100% | 2.05 | 3.70 |
| Adam, bias-corrected | -0.959 | +0.832 | -0.905 | 99% | 1.94 | 2.02 | -0.983 | +0.982 | 100% | 2.28 | 2.84 |
| AdamW | -0.972 | +0.872 | -0.753 | 94% | 1.92 | 2.00 | -0.996 | +0.994 | 100% | 2.09 | 3.69 |
| Adam + | +0.677 | +0.942 | -0.128 | 1% | 1.76 | 1.87 | -0.119 | +0.844 | 3% | 1.94 | 2.90 |
| Padam ( ) | -0.982 | +0.930 | -0.957 | 100% | 1.97 | 2.01 | -1.000 | +0.999 | 100% | 2.00 | 2.06 |
| full-set gradient | minibatch gradient | ||||||
|---|---|---|---|---|---|---|---|
| model | reversal | reversal | |||||
| GPT-2 | 0 | full | -0.98 | 93% | -0.98 | 93% | +1.00 |
| GPT-2 | 0 | 1/2 | -0.94 | 90% | -0.95 | 89% | +0.95 |
| GPT-2 | 0 | 1/4 | -0.79 | 75% | -0.75 | 70% | +0.90 |
| GPT-2 | 0 | 1/8 | -0.44 | 43% | -0.36 | 30% | +0.62 |
| GPT-2 | 0.9 | full | -0.93 | 65% | -0.93 | 65% | +1.00 |