On the Validity of Modeling SGD with Stochastic Differential Equations (SDEs)

1 Introduction . Training with Stochastic Gradient Gescent ( SGD ) ( 1 ) and finite learning rate ( LR ) is largely considered essential for getting best performance out of deep nets : using infinitesimal LR ( which turns the process into Gradient Flow ( GF ) ) or finite LR with full gradients results in noticeably worse test error despite sometimes giving better training error ( Wu et al. , 2020 ; Smith et al. , 2020 ; Bjorck et al. , 2018 ) . Mathematical explorations of the implicit bias of finite-LR SGD toward good generalization have focused on the noise arising from gradients being estimated from small batches . This has motivated modeling SGD as a stochastic process and , in particular , studying Stochastic Differential Equations ( SDEs ) to understand the evolution of net parameters . Early attempts to analyze the effect of noise try to model it as as a fixed Gaussian ( Jastrzebski et al. , 2017 ; Mandt et al. , 2017 ) . Current approaches approximate SGD using a parameter-dependent noise distribution that match the first and second order moments of of the SGD ( Equation ( 2 ) ) . It is important to realize that this approximation is heuristic for finite LR , meaning it is not known whether the two trajectories actually track each other closely . Experimental verification seems difficult because simulating the ( continuous ) SDE requires full gradient/noise computation over suitably fine time intervals . Recently , Li et al . ( 2017 , 2019a ) ; Feng et al . ( 2017 ) ; Hu et al . ( 2019 ) provided rigorous proofs that the trajectories are arbitrarily close in a natural sense , but the proof needs the LR of SGD to be an unrealistically small ( unspecified ) constant so the approximation remains heuristic . In the worst case , the LR needs to be exponentially small , i.e. , e ⌦ ( T ) , where T is the continuous training time . Furthermore , noise plays no role in these approximation analyses and the same analysis indeed shows GD , SGD and SDE all converge weakly to GF at the same rate . Thus whenever their requirements for LR are met , there should be no performance difference between SGD and full-batch GD . However , for a common practical LR choice we observe some difference in Figure 1 , indicating that the LR is usually outside of the regime their result requires . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Setting aside the issue of correctness of the SDE approximation , there is no doubt it has yielded important insights of practical importance , especially the linear scaling rule ( LSR ; see Definition 2.1 ) relating batch size and optimal LR , which allows much faster training using high parallelism ( Krizhevsky , 2014 ; Goyal et al. , 2017 ) . However , since the scaling rule depends upon the validity of the SDE approximation , it is not mathematically understood when the rule fails . ( Empirical investigation , with some intuition based upon analysis of simpler models , appears in ( Goyal et al. , 2017 ; Smith et al. , 2020 ) . This paper casts new light on the SDE approximation via the following contributions : 1 . A new and efficient numerical method , Stochastic Variance Amplified Gradient ( SVAG ) , to test if the trajectories of SGD and its corresponding SDE are close for a given model , dataset , and hyperparameter configuration . In Theorem 4.3 , we prove ( using ideas similar to Li et al . ( 2019a ) ) that SVAG provides an order-1 weak approximation to the corresponding SDE . ( Section 4 ) 2 . Empirical testing showing that the trajectory under SVAG converges and closely follows SGD , suggesting ( in combination with the previous result ) that the SDE approximation can be a meaningful approach to understanding the implicit bias of SGD in deep learning . 3 . New theoretical insight into the observation in ( Goyal et al. , 2017 ; Smith et al. , 2020 ) that linear scaling rule fails at large LR/batch sizes ( Section 5 ) . It applies to networks that use normalization layers ( scale-invariant nets in Arora et al . ( 2019b ) ) , which includes most popular architectures . We give a necessary condition for the SDE approximation to hold : at equilibrium , the squared gradient norm must be smaller than its variance . 2 Preliminaries and Overview . We use | · | to denote the ` 2 norm of a vector and ⌦ to denote the tensor product . Stochastic Gradient Descent ( SGD ) is often used to solve optimization problems of the form minx2Rd L ( x ) : = E L ( x ) where { L : 2 } is a family of functions from Rd to R and is a -valued variable , e.g. , denoting a random batch of training data . We consider the general case of an expectation over arbitrary index sets and distributions . xk+1 = xk ⌘rL k ( xk ) , ( SGD ) ( 1 ) where each k is an i.i.d . random variable with the same distribution as . Taking learning rate ( LR ) ⌘ toward 0 turns SGD into ( deterministic ) Gradient Descent ( GD ) with infinitesimal LR , also called Gradient Flow . Infinitesimal LR is more compatible with traditional calculus-based analyses , but SGD with finite LR yields the best generalization properties in practice . Stochastic processes give a way to ( heuristically ) model SGD as a continuous-time evolution ( i.e. , stochastic differential equation or SDE ) without ignoring the crucial role of noise . Driven by the intuition that the benefit of SGD depends primarily on the covariance of noise in gradient estimation ( and not , say , the higher moments ) , researchers arrived at following SDE for parameter vector Xt : dXt = rL ( Xt ) dt+ ( ⌘⌃ ( Xt ) ) 1/2dWt ( SDE approximation ) ( 2 ) where Wt is Wiener Process , and ⌃ ( X ) : = E [ ( rL ( X ) rL ( X ) ) ( L ( X ) rL ( X ) ) > ] is the covariance of the gradient noise . When the gradient noise is modeled by white noise as above , it is called an Itô SDE . Replacing Wt with a more general distribution with stationary and independent increments ( i.e. , a Lévy process , described in Definition A.1 ) yields a Lévy SDE . The SDE view—specifically , the belief in key role played by noise covariance—motivated the famous Linear Scaling Rule , a rule of thumb to train models with large minibatch sizes ( e.g. , in highly parallel architectures ) by changing LR proportionately , thereby preserving the scale of the gradient noise . Definition 2.1 ( Linear Scaling Rule ( LSR ) ) . ( Krizhevsky , 2014 ; Goyal et al. , 2017 ) When multiplying the minibatch size by > 0 , multiply the learning rate ( LR ) also by . If the SDE approximation accurately captures the SGD dynamics for a specific training setting , then LSR should work ; however , LSR can work even when the SDE approximation fails . We hope to ( 1 ) understand when and why the SDE approximation can fail and ( 2 ) provide provable and practically applicable guidance on when LSR can fail . Experimentally verifying if the SDE approximation is valid is computationally challenging , because it requires repeatedly computing the full gradient and the noise covariance at very fine time intervals , e.g . the Euler-Maruyama method Equation ( 16 ) . We are not aware of any empirical verification using conventional techniques , which we discuss in more detail in Appendix A.1 . Section 4 gives a new , tractable simulation algorithm , SVAG , and presents theory and experiments suggesting it is a reasonably good approximation to both the SDE and SGD . Formalizing closeness of two stochastic processes . Two stochastic processes ( e.g. , SGD and SDE ) track each other closely if they lead to similar distributions on outcomes ( e.g. , trained nets ) . Mathematics formulates closeness of distributions in terms of expectations of suitable classes of test functions1 ; see Section 4.2 . The test functions of greatest interest for ML are of course train and test error . These do not satisfy formal conditions such as differentiability assumed in classical theory but can be still used in experiments ( see Figure 4 ) . Section 5 uses test functions such as weight norm |xt| , gradient norm |rL ( xt ) | and trace of noise covariance Tr [ ⌃ ( xt ) ] and proves a sufficient condition for the failure of SDE approximation . Mathematical analyses of closeness of SGD and SDE will often consider the discrete process x̂k+1 = x̂k ⌘rL ( x̂k ) + ⌘⌃ 1 2 ( x̂k ) zk , ( Noisy Gradient Descent/NGD ) ( 3 ) where zk i.i.d . ⇠ N ( 0 , Id ) . A basic step in analysis will be the following Error Decomposition : Eg ( X⌘k ) Eg ( xk ) = ( Eg ( X⌘k ) ) Eg ( x̂k ) ) | { z } Discretization Error + ( Eg ( x̂k ) Eg ( xk ) ) | { z } Gap due to non-Gaussian noise ( 4 ) Understanding the failure caused by discretization error : In Section 5 , a testable condition of SDE approximation is derived for scale-invariant nets ( i.e . nets using normalization layers ) . This condition only involves the Noise-Signal-Ratio , but not the shape of the noise . We further extend this condition to LSR and develops a method predicting the largest batch size at which LSR succeeds , which only takes a single run with small batch size . 2.1 Understanding the Role of Non-Gaussian Noise . Some works have challenged the traditional assumption that SGD noise is Gaussian . Simsekli et al . ( 2019 ) ; Nguyen et al . ( 2019 ) suggested that SGD noise is heavy-tailed , which Zhou et al . ( 2020 ) claimed causes adaptive gradient methods to generalize better than SGD . Xie et al . ( 2021 ) argued that the experimental evidence in ( Simsekli et al. , 2019 ) made strong assumptions on the nature of the gradient noise , and we furthermore prove in Appendix B.3 that their measurement method could flag Gaussian distributions as non-Gaussian . Below , we clarify how the Gaussian noise assumption interacts with our findings . Non-Gaussian noise is not essential to SGD performance . We provide experimental evidence in Figure 1 and Appendix F.3 that SGD ( 1 ) and NGD ( 3 ) with matching covariances achieve 1The discriminator net in GANs is an example of test function in machine learning . similar test performance on CIFAR10 ( ⇠ 89 % ) , suggesting that even if the gradient noise in SGD is non-Gaussian , modeling it by a Gaussian estimation is sufficient to understand generalization properties . Similar experiments were conducted in ( Wu et al. , 2020 ) but used SGD with momentum and BatchNorm , which prevents the covariance of NGD noise from being equal to that of SGD . These findings confirm the conclusion in ( Cheng et al. , 2020 ) that differences in the third-and-higher moments in SGD noise don ’ t affect the test accuracy significantly , though differences in the second moments do . LSR can work when SDE approximation fails . We note that ( Smith et al. , 2020 ) derives LSR ( Definition 2.1 ) by assuming the Itô SDE approximation ( 2 ) holds , but in fact the validity of the SDE approximation is a sufficient but not necessary condition for LSR to work . In Section B.1 , we provide a concrete example where LSR holds for all LRs and batch sizes , but the dynamics are constantly away from the Itô SDE limit . This example also illustrates that the failure of the SDE approximation can be caused solely by non-Gaussian noise , even when there is no discretization error ( i.e. , the loss landscape and noise distribution are parameter-independent ) . SVAG does not require Gaussian gradient noise . In Section 4 , we present an efficient algorithm SVAG to simulate the Itô SDE corresponding to a given training setting . In particular , Theorem 4.3 reveals that SVAG simultaneously causes the discretization error and the gap by non-Gaussian noise to disappear as it converges to the SDE approximation . From Figure 4 and Appendix F.1 , we can observe that for vision tasks , the test accuracy of deep nets trained by SGD in standard settings stays the same when interpolating towards SDE via SVAG , suggesting that neither the potentially non-Gaussian nature of SGD noise nor the discrete nature of SGD dynamics is an essential ingredient of the generalization mystery of deep learning .

Summary: The paper has two main contributions. Firstly, it presents a practical simulation algorithm, SVAG, which is shown to converge weakly to the Itô SDE approximation of SGD over finite time horizons (with no restrictions on the learning rate $\eta$). Experimental evidence on deep nets is given to support the claim that SVAG can be used as a diagnostic for empirically testing if SGD and its SDE approximation have similar behaviour. The second main contribution provides sufficient conditions for the failure of the SDE approximation and linear scaling rule when the loss function is scale invariant. This theory is further supported by experiments on deep nets.