Edge of Stability Selectively Shapes Learning Across the Data Distribution

Shauna Kwag, Anakha Ganesh, Tomaso Poggio, Pierfrancesco Beneventano

Keywords: edge of stability; training dynamics; implicit bias; curvature; data geometry; generalization

TL;DR: At the edge of stability, optimization shapes not just how fast a network learns, but which examples it prioritizes, favoring subsets aligned with the sharpest direction and resistant to saturation.

Abstract: 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.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 14