Edge of Stability Selectively Shapes Learning Across the Data Distribution
For machine learning practitioners, this reveals that EoS can cause uneven learning across data subgroups, which is important for understanding and mitigating bias in training dynamics.
The paper shows that the edge of stability (EoS) in optimization is not just a global property but selectively redistributes learning across data subgroups, amplifying progress on some while suppressing others. Using causal interventions, they identify two conditions for a group to benefit: alignment of its aggregate gradient with the top Hessian eigenvector and sustained non-vanishing gradient magnitude.
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.