Generalization at the Edge of Stability

TL;DR

This paper develops a generalization theory for neural networks trained in the Edge of Stability (EoS) regime — where the top Hessian eigenvalue exceeds the classical stability threshold 2/η and training trajectories become chaotic. The authors model stochastic optimizers as random dynamical systems (RDS) that settle onto a random pullback attractor, and introduce a new complexity measure — the Sharpness Dimension (SD) — derived from the full Hessian spectrum that bounds worst-case generalization on that attractor. Unlike trajectory-based PH-dimension or single-point sharpness/trace measures, SD depends on the spectrum’s structure of partial determinants, captures the effective dimensionality of expanding-vs-contracting directions, and empirically correlates with the generalization gap on MLPs, GPT-2, and grokking experiments — exactly where flatness-based proxies fail.

Key claims

Cohen et al.’s Edge of Stability — top Hessian eigenvalue oscillating near and above 2/η — implies chaotic training dynamics with a positive top Lyapunov exponent, under which trajectories diverge exponentially even though the optimizer remains bounded [§4, §3].
Trace-based and pointwise sharpness flatness criteria are neither necessary nor sufficient for good generalization, so pointwise analysis of individual solutions is fundamentally inadequate in the EoS regime [§1, §2].
Stochastic optimization at fixed learning rate (e.g. SGD with persistent minibatch noise) does not converge to a point but settles on a compact random pullback attractor whose existence is guaranteed by a pullback-absorbing-set argument [§3 Dfn. 3.2, §4 Prop. 4.1].
The proposed Sharpness Dimension SD is the largest order k such that the global sharpness index σ_k ≥ 0; it measures the effective number of expanding directions before global volume contraction dominates, and in the EoS regime is strictly smaller than the ambient parameter dimension [§4, Dfn. 4].
Main theorem: the worst-case generalization gap over the random pullback attractor is upper-bounded by an expression in SD, mediated by the Minkowski dimension of the attractor; the proof covers the attractor with ellipsoids whose axes are determined by σ_k and reduces to Dupuis et al.’s fractal-dimension covering bound [§4 Thm 4, App. B].
SD depends on the complete Hessian spectrum and the structure of its partial determinants — a complexity that trace and spectral norm cannot capture [Abstract, §4].
The Hessian spectrum is estimated at scale via Stochastic Lanczos Quadrature (SLQ) with operator resampling across runs — minibatch Hessian–vector products with a fresh batch per Lanczos run — yielding SD-SLQ directly in expectation rather than via post-hoc averaging [§5].
On small 3-layer MLPs on MNIST with exact SVD, SD and SD-SLQ show stronger Kendall correlation with the generalization gap than Euclidean and data-dependent PH-Dim, α-weighted lifetime sums, or top RDS-Sharpness alone [§5.1, Fig. 2].
On 5-layer 200-width MLPs (278K params), SD-SLQ ties backward PH-Dim on granulated Kendall but outperforms all baselines on classical Kendall — and unlike trajectory-based PH-Dim does not require recording the entire training trajectory [§5.1, Fig. 3].
Empirical correlation with generalization extends to GPT-2 transformers and the grokking setting (delayed/sudden generalization), where standard sharpness and trace measures are known to fail [§1 contributions, §5].

Method

The framework casts SGD with constant learning rate as a discrete-time random dynamical system (θ, φ), where θ is an ergodic shift on infinite minibatch-index sequences and φ(n, ω) is the cocycle that composes n update steps under the noise realization ω. Because fresh noise is injected at every step, asymptotic behavior must be analyzed in the pullback sense: fix a noise realization ω, push every bounded initial set forward from the distant past, and look at where they land at the present time. The resulting random pullback attractor A(ω) is a compact, noise-conditioned snapshot of the asymptotic state — and the object the bound is stated over.

The complexity measure is built from global sharpness indices σ_k(ω), defined as long-run logarithmic stretching rates of k-dimensional volumes along the principal directions of the one-step Jacobian. SD is the largest k with σ_k ≥ 0 — the effective volume-preserving dimension of the attractor. In the EoS regime, the top Lyapunov exponent is positive (chaos) but global dissipativity bounds composite volumes, so SD < d strictly. The main theorem (Thm 4) shows that under boundedness + Lipschitz + integrability assumptions, the worst-case generalization gap is O(√(SD · log n / n)) (modulo a mutual-information term that the companion result Thm B.5 removes via Tuci et al.’s set-stability framework). The proof bounds the Minkowski dimension of A(ω) by SD via an ellipsoidal cover whose axes are calibrated by σ_k, then plugs into Dupuis et al.’s existing fractal-dimension generalization bound.

For practical computation, the authors deploy Stochastic Lanczos Quadrature with operator resampling: each of R Lanczos runs uses a fresh minibatch for its Hessian–vector products, the resulting per-run quadrature nodes are smoothed (histogram or KDE) into a spectral density, and the integral SD = ∫ ... dλ is computed by the trapezoid rule. This avoids the quadratic-memory cost of materializing the Hessian and the failure of Krylov methods in the near-zero-mass regime that dominates DNN spectra.

Results

3-layer MLP (12,960 params) on MNIST, exact SVD: SD-based measures dominate. Backward PH-Dim and forward α-weighted lifetime sums show only weak positive correlation with the generalization gap; top RDS-sharpness alone shows weak negative correlation. Both exact SD and SD-SLQ give substantially higher Kendall and granulated-Kendall coefficients, with exact SD strongest [Fig. 2].
5-layer MLP (278,800 params) on MNIST: exact SVD is intractable. Backward PH-Dim, forward α ratios, SD-SLQ, and SD-KDE reach similar GKC averages; on classical Kendall coefficient SD-SLQ outperforms all other indicators [Fig. 3].
The paper extends empirically to GPT-2 transformers and to grokking — where sharpness and Hessian trace are well-known to fail — and reports that SD continues to track the generalization gap [§1, §5].
Code is promised at https://circle-group.github.io/research/GATES.

Why it’s interesting

The wiki has a thick training-stability cluster (Training stability at scale, Hyperparameter scaling laws) framed almost entirely around preventing loss spikes — QK-Norm, spectral-sphere optimizers, Birkhoff-polytope residual mixing. This paper inverts the framing: it accepts that modern training operates in the chaotic EoS regime by design (Andreyev & Beneventano observe expected mini-batch Hessian concentrates near the threshold), and asks what generalization theory still applies. The answer — a fractal-dimension bound on the entire random attractor rather than on any single solution — provides theoretical grounding for why architectural stability fixes that suppress symptoms (outliers, signal-gain explosions) without leaving the EoS regime can still produce networks that generalize well. It also connects to Preconditioned Inexact Stochastic ADMM for Deep Models (PISA / SISA / NSISA)‘s PISA convergence-theory work in the optimizer-side stability fork, but from the orthogonal direction of generalization rather than convergence. The SLQ-based numerical recipe for spectrum estimation is also directly useful for any future Luma work that wants Hessian-spectrum diagnostics at DiT/MMDiT scale.