Strength of Minibatch Noise in SGD

1 INTRODUCTION . Stochastic gradient descent ( SGD ) is the simple and efficient optimization algorithm behind the success of deep learning ( Allen-Zhu et al. , 2019 ; Xing et al. , 2018 ; Zhang et al. , 2018 ; Wang et al. , 2020 ; He and Tao , 2020 ; Liu et al. , 2021 ; Simsekli et al. , 2019 ; Wu et al. , 2020 ) . Minibatch noise , also known as the SGD noise , is the primary type of noise in the learning dynamics of neural networks . Practically , minibatch noise is unavoidable because a modern computer ’ s memory is limited while the size of the datasets we use is large ; this demands the dataset to be split into “ minibatches '' for training . At the same time , using minibatch is also a recommended practice because using a smaller batch size often leads to better generalization performance ( Hoffer et al. , 2017 ) . Therefore , understanding minibatch noise in SGD has been one of the primary topics in deep learning theory . Dominantly many theoretical studies take two approximations : ( 1 ) the continuous-time approximation , which takes the infinitesimal step-size limit ; ( 2 ) the Hessian approximation , which assumes that the covariance matrix of the SGD noise is equal to the Hessian H . While these approximations have been shown to provide some qualitative understanding , the limitation of these approximations is not well understood . For example , it is still unsure when such approximations are valid , which hinders our capability to assess the correctness of the results obtained by approximations . In this work , we fill this gap by deriving analytical formulae for discrete-time SGD with arbitrary learning rates and exact minibatch noise covariance . In summary , the main contributions are : ( 1 ) we derive the strength and the shape of the minibatch SGD noise in the cases where the noise for discrete-time SGD is analytically solvable ; ( 2 ) we show that the SGD noise takes a different form in different kinds of minima and propose general and more accurate approximations . This work is organized as follows : Sec . 2 introduces the background . Sec . 3 discusses the related works . Sec . 4 outlines our theoretical results . Sec . 5 derives new approximation formulae for SGD noises . In Sec . 6 , we show how our results can provide practical and theoretical insights to problems relevant to contemporary machine learning research . For reference , the relationship of this work to the previous works is shown in Table 1 . 2 BACKGROUND . In this section , we introduce the minibatch SGD algorithm . Let { xi , yi } Ni=1 be a training set . We can define the gradient descent ( GD ) algorithm for a differentiable loss function L as wt =wt−1 − λ∇wL ( w , { x , y } ) , where λ is the learning rate and w ∈ RD is the weights of the model . We consider an additive loss function for applying the minibatch SGD . Definition 1 . A loss function L ( { xi , yi } Ni=1 , w ) is additive if L ( { xi , yi } Ni=1 , w ) = 1 N ∑ N i=1 ` ( xi , yi , w ) for some differentiable , non-negative function ` ( ⋅ ) . This definition is quite general . Most commonly studied and used loss functions are additive , e.g. , the mean-square error ( MSE ) and cross-entropy loss . For an additive loss , the minibatch SGD with momentum algorithm can be defined . Definition 2 . The minibatch SGD with momentum algorithm by sampling with replacement computes the update to the parameter w with the following set of equations : ⎧⎪⎪⎪⎪ ⎨ ⎪⎪⎪⎪⎩ ĝt = 1 S ∑i∈Bt ∇ ` ( xi , yi , wt−1 ) ; mt = µmt−1 + ĝt ; wt =wt−1 − λmt . ( 1 ) where µ ∈ [ 0,1 ) is the momentum hyperparameter , S ∶= ∣Bt∣ is the minibatch size , and the set Bt = { i1 , ... iS } are S i.i.d . random integers sampled uniformly from [ 1 , N ] . One can decompose the gradient into a deterministic plus a stochastic term . Note that EB [ ĝt ] = ∇L is equal to the gradient for the GD algorithm . We use EB ( ⋅ ) to denote the expectation over batches , and use Ew ( ⋅ ) to denote the expectation over the stationary distribution of the model parameters . Therefore , we can write ĝt = EB [ ĝt ] + ηt , where ηt ∶= 1S ∑i∈Bt ∇ ` ( xi , yi , wt−1 ) − EB [ ĝt ] is the noise term ; the noise covariance is C ( wt ) ∶= cov ( ηt , ηt ) . Of central importance to us is the averaged asymptotic noise covariance C ∶= limt→∞Ewt [ C ( wt ) ] . Also , we consider the asymptotic model fluctuation Σ ∶= limt→∞ cov ( wt , wt ) . Σ gives the strength and shape of the fluctuation of w around a local minimum and is another quantity of central importance to this work . Throughout this work , C is called the “ noise '' and Σ the “ fluctuation '' . 3 RELATED WORKS . Noise and Fluctuation in SGD . Deep learning models are trained with SGD and its variants . To understand the parameter distribution in deep learning , one needs to understand the stationary distribution of SGD ( Mandt et al. , 2017 ) . Sato and Nakagawa ( 2014 ) describes the stationary distribution of stochastic gradient Langevin dynamics using discrete-time Fokker-Planck equation . Yaida ( 2019 ) connects the covariance of parameter Σ to that of the noise C through the fluctuationdissipation theorem . When Σ is known , one may obtain by Laplace approximation the stationary distribution of the model parameter around a local minimum w∗ as N ( w∗ , Σ ) . Therefore , knowing Σ can be of great practical use . For example , it has been used to estimate the local minimum escape efficiency ( Zhu et al. , 2019 ; Liu et al. , 2021 ) and argue that SGD prefers a flatter minimum ; it can also be used to assess parameter uncertainty and prediction uncertainty when a Bayesian prior is specified ( Mandt et al. , 2017 ; Gal and Ghahramani , 2016 ; Pearce et al. , 2020 ) . Empirically , both the fluctuation and the noise are known to crucially affect the generalization of a deep neural network . Wu et al . ( 2020 ) shows that the strength and shape of the Σ due to the minibatch noise lead to better generalization of neural networks in comparison to an artificially constructed noise . Hessian Approximation of the Minibatch Noise . However , it is not yet known what form C and Σ actually take for SGD in a realistic learning setting . Early attempts assume isotropic noise in the continuous-time limit ( Sato and Nakagawa , 2014 ; Mandt et al. , 2017 ) . In this setting , the noise is an isotropic Gaussian with C ∼ ID , and Σ is known to be proportional to the inverse Hessian H−1 . More recently , the importance of noise structure was realized ( Hoffer et al. , 2017 ; Jastrzebski et al. , 2018 ; Zhu et al. , 2019 ; HaoChen et al. , 2020 ) . “ Hessian approximation '' , which assumes C ≈ c0H for some constant c0 , has often been adopted for understanding SGD ( see Table 1 ) ; this assumption is often motivated by the fact that C = Jw ≈H , where Jw is the Fisher information matrix ( FIM ) ( Zhu et al. , 2019 ) ; the fluctuation can be solved to be isotropic : Σ ∼ ID . However , it is not known under what conditions the Hessian approximation is valid , while previous works have argued that it can be very inaccurate ( Martens , 2014 ; Liu et al. , 2021 ; Thomas et al. , 2020 ; Kunstner et al. , 2019 ) . However , Martens ( 2014 ) and Kunstner et al . ( 2019 ) only focuses on the natural gradient descent ( NGD ) setting ; Thomas et al . ( 2020 ) is closest to ours , but it does not apply to the case with momentum , a matrix learning rate , or regularization . Discrete-time SGD with a Large Learning Rate . Recently , it has been realized that networks trained at a large learning rate have a dramatically better performance than networks trained with a vanishing learning rate ( lazy training ) ( Chizat and Bach , 2018 ) . Lewkowycz et al . ( 2020 ) shows that there is a qualitative difference between the lazy training regime and the large learning rate regime ; the performance features two plateaus in testing accuracy in the two regimes , with the large learning rate regime performing much better . However , the theory regarding discrete-time SGD at a large learning rate is almost non-existent , and it is also not known what Σ may be when the learning rate is non-vanishing . Our work also sheds light on the behavior of SGD at a large learning rate . Some other works also consider discrete-time SGD in a similar setting ( Fontaine et al. , 2021 ; Dieuleveut et al. , 2020 ; Toulis et al. , 2017 ) , but the focus is not on deriving analytical formulae or does not deal with the stationary distribution . 4 SGD NOISE AND FLUCTUATION IN LINEAR REGRESSION . This section derives the shape and strength of SGD noise and fluctuation for linear regression ; concurrent to our work , Kunin et al . ( 2021 ) also studies the same problem but with continuous-time approximation ; our result is thus more general . To emphasize the message , we discuss the label noise case in more detail . The other situations also deserve detailed analysis ; we delay such discussion to the appendix due to space constraints . Notation : S denotes the minibatch size . w ∈ RD is the model parameter viewed in a vectorized form ; λ ∈ R+ denotes a scalar learning rate ; when the learning rate takes the form of a preconditioning matrix , we use Λ ∈ RD×D . A ∈ RD×D denotes the covariance matrix of the input data . When a matrix X is positive semi-definite , we write X ≥ 0 ; throughout , we require Λ ≥ 0. γ ∈ R denotes the weight decay hyperparameter ; when the weight decay hyperparameter is a matrix , we write Γ ∈ RD×D . µ is the momentum hyperparameter in SGD . For two matrices X , Y , the commutator is defined as [ X , Y ] ∶= XY − Y X . Other notations are introduced in the context.1 The results of this section are numerically verified in Sec . A . 4.1 KEY PREVIOUS RESULTS . When N ≫ S , the following proposition is well-known and gives the exact noise due to minibatch sampling . See Appendix E.1 for derivation . Proposition 1 . The noise covariance of SGD as defined in Definition 2 is C ( w ) = 1 SN N ∑ i=1 ∇ ` i ( w ) ∇ ` i ( w ) T − 1 S ∇L ( w ) ∇L ( w ) T , ( 2 ) where the notations ` i ( w ) ∶= l ( xi , yi , w ) and L ( w ) ∶= L ( { xi , yi } Ni=1 , w ) are used . This gradient covariance matrix C is crucial to understand the minibatch noise . The standard literature often assumes C ( w ) ≈H ( w ) ; however , the following well-known proposition shows that this approximation can easily break down . Proposition 2 . Let w∗ be the solution such that L ( w∗ ) = 0 , then C ( w∗ ) = 0 . Proof . Because ` i is non-negative for all i , L ( w∗ ) = 0 implies that ` i ( w∗ ) = 0 . The differentiability in turn implies that each ∇ ` i ( w∗ ) = 0 ; therefore , C = 0 . ◻ This proposition implies that there is no noise if our model can achieve zero training loss ( which is achievable for an overparametrized model ) . This already suggests that the Hessian approximation C ∼ H is wrong since the Hessian is unlikely to vanish in any minimum . The fact that the noise strength vanishes at L = 0 suggests that the SGD noise might at least be proportional to L ( w ) , which we will show to be true for many cases . The following theorem relates C and Σ of the discrete-time SGD algorithm with momentum for a matrix learning rate . Theorem 1 . ( Liu et al. , 2021 ) Consider running SGD on a quadratic loss function with Hessian H , learning rate matrix Λ , momentum µ . Assuming ergodicity , then ( 1 − µ ) ( ΛHΣ +ΣHΛ ) − 1 + µ2 1 − µ2 ΛHΣHΛ + µ 1 − µ2 ( ΛHΛHΣ +ΣHΛHΛ ) = ΛCΛ . ( 3 ) Propostion 1 and Theorem 1 allow one to solve C and Σ . Equation ( 3 ) can be seen as a general form of the Lyapunov equation ( Lyapunov , 1992 ) and is hard to solve in general ( Hammarling , 1982 ; Ye et al. , 1998 ; Simoncini , 2016 ) . Solving this analytical equation in settings of machine learning relevance is one of the main technical contributions of this work . 1We use the word global minimum to refer to the global minimum of the loss function , i.e. , where L = 0 and a local minimum refers to a minimum that has a non-negative loss , i.e. , L ≥ 0 .

Summary: This paper studies the properties of the gradient noise in mini-batch SGD using discrete-time analysis for a fixed learning rate. Their analysis is more general than prior work and considers SGD with momentum, a learning rate matrix (that subsumes preconditioning methods), and regularization. They provide closed form expressions for the noise covariance for various machine learning problems including linear regression with label noise (Theorem 3), with regularization (Proposition 3), and more general models under reasonable assumptions (Theorem 5). Their theory matches the empirical observations about the effect of mini-batch size and learning rate on the gradient noise for small batch sizes and large learning rates where there was previously a gap in theory. Section 6 illustrates the significance of the results as it is able to explain various observations in machine learning theory and deep learning practice as well as providing ideas for practitioners.