Training Instabilities Induce Flatness Bias in Gradient Descent

TL;DR

The paper investigates how training instabilities in gradient descent can improve generalization by biasing optimization toward flatter regions of the loss landscape. It introduces the Rotational Polarity of Eigenvectors (RPE), a geometric mechanism in which leading Hessian eigenvectors rotate during unstable training, promoting exploration and flattening; this effect extends to SGD and can be harnessed in adaptive optimizers via Clipped-Ada. By modeling higher-order curvature moments with a coupled instability-dynamics framework, the authors prove an implicit flatness bias and demonstrate empirical flattening and generalization gains on CIFAR10 and Fashion-MNIST. The work argues for embracing certain instability regimes (a Goldilocks zone) rather than avoiding them, with practical implications for learning-rate schedules and optimizer design in deep networks.

Abstract

Classical analyses of gradient descent (GD) define a stability threshold based on the largest eigenvalue of the loss Hessian, often termed sharpness. When the learning rate lies below this threshold, training is stable and the loss decreases monotonically. Yet, modern deep networks often achieve their best performance beyond this regime. We demonstrate that such instabilities induce an implicit bias in GD, driving parameters toward flatter regions of the loss landscape and thereby improving generalization. The key mechanism is the Rotational Polarity of Eigenvectors (RPE), a geometric phenomenon in which the leading eigenvectors of the Hessian rotate during training instabilities. These rotations, which increase with learning rates, promote exploration and provably lead to flatter minima. This theoretical framework extends to stochastic GD, where instability-driven flattening persists and its empirical effects outweigh minibatch noise. Finally, we show that restoring instabilities in Adam further improves generalization. Together, these results establish and understand the constructive role of training instabilities in deep learning.

Figures (22)

• Figure 1: Stable and unstable phases of training. The magnitude of parameter oscillations, measured by $\|{\bm{\theta}}-{\bm{\theta}}^*\|_2$, cleanly separates stable and unstable phases in an MLP trained on fMNIST with GD. Sharpness increases until $\lambda_{\max}\!\approx\!2/\eta$ (progressive sharpening), after which loss and curvature exhibit spikes corresponding to valley jumping.
• Figure 1: Large learning rates improve generalization on CIFAR10. Validation accuracy across architectures and datasets.
• Figure 2: Loss and sharpness through a training instability. We plot the loss (top) and curvature (bottom; leading Hessian eigenvalue estimated via Hutchinson’s trick hutchinson1989stochastic) across representative snapshots. Fluctuations in the curve of curvature coincide with strong reorientations of leading eigenvectors, via rotations, that precede rapid flattening and the restoration of stability.
• Figure 2: Aggregate correlation statistics between $|\gamma_{U_t}|$ and $\lambda_{1,t}$. Strong correlations across seeds affirm C3 of Condition \ref{['cond:flat:standing']}.
• Figure 3: Optimization trajectories in a $2$-parameter DLN with varying $\eta$. Left: the regimes of $\gamma_\beta$ vs learning rate. Middle: the loss landscape. Right: rotation of the leading Hessian eigenvector ${\bm{v}}_1$. Stable regimes (orange/red/green) show convergence to sharper axis; unstable regimes (blue) show divergence.

Theorems & Definitions (56)

• Definition 5.1: Instability surplus
• Theorem 5.4: Point-mass contraction (PMC)
• Corollary 5.5: Mass conditions for median contraction
• Theorem 5.6: Upper-envelope drift
• Lemma 5.9: Two-sided monotonicity
• Lemma 5.10: CLT-controlled monotone tails
• Proposition 5.11: Median dominates mode
• Theorem 5.12: Even-time median drift
• Definition F.1
• Lemma F.2: Two–step adjacent cancellations are negative