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

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.

• The paper demonstrates that transient learning dynamics enable SGD to escape sharp valleys by leveraging an early noise-driven exploration phase.
• It combines detailed empirical mapping with a statistical physics model to show that increased noise prolongs exploration and biases convergence toward flatter minima.
• The findings imply that tuning hyperparameters like learning rate and batch size can strategically enhance generalization in deep learning models.

Transient Dynamics and Valley Selection in Stochastic Gradient Descent

Overview

The study "Transient learning dynamics drive escape from sharp valleys in Stochastic Gradient Descent" (2601.10962) presents a comprehensive dissection of how stochastic gradient descent (SGD) selects among local minima during neural network training. By combining detailed empirical analysis with an analytically tractable statistical physics model, the authors provide strong evidence that SGD's bias toward flatter, more generalizable minima is fundamentally a dynamical, nonequilibrium effect, governed by an early-stage exploration and subsequent transient "freezing" of the parameter trajectory.

Empirical Analysis of Loss Landscape and SGD Trajectories

The authors initiate their investigation by systematically mapping the loss landscape navigated by SGD in a canonical neural setting—a fully connected MLP trained on a subset of MNIST. Upon varying SGD hyperparameters (learning rate, batch size), and using identical initialization, they observe that final solutions cluster into distinct regions in parameter space, each corresponding to a local minimum ("valley") with full training accuracy, but separated by significant loss barriers.

Figure 1: Multiple distinct solution valleys emerge in neural loss landscapes, as revealed by PCA-projected SGD trajectories and Jaccard-based functional similarity.

Pairwise analysis using Jaccard similarity of test error patterns, which is invariant to weight permutations, reveals that the heights of barriers separating minima are strongly anticorrelated with the similarity of solutions—functionally dissimilar solutions are separated by higher barriers. These empirical findings underscore a salient structural property: neural loss surfaces at scale admit multiple well-isolated, functionally meaningful valleys.

SGD Noise, Valley Flatness, and Generalization

Adjusting the noise properties of SGD via batch size and learning rate inherently modulates the diversity of solutions discovered. Systematic grid sweeps over batch sizes and learning rates show that increased noise (i.e., smaller $B$, larger $\eta$) produces broader exploration and, crucially, a marked increase in the frequency of convergence to flatter minima with higher test accuracy.

Figure 2: Higher SGD noise broadens exploration and enhances the likelihood of convergence to flatter, more generalizable minima.

Analysis of Hessian spectra demonstrates that these flatter solutions correspond to valleys with lower top eigenvalues, confirming an established—yet now explicitly dynamical—link between flatness and generalization.

Early Transient Phase: Exploration and Freezing

Temporal evolution of the parameter trajectory during training reveals a well-defined "transient phase" characterized by large loss gradients and frequent transitions ("jumps") between different valleys. The termination of this exploratory phase, termed the freezing time ($t_\mathrm{freeze}$), coincides with entry into the basin of a particular minimum, after which further inter-valley transitions are suppressed by growing barriers.

Notably, larger SGD noise systematically extends the duration of the transient phase, as evidenced by delayed $t_\mathrm{freeze}$ at fixed learning rate across batch size reductions.

Figure 3: Larger SGD noise increases the duration of the transient search phase, facilitating discovery of flatter valleys.

This prolongation is not simply temporal but is directly leveraged: flatness and generalization of the final solution monotonically increase as a function of the length of the exploratory period, a relationship validated by strong empirical correlation.

Continuation Method: Direct Evidence of Valley Hopping

By devising a "continuation training" experiment—branching the training process at intermediate times to deterministic full-batch GD—the authors provide direct evidence for the mechanism of valley hopping: trajectories switched to GD early in training become trapped in functionally poor minima, whereas later switching exploits SGD's exploration to discover increasingly flatter, higher-quality valleys.

Figure 4: Continuation experiments reveal valley-jumping dynamics during early training and establish a temporal link between exploration and final solution quality.

Tracking test loss, flatness, and output similarity, the continuation experiments confirm that valley selection occurs nearly exclusively in the early, noise-driven transient phase, and is essentially "frozen in" once the optimizer descends below a threshold loss.

Minimal Model: Anisotropic Noise and Nonequilibrium Selection

To explain and quantitatively capture the observed dynamics, the paper introduces a minimal two-valley model. It features a bifurcating loss landscape with valleys of equal depth but unequal flatness, and models SGD as a Langevin process with landscape-dependent, anisotropic noise proportional to the local Hessian.

Figure 5: A minimal two-valley model with anisotropic, curvature-dependent noise reproduces empirical learning dynamics and the phase diagram of valley selection.

Simulations with varying noise and learning rates reproduce all key empirical findings: increased noise both enhances the probability of converging to the flatter minimum and delays the effective freezing time. Analytical treatment using Fokker-Planck formalism and Kramers' escape theory demonstrates that anisotropic (curvature-proportional) noise distorts the effective potential, explicitly biasing dynamics toward flatter regions.

Analytical Theory: Transient Freezing as the Origin of Flatness Bias

Crucially, the paper's theory demonstrates that the observed preference for flat minima is not a property of the eventual steady-state distribution of SGD, but is intrinsically a result of the freezing of exploration at a dynamically determined point in training. While higher SGD noise reduces selectivity for flatness at any fixed loss—contrary to some prior interpretations—the longer transient allowed by high noise ensures that freezing occurs at a flatter and higher-quality valley.

Figure 6: Analytical results reveal that the transient freezing mechanism due to anisotropic noise explains the bias toward flatter minima.

Quantitative expressions predict, and simulations confirm, that the probability of final convergence to a flat minimum increases with both noise strength and flatness contrast.

Implications and Future Directions

This work has both theoretical and practical implications:

• Theoretical Insight: The unification of loss landscape geometry, stochastic learning dynamics, and generalization quality under a nonequilibrium statistical physics framework advances our understanding of why SGD generalizes well. The identification of a transient, noise-regulated exploration phase as the locus of generalization selection provides a concrete, mechanistic explanation beyond prior post-hoc or instability-based theories.
• Optimization Strategy: It follows that modulation of learning rate/batch size schedules—interpreted as tools to regulate the length and noise properties of the transient exploration phase—can be systematically designed to enhance generalization, rather than for mere convergence or training speed. Explicitly delaying freezing or shaping anisotropic noise properties could enable access to even flatter, more robust solution basins.
• Modeling Generality: While based on a tractable two-valley model and canonical neural architectures/datasets, both the multi-valley structure and transient exploration/freezing phenomena are expected to be generic across large-scale deep networks and datasets.

Conclusion

This paper provides a rigorous, mechanistically grounded account of SGD's implicit bias toward flat, generalizable minima, reconciling empirical observations, simulation, and theory. The identification of a dynamical, transient "freezing" process under anisotropic noise not only clarifies the trajectory-level origins of generalization in deep learning but also motivates new classes of optimization and regularization strategies for advancing the field.