Convergence Rates for Gradient Descent on the Edge of Stability for Overparametrised Least Squares

Classical optimisation theory guarantees monotonic objective decrease for gradient descent (GD) when employed in a small step size, or stable, regime. In contrast, gradient descent on neural networks is frequently performed in a large step size regime called the edge of stability, in which the objective decreases non-monotonically with an observed implicit bias towards flat minima. In this paper, we take a step toward quantifying this phenomenon by providing convergence rates for gradient descent with large learning rates in an overparametrised least squares setting. The key insight behind our analysis is that, as a consequence of overparametrisation, the set of global minimisers forms a Riemannian manifold $M$, which enables the decomposition of the GD dynamics into components parallel and orthogonal to $M$.