📊 statistics

Edge of Stability Selectively Shapes Learning Across the Data Distribution

This paper demonstrates that the Edge of Stability is not merely a global optimization boundary but a selective mechanism that redistributes learning across the data distribution by amplifying progress on groups with top-Hessian-aligned, non-saturating gradients while suppressing others.

Original authors: Shauna Kwag, Anakha Ganesh, Tomaso Poggio, Pierfrancesco Beneventano

Published 2026-06-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Shauna Kwag, Anakha Ganesh, Tomaso Poggio, Pierfrancesco Beneventano

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 running a massive training camp for a team of deep neural networks (AI models). The goal is to teach the model to recognize things, like distinguishing between cars and trucks.

Usually, we think of training as a smooth, steady march toward a solution. But this paper reveals that when the training gets "edgy"—specifically, when the learning rate is high enough to push the system to the very edge of stability—it stops being a smooth march. Instead, it becomes a selective filter that decides who gets to learn and who gets left behind.

Here is the breakdown of their discovery using simple analogies:

1. The "Edge of Stability" is a Tightrope

Imagine the AI is walking on a tightrope. If it walks too slowly (low learning rate), it's safe but slow. If it walks too fast, it falls. The "Edge of Stability" (EoS) is the sweet spot where the AI walks so fast it starts to wobble and oscillate back and forth, but just barely stays on the rope.

For a long time, scientists thought this wobbling was just a weird side effect. This paper says: No, the wobbling is actually a powerful tool. It acts like a spotlight that shines on specific students in the class while ignoring others.

2. The Selective Spotlight: Who Gets Helped?

The researchers set up a "fork in the road" experiment. They trained two identical AI models on the exact same data.

Model A kept walking on the tightrope (staying at the Edge of Stability).
Model B slowed down and stepped off the tightrope (exiting the Edge of Stability).

The Surprise: Model A didn't just get better overall. It got much better at recognizing specific types of examples, while Model B got better at different ones. The tightrope didn't help everyone equally; it redistributed the learning effort.

The Winners: Examples that were "outliers" (strange, unusual inputs) or had labels that didn't quite fit (output-outliers) got a massive boost. They learned faster on the tightrope.
The Losers: The "normal" examples (inliers) and the ones right on the boundary between categories actually learned slower on the tightrope compared to the model that stepped off.

3. The Two Rules of the Spotlight

Why did the tightrope help some and hurt others? The paper identifies two strict rules. For a group of data to get the "Edge of Stability" boost, they must pass a two-part test:

Rule #1: The "Directional Alignment" (The Push)

Imagine the AI is being pushed by a giant wind. The "Edge of Stability" only helps if the wind is blowing in a very specific direction.

The Test: The group's "gradient" (the direction they need to learn) must align perfectly with the "top Hessian eigenvector."
The Analogy: Think of a group of people trying to push a heavy boulder. If they all push in the same direction, the boulder moves fast. If they push in random directions, they cancel each other out.
The Finding: The paper proved that if you take the "outlier" data and scramble the direction they push (while keeping the distance the same), they lose their advantage. They must push coherently in the same direction to get the boost.

Rule #2: The "Persistence" (The Stamina)

The second rule is about not giving up.

The Test: The group must keep "pushing" (having a non-zero gradient) throughout the training.
The Analogy: Imagine a race. If a runner gets confident and stops running because they think they've won, they stop making progress.
The Finding: In some loss functions (like Cross-Entropy), once the AI gets confident about a "normal" example, it stops trying to learn from it (the gradient vanishes). But "outliers" or "mislabelled" examples keep confusing the AI, so the AI keeps trying to learn from them. Because they keep pushing, they get the boost. If you switch to a loss function where confident examples stop pushing, the advantage shifts to the ones that keep pushing.

4. The Big Picture: Geometry Dictates the Winner

The most fascinating part is that which group wins depends entirely on the shape of the data.

