On the different regimes of stochastic gradient descent

Significance

The success of deep learning contrasts with its limited understanding. One example is stochastic gradient descent, the main algorithm used to train neural networks. It depends on hyperparameters whose choice has little theoretical guidance and relies on expensive trial-and-error procedures. In this work, we clarify how these hyperparameters affect the training dynamics of neural networks, leading to a phase diagram with distinct dynamical regimes. Our results explain the surprising observation that these hyperparameters strongly depend on the number of data available. We show that this dependence is controlled by the difficulty of the task being learned.

Abstract

Modern deep networks are trained with stochastic gradient descent (SGD) whose key hyperparameters are the number of data considered at each step or batch size $B$, and the step size or learning rate $\eta$. For small $B$ and large $\eta$, SGD corresponds to a stochastic evolution of the parameters, whose noise amplitude is governed by the “temperature” $T\equiv \eta /B$. Yet this description is observed to break down for sufficiently large batches $B\ge B^{\ast }$, or simplifies to gradient descent (GD) when the temperature is sufficiently small. Understanding where these cross-overs take place remains a central challenge. Here, we resolve these questions for a teacher-student perceptron classification model and show empirically that our key predictions still apply to deep networks. Specifically, we obtain a phase diagram in the $B$-$\eta$ plane that separates three dynamical phases: i) a noise-dominated SGD governed by temperature, ii) a large-first-step-dominated SGD and iii) GD. These different phases also correspond to different regimes of generalization error. Remarkably, our analysis reveals that the batch size $B^{\ast }$ separating regimes (i) and (ii) scale with the size $P$ of the training set, with an exponent that characterizes the hardness of the classification problem.

Stochastic gradient descent (SGD), with its variations, has been the algorithm of choice to minimize the loss and train neural networks since the introduction of back-propagation (1–3). When minimizing an empirical loss on a training set of size $P$, SGD consists in estimating the loss gradient using a mini-batch of the data selected randomly at each step. When the size $B$ of the batch is small, SGD becomes a noisy version of gradient descent (GD) and the magnitude of this noise is controlled by the “temperature” $T=\eta /B$, with $\eta$ the learning rate (4). This temperature scale is not a thermodynamic quantity, since in general no thermal equilibrium is reached before the dynamics. However, understanding the scale of hyperparameters where this description holds and SGD noise matters has remained a challenge. Specific questions include i) below which temperature $T_c$ noise becomes irrelevant and the dynamics corresponds to gradient descent? ii) What determines the critical batch size $B^{\ast }$, beyond which SGD is observed not to be controlled by temperature in a variety of settings (5, 6)? This question is of practical importance: after searching for an optimal temperature, practitioners can maximize the batch size up to $B^{\ast }$ while keeping temperature fixed, as large batches can lead to faster computations in practice (7, 8). iii) It was observed that the variation of the network weights during training increases as power laws of both $T$ and $P$, both for deep nets and for simple models like the perceptron (9). Yet, a quantitative explanation of this phenomenon is lacking.

SGD noise has attracted significant research interest, in particular for its connections to loss landscape properties and performance. Recent works describing SGD via a stochastic differential equation (SDE) have emphasized several effects of this noise, including its ability to escape saddles (10–16), to bias the dynamics toward broader minima (17–21) or toward sparser distributions of the network parameters (22–26). Yet, most of these works assume that fresh samples are considered at each time step, and are thus unable to explain the dependence of SGD on the finite size of the training set $P$. Here, we study this dependence on classification problems and we identify three regimes of SGD in deep learning. Following ref. 9, we first consider a perceptron setting where classes are given by a teacher $y(\mathbf{x})=\mathrm{sign}({\mathbf{w}}^{\ast }·\mathbf{x})$, where ${\mathbf{w}}^{\ast }$ is a unit vector, and are learned by a student $f(\mathbf{x})\propto \mathbf{w}·\mathbf{x}$. We solve this problem analytically for the hinge loss, which vanishes if all data are fitted by some margin $\kappa$. Our central results, summarized in Fig. 1A and the phase diagrams therein, are as follows.

Fig. 1.

SGD phase diagrams for different data and architectures. The different datasets considered are a teacher perceptron model with $P=8,192$, dimension $d=128$ and data distribution of Eq. 1 with $\chi =1$ (A.1), $P=32,768$ images of MNIST (B.1) and $P$ = 16,384 images of CIFAR 10 (C.1). The different neural network architectures trained on these datasets correspond respectively to (A.2) a perceptron, for which the output is linear in both the input $\mathbf{x}$ and the weights $\mathbf{w}$, and trained with hinge loss margin $\kappa =2^{-7}$; (B.2) a fully connected network with five hidden layers, $128$ hidden neurons per layer and margin $\kappa =2^{-15}$; (C.2) a CNN made by several blocks composed of depth-wise, point-wise and standard convolutions plus residual connections (more details in SI Appendix, section 1), with margin $\kappa =2^{-15}$. Panels (A.3), (B.3), and (C.3) display the alignment after training in the $\eta ,B$ phase diagram. The black dots correspond to diverging trainings where the algorithm does not converge. We can distinguish the noise-dominated SGD regime, for which the alignment is constant along the diagonals $\frac{\eta }{B}=T$. Within the first-step-dominated SGD, instead, the alignment is constant at constant $\eta$. For small $\eta$, one enters in the gradient descent (GD) regime where the alignment does not depend on $\eta$ and $B$. Taking this value $m_{\mathit{GD}}$ of the alignment as a reference, the black dashed line ${\eta }_c(B)$ delimiting the GD region corresponds to the alignment taking value $2m_{\mathit{GD}}$. The vertical black dashed line guides the eye to indicate the critical batch size $B^{\ast }$. Panels (A.4), (B.4), and (C.4) display the test error again as a function of $\eta ,B$. As expected, the test error is constant along the diagonals $\frac{\eta }{B}=T$ for noise-dominated SGD and constant in the GD regime. For first-step-dominated SGD, the test error can be affected by both $\eta$ and $B$, and can improve at very large batches (Discussion in Sec. 2.A).

