On the Stability of Gradient Descent for Large Learning Rate
Alexandru Crăciun Debarghya Ghoshdastidar
There currently is a significant interest in understanding the Edge of
Stability (EoS) phenomenon, which has been observed in neural networks
training, characterized by a non-monotonic decrease of the loss function over
epochs, while the sharpness of the loss (spectral norm of the Hessian)
progressively approaches and stabilizes around 2/(learning rate). Reasons for
the existence of EoS when training using gradient descent have recently been
proposed -- a lack of flat minima near the gradient descent trajectory together
with the presence of compact forward-invariant sets. In this paper, we show
that linear neural networks optimized under a quadratic loss function satisfy
the first assumption and also a necessary condition for the second assumption.
More precisely, we prove that the gradient descent map is non-singular, the set
of global minimizers of the loss function forms a smooth manifold, and the
stable minima form a bounded subset in parameter space. Additionally, we prove
that if the step-size is too big, then the set of initializations from which
gradient descent converges to a critical point has measure zero.