Flatland: The Adventures of Gradient Descent with Large Step Sizes

The training of neural networks often entails objective functions that are not globally $L$-smooth. For these functions, it is both theoretically and practically difficult to reply to the question: what is the largest possible step size that ensures the convergence of gradient descent (GD)? We address this longstanding open question in deep learning by providing a unifying definition of "large'' step sizes that requires only local Lipschitz (or even Hölder) continuity of the gradient. We design first-order adaptive methods that provably yield large step sizes and show that they operate at the edge of stability (EoS) right from the start of the training. In particular, the loss decreases nonmonotonically and the product between the step size and sharpness, i.e., the largest eigenvalue of the hessian, stays above the EoS threshold of 2 throughout training. Using our method, we are also able to minimize the sharpness all the way down to its global minimum. Contrary to expectation, we find that encountering globally-flat regions too early in the training may both slow down convergence and jeopardize the generalization ability of the network. Exploiting a self-stabilization argument, we allow GD to enter slightly sharper valleys and turn unsuccessful training runs into very successful ones.

Training a neural network depends heavily on choosing the correct learning rate. Our work studies how to safely use very large learning rates for gradient descent. We give a new practical meaning to "large" learning rates that does not rely on unrealistic assumptions about the whole loss landscape. We also design adaptive first-order methods that automatically find such large learning rates during training. These methods make training operate at the "edge of stability," where the loss no longer decreases smoothly but instead oscillates while still making progress. Surprisingly, this behavior can push the network into extremely flat regions of the loss landscape. While flatness is often thought to be good for generalization, our work shows that becoming too flat too early can actually hurt training, as the network may get stuck before it learns useful features. Finally, we show that slightly limiting these large learning rates can help avoid these overly flat saddle regions. This lets gradient descent move through somewhat sharper but more useful valleys, improving training success and test accuracy. Overall, the work gives both theory and a practical method for choosing very large learning rates safely.