Transient learning dynamics drive escape from sharp valleys in Stochastic Gradient Descent

TL;DR

The paper addresses why stochastic gradient descent (SGD) tends to find flat, generalizable minima by analyzing its early transient dynamics and introducing a minimal two-valley model with anisotropic, landscape-dependent noise. It shows that the noise reshapes the loss into an effective nonequilibrium potential $\mathcal{L}_{\mathrm{eff}}^{\pm}$ and that a transient freezing point $y_{\mathrm{freeze}}$ controls final valley selection, with larger noise extending the exploratory phase and biasing toward flatter valleys (higher $P_{\mathrm{flat}}^{\mathrm{tr,SGD}}$). The results unify learning dynamics, geometry of the loss landscape, and generalization, and suggest optimization strategies that actively regulate early transient dynamics, noise anisotropy, and freeze timing. This framework provides principled guidance for designing stochastic optimizers that enhance access to flatter, more robust minima in high-dimensional networks.

Abstract

Stochastic gradient descent (SGD) is central to deep learning, yet the dynamical origin of its preference for flatter, more generalizable solutions remains unclear. Here, by analyzing SGD learning dynamics, we identify a nonequilibrium mechanism governing solution selection. Numerical experiments reveal a transient exploratory phase in which SGD trajectories repeatedly escape sharp valleys and transition toward flatter regions of the loss landscape. By using a tractable physical model, we show that the SGD noise reshapes the landscape into an effective potential that favors flat solutions. Crucially, we uncover a transient freezing mechanism: as training proceeds, growing energy barriers suppress inter-valley transitions and ultimately trap the dynamics within a single basin. Increasing the SGD noise strength delays this freezing, which enhances convergence to flatter minima. Together, these results provide a unified physical framework linking learning dynamics, loss-landscape geometry, and generalization, and suggest principles for the design of more effective optimization algorithms.

Figures (6)

• Figure 1: The loss landscape of a neural network is composed of multiple distinct solution valleys separated by loss barriers. (A) A 2D Principal Component Analysis (PCA) projection of the weight training trajectories on subset of MNIST. All runs start from an identical initialization (star) but use different SGD hyperparameters (learning rate $\eta$ and batch size $B$), indicated by different colors and marker styles, respectively. The trajectories shown here all converge, achieving 100% training accuracy. (B) The Jaccard similarity ($\mathrm{Sim}_J$) matrix between pairs of solutions, calculated based on the overlap of misclassified test samples. (C, D) The training loss along the linear interpolation path between two final solutions. A negligible loss barrier exists for highly similar solutions ($\mathrm{Sim}_J = 0.95$) (C), while a significant barrier separates dissimilar solutions ($\mathrm{Sim}_J = 0.625$) (D). The color scale represents the normalized training time $t/t_\mathrm{max}$. (E) The height of the loss barrier $\Delta \mathcal{L}_\mathrm{barrier}$ exhibits a strong negative correlation with the Jaccard similarity $\mathrm{Sim}_J$ between solution pairs (linear fit, $R^2=0.64$).
• Figure 2: Increased SGD noise promotes broader exploration and convergence to flatter, more generalizable valleys. (A) Heatmap of average Jaccard similarity $\langle \mathrm{Sim}_J \rangle$ between final solutions across learning rates $\eta$ and batch sizes $B$ ($B=1000$ corresponds to full-batch Gradient Descent). (B, C) PCA-projected training trajectories for low-noise ($B=50, \eta=0.01$, ① in A) and high-noise ($B=50, \eta=0.05$, ② in A) conditions. (D, F) Heatmaps of maximum test accuracy $\max(Acc_{\mathrm{test}})$ (D) and maximum solution flatness $\max(F)$ (F), where flatness is quantified as the inverse geometric mean of the top Hessian eigenvalues. (E, G) Violin plots showing the distributions of test accuracy (E, for $\eta=0.05$, gray box in D) and flatness (G) across batch sizes. Gray regions in the heatmaps (A, D, F) indicate divergent settings where no runs achieve 100% training accuracy and are therefore excluded from comparisons with other conditions. See Sec. I A in Supplemental Material for details and Fig. S2 for additional numerical results.
• Figure 3: SGD noise extends the transient phase and facilitates the discovery of flatter solutions. (A) Representative training trajectory showing training loss (magenta) and accuracy (cyan) for $B=50$ and $\eta=0.05$. The freezing time $t_{\mathrm{freeze}}$ is defined as the first iteration where the training loss drops below $\mathcal{L}_c=0.1$. (B, D) Mean training loss $\langle \mathcal{L}_{\mathrm{train}} \rangle$ (B) and mean flatness $\langle F \rangle$ (D) across multiple runs for different batch sizes $B$ at fixed learning rate $\eta=0.02$. (C) Heatmap of normalized mean freezing time $\eta \langle t_{\mathrm{freeze}} \rangle$ across learning rates and batch sizes, averaged over all runs that reached $\mathcal{L}_c$. Gray regions in heatmaps indicate settings where no runs reach 100% training accuracy.
• Figure 4: Continuation training reveals the valley-jumping dynamics of SGD during the early transient phase. (A) A 2D PCA projection of the main SGD trajectory (blue solid line; reference setting $B=50$, $\eta=0.05$) and trajectories obtained by switching to full-batch GD ($B=1000$, $\eta=0.05$) at different branching times $t_c$ (colored lines with circle markers). (B, C) Evolution of the test loss $\mathcal{L}_\mathrm{test}$ (B) and solution flatness $F$ (C) for the main SGD trajectory and the GD-continued trajectories branched at different $t_c$. Insets show the final values of $\mathcal{L}_{\mathrm{test}}$ and $F$ for each continuation trajectory (evaluated at $t=2000$) as a function of their branching time $t_c$. (D) Final flatness $F$ of the continued solutions exhibits a strong negative correlation with their final test loss $\mathcal{L}_\mathrm{test}$ (linear fit, $R^2=0.76$). (E) Jaccard similarity $Sim_J$ among solutions obtained from different branching times $t_c$, with the red square marking the region where $t_c \geq t_{\mathrm{freeze}}$.
• Figure 5: A minimal two-valley model with anisotropic, landscape-dependent noise recapitulates the empirically observed training dynamics. (A) Conceptual schematic of key characteristics of empirical SGD dynamics. (B) The constructed 2D loss landscape $\mathcal{L}(x, y)$ with bifurcating sharper and flatter valleys of equal depth, as defined in Eq. \ref{['eq:loss_toy_model']}. (C and D) Heatmaps showing the probability of converging to the flatter valley $P_{\mathrm{flat}}$ (C) and the normalized mean freezing time $\eta \langle t_{\mathrm{freeze}} \rangle$ (D), as functions of learning rate $\eta$ and noise strength $\sigma$. Statistics were collected from 2000 simulations, half initialized on the sharper side and half on the flatter side for each hyperparameter setting. Full simulation procedures and landscape parameters are provided in the Supplemental Material, Sec. II A.