Noise-Dominated SGD. For small batches and large learning rates, the dynamics is well described by an SDE with a noise scale $T=\eta /B$. We show that this noise controls the components ${\mathbf{w}}_{\perp }$ of the student weights orthogonal to the teacher, such that at the end of training $\Vert {\mathbf{w}}_{\perp }\Vert \sim T$. Using this result, together with considerations on the angle of the predictor that can fit the entire training set, implies that the weight magnitude after training indeed depends both on $T$ and $P$, as $\Vert \mathbf{w}\Vert \sim T{P}^{\gamma }$, thus explaining the observation of Sclocchi et al. (9). The exponent $\gamma$ characterizes the difficulty of the classification problem.

Gradient Descent. This regime breaks down when the temperature, and therefore the magnitude of orthogonal weights $\Vert {\mathbf{w}}_{\perp }\Vert$, is too small. Then, fitting all data no longer corresponds to a constraint on the angle of the predictor $\Vert {\mathbf{w}}_{\perp }\Vert /\Vert \mathbf{w}\Vert$, but instead to a constraint on $\Vert \mathbf{w}\Vert$ to satisfy the margin. This latter constraint is the most stringent one for $T\le T_c\sim \kappa$. In that regime, temperature can be neglected and the dynamics corresponds to gradient descent, thus answering point (i) above.

First-Step-Dominated SGD. The noise-dominated SGD regime also breaks down when the batch and learning rates are increased at fixed $T$. Indeed, the first step will increase the weight magnitude to $\Vert \mathbf{w}\Vert \sim \eta$, which is larger than the noise-dominated prediction $\Vert \mathbf{w}\Vert \sim T{P}^{\gamma }$ if $B\ge B^{\ast }\equiv P^{\gamma }$, answering (ii) above.

Our second central result is empirical: We find this phase diagram also in deep neural networks. We train using the hinge loss as in the perceptron case (using the cross-entropy loss with early stopping leads to similar performance and weights increase (9, 27)). In this case, it is useful to introduce the quantity ${⟨y(\mathbf{x})f(\mathbf{x})⟩}_{\mathbf{x}}$ characterizing the magnitude with which the output aligns to the task (for the perceptron, it simply corresponds to the quantity $\mathbf{w}·{\mathbf{w}}^{\ast }$). Fig. 1 shows this alignment in the diagram $\eta ,B$ for fully connected (B) and Convolutional networks (C), revealing in each case the three regimes also obtained for the perceptron. As shown below, in both cases, $B^{\ast }$ depends as a power law of the training set size, as predicted.

Related Works.

SGD descriptions. In the high-dimensional setting, several works have analyzed online-SGD in teacher-student problems (28–32). It has been rigorously shown (15) that, for online-SGD in high dimensions, the effective dynamics of summary statistics obey an ordinary differential equation with correction terms given by SGD stochasticity. Our analysis of the perceptron falls into this regime, with fixed points of the dynamics depending explicitly on $T=\eta /B$. We further introduce the effects of a finite training set $P$ by studying when this online description breaks down. Other studies (33, 34) have analyzed the correlations of SGD noise with the training set using dynamical mean-field theory. They consider the regime where the batch size $B$ is extensive in the training set size $P$, while our online description considers the regime where $B$ is much smaller than $P$.

In finite-dimensional descriptions of SGD, theoretical justification for approximating SGD with an SDE has generally relied upon vanishing learning rates (13, 35). This condition is too restrictive for standard deep learning settings and (36) has shown empirically that the applicability of the SDE approximation can be extended to finite learning rates. Other SDE approximations assume isotropic and homogeneous noise (10, 37–39) or a noise covariance proportional to Hessian of the loss (4, 14, 19, 40), while in the online limit of the perceptron model, we can compute the exact dependence of the noise covariance on the model weights. Other studies have observed a link between the effects of SGD noise and the size of the training set (41). In our work, we show that these effects do not depend on a direct increase of the noise scale with the training set size, but rather on the fact a larger $P$ implies that larger weights are needed to fit the data, a process that depends on SGD noise.

Critical batch size. An empirical study at fixed training set size of large-batch training has observed in several learning settings that $B^{\ast }$ is inversely proportional to the signal-to-noise ratio of loss gradients across training samples (8). Our work is consistent with these findings, but most importantly further predicts and tests a nontrivial power law dependence of $B^{\ast }$ with $P$.

1. Noise-Dominated Regime

A. Perceptron Model.

We consider data $\mathbf{x}\in {\mathbb{R}}^d$ with $d\gg 1$ that are linearly separable. Without loss of generality, we chose that the class label $y(\mathbf{x})=\pm 1$ is given by the sign of the first component $y(\mathbf{x})=\text{sign}(x_1)$. ${\{{\mathbf{e}}_i\}}_{i=1,...,d}$ is the canonical basis and ${\mathbf{e}}_1$ corresponds to the teacher direction. The informative component $x_1$ of each datum is independently sampled from the probability distribution

