🤖 machine learning

Dynamics of Gradient Descent with Large Step Size Near a Manifold of Flat Minima

This paper extends the theory of gradient descent with large step sizes from isolated flat minima to manifolds of flat minima in overparametrised least-squares with vector-valued outputs, establishing a generalized normal form and convergence results that reveal the fibre bundle structure of flat minima in deep matrix factorisation.

Original authors: Lachlan Ewen MacDonald, René Vidal

Published 2026-07-10

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Lachlan Ewen MacDonald, René Vidal

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find the lowest point in a vast, foggy landscape. This landscape represents the "loss" of a deep neural network—a measure of how wrong the computer is. The goal is to get to the very bottom.

For a long time, computer scientists thought the best way to do this was to take tiny, careful steps downhill. They believed that if you took a step that was too big, you would overshoot the bottom, bounce back up, and eventually fly off into the sky (diverge). They had a strict rule: your step size had to be smaller than a specific limit determined by how "sharp" the bottom of the valley was. If the bottom was a needle point (sharp), you had to take microscopic steps. If it was a wide, flat bowl, you could take slightly larger steps.

But here's the plot twist: in the real world of training AI, people started taking huge steps. Instead of crashing, the AI got smarter, faster, and surprisingly, it seemed to prefer landing in those wide, flat bowls rather than the sharp needle points. This behavior was a mystery. It was like watching a skier take a massive jump, land perfectly in a soft snowbank, and glide to a stop, defying all the physics textbooks that said they should have crashed.

This paper is the team of detectives who finally figured out why the skier didn't crash, and exactly how they landed.

The Big Discovery: The "Flip" and the "Slide"

The authors, Lachlan MacDonald and René Vidal, took the old, simple theories and upgraded them to handle the messy, high-dimensional reality of modern AI. They discovered that when you take a large step, the AI doesn't just wander randomly. It actually splits its behavior into two distinct modes, like a car that can drive forward and also spin in place.

1. The "Flip" (The Bouncing Ball):
Imagine the AI is bouncing on a trampoline. If the step size is just right, the AI bounces up and down in a very specific rhythm. It doesn't stop bouncing immediately; instead, it settles into a stable, repeating pattern of bouncing back and forth. The paper proves that if the step size is slightly too big (but not too big), the AI will oscillate in a perfect, predictable loop right above the flat bottom. It's not a mistake; it's a feature.

2. The "Slide" (The River):
While the AI is bouncing up and down (the flip), it is also slowly sliding sideways along the flat bottom of the valley. The authors show that the AI is essentially performing a special kind of "Riemannian gradient descent" along this flat surface. Think of it like a river flowing along a flat plain. The river doesn't care about the tiny bumps in the sand; it just flows smoothly toward the flattest, most stable part of the landscape.

The Shape of the "Flat"

One of the most exciting parts of this paper is how they describe the "flat minima" (the wide, safe valleys). Previous theories treated these flat spots as isolated islands. But the authors show that in complex problems like matrix factorization (a key technique in deep learning), these flat spots aren't just single points. They are actually a fibre bundle over a product of spheres.

Let's translate that into plain English: Imagine the flat valley isn't a single room, but a giant, multi-layered structure. The base of this structure is a collection of spheres (like the surface of a ball). On top of every point on these spheres, there is a whole "fiber" or a small room of solutions. The AI doesn't just find one solution; it finds a whole connected, smooth family of solutions that are all equally good. The paper proves that the "sharpness" (how steep the sides are) changes smoothly as you move around this structure, behaving in a very orderly way called "Morse-Bott."

What They Proved vs. What They Simulated

The authors didn't just guess this; they built a rigorous mathematical framework to prove it.

The Theory: They proved that for a wide range of step sizes, the AI's behavior splits into three distinct regimes:
- Subcritical (Safe but slow): If the step is small enough, the AI converges exponentially fast to a "sub-optimally flat" minimum. It's safe, but maybe not the absolute flattest spot.
- Critical (The Sweet Spot): If the step size hits a specific threshold (exactly $2/\lambda_1$ , where $\lambda_1$ is the sharpness), the AI converges to the flat minimum at a rate of $t^{-1/2}$ . This is a specific, proven mathematical rate. They showed this happens in simulations of matrix factorization.
- Supercritical (The Dance): If the step is slightly larger than the threshold, the AI doesn't stop at the bottom. Instead, it converges exponentially to a period-2 orbit. This means it settles into a stable, repeating cycle of two points, oscillating back and forth. The paper proves this cycle exists and is stable.
The Simulations: To back up their math, they ran experiments on matrix factorization problems (specifically 3-layer, $2 \times 2$ matrix factorization). In these simulations, they watched the AI take large steps. The graphs showed exactly what the math predicted: the AI bounced, then settled into the $t^{-1/2}$ slide, or locked into the period-2 dance.