If your data has a specific shape where "outliers" are the ones pushing in the right direction, the Edge of Stability will make the AI an expert on outliers (which can help with robustness against weird attacks).
If you change the data shape so that "boundary" examples are the ones pushing in the right direction, the Edge of Stability will suddenly make the AI an expert on the boundaries instead.

The Takeaway:
The Edge of Stability isn't just a safety limit; it's a mechanism that allocates learning. It acts like a manager who decides, "We are going to focus our energy on fixing these specific problems right now, and we will ignore the others."

This changes how we view AI training. It's not just about finding the "flattest" or "best" solution globally. It's about understanding that the training process naturally prioritizes certain parts of the data based on their geometry and how long they keep "fighting" to be learned. The "Edge of Stability" is the tool that enforces this priority.

Technical Summary: Edge of Stability Selectively Shapes Learning Across the Data Distribution

1. Problem Statement

Existing literature on the Edge of Stability (EoS) treats it primarily as a global property of optimization, where the top Hessian eigenvalue self-stabilizes near the threshold $2/\eta$ (where $\eta$ is the learning rate). While the dynamics of this regime—characterized by sign-alternating oscillations and a slow corrective drift—are well-documented, the functional consequences of operating at this threshold remain unclear. Specifically, it is unknown whether the stability constraint provides a uniform benefit or if it differentially shapes learning across different subsets of the training data. The central question is: Does the EoS regime induce a selective learning bias, and if so, what geometric and dynamic properties determine which data subsets benefit?

2. Methodology

2.1 Experimental Setup and Prototype Taxonomy

The authors instantiate a controlled experimental framework using a binary classification task on CIFAR-10 (Automobile vs. Truck) with a two-hidden-layer MLP trained via full-batch gradient descent. To isolate the effects of data geometry, they define a prototype taxonomy partitioning the training distribution into four groups based on their position relative to class centroids ( $\mu_c$ ) and label consistency, independent of the trained model:

Inliers: High-density points near the class centroid with correct labels.
Boundary: In-distribution points near the inter-class boundary (high label ambiguity).
Input-outliers: Geometrically atypical inputs (far from $\mu_c$ ) with correct labels.
Output-outliers: High-density inputs assigned incorrect labels (synthetically constructed).

2.2 Branching Intervention (Causal Identification)

To causally identify the effect of remaining in the EoS regime versus exiting it, the authors employ a branching intervention:

Baseline Branch: Continues training at the original learning rate $\eta$ , remaining in the EoS regime.
Exit Branch: At the onset of EoS (detected when the top Hessian eigenvalue $\lambda_1$ reaches $2/\eta$ ), the learning rate is halved ( $\eta \to \eta/2$ ). This raises the stability threshold to $4/\eta$ , allowing the optimizer to exit the EoS regime.
Comparison: By comparing the divergence in per-group losses ( $\ell_k$ ) between the two branches post-intervention, the authors isolate the causal impact of the EoS stability constraint.

2.3 Mechanism Isolation

The authors test two hypothesized necessary conditions for a subset to benefit from EoS:

Directional Alignment: The aggregate gradient of the subset must align with the top Hessian eigenvector ( $v_1$ ). This is tested by constructing "random-direction" input-outliers where displacement distance is preserved, but the direction is randomized, destroying coherent alignment.
Gradient Persistence: The subset must sustain a non-vanishing gradient magnitude. This is tested by comparing Mean Squared Error (MSE) loss (gradients persist) against Cross-Entropy (CE) loss (gradients vanish for confidently classified points).

3. Key Results

3.1 Selective Learning Trade-off

The branching intervention reveals that EoS is selective, not global. The stability constraint redistributes optimization effort unevenly:

Beneficiaries: Subsets whose gradients align with $v_1$ and persist (e.g., input-outliers under MSE, output-outliers under CE) show accelerated loss reduction in the EoS branch compared to the exit branch.
Suppressors: Subsets with orthogonal gradients or vanishing magnitudes (e.g., inliers and boundary points under specific conditions) show suppressed progress in the EoS branch relative to the exit branch.
This trade-off replicates across architectures (CNN, ResNet), optimizers (SGD, Momentum), and class pairs.

