# Edge of Stability Selectively Shapes Learning Across the Data Distribution

arXiv: https://arxiv.org/abs/2606.04212 · HTML: https://mlsanity.com/paper/2606.04212/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 38 of its 43 references that are in the map.
- First posted on arXiv 2 Jun 2026.

## Builds on (most cited first, 30 of 38)

- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (2017)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [Understanding deep learning requires rethinking generalization](https://mlsanity.com/paper/1611.03530/) (2016)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [An Empirical Study of Example Forgetting during Deep Neural Network Learning](https://mlsanity.com/paper/1812.05159/) (2018)
- [Exploring Generalization in Deep Learning](https://mlsanity.com/paper/1706.08947/) (2017)
- [Deep Learning on a Data Diet: Finding Important Examples Early in Training](https://mlsanity.com/paper/2107.07075/) (2021)
- [Beyond neural scaling laws: beating power law scaling via data pruning](https://mlsanity.com/paper/2206.14486/) (2022)
- [Fantastic Generalization Measures and Where to Find Them](https://mlsanity.com/paper/1912.02178/) (2019)
- [What Neural Networks Memorize and Why: Discovering the Long Tail via Influence Estimation](https://mlsanity.com/paper/2008.03703/) (2020)
- [Sharp Minima Can Generalize For Deep Nets](https://mlsanity.com/paper/1703.04933/) (2017)
- [Does Learning Require Memorization? A Short Tale about a Long Tail](https://mlsanity.com/paper/1906.05271/) (2019)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [Dataset Cartography: Mapping and Diagnosing Datasets with Training Dynamics](https://mlsanity.com/paper/2009.10795/) (2020)
- [Three Factors Influencing Minima in SGD](https://mlsanity.com/paper/1711.04623/) (2017)
- [Gradient Starvation: A Learning Proclivity in Neural Networks](https://mlsanity.com/paper/2011.09468/) (2020)
- [The large learning rate phase of deep learning: the catapult mechanism](https://mlsanity.com/paper/2003.02218/) (2020)
- [Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability](https://mlsanity.com/paper/2209.15594/) (2022)
- [The Break-Even Point on Optimization Trajectories of Deep Neural Networks](https://mlsanity.com/paper/2002.09572/) (2020)
- [Understanding Gradient Descent on Edge of Stability in Deep Learning](https://mlsanity.com/paper/2205.09745/) (2022)
- [Adaptive Gradient Methods at the Edge of Stability](https://mlsanity.com/paper/2207.14484/) (2022)
- [On the Relation Between the Sharpest Directions of DNN Loss and the SGD Step Length](https://mlsanity.com/paper/1807.05031/) (2018)
- [Understanding the unstable convergence of gradient descent](https://mlsanity.com/paper/2204.01050/) (2022)
- [Understanding the Generalization Benefit of Normalization Layers: Sharpness Reduction](https://mlsanity.com/paper/2206.07085/) (2022)
- [Towards Understanding Generalization of Deep Learning: Perspective of Loss Landscapes](https://mlsanity.com/paper/1706.10239/) (2017)
- [A Walk with SGD](https://mlsanity.com/paper/1802.08770/) (2018)
- [Understanding Edge-of-Stability Training Dynamics with a Minimalist Example](https://mlsanity.com/paper/2210.03294/) (2022)
- [Learning threshold neurons via the "edge of stability"](https://mlsanity.com/paper/2212.07469/) (2022)

## Abstract (arXiv)

Existing analyses of the edge of stability (EoS) treat it as a global property of optimization. We show that it is also selective: the stability constraint redistributes learning across subsets of the training distribution, amplifying progress on some groups while suppressing progress on others. Using a branching intervention that enters or exits the EoS regime from the same training state, we causally demonstrate this trade-off and identify two necessary conditions for a group to benefit. First, its aggregate gradient must align with the top Hessian eigenvector. We isolate this mechanism with a controlled perturbation that preserves distance but randomizes direction, destroying alignment and eliminating the advantage. Second, the group must sustain non-vanishing gradient magnitude over time. Under cross-entropy loss, gradient saturation decouples confidently classified groups, shifting the advantage to output-outliers, whose gradients persist. Together, these results show that EoS functions not only as a stability boundary, but as a mechanism governing the allocation of learning across the data distribution.