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.
Figures & tables
Figure 1. The left panel illustrates the evolution of the regularized loss functions across networks with deterministic step sizes as specified in ( 27 ) . The right panel presents the Frobenius norms of individual weight matrices for networks with depths 3 and 4 during training.
Figure 2. The left panel displays the evolution of the regularized losses across networks with deterministic stepsizes that do not meet requirements in Theorem 13 . The right panel shows the Frobenius norms of individual matrices for networks of depths 2, 3, and 4 during training.
Figure 3. The left panel illustrates the evolution of the regularized loss across networks with adaptive stepsizes ( 28 ) . The right panel presents the Frobenius norms of individual weight matrices for networks of depth 3 and 4.
Figure 4. The left panel shows the evolution of the regularized loss across networks with the modified adaptive step sizes ( 29 ) . The right panel displays the Frobenius norms of individual weight matrices for networks of depth 2, 3, and 4 during training with the same adaptive step sizes.
Figure 5. Convergence comparison for deterministic versus adaptive stepsizes on a linear network of depth 3 with parameters β=5 , λ=5⋅10−5 , r=5 , and mini-batch size T=30 .
Figure 6. Training nonlinear neural networks: evolution of the regularized loss functions across networks under deterministic step sizes in ( 27 ) on the left panel. Frobenius norms of individual weight matrices of the corresponding networks for depth 3, and 4 on the right panel.
Figure 7. Training nonlinear neural networks: evolution of the regularized loss functions across networks under adaptive stepsizes in ( 28 ) on the left panel. Frobenius norms of individual weight matrices of the corresponding networks for depth 3 and 4 on the right panel.
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
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.
Guillaume Corlouer, Avi Semler, Alexander Strang +1
Moirai · University of Oxford · University of California, Berkeley +1
Modern deep learning has been shown to operate using learning rates far larger than those justified by classical optimization theory. Most prior analyses of this phenomenon focus on deterministic gradient descent, leaving the stochastic setting largely unexplored. In this work, we provide sharp convergence guarantees for Stochastic Gradient Descent (SGD) applied to the multiclass cross-entropy loss, for both linear classifiers and two-layer neural networks. We show that the stochasticity of SGD may cause the dynamics to alternate between an edge-of-stability regime that is dominated by curvature-driven oscillations, and a stable regime in which the expected loss decreases at a controlled rate. Despite that, we prove that SGD stabilizes the dynamics, ensuring that the iterates return to stability in a fixed number of iterations and allowing convergence in the best-iterate sense even with large learning rates. Experiments validate our theoretical findings and illustrate the benefits of SGD in the large-stepsize regime.