3.2 Mechanism: Alignment $\times$ Persistence

Two conditions are jointly necessary for a group to capture the EoS advantage:

Alignment: The group's aggregate gradient must align with the top Hessian eigenvector ( $v_1$ $v_{1}$ ).
- Evidence: When input-outliers are displaced in random directions (preserving distance but destroying alignment), their curvature influence $(\nabla \ell_k \cdot v_1)^2$ collapses by an order of magnitude, and the EoS advantage disappears.
Persistence: The group must maintain a non-vanishing gradient magnitude.
- Evidence: Under Cross-Entropy loss, gradients for confidently classified input-outliers saturate (vanish). Consequently, the EoS advantage shifts to output-outliers, whose gradients remain active due to label inconsistency, despite potentially lower alignment.

3.3 Geometry Determines the Beneficiary

The subset that benefits from EoS is determined by the geometric composition of the data distribution:

In the standard CIFAR-10 setup, input-outliers (furthest from centroids) dominate $v_1$ and benefit.
In a harder class pair (Cat vs. Dog) or with different outlier construction parameters, boundary points can become the most distant group, dominating $v_1$ and becoming the primary beneficiary.
Natural Data: On the natural CIFAR-10 distribution, a positive correlation emerges at EoS onset between centroid distance and alignment with $v_1$ , confirming that geometric outliers naturally drive the unstable mode.

3.4 Generalization and Robustness Implications

The paper presents preliminary evidence that the training-time selectivity transfers to test-time behavior:

Adversarial Robustness: EoS improves robustness on boundary points only when boundary points dominate $v_1$ . If input-outliers dominate, the EoS branch performs worse on boundary robustness because optimization budget is diverted.
Out-of-Distribution (OOD) Generalization: EoS improves OOD performance on extrapolated inputs only when those inputs dominate $v_1$ during training.

4. Theoretical Framework

The authors extend the self-stabilization framework (Damian et al., 2023) from global loss to per-subset losses. They derive that the cycle-averaged drift in the EoS regime includes an additional term proportional to $-\nabla S$ (where $S$ is sharpness). The differential effect on a subset $k$ is governed by the selector:
$Q_k = \langle \nabla \ell_k, \nabla S \rangle$
Under a single-mode approximation, this selector is proportional to the squared projection of the subset gradient onto the top eigenvector:
$C_k = (\nabla \ell_k \cdot v_1)^2 = \|\nabla \ell_k\|^2 \cos^2 \theta_k$
This decomposition confirms that both directional alignment ( $\cos^2 \theta_k$ ) and gradient magnitude persistence ( $\|\nabla \ell_k\|^2$ ) are required for a subset to experience the EoS drift.

5. Significance and Claims

The paper claims to shift the understanding of EoS from a scalar stability boundary to a mechanism governing the allocation of learning across the data distribution.

Reframing Flatness: The authors argue that "flatness" is not a scalar property of the solution but a directional property determined by data geometry. The functional benefit of low curvature depends entirely on which subset dominates the top Hessian eigendirection.
Inductive Bias: EoS acts as an inductive bias that selects specific subsets for optimization based on their geometric relationship to the unstable mode, rather than merely selecting among equivalent minimizers.
Unifying Perspectives: The work connects implicit regularization in parameter space (curvature constraints) with inductive bias over the data distribution (which examples are learned).
Modest Scope: The authors acknowledge limitations, noting that experiments are restricted to small models and full-batch training due to the computational cost of Hessian-vector products. They state that whether these selectivity effects persist for larger datasets, mini-batch optimizers, or representation-space prototypes remains an open question.

In conclusion, the paper demonstrates that the Edge of Stability is a selective regime where the stability constraint amplifies learning for subsets that align with and persist along the top Hessian eigendirection, while suppressing others, thereby fundamentally shaping the model's functional properties.