$$\begin{array}{c}\hfill \rho (x_1)={|x_1|}^{\chi }{e}^{-x_1^2/2}/Z,\end{array}$$ [1]

where $Z$ is a normalization constant and $\chi >-1$ (42). The other $d-1$ components ${\mathbf{x}}_{\perp }={[x_i]}_{i=2,...,d}$ are distributed as standard multivariate Gaussian numbers, i.e., ${\mathbf{x}}_{\perp }\sim \mathcal{N}(\mathbf{0},{\mathbf{I}}_{d-1})$, being ${\mathbf{I}}_{d-1}$ the $(d-1)$ identity matrix. The parameter $\chi$ controls the data distribution near the decision boundary ${\mathbf{x}}_1=0$ and how “hard” the classification problem is. Smaller $\chi$ corresponds to harder problems, as more points lie close to the boundary. The case $\chi =0$ corresponds to the Gaussian case.^*

As an architecture, we chose the perceptron $f(\mathbf{w},\mathbf{x})=\mathbf{w}·\mathbf{x}/\sqrt{d}$. The weights are trained by minimizing the hinge loss $L(\mathbf{w})=\frac{1}{P}\sum _{\mu =1}^P{(\kappa -y^{\mu }f(\mathbf{w},{\mathbf{x}}^{\mu }))}^{+}$, where ${(x)}^{+}=max(0,x)$, $\kappa >0$ is the margin, ${\{({\mathbf{x}}^{\mu },y^{\mu }=y({\mathbf{x}}^{\mu }))\}}_{\mu =1,...,P}$ is the training set with size $P\gg d$. We denote ${\mathbf{w}}^t$ the weights obtained at time $t$, defined as the number of training steps times the learning rate, starting from the initialization ${\mathbf{w}}^0=\mathbf{0}$.

The hinge loss is minimized with a SGD algorithm, in which the batch is randomly selected at each time step among all the $P$ data. The learning rate $\eta$ is kept constant during training. The end of training is reached at time $t^{\ast }$ when $L({\mathbf{w}}^{t^{\ast }})=n({\mathbf{w}}^{t^{\ast }})=0$, where $n({\mathbf{w}}^t)$ indicates the fraction of data not fitted at time $t$: $n(\mathbf{w})=\frac{1}{P}\sum _{\mu =1}^P\theta (\kappa -y^{\mu }f(\mathbf{w},{\mathbf{x}}^{\mu }))$, with $\theta (\cdots )$ the Heaviside step function.

B. Stochastic Differential Equation.

In the limit of small batches $B\ll P$, SGD can be described with a stochastic differential equation (SDE) (35):

$$\begin{array}{c}\hfill d{\mathbf{w}}^t=-dt\mathrm{\nabla }L({\mathbf{w}}^t)+\sqrt{T}\sqrt{\mathbf{\Sigma }\left({\mathbf{w}}^t\right)}d{\mathbf{W}}^t,\end{array}$$ [2]

where ${\mathbf{W}}^t$ is a $d$-dimensional Wiener process (Ito’s convention) and $\mathbf{\Sigma }\left(\mathbf{w}\right)/B$ is the covariance of the mini-batch gradients (43), given in SI Appendix, section 1.A. We first make the approximation of substituting the empirical gradient $\mathrm{\nabla }L({\mathbf{w}}^t)$ and this covariance with their population averages $\mathrm{\nabla }L({\mathbf{w}}^t)\approx {\mathbf{g}}^t={\mathbb{E}}_{\mathbf{x}}\left[-{\mathrm{\nabla }}_{\mathbf{w}}l({\mathbf{w}}^t,\mathbf{x})\right]$ and ${\mathbf{\Sigma }}^t\approx {\mathbb{E}}_{\mathbf{x}}\left[{\mathrm{\nabla }}_{\mathbf{w}}l({\mathbf{w}}^t,\mathbf{x})\otimes {\mathrm{\nabla }}_{\mathbf{w}}l({\mathbf{w}}^t,\mathbf{x})\right]-{\mathbf{g}}^{\mathbf{t}}\otimes {\mathbf{g}}^{\mathbf{t}}$. This approximation does not include finite training set effects. Our strategy below is to estimate the time where such a simplified description breaks down; the solution of the SDE at this time provides the correct asymptotic behaviors for the network at the end of training, as we observe experimentally.

B.1. Asymptotic online-SDE dynamics.

We decompose the student weights as $\mathbf{w}=w_1{\mathbf{e}}_1+{\mathbf{w}}_{\perp }$ and we study the dynamics of both the scalar variables $w_1^t$ (the “overlap” between the student and teacher (28)), and the magnitude of the weights in the perpendicular directions $\Vert {\mathbf{w}}_{\perp }^t\Vert$. Given that we consider a data distribution that is radially symmetric in ${\mathbf{x}}_{\perp }$, but not in $\mathbf{x}$, $w_1$, and $\Vert {\mathbf{w}}_{\perp }\Vert$ are the natural order parameters in our problem. Therefore, the $d$-dimensional dynamics of $\mathbf{w}$ can be studied through a two-dimensional summary statistics.

One finds that the online-SDE dynamics of $w_1^t$ and $\Vert {\mathbf{w}}_{\perp }^t\Vert$, using Ito’s lemma (SI Appendix, section 6.B), can be written as