What They Explicitly Rule Out

It is important to know what this paper says is not happening.

It is not random chaos: The paper explicitly rules out the idea that large steps lead to random, unpredictable behavior in these specific regimes. The oscillation is a structured, stable period-2 orbit, not a chaotic mess.
It is not just about single points: The paper argues against the idea that flat minima are isolated. They prove that in these systems, the minima form a continuous, smooth manifold (a connected surface), not a scattered collection of dots.
It is not a global guarantee: The authors are careful to state that their proofs are local. They prove what happens near the flat minima. They do not claim to have solved the mystery of how the AI finds the flat minima from a random starting point far away (the "progressive sharpening" phase). They only explain what happens once the AI is already in the neighborhood of the flat spot.

The "Edge of Stability"

The paper connects this to a phenomenon called the "Edge of Stability." This is the regime where the AI teeters on the edge of crashing but doesn't fall. The authors show that this isn't a bug; it's a specific dynamical state where the AI implicitly performs a "Riemannian gradient descent" on the sharpness itself. It's as if the AI is using the bounce to feel out the landscape and slide toward the flattest possible spot.

The Bottom Line

This paper takes a complex, high-dimensional problem and gives us a clear map. It shows that when we take big steps in training AI, we aren't just guessing. We are engaging in a sophisticated dance where the AI bounces in a stable rhythm while simultaneously sliding along a smooth, connected surface of perfect solutions.

They proved that for matrix factorization, this surface is a beautiful geometric structure (a fibre bundle over spheres), and they proved that the AI's movement on this surface follows strict, predictable laws. While they haven't solved the entire mystery of deep learning (like how to get from the start to the finish line), they have provided the first rigorous, mathematical explanation for why taking huge steps works so well once you are close to the solution.

In short: The AI isn't crashing; it's dancing. And thanks to this paper, we finally have the sheet music.

Technical Summary: Dynamics of Gradient Descent with Large Step Size Near a Manifold of Flat Minima

Problem Statement

Classical convergence theory for Gradient Descent (GD) relies on the stability condition $\eta < 2/\lambda_1(x)$ , where $\eta$ is the step size and $\lambda_1(x)$ is the largest eigenvalue (sharpness) of the objective Hessian. In deep learning, this condition is frequently violated; yet, GD with large step sizes often converges, exhibiting phenomena such as "edge of stability" dynamics and implicit bias toward "flat minima" (minima with small sharpness).

Previous work by MacDonald et al. [39] provided a rigorous dynamical systems analysis for large-step GD in codimension 1 least-squares problems (single scalar output) near an isolated flat minimum. However, this setting is practically restrictive, failing to cover:

Higher codimension problems: Overparametrised regression with vector-valued outputs and arbitrary numbers of observations.
Manifolds of minima: Many deep learning problems, such as matrix factorisation, possess a continuous manifold of flat minima rather than isolated points.

This paper extends the theory of [39] to these broader settings, addressing the technical challenges of higher-dimensional bifurcations and singular partial differential equations (PDEs) arising along manifolds.

Methodology

1. Geometric Framework and Assumptions

The authors consider overparametrised least-squares objectives $\ell(w) = \frac{1}{2}\|\tau - f(w)\|^2$ where $f: \mathbb{R}^p \to \mathbb{R}^q$ ( $p > q$ ). They establish a hierarchy of geometric assumptions to ensure the existence of a normal form:

Assumption 3.1: The target $\tau$ is a regular value, ensuring the solution set $M = f^{-1}(\tau)$ is a smooth submanifold.
Assumption 3.3: The top eigenvalue $\lambda_1$ of the Hessian is simple (distinct from $\lambda_2$ ) on a non-empty open subset, allowing for a well-defined top eigenvector field $\nu_1$ .
Assumption 3.4: There exists a submanifold $F \subset M$ of local minima for $\lambda_1$ where $\lambda_1$ is Morse-Bott (strongly convex in directions normal to $F$ within $M$ ).
Assumption 3.5: A regularity condition on a coefficient $\alpha_\eta$ derived from higher-order derivatives of $\ell$ , ensuring the solvability of a singular PDE required for the normal form.

2. Normal Form Derivation

The core technical contribution is the derivation of a normal form for GD near the manifold $F$ . The authors perform a sequence of coordinate transformations:

Tubular Coordinates: Mapping the neighborhood of $M$ to $M \times \mathbb{R}^q$ .
Center Manifold Reduction: Using a center manifold theorem to reduce the $q$ -dimensional orthogonal dynamics to a 1-dimensional system corresponding to the top eigendirection $\nu_1$ .
Solving a Singular PDE: A critical step involves solving a PDE of the form $Xu = \alpha u + \beta$ where the vector field $X$ vanishes along the manifold $F$ . Unlike the isolated point case in [39], the singularity here is along a manifold. The authors develop novel techniques involving logarithmic expansions and invariant foliations to prove the existence of a $C^1$ solution, correcting a flawed proof in [39] regarding the coupling between the manifold coordinates and the orthogonal dynamics.
Strong-Stable Foliation: Disentangling the contracting directions from the dynamics on the center manifold.

