Adaptive Gradient Methods at the Edge of Stability
This work provides foundational insights into the behavior of adaptive gradient methods, which is crucial for the deep learning community to better understand and optimize training algorithms.
The paper investigates the training dynamics of adaptive gradient methods like Adam in deep learning, showing that during full-batch training, the maximum eigenvalue of the preconditioned Hessian equilibrates at a stability threshold (e.g., 38/η for Adam with β₁=0.9), and that these methods can advance into high-curvature regions while adapting the preconditioner, unlike non-adaptive methods.
Very little is known about the training dynamics of adaptive gradient methods like Adam in deep learning. In this paper, we shed light on the behavior of these algorithms in the full-batch and sufficiently large batch settings. Specifically, we empirically demonstrate that during full-batch training, the maximum eigenvalue of the preconditioned Hessian typically equilibrates at a certain numerical value -- the stability threshold of a gradient descent algorithm. For Adam with step size $η$ and $β_1 = 0.9$, this stability threshold is $38/η$. Similar effects occur during minibatch training, especially as the batch size grows. Yet, even though adaptive methods train at the ``Adaptive Edge of Stability'' (AEoS), their behavior in this regime differs in a significant way from that of non-adaptive methods at the EoS. Whereas non-adaptive algorithms at the EoS are blocked from entering high-curvature regions of the loss landscape, adaptive gradient methods at the AEoS can keep advancing into high-curvature regions, while adapting the preconditioner to compensate. Our findings can serve as a foundation for the community's future understanding of adaptive gradient methods in deep learning.