$$\begin{array}{c}\hfill \begin{array}{c}\hfill \left\{\begin{array}{cc}d{w}_1^t\hfill & =dt{g}_1^t+\sqrt{\frac{T}{d}}{\sigma }_1^{t}d{\tilde{W}}_1^t\hfill \\ d\Vert {\mathbf{w}}_{\perp }^t\Vert \hfill & =dt\left[g_{\perp }^t+\frac{T}{2\Vert {\mathbf{w}}_{\perp }^t\Vert }{n}^t\right]+\sqrt{\frac{T}{d}}{\sigma }_2^{t}d{\tilde{W}}_2^t\hfill \end{array}\right.,\end{array}\end{array}$$ [3]

where ${\tilde{W}}_1^t$ and ${\tilde{W}}_2^t$ are Wiener processes and the expressions for $g_1^t$, $g_{\perp }^t$, $n^t$, ${\sigma }_1^t$, ${\sigma }_2^t$, reported in SI Appendix, section 6.A, are functions of $w_1^t$ and $\Vert {\mathbf{w}}_{\perp }^t\Vert$ through the time-dependent ratios:

$$\begin{array}{c}\hfill \lambda =\frac{w_1^t}{\Vert {\mathbf{w}}_{\perp }^t\Vert },r=\frac{\kappa \sqrt{d}}{\Vert {\mathbf{w}}_{\perp }^t\Vert }.\end{array}$$ [4]

The quantity $\lambda$ measures the angle $\theta$ between the student and the teacher directions, since $\theta =arctan{\lambda }^{-1}$. The ratio $r$ compares the hinge loss margin $\kappa$ with the magnitude of the orthogonal components $\Vert {\mathbf{w}}_{\perp }^t\Vert$.

In the limit $d\gg 1$, the stochastic part of Eq. 3 is negligible, as we show in SI Appendix, section 4.B. The variables $w_1^t$ and $\Vert {\mathbf{w}}_{\perp }^t\Vert$, therefore, have a deterministic time evolution given by the deterministic part of Eq. 3. This result has been proved more generally for the summary statistics of different learning tasks in high dimensions (15). Note that, even if the stochastic fluctuations are negligible, the SGD noise affects the dynamics through the term $\frac{T}{2\Vert {\mathbf{w}}_{\perp }^t\Vert }{n}^t$ in the evolution of $\Vert {\mathbf{w}}_{\perp }^t\Vert$.

The term $g_1^t$ (SI Appendix, section 6.A) is always positive and vanishes in the limit $\lambda \to \infty$, which determines the fixed point of Eq. 3. Therefore, we consider its vicinity

$$\begin{array}{c}\hfill \lambda =\frac{w_1^t}{\Vert {\mathbf{w}}_{\perp }^t\Vert }\to \infty ,\end{array}$$ [5]

corresponding to a vanishing angle $\theta =arctan{\lambda }^{-1}\sim {\lambda }^{-1}$ between the student and teacher directions. In addition, we consider the limit

$$\begin{array}{c}\hfill r=\frac{\kappa \sqrt{d}}{\Vert {\mathbf{w}}_{\perp }^t\Vert }\to 0,\end{array}$$ [6]

and argue below why this limit corresponds to the noise-dominated regime of SGD.

Under the conditions of Eqs. 5 to 6, the deterministic part of Eq. 3 reads:

$$\begin{array}{c}\hfill \begin{array}{c}\hfill \left\{\begin{array}{cc}d{w}_1^t\hfill & =dt{\lambda }^{-\chi -2}\frac{c_1}{\sqrt{d}}\hfill \\ d\Vert {\mathbf{w}}_{\perp }^t\Vert \hfill & =dt{\lambda }^{-\chi -1}\left[-\frac{1}{\sqrt{d}\sqrt{2\pi }}+\frac{T}{2\Vert {\mathbf{w}}_{\perp }^t\Vert }{c}_n\right]\hfill \end{array}\right.,\end{array}\end{array}$$ [7]

with constants $c_1$, $c_n$ (SI Appendix, section 6.A). Solving Eq. 7 gives

$$\begin{array}{c}\hfill \Vert {\mathbf{w}}_{\perp }^t\Vert \propto T\sqrt{d},w_1^t\sim T\sqrt{d}{\left(\frac{t}{\mathit{Td}}\right)}^{\frac{1}{3+\chi }}.\end{array}$$ [8]

Therefore, the orthogonal component $\Vert {\mathbf{w}}_{\perp }^t\Vert$ tends to a steady state proportional to the SGD temperature $T$, while the informative one $w_1^t$ grows as a power law of time. Eq. 8 implies that $\lambda \gg 1$ is obtained in the large time limit $t\gg Td$, while $r\ll 1$ corresponds to $T\gg \kappa$, and therefore holds at sufficiently high temperatures.

B.2. SDE breakdown and solution obtained at the end of training.

Due to the finiteness of $P$, the online solution of Eq. 8 is expected to be no longer valid at some time $\widehat{t}<t^{\ast }$. In SI Appendix, section 5, we provide an argument, confirmed empirically, predicting that $\widehat{t}$ is reached when the number of training points contributing to the hinge loss gradient is of order $\mathcal{O}(d)$. This is obtained by assuming that the student perceptron follows the dynamics Eq. 8 and applying the central limit theorem to the empirical gradient $\mathrm{\nabla }L({\mathbf{w}}^t)$. One finds that the magnitude of its population average follows $\Vert {\mathbb{E}}_{\mathbf{x}}\left[\mathrm{\nabla }L({\mathbf{w}}^t)\right]\Vert \sim \frac{1}{\sqrt{d}}n({\mathbf{w}}^t)$ while that of its finite-$P$ fluctuations is given by $\frac{1}{\sqrt{P}}\sqrt{n({\mathbf{w}}^t)}$. For $P$ finite but much larger than $d$, the average gradient is much larger than its fluctuations as long as the fraction of unfitted training points $n({\mathbf{w}}^t)$ satisfies $n({\mathbf{w}}^t)\gg \frac{d}{P}$. As the training progresses, more data points get fitted and $n({\mathbf{w}}^t)$ decreases until reaching the condition

$$\begin{array}{c}\hfill n({\mathbf{w}}^{\widehat{t}})\sim \mathcal{O}\left(\frac{d}{P}\right).\end{array}$$ [9]

After this point, the empirical gradient and the population one greatly differ, and the online dynamics is no longer a valid description of the training dynamics. This is the time beyond which test and train errors start to differ. At leading order in $\lambda \to \infty$ as $r\to 0$, we have that $n({\mathbf{w}}^t)\sim {\lambda }^{-\chi -1}$. In fact, for a density of points $\rho (x_1)\sim x_1^{\chi }$ at small $x_1$, $n({\mathbf{w}}^t)\sim {\theta }^{\chi +1}\sim {\lambda }^{-\chi -1}$ is the fraction of points with coordinate $x_1$ smaller than $\theta \sim {\lambda }^{-1}$.^†

Therefore, using Eq. 8,

$$\begin{array}{c}\hfill n({\mathbf{w}}^t)\sim {\lambda }^{-\chi -1}\sim {\left(\frac{t}{\mathit{Td}}\right)}^{-\frac{\chi +1}{\chi +3}}.\end{array}$$ [10]

Note that, for $t\ll \widehat{t}$, $n({\mathbf{w}}^t)$ corresponds to the test error. Neglecting the dependence in $d$, we finally obtain:

$$\begin{array}{c}\hfill \widehat{t}\sim T{P}^b,\Vert {\mathbf{w}}^{\widehat{t}}\Vert \sim w_1^{\widehat{t}}\sim T{P}^{\gamma },\end{array}$$ [11]

where $b=1+\frac{2}{1+\chi }$ and $\gamma =\frac{1}{1+\chi }$.

In SI Appendix, Fig. S9, we show experimentally that the asymptotic solution of the online SDE is a valid description of SGD up to $\widehat{t}$, as predicted by Eq. 9. In SI Appendix, Figs. S10 and S11, we observe empirically that the power law scaling in $T$ and $P$ of the stopping time $t^{\ast }$ and of $\Vert {\mathbf{w}}^{t^{\ast }}\Vert$ are the same as those of $\widehat{t}$ and $\Vert {\mathbf{w}}^{\widehat{t}}\Vert$, therefore Eq. 11 also hold to characterize the end of training.

C. T Condition for the Noise-Dominated Regime.

To fit a data point $({\mathbf{x}}^{\mu },y^{\mu })$, the perceptron weights must satisfy the condition $y^{\mu }\frac{w_1{x}_1^{\mu }}{\sqrt{d}}+y^{\mu }\frac{{\mathbf{w}}_{\perp }·{\mathbf{x}}_{\perp }^{\mu }}{\sqrt{d}}\ge \kappa$, which corresponds to ref. 9:

$$\begin{array}{c}\hfill \frac{w_1}{\Vert {\mathbf{w}}_{\perp }\Vert }\ge \frac{1}{|x_1^{\mu }|}\left(\frac{\kappa \sqrt{d}}{\Vert {\mathbf{w}}_{\perp }\Vert }+c^{\mu }\right),\end{array}$$ [12]

where we have defined the random quantity $c^{\mu }=-y^{\mu }\frac{{\mathbf{w}}_{\perp }}{\Vert {\mathbf{w}}_{\perp }\Vert }·{\mathbf{x}}_{\perp }^{\mu }=\mathcal{O}(1)$, and where $\frac{\kappa \sqrt{d}}{\Vert {\mathbf{w}}_{\perp }\Vert }\propto \frac{\kappa }{T}$ in the noise-dominated regime where Eq. 8 applies.

For $T\gg \kappa$, $\frac{\kappa \sqrt{d}}{\Vert {\mathbf{w}}_{\perp }\Vert }$ is negligible with respect to $c^{\mu }=\mathcal{O}(1)$ and Eq. 12 now becomes $\frac{w_1}{\Vert {\mathbf{w}}_{\perp }\Vert }\ge \frac{c^{\mu }}{|x_1^{\mu }|}$, independent of the margin $\kappa$. In this case, therefore, fitting $({\mathbf{x}}^{\mu },y^{\mu })$ is constrained by the SGD noise which inflates the noninformative weight component $\Vert {\mathbf{w}}_{\perp }\Vert$. For $T\ll \kappa$, instead, fitting $({\mathbf{x}}^{\mu },y^{\mu })$ is constrained by the margin $\kappa$ and the SGD noise is negligible in Eq. 12, implying that the temperature delimiting the noise-dominated regime of SGD follows:

$$\begin{array}{c}\hfill T_c\propto \kappa .\end{array}$$ [13]

For $T\ll T_c$, the final magnitude of $\Vert {\mathbf{w}}_{\perp }\Vert$ is independent of $T$ as expected. Instead, one finds that it is proportional to $\kappa$ (SI Appendix, Fig. S11B).

2. Critical Batch Size $B^{\ast }$

In the small batch regime $B\ll B^{\ast }$, for a given temperature $T=\eta /B$, different SGD simulations converge to the asymptotic trajectory given by the online-SDE, as observed in Fig. 2 showing SGD trajectories in the $(w_1,w_{\perp })$ plane, at fixed $T$ and varying batch size. This is not the case in the large batch regime $B\gg B^{\ast }$. In fact, for fixed $T$, a larger batch size corresponds to a larger learning rate and therefore a larger initial step. If the initial step is too large, SGD would converge to a zero-loss solution before being able to converge to the online-SDE asymptotic trajectory. This situation is vividly apparent in Fig. 2. Therefore, we can estimate the scale of the critical batch-size $B^{\ast }$ by comparing the first step $w_1^{\eta }\propto \eta$ with the final value $w_1^{t^{\ast }}\sim \frac{\eta }{B}{P}^{\gamma }$ in the noise-dominated regime. When $w_1^{\eta }\ll w_1^{t^{\ast }}$, SGD converges to the asymptotic online-SDE and the final minimum depends on $T=\eta /B$. When $w_1^{\eta }\gg w_1^{t^{\ast }}$, SGD does not converge to the asymptotic online-SDE and the final weight magnitude depends only on $w_1^{\eta }$, as shown in Fig. 3. The condition $w_1^{\eta }\sim w_1^{t^{\ast }}$, that is $\eta \sim \frac{\eta }{B^{\ast }}{P}^{\gamma }$ gives

$$\begin{array}{c}\hfill B^{\ast }\sim P^{\gamma },\end{array}$$ [14]

Fig. 2.

Dynamical trajectories of SGD in the $(w_1,w_{\perp })$ plane, at fixed $T$ and varying batch size as indicated in caption. Black circles indicate the first step of SGD, black stars indicate the last one. For small enough batches (and therefore small learning rates), trajectories converge to the online SDE solution (black dashed line). For large batches, this is not true anymore, and the final magnitude of the weights increases with batch size. The location of stopping weights corresponds to zero loss, which can be approximately determined by measuring the hinge loss values $L_{\mathit{train}}(w_1,w_{\perp })$ (shown in color) computed as a function of the perceptron weights $\mathbf{w}=w_1{\mathbf{e}}_1+w_{\perp }\mathit{\xi }$. Here, $\mathit{\xi }$ is a $(d-1)$-dimensional Gaussian random vector. The white area corresponds to interpolating solutions $L_{\mathit{train}}=0$ in this simplified set-up. For full-batch, we observe that $\mathbf{w}$ can land directly in the white area and therefore fit the data with at most few steps. This behavior affects the test error when $\eta$ is large (Fig. 1A4). Data correspond to $P$ = 16,384, $d=128$, $\kappa =0.01$, $\chi =1$, $T=2$.

Fig. 3.

For large learning rates, the dynamics of the alignment is different in the small batch and large batch regimes. (Main panel) Perceptron: in this case, the alignment is proportional to the student component $w_1$; data for fixed $\eta =512$, same setting as Fig. 1. For small $B$, $w_1$ grows during the training dynamics, while, for large $B$, its final value is reached after a single training step. (Inset) Fully connected network on MNIST, small margin ($\kappa =2^{-15}$), fixed $\eta =16$, same setting as Fig. 1: For small and large batch, the alignment shows a similar dynamics to the perceptron case, although for large batch it reaches its final value after some training steps (and not just a single step).

with $\gamma =\frac{1}{1+\chi }$ depending on the data distribution. The relationship of Eq. 14 is well verified from the empirical data of Fig. 4A.

Fig. 4.

The critical batch size $B^{\ast }$ depends on the size of the training set as predicted by Eq. 14. (A) $w_1$ at the end of training for the perceptron. Inset: $w_1$ depends on $\eta$, $B$, and $P$ for small $B$, while it only depends on $\eta$ for large $B$. We observe that the cross-over between small and large $B$ depends on $P$. Main panel: the curves collapse in one curve by rescaling $w_1$ by $\eta$ and $B$ by $B^{\ast }\propto P^{\gamma }$, $\gamma =\frac{1}{1+\chi }$. This is consistent with $w_1\sim \frac{\eta }{B}{P}^{\gamma }$ for small $B$ (Eq. 11) and $w_1\sim \eta$ for large $B$ (Section 2). (B) Fully connected network on parity MNIST. Alignment ${⟨y(\mathbf{x})f(\mathbf{x})⟩}_{\mathbf{x}}$ at the end of training (measured as in Eq. 15). Inset: as for $w_1$ in panel (A), ${⟨y(\mathbf{x})f(\mathbf{x})⟩}_{\mathbf{x}}$ depends on $\eta$, $B$, and $P$ for small $B$, and it only depends on $\eta$ for large $B$. Main panel: rescaling $B$ by $P^{0.2}$ aligns the cross-over batch size at different $P$s, suggesting a dependence $B^{\ast }\propto P^{0.2}$. The curves are approximately collapsed by the rescaling of the y-axis as ${⟨y(\mathbf{x})f(\mathbf{x})⟩}_{\mathbf{x}}/\eta$.

This finding is consistent with $B^{\ast }$ being inversely proportional to the signal-to-noise ratio of loss gradients across training samples as proposed in ref. 8. In fact, in our setting, increasing the gradient noise scale by a factor $\sigma$ corresponds to the substitution $T\to \sigma T$, which transforms Eq. 14 into $B^{\ast }\sim \sigma P^{\gamma }$. Moreover, at large $P$, increasing $\chi$ reduces the exponent $\gamma$ and $B^{\ast }$. In fact, larger $\chi$ corresponds to fewer training points close to the decision boundary and therefore a larger gradients’ signal-to-noise ratio.

A. Performance in Large Step SGD.

The first training step leads to $w_1^{\eta }=\frac{\eta }{B}\sum _{\mu \in {\mathbb{B}}_1}\frac{|x_1^{\mu }|}{\sqrt{d}}$ and ${\mathbf{w}}_{\perp }^{\eta }=\frac{\eta }{B}\sum _{\mu \in {\mathbb{B}}_1}{y}^{\mu }\frac{{\mathbf{x}}_{\perp }^{\mu }}{\sqrt{d}}$, where ${\mathbb{B}}_1$ is the first sampled batch. ${\mathbf{w}}_{\perp }^{\eta }$ is a zero-mean random vector of norm $\Vert {\mathbf{w}}_{\perp }^{\eta }\Vert =\mathcal{O}(\frac{\eta }{\sqrt{B}})$, while $w_1^{\eta }=\mathcal{O}(\eta )$. We can distinguish several regimes:

• (i)If $w_1^{\eta }|x_1^{\mu }|\sim \eta |x_1^{\mu }|\gg \kappa \forall \mu$, then the margin $\kappa$ can be neglected. From extreme value statistics $\underset{\mu }{min}|x_1^{\mu }|=\mathcal{O}(P^{-1/(1+\chi )})$, thus this condition is saturated when $\eta \sim \kappa P^{1/(1+\chi )}$. It corresponds to the horizontal dashed line in the diagrams of Fig. 1.
• (ii)Above this line in the regime $\eta \gg \kappa P^{1/(1+\chi )}$, the fraction of unfitted points as well as the test error after one step are given by the angle of the predictor ${\lambda }^{\eta }=\frac{w_1^{\eta }}{\Vert {\mathbf{w}}_{\perp }^{\eta }\Vert }\sim \mathcal{O}(\sqrt{B})$. Following Eq. 10, $n({\mathbf{w}}^{\eta })\sim {({\lambda }^{\eta })}^{-\chi -1}\sim B^{-(\chi +1)/2}$. If $n({\mathbf{w}}^{\eta })\ll 1/P$ or equivalently $B\gg P^{2/(1+\chi )}$, then with high probability all points are fitted after one step and the dynamics stops; the final test error is then of order $ϵ\sim B^{-(\chi +1)/2}$. Since $B\le P$, this condition can only occur for $\chi \ge 1$. Note that the error is smaller than in the noise-dominated regime where $ϵ\sim 1/P$.

For the marginal case $\chi =1$ in Fig. 2, we observe that the full batch case $B$ = 16,384 reaches (nearly) zero training loss after the first training step. Correspondingly, in the phase diagram of Fig. 1A4, the lowest test error is achieved for full batch in the first-step-dominated regime.

B. Notes on Variations of SGD.

Our analysis does not consider variations of SGD sometimes used in practice, such as the addition of momentum, adaptive learning rates, or weight decay. A detailed extension of the present approach to these cases is an interesting work left for the future. In SI Appendix, section 2, we provide evidence supporting that adding momentum is akin to rescaling the learning rate by a momentum-dependent factor. By contrast, adding weight decay in the noise-dominated regime is akin to stopping the dynamics early, and it is therefore lowering performance in the case of the perceptron with no label noise. For large batch sizes, the first-step-dominated dynamics eventually converge to the noise-dominated one at a time-scale depending on the weight-decay coefficient (SI Appendix, section 2.B).

3. Experiments for Deep Networks

We consider a $5$-hidden-layer fully connected architecture with GELU activation functions (44) that classifies the parity MNIST dataset (even vs. odd digits) and a deep convolutional architecture (MobileNet) which classifies the CIFAR10 images (animals vs. the rest). As before, training corresponds to minimizing the hinge loss with SGD at constant $\eta$, until reaching zero training loss, see SI Appendix, section 1 for further details.

For deep neural networks, the alignment between the network output $f(\mathbf{x})$ and the data labels $y(\mathbf{x})\in \{+1,-1\}$ cannot simply be written as a projection $w_1$ of the weights in some direction. Instead, we define it as

$$\begin{array}{c}\hfill {⟨y(\mathbf{x})f(\mathbf{x})⟩}_{\mathbf{x}}=\frac{1}{|S_{\mathit{test}}|}{\displaystyle \sum _{{\mathbf{x}}^{\nu }\in S_{\mathit{test}}}}y({\mathbf{x}}^{\nu })f({\mathbf{x}}^{\nu }),\end{array}$$ [15]

where the average is made over the test set $S_{\mathit{test}}$. For the perceptron, $w_1$ and ${⟨y(\mathbf{x})f(\mathbf{x})⟩}_{\mathbf{x}}$ are proportional.

A. Small Margin $\kappa$.

An interesting parameter to vary is the margin $\kappa$, which is equivalent to changing the initialization scale of the network. For very small $\kappa$ and gradient flow, tiny changes of weights can fit the data, corresponding to the lazy regime where neural nets behave as kernel methods (45–47). However, by cranking up the learning rate one escapes this regime (48), which is the case in our set-up (Fig. S3 in SI Appendix, section 3). Our central results are as follows:

• (i)The phase diagram predicted for the perceptron holds true, as shown Fig. 1B3 and C3 representing the alignment in the ($B$, $\eta$) plane. We observe a noise-dominated SGD where the alignment depends on the ratio $T=\eta /B$; a first-step-dominated SGD where the alignment is constant at a given $\eta$; and a GD regime at small $\eta$ where the alignment does not depend on $\eta$ and $B$.
• (ii)These three regimes also delineate different behaviors of the test error, as shown in Fig. 1B4 and C4.
• (iii)Fig. 3, Inset confirms the prediction that in the first-step-dominated SGD regime, the alignment builds up in a few steps. By contrast, in the noise-dominated SGD regime, it builds up slowly over time.
• (iv)The critical batch $B^{\ast }$ separating these two regimes indeed depends on a power law of the training set $P$. This result is revealed in Fig. 4B, reporting the alignment at the end of training as a function of the batch size $B$. The Inset shows that the alignment depends on the batch size for small $B$, while it only depends on $\eta$ for large $B$. The cross-over $B^{\ast }$ between the two regimes depends on $P$. It is estimated in the Main panel by rescaling the x-axis by $P^{-0.2}$ so as to align these cross-overs, indicating a dependence $B^{\ast }\sim P^{0.2}$. The same phenomenology is observed for the perceptron in Fig. 4A, with the relationship $B^{\ast }\sim P^{\frac{1}{1+\chi }}$ given by Eq. 14. We can use this measurement of $B^{\ast }$ in neural networks to estimate a value of $\chi$ for the corresponding dataset. From the data of Fig. 4B, we obtain ${\chi }_{\mathit{MNIST}}\approx 4$ for the MNIST dataset, while from SI Appendix, Fig. S8A ${\chi }_{\mathit{CIFAR}}\approx 1.5$ for CIFAR10. These estimates are relatively close to the values $6$ and $1.5$, for MNIST and CIFAR10 respectively, independently measured in a study of kernel ridge regression (42). In that context, the value of $\chi$ directly affects the kernel eigenvalue spectrum and therefore the task difficulty. These estimates suggest that our toy model for $\chi >0$ captures some properties of image classification tasks.

B. Large Margin $\kappa$.

In the alternative case where the margin is large, the predictor has to grow in time or “inflate” in order to become of the order of the margin and fit the data (45, 47, 49). As a result, the weights and the magnitude of the gradients initially grow in time. In our set-up, this regime can be obtained choosing $\kappa =1$, because at initialization the output function is small.^‡ Such an inflation is not captured by the perceptron model. As a consequence, reasoning on the magnitude of the first gradient step may not be appropriate, and the finding (iii) above is not observed as shown in SI Appendix, Fig. S7A. Yet, after inflation occurs, as for the perceptron the output still needs to grow even further in the good direction to overcome the noise from SGD, the more so the larger the training set. This effect is observed both for CNNs and fully-connected nets on different datasets and, in the latter case, strongly correlates with performance (9).

For fully connected nets learning MNIST, observations are indeed similar to the small margin case: (i,ii) a phase diagram suggesting three regimes affecting performance is shown in SI Appendix, Fig. S5 and (iv) a critical batch size $B^{\ast }$ that again appears to follow a power law of the training set size, as $B^{\ast }\sim P^{0.4}$ as shown in SI Appendix, Fig. S7B. Interestingly, using early-stopping or not does not affect performance, as shown in SI Appendix, Fig. S5D. It is consistent with our analysis of weight changes, which for the perceptron are similar at early stopping or at the end of training. For CNNs learning CIFAR10, the picture is different. Although three learning regimes may be identifiable in SI Appendix, Fig. S6, most of the dependence of performance with $\eta ,B$ is gone when using early stopping as shown in SI Appendix, Fig. S6D. It suggests that effects other than the growth of weights induced by SGD control performance in this example, as discussed below.

Conclusion

In deep nets, the effects of SGD on learning are known to depend on the size of the training set. We used a simple toy model, the perceptron, to explain these observations and relate them to the difficulty of the task. SGD noise increases the dependence of the predictor along incorrect input directions, a phenomenon that must be compensated by aligning more strongly toward the true function. As a result, alignment and weight changes depend on both $T$ and $P$. If temperature is too small, this alignment is instead fixed by overcoming the margin and SGD is equivalent to GD. If the batch size is larger than a $P$-dependent $B^{\ast }$, the weight changes are instead governed by a the first few steps of gradient descent. As one would expect, the alignment magnitude correlates with performance. It was observed in several cases in ref. 9 and is also reflected by our observation that different alignment regimes correspond to different regimes of performance. In SI Appendix, section 3, we discuss why it is so in the simple example of a teacher perceptron learned by a multi-layer network. As shown in SI Appendix, Fig. S4, as weights align more strongly in the true direction as SGD noise increases, the network becomes less sensitive to irrelevant directions in input space, thus performing better. Yet, depending on the data structure, a strong alignment could be beneficial or not- a versatility of outcome that is observed (5, 9, 40, 50). Obviously, other effects of SGD independent on training set size may affect further performance, such as escaping saddles (10–16), biasing the dynamics toward broader minima (17–21, 43, 50–58) or finding sparser solutions (22, 24–26). We have shown examples where the alignment effects appear to dominate, and others where SGD instead is akin to a regularization similar to early stopping, which is the case for the CNN in SI Appendix, Fig. S6D—a situation predicted in some theoretical approaches, see e.g., ref. 24. Determining which effect of SGD most strongly affects performance given the structure of the task and the architecture is an important practical question for the future. Our work illustrates that subtle aspects enter this equation—such as the size of the training set or the density of points near the decision boundary.

Supplementary Material

Appendix 01 (PDF)

Data, Materials, and Software Availability

Experiments’ code is available at https://github.com/pcsl-epfl/regimes_of_SGD (59).