Deep linear networks (DLNs) are used as an analytically tractable model of the training dynamics of deep neural networks. While gradient descent in DLNs is known to exhibit saddle-to-saddle dynamics, the impact of stochastic gradient descent (SGD) noise on this regime remains poorly understood. We investigate the dynamics of SGD during training of DLNs in the saddle-to-saddle regime. We model the training dynamics as stochastic Langevin dynamics with anisotropic, state-dependent noise. Under the assumption of aligned and balanced weights, we derive an exact decomposition of the dynamics into a system of one-dimensional per-mode stochastic differential equations. This establishes that the maximal diffusion along a mode precedes the corresponding feature being completely learned. We also derive the stationary distribution of SGD for each mode: in the absence of label noise, its marginal distribution along specific features coincides with the stationary distribution of gradient flow, while in the presence of label noise it approximates a Boltzmann distribution. Finally, we confirm experimentally that the theoretical results hold qualitatively even without aligned or balanced weights. These results establish that SGD noise encodes information about the progression of feature learning but does not fundamentally alter the saddle-to-saddle dynamics.
Figures & tables
Figure 1: Predicting when modes are learned using the predicted time of maximum diffusion, in a 4-layer linear network trained with SGD. (a) The modes are learned in order of magnitude, with the time of learning being predicted by the time of maximum diffusion. (b) The diffusion along a mode peaks while a mode is being learned, and our theoretical prediction (see Equation 7 ) matches what is observed. The vertical lines in (a) correspond to the peaks (marked with ★ in (b)) of the theoretical prediction for the diffusion (see Equation 9 ).
Figure 2: Saddle-to-saddle dynamics with different optimizers in a depth-6 linear network, with the sharp changes corresponding to increases in the numerical rank of the network. (a) Train loss plateaus with discrete optimizers (b) Train loss plateaus with numerical simulation of their continuous counterparts (c) Mode growth over training with gradient descent, showing that the 5 singular values of the teacher matrix are learned in descending order of magnitude (d) Mode growth over training for stochastic gradient descent, showing that the stagewise dynamics are retained, but modes take longer to be learned. (See Appendix H for details of the numerical simulations and experiments.)
Figure 3: Comparison of diffusion along modes with mode amplitude , for (a) SGD and (b) an Euler-Maruyama simulation of anisotropic Langevin dynamics. In each column, the empirical value of ηaα(θ)⊤Σ(θ)aα(θ) is shown on the left and Proposition 3.5 ’s theoretical prediction for Dα(θ) is shown on the right. The shaded bands show the time that the corresponding mode is learned. In agreement with the theoretical prediction, both have peaks in the diffusion before the mode is learned.
Figure 4: Comparison of the empirical distributions of the amplitude of the first mode w0 at the end of training for SGD, anisotropic Gaussian, and isotropic Gaussian noise. In the absence of label noise, SGD (a) concentrates entirely on the value of the top singular value of the teacher matrix, but anisotropic noise (b) does not have this behavior. However, the variance of the distribution for anisotropic noise is lower than for isotropic noise (c), and is thus more similar to SGD.
Figure 5: End-of-training distribution in the presence of label noise (variance 0.1) of the amplitude of the first mode w0 for SGD. The distribution changes from being concentrated at a point to having greater variance.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Measuring magnitude of cross modes over SGD training . For both a depth-2 and a depth-4 linear network, we observe that the majority of cross modes are small for most of training, except for a small number of cross modes which peak at times corresponding to modes being learned.
Figure 7: Measuring balance from balanced and unbalanced initialization . (a) shows that from an unbalanced Gaussian initialization, there is not strong balance at the start of training, but as training continues balance increases. (b) When we enforce balance at initialization, it is approximately maintained, and for all of training the weights are significantly more balanced than at any point in the unbalanced-initialization trajectory. (c) and (d) show that this holds for deeper (depth-4) linear networks. The unnormalized lines are the numerator of the expression for rl , and show little change.
Figure 8: Frobenius norms of weight matrices over SGD training . Frobenius norms of layers increase over training, accounting for most of the decrease in normalized balance measure rl from standard initializations.
Figure 9: Learning rate versus maximum diffusion for each mode. We observe a linear relationship between learning rate and the diffusion along each mode.
Figure 10: Batch size versus maximum diffusion for each mode. We observe a linear relationship between the reciprocal of batch size and the diffusion along each mode.
Figure 11: Empirical versus theoretical modewise diffusion for a non-rectangular DLN. The structure of the diffusion, including the location of the peaks and tending to zero once the mode is learned is maintained.
Figure 12: Variance of the end-of-training distribution of the first mode versus simulation fineness. We vary the fineness (i.e., the Δt parameter) of the Euler-Maruyama simulation of the stochastic gradient flow SDE ( Equation 1 ). For finer simulations, the variance reduces.
One crucial factor behind the success of deep learning lies in the implicit bias induced by noise inherent in gradient-based training algorithms. Motivated by empirical observations that training with noisy labels improves model generalization, we delve into the underlying mechanisms behind stochastic gradient descent (SGD) with label noise. Focusing on a two-layer over-parameterized linear network, we analyze the learning dynamics of label noise SGD, unveiling a two-phase learning behavior. In Phase I, the magnitudes of model weights progressively diminish, and the model escapes the lazy regime; enters the rich regime. In Phase II, the alignment between model weights and the ground-truth interpolator increases, and the model eventually converges. Our analysis highlights the critical role of label noise in driving the transition from the lazy to the rich regime and minimally explains its empirical success. Furthermore, we extend these insights to Sharpness-Aware Minimization (SAM), showing that the principles governing label noise SGD also apply to broader optimization algorithms. Extensive experiments, conducted under both synthetic and real-world setups, strongly support our theory. Our code is released at https://github.com/a-usually/Label-Noise-SGD.
Tongcheng Zhang, Zhanpeng Zhou, Mingze Wang +4
Sch. of Computer Science and Zhiyuan College, Shanghai Jiao Tong University · Peking University · University of Sydney +4
Stochastic gradient descent (SGD) is the cornerstone of modern deep learning, yet its theoretical convergence properties remain poorly understood in non-convex settings. To address this gap, we study SGD dynamics in deep linear neural networks - simplified models that retain key non-convex characteristics while remaining analytically tractable. Building on recent results for gradient descent, we extend convergence guarantees to SGD under regularized squared loss. We establish almost sure boundedness of SGD iterates. Using the framework of asymptotic pseudotrajectories from stochastic approximation theory, our analysis establishes then, for both deterministic and adaptive step-size choices, convergence of SGD to the set of critical points under mild conditions (without relying on Lipschitz properties or whitened data). To our knowledge, this is the first work providing almost sure convergence guarantees for SGD in deep linear networks, bridging the gap between theoretical optimization and practical stochastic training.
Gabin Maxime Nguegnang, Bubacarr Bah, Holger Rauhut +1
Understanding the behavior of stochastic gradient methods is a central problem in modern machine learning. Recent work has highlighted diagonal linear networks as a simplified yet expressive setting for analyzing the optimization and generalization properties of neural models. In this work, we show that in the high-dimensional regime, stochastic gradient descent on diagonal linear networks is well-approximated by continuous dynamics governed by a stochastic differential equation (SDE), which explicitly decouples the drift from the gradient noise. We further derive a deterministic partial differential equation whose solution propagates the relevant state of the iterates and characterizes the time evolution of a broad class of observable statistics, including the risk, curvature, and other metrics for optimality. Finally, we show that, under a suitable parametrization, the stochastic dynamics are globally well posed and converge exponentially fast to zero risk with high probability, yielding a fully explicit non-asymptotic description of their long-time behavior. Numerical simulations corroborate our theoretical findings.
Begoña García Malaxechebarría, Courtney Paquette, Maryam Fazel +1
University of Washington · McGill University · Google DeepMind +2