The resulting normal form decomposes the GD dynamics into:

Tangent to $M$ ( $x$ ): Riemannian Gradient Descent (RGD) on the sharpness $\lambda_1$ along $M$ , with a step size controlled by the square of the orthogonal distance $y_1^2$ .
Top Orthogonal Direction ( $y_1$ ): A flip bifurcation dynamic (similar to a period-doubling bifurcation) governed by $(1 - \eta\lambda_1)y_1 + y_1^3$ .
Other Orthogonal Directions ( $y_{2:q}$ ): Exponentially contracting stable directions.

3. Case Study: Deep Matrix Factorisation

The authors verify that deep matrix factorisation fits this framework. They prove that the set of flat minima $F$ forms a smooth fibre bundle over a product of spheres. Furthermore, they demonstrate that the sharpness $\lambda_1$ is Morse-Bott along this manifold, satisfying the necessary geometric hypotheses.

Key Contributions and Results

1. Generalised Normal Form

The paper establishes a normal form for GD with large step sizes in the neighborhood of a manifold of flat minima for arbitrary codimension. This generalises the results of [39] from isolated minima in codimension 1 to continuous manifolds in arbitrary dimensions.

2. Convergence Theorems

The authors prove three convergence regimes based on the step size $\eta$ relative to the stability threshold $2/\lambda_1|_F$ :

Subcritical Regime ( $\eta < 2/\lambda_1|_F$ ):
After an initial transient, iterates converge exponentially to a suboptimally flat global minimum. The step size decays exponentially as the iterates approach the manifold.
Critical Regime ( $\eta = 2/\lambda_1|_F$ ):
Iterates converge non-monotonically with a rate of $\Theta(t^{-1/2})$ to a point in the flat minima manifold $F$ $F$ .
- Note: The proof for the general case relies on a conjecture (Conjecture D.7) regarding the existence of a stable foliation for normally parabolic dynamical systems, analogous to the strong-stable foliation theorem for normally hyperbolic manifolds. The result is proven rigorously if the Hessian of $\lambda_1$ is a scalar multiple of the identity along $F$ , or if the conjecture holds.
Supercritical Regime ( $\eta > 2/\lambda_1|_F$ ):
Iterates converge exponentially to a stable period-2 orbit along the span of the top eigenvector field $\nu_1|_F$ . The amplitude of this oscillation scales as $\Theta(\sqrt{\eta\lambda_1|_F - 2})$ .

3. Structural Results for Matrix Factorisation

For deep matrix factorisation, the paper provides novel structural insights:

The flat minima manifold $F$ is a fibre bundle over a product of spheres.
The sharpness $\lambda_1$ is Morse-Bott along $F$ .
The eigenvalues of the Hessian of $\lambda_1$ restricted to the normal bundle of $F$ are strictly positive, with specific lower bounds derived (e.g., for $L=2$ , the lower bound is $4\sigma_1/(\sigma_1 - \sigma_{d_L})$ ).

4. Numerical Verification

The theoretical predictions are verified via numerical experiments on 3-layer, $2 \times 2$ matrix factorisation problems. The experiments confirm:

Exponential convergence in the subcritical regime.
$t^{-1/2}$ convergence in the critical regime.
Exponential convergence to a period-2 cycle in the supercritical regime.

Significance and Claims

The paper claims to bridge a significant gap in the theoretical understanding of deep learning optimization:

Beyond Codimension 1: It moves the theory of large-step GD from the impractical "single datum" setting to general overparametrised regression with vector outputs.
Manifolds of Minima: It addresses the reality that deep learning loss landscapes often contain manifolds of minima, showing that the "edge of stability" dynamics persist in these higher-dimensional, continuous settings.
Technical Novelty: The solution of the singular PDE along a manifold and the correction of the proof in [39] regarding the coupling of coordinates are presented as independent technical contributions of interest to dynamical systems theory.
Implicit Bias: The results rigorously formalize the observation that large-step GD implicitly performs Riemannian gradient descent on the sharpness along the solution manifold, providing a mechanistic explanation for the bias toward flat minima.

The authors remain modest regarding the critical regime, explicitly stating that the general convergence proof relies on a conjecture about parabolic foliations, which lies outside the scope of the current work but is argued to be plausible based on analogies to hyperbolic systems. They also note that their results are local (near the flat minima manifold) and do not yet explain the global "progressive sharpening" phase observed in training.