Physics inspired optimizers improve training beyond stability edge
Physics-Inspired Optimizers beyond the Edge of Stability: Period-Two Orbits and Delayed Instability
Machine Learning
Summary
Training deep learning models often involves pushing the training process to a limit called the edge of stability, where changes happen in a delicate balance. The authors study special momentum-based optimization methods inspired by physics, which adjust how the training steps behave beyond this edge. They discover patterns like repeating cycles and delayed instability that explain how these optimizers keep learning stable even with large learning rates. Their experiments with neural networks show that these new insights help understand and potentially improve training dynamics.
What this means in practice
- •For deep learning engineers: Use physics-inspired optimizers to maintain stable training at large learning rates beyond conventional stability limits.
- •For machine learning framework developers: Incorporate kinetic map-based optimizers to better handle delayed instability and oscillations in neural network training.
Authors
Alexander Itin
Abstract
Gradient descent with momentum can be viewed as a particle moving with friction in the loss landscape. In deep learning the learning rate is usually large, so training is a discrete dynamical system: in many full-batch experiments the sharpness grows to the optimizer's stability limit, the edge of stability (EoS), and training continues there. Physics-inspired optimizers change this motion; Velocity-Regularized Adam (VRAdam), for example, shrinks the step at high speed. We study kinetic maps, in which the step is the gradient of a kinetic energy of the momentum, and ask what they do beyond the EoS. (i) In a reduced model, a kinetic map creates a symmetric period-two orbit beyond the edge, strictly linearly stable if and only if the map's elasticity (the log-log slope of step length versus momentum) lies between 0 and 1; on it, training sits at an effective edge of stability. VRAdam's damper fails above twice the critical learning rate, where its kinetic energy is not convex; the Legendre dual of VRAdam's own Lagrangian is convex and passes at every learning rate. (ii) Because the sharpness rises gradually, the loss of stability is delayed, as in Neishtadt's theory of slowly swept bifurcations. In a network with a controlled sweep, the dynamics match this theory without fitted parameters, including the memory of where the sweep started and its erasure by noise. (iii) With momentum, reaching the oscillation also requires elasticity at large momentum: maps that bound the step have competing period-four and chaotic attractors already at beta1 = 0.9, while the cubic-Legendre map reached the oscillation from every sampled start up to beta1 = 0.99. Experiments on an exact quadratic problem, a full-batch MLP and ResNet-32 support the theory; on ResNet-18, much of the apparent benefit of the maps is reproduced by matching the effective learning rate.