🤖 machine learning

Flatland: The Adventures of Gradient Descent with Large Step Sizes

This paper addresses the open question of convergence for gradient descent with large step sizes on non-L-smooth functions by proposing adaptive methods that operate at the edge of stability, demonstrating that allowing training to enter slightly sharper valleys via self-stabilization can improve convergence and generalization compared to prematurely encountering flat regions.

Original authors: Leonardo Galli, Curtis Fox, Wiebke Bartolomaeus, Mark Schmidt, Holger Rauhut

Published 2026-06-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Leonardo Galli, Curtis Fox, Wiebke Bartolomaeus, Mark Schmidt, Holger Rauhut

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 mountain range. This mountain range represents the "loss landscape" of a neural network. Your goal is to get to the very bottom (the best solution) as quickly as possible.

In the world of deep learning, the tool you use to walk down the mountain is called Gradient Descent. It's like a hiker who always takes a step in the direction that goes steepest downhill. The size of that step is the learning rate (or step size).

For a long time, researchers believed that taking small, careful steps was the safest way to ensure you didn't trip and fall off a cliff. However, a recent observation showed that sometimes, taking huge, reckless steps actually helps you find a better, flatter valley at the bottom, which leads to a smarter AI. But there's a catch: if your steps are too big, you might start bouncing wildly and never stop, or you might get stuck in a weird, flat spot that looks like the bottom but isn't.

This paper, titled "Flatland: The Adventures of Gradient Descent with Large Step Sizes," is like a guidebook for hikers who want to take those giant steps without falling off the edge of the world.

The Problem: The "Edge of Stability"

The authors noticed that when you take steps that are just the right size to be "large," the hiker starts to oscillate. They bounce up and down the sides of the valley instead of walking straight down. This is called the Edge of Stability (EoS).

Think of it like a skateboarder on a half-pipe. If they go too slow, they stop in the middle. If they go just right, they ride back and forth, gaining momentum and eventually finding the smoothest, flattest part of the bottom. The paper argues that staying on this "edge" is actually good for training AI.

The Solution: A Smart Compass (The Algorithm)

The big question was: How do we know how big a step is "too big" without crashing?

The authors created a new "compass" (an algorithm) that automatically adjusts the step size. Instead of guessing, the compass checks the terrain right in front of the hiker.

If the ground is smooth, it says, "Take a huge step!"
If the ground is jagged, it says, "Slow down, but don't stop."

This compass allows the hiker to stay on the Edge of Stability from the very first step, ensuring they move fast but don't fly off the mountain.

The Surprise: The "Flatland" Trap

Here is the most interesting part of the story. The authors discovered that if you take steps that are too large, you might accidentally stumble into a place called Flatland.

Imagine you are hiking and you reach a spot that looks perfectly flat. You think, "Great, I've found the bottom!" But actually, you've landed on a saddle point (like the middle of a horse's saddle). It's flat in one direction but slopes down in another.

The Trap: In these "Flatland" regions, the AI stops learning. It thinks it's done because the ground is flat, but it's actually just stuck in a dead end where it can't figure out how to solve the problem. The paper calls these "unsuccessfully-flat" runs. The AI is essentially guessing randomly and getting stuck there.

The Fix: Don't Go Too Flat

The paper's main advice is counter-intuitive: Don't aim for the absolute flattest spot.

The authors found that the best results come from a "successfully-flat" run. This is where the AI stays in a valley that is mostly flat but still has a tiny bit of slope.

The Metaphor: Imagine you are looking for the best spot to set up a tent.
- The Trap: You find a perfectly flat, featureless plain. It looks perfect, but there's no wind to help you dry your clothes, and you can't tell which way is North. You get stuck.
- The Solution: You find a spot that is almost flat, but has a gentle, slight slope. This slight slope helps you orient yourself and keeps you moving toward the best possible view.

