Edge of Stability has been a thing for a while now. We can show that most training happens at the edge of stability. This is very surprising because of a "negative " fact: Most theory does not apply (e.g., Descent Lemma).
The question becomes:
Does the fact that training happens at the EoS impact performance? In what way? With what mechanisms? How is that affected by the architecture?
To our knowledge, these questions they are unanswered to date and our new paper is the about the first step towards this!
EoS picks what part of the distribution to learn and at what speed!
arxiv.org/abs/2606.04212
In particular, the consequence is that for MLPs:
EoS can improve robustness or OOD behavior, but only when the relevant subset is the one selected by EoS.
If boundary points dominate, EoS helps near-boundary robustness.
If distributional outliers dominate, EoS helps extrapolate toward the tail.
Precisely:
1. We make EoS causal. We fork training from the same state at EoS onset:
one branch stays at EoS, the other exits by lowering the learning rate. Same everything, only the stability constraint changes.
2. EoS is selective:
it learns some parts of the data distribution faster, while slowing others.
3. The selector is surprisingly simple. A group benefits from EoS when its gradient is big and
has high alignment with v_1, the top Hessian eigenvector.
Translation: to benefit from EoS, a subset must
(i) point in the sharpest direction, and
(ii) keep a non-vanishing gradient.
This gives concrete answers to the questions above:
Does EoS affect performance?
Yes, but not uniformly. It reallocates learning.
In what way?
It prioritizes the subset with largest curvature influence.
By what mechanism?
Self-stabilization at the sharpness boundary couples training to v_1; only groups aligned with v_1, and whose gradients persist, get the extra progress.
How does architecture matter?
Architecture changes the map from input geometry to gradient geometry. Thus it can change which subset dominates, but the same predictor remains:
largest curvature influence wins.
With the amazing Shauna Kwag,
@anakha_g, and
@TomasoPoggio!