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.