The authors propose a simple rule to avoid the trap: Limit your step size so it never gets bigger than the number of categories you are trying to learn. (For example, if you are teaching an AI to recognize 100 types of animals, don't let your step size exceed 100). This keeps the AI in the "slightly sloped" valley where it can actually learn, rather than getting stuck in the "perfectly flat" trap.

Summary

Big Steps are Good: Taking large steps helps AI learn faster and find better solutions, provided you stay on the "Edge of Stability."
The Trap: If you take steps that are too big, you might get stuck in a "Flatland" (a saddle point) where the AI stops learning and just guesses randomly.
The Fix: Use a smart method to adjust steps, but put a "speed limit" on them. Keep the steps large enough to be fast, but small enough to avoid the perfectly flat trap. This leads to a smarter, more accurate AI.

In short, the paper teaches us that in the race to train AI, sometimes you need to run fast, but you must be careful not to run so fast that you trip over a perfectly flat rock and stop moving entirely.

Technical Summary: Flatland: The Adventures of Gradient Descent with Large Step Sizes

Problem Statement

The training of neural networks frequently involves objective functions that are not globally $L$ -smooth. For such functions, a longstanding open question in deep learning remains: what is the largest possible step size that ensures the convergence of Gradient Descent (GD)? While previous work (e.g., Cohen et al., 2021) observed that GD with fixed step sizes often enters an "Edge of Stability" (EoS) phase—where the loss decreases non-monotonically and the product of step size and sharpness stabilizes around 2—there is no unified definition of "large" step sizes that guarantees convergence without relying on second-order information. Furthermore, it is unclear whether GD can consistently operate at the EoS from the start of training or if it risks diverging due to uncontrolled oscillations.

Methodology

The authors propose a unifying framework for defining and finding "large" step sizes based on local segment smoothness rather than global Lipschitz continuity.

Segment Smoothness: The paper introduces the concept of local segment smoothness, assuming the objective function $f$ is continuously differentiable and that the gradient is Lipschitz (or Hölder continuous) only along the anti-gradient direction between iterates $w$ and $w - \eta \nabla f(w)$ .
Equality Line Search: To find step sizes that satisfy the definition of "large" (specifically, step sizes $\eta_k$ such that $\eta_k \geq 1/l(w_k, \eta_k)$ and occasionally exceed $2/l(w_k, \eta_k)$ ), the authors design a new Equality Line Search algorithm. This method combines backtracking and extrapolation steps. If the standard non-monotone line search condition is initially satisfied, the step size is doubled; if not, it is halved. The process continues until the condition is satisfied as an equality, approximating the maximal step size allowed by the local geometry.
Theoretical Guarantees: The authors prove that this approach yields step sizes that are "Lipschitz-aware" and "Hölder-aware." Crucially, they show that as the gradient vanishes near a stationary point, the step size is lower-bounded by the reciprocal of the sharpness (the largest eigenvalue of the Hessian), effectively making the method "sharpness-aware" without explicit Hessian computation.
Avoiding Flat Traps: The authors identify a phenomenon where GD with large step sizes can get "stuck" in globally-flat regions (saddle points) where sharpness reaches a global minimum of $2/K$ (for $K$ classes). To mitigate this, they propose NLS-ub, a variant that caps the step size below $K$ , forcing the optimizer to traverse slightly sharper valleys.

Key Contributions

Unified Definition: A formal definition of "large" step sizes that ensures GD convergence for non-globally-smooth functions, relying only on local segment smoothness.
First-Order Sharpness Awareness: An algorithm (Equality Line Search) that provably yields step sizes operating at the Edge of Stability from the very first iteration, without requiring second-order information.
Characterization of Globally-Flat Regions: The identification of specific saddle points where the sharpness achieves a global minimum of $2/K$ . The authors demonstrate that encountering these points too early in training leads to "unsuccessfully-flat" runs where the network fails to learn meaningful features, resulting in poor convergence and generalization.
Self-Stabilization Mechanism: A practical strategy (NLS-ub) to limit step sizes below the number of classes $K$ , which successfully converts "unsuccessfully-flat" runs into "successfully-flat" ones that achieve lower training loss and higher test accuracy.

Experimental Results

The authors evaluated their methods on various architectures (MLP, CNN, ResNet, VGG, DenseNet, Vision Transformers) and datasets (CIFAR10/100, SVHN, EMNIST) using both Mean Squared Error (MSE) and Cross-Entropy losses.

Edge of Stability: The proposed Equality Line Search (NLS) and its deterministic variant (PoNLS) consistently operate at the EoS, with the product of step size and sharpness remaining above the threshold of 2 throughout training.
Globally-Flat Phenomenon: In several experiments (e.g., ResNet on SVHN), NLS drove the sharpness to exactly $2/K$ very early in training. In these cases, the training loss stagnated at the level of random guessing, and the gradient norms of hidden layers collapsed to near zero, indicating the optimizer had reached a saddle point.
Success vs. Failure: Runs that avoided hitting the exact $2/K$ minimum (staying slightly above it) were labeled "successfully-flat." These runs achieved significantly better training loss and test accuracy.
NLS-ub Efficacy: By capping the step size below $K$ , NLS-ub prevented the optimizer from entering the globally-flat saddle regions. This modification consistently improved convergence and test accuracy compared to standard NLS and other baselines like SAM (Sharpness Aware Minimization) and CDAT.
Loss Function Sensitivity: The phenomenon of hitting global sharpness minima was observed with MSE loss (where the minimum is $2/K$ ). With Cross-Entropy loss, the sharpness minimum is 0, and similar "stuck" behaviors were observed when sharpness approached 0, though the dynamics differed slightly.

Significance and Claims

The paper claims to provide a concrete answer to the open question of how to select the largest possible learning rates while ensuring GD convergence. It challenges the prevailing intuition that "flatter" minima are always preferable. The authors demonstrate that finding the globally-flat regions (global minima of sharpness) is not beneficial for convergence speed or generalization; in fact, it can be detrimental, causing the network to stall at saddle points.

The significance lies in:

Operationalizing Large Steps: Providing a practical, first-order method to automatically select step sizes that operate at the Edge of Stability, accelerating training without manual tuning.
Reframing Flatness: Showing that the pursuit of the absolute flattest minima can be a trap, and that slightly sharper valleys (avoiding the $2/K$ or 0 sharpness floor) yield superior performance.
Theoretical Foundation: Establishing convergence guarantees for GD under local segment smoothness, a weaker assumption than global $L$ -smoothness, which is often violated in deep learning.

The authors remain modest regarding the scope, noting that while their method works well for full-batch GD and preliminary stochastic experiments, the theoretical convergence of Stochastic Gradient Descent (SGD) under these specific assumptions requires further investigation. They also acknowledge that the mechanism causing GD to enter these globally-flat regions so efficiently is not yet fully understood and may be related to the neural collapse phenomenon.