Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Lachlan Ewen MacDonald, Hancheng Min, Leandro Palma, Salma Tarmoun, Ziqing Xu, René Vidal

TL;DR

This work provides a rigorous framework for understanding gradient descent with large learning rates in overparametrised least-squares problems by exploiting a solution manifold $M=f^{-1}igl\\{y\bigr\ brace$. Using tubular coordinates, the authors derive a normal form that separates GD into a Riemannian gradient descent on the sharpness along $M$ and an orthogonal, bifurcation-type dynamics, enabling quantitative convergence results across three regimes: subcritical (linear convergence after a finite transient to a suboptimal flat minimum), critical (parabolic, power-law convergence toward the optimally flat minimum), and supercritical (linear convergence to a period-two orbit orthogonal to $M$). The theory is validated on multilayer scalar factorisation, illustrating how overparametrisation can induce a bias toward flatter minima and providing explicit convergence rates. This advances the theoretical understanding of the edge of stability in non-convex, overparametrised settings and offers principled explanations for the practical success of GD with large learning rates in deep learning contexts.

Abstract

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$. The parallel component corresponds to Riemannian gradient descent on the objective sharpness, while the orthogonal component is a bifurcating dynamical system. This insight allows us to derive convergence rates in three regimes characterised by the learning rate size: (a) the subcritical regime, in which transient instability is overcome in finite time before linear convergence to a suboptimally flat global minimum; (b) the critical regime, in which instability persists for all time with a power-law convergence toward the optimally flat global minimum; and (c) the supercritical regime, in which instability persists for all time with linear convergence to an orbit of period two centred on the optimally flat global minimum.

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

TL;DR

This work provides a rigorous framework for understanding gradient descent with large learning rates in overparametrised least-squares problems by exploiting a solution manifold . Using tubular coordinates, the authors derive a normal form that separates GD into a Riemannian gradient descent on the sharpness along and an orthogonal, bifurcation-type dynamics, enabling quantitative convergence results across three regimes: subcritical (linear convergence after a finite transient to a suboptimal flat minimum), critical (parabolic, power-law convergence toward the optimally flat minimum), and supercritical (linear convergence to a period-two orbit orthogonal to ). The theory is validated on multilayer scalar factorisation, illustrating how overparametrisation can induce a bias toward flatter minima and providing explicit convergence rates. This advances the theoretical understanding of the edge of stability in non-convex, overparametrised settings and offers principled explanations for the practical success of GD with large learning rates in deep learning contexts.

Abstract

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 , which enables the decomposition of the GD dynamics into components parallel and orthogonal to . The parallel component corresponds to Riemannian gradient descent on the objective sharpness, while the orthogonal component is a bifurcating dynamical system. This insight allows us to derive convergence rates in three regimes characterised by the learning rate size: (a) the subcritical regime, in which transient instability is overcome in finite time before linear convergence to a suboptimally flat global minimum; (b) the critical regime, in which instability persists for all time with a power-law convergence toward the optimally flat global minimum; and (c) the supercritical regime, in which instability persists for all time with linear convergence to an orbit of period two centred on the optimally flat global minimum.
Paper Structure (22 sections, 35 theorems, 248 equations, 4 figures)

This paper contains 22 sections, 35 theorems, 248 equations, 4 figures.

Table of Contents

  1. Introduction
  2. Related work
  3. Theoretical setting
  4. Notation
  5. Problem setting
  6. A normal form for gradient descent about the solution manifold
  7. Convergence theorems
  8. Subcritical regime
  9. Critical regime
  10. Supercritical regime
  11. Limitations, discussion and conclusion
  12. Differential geometry background
  13. Riemannian metrics, connections and geodesics
  14. Riemannian geometry of submanifolds and hypersurfaces in Euclidean space
  15. Case study: deep scalar factorisation
  16. ...and 7 more sections

Key Result

Theorem 4.1

There is an analytic change of coordinates $(\eta,\theta^{\parallel},\theta^{\perp})\mapsto (\eta,\theta)$ for a tubular neighbourhood $N$ of $M$ in which the gradient descent map $\mathrm{GD}^{\ell}:(\eta, \theta) \mapsto (\eta,\theta-\eta\nabla\ell(\theta^{\parallel}))$ takes the form $(\eta,\thet with $c(\eta,\theta^{\parallel})$ defined in Assumption ass:generic.

Figures (4)

  • Figure 1: $M = \{(x,y):xy=1\}$ with tubular neighbourhood $N$ (shaded) and line $\mathcal{L}$ (dotted). Inside $N$, any point $\theta$ is closest to a unique point $\theta^{\parallel}$ on $M$, with $\theta-\theta^{\parallel} = \theta^{\perp}\,n(\theta^{\parallel})$ orthogonal to $M$ for some $\theta^{\perp}\in\mathbb{R}$. Assumption \ref{['ass:orthogonal']} says that $\nabla\ell(\theta')$ is parallel to $\mathcal{L}$ at any point $\theta'\in\mathcal{L}$.
  • Figure 2: Log-scale plots of distance from $\theta^{\parallel}_t$ to $\theta^{\parallel}_*$ (left), magnitude of $\theta^{\perp}_t$ (centre) and sharpness suboptimality gap (right) for gradient descent on depth 5 scalar factorisation in the subcritical regime. Trajectories from five different initialisations shown. Initial instability in $|\theta^{\perp}_t|$ (top) is overcome in finite time with rapid convergence to a suboptimally flat global minimum (bottom).
  • Figure 3: Log-$y$-scale plots of distance from $\theta^{\parallel}_t$ to $\theta^{\parallel}_*$ (left), magnitude of $\theta^{\perp}_t$ (centre) and sharpness suboptimality gap (right) on depth 5 scalar factorisation in the critical regime. Trajectories from five different initialisations shown. The iterates $|\theta^{\perp}_t|$ may or may not initially increase; in both cases, asymptotic, power law decrease of $|\theta^{\perp}_t|$ and $\|\theta^{\parallel}_t-\theta^{\parallel}_*\|$ to zero indicates power law convergence to the optimally flat global minimum.
  • Figure 4: Log-scale plots of distance from $\theta^{\parallel}_t$ to $\theta^{\parallel}_*$ (left), magnitude of $\theta^{\perp}_t$ (centre) and sharpness suboptimality gap (right) on depth 5 scalar factorisation in the supercritical regime. Trajectories from five different initialisations shown. The iterates $|\theta^{\perp}_t|$ converge to a stable, period-two orbit driving linear convergence of $\theta^{\parallel}_t$ to the optimally flat global minimum $\theta^{\parallel}_*$.

Theorems & Definitions (66)

  • Theorem 4.1: Normal form for GD about $M$
  • proof : Outline of proof
  • Theorem 5.1
  • Theorem 5.2
  • Theorem 5.3
  • Lemma A.1
  • proof
  • Lemma A.2
  • proof
  • Lemma A.3
  • ...and 56 more