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: This paper validates the SDE approximation to SGD under finite (non-infinitesimal) learning rate by proposing a novel algorithm SVAG. This paper theoretically proves that SVAG is an order-1 weak approximation to the SDE, and empirical experiments show that SVAG behaves similarly to SGD. Using SVAG as a theoretical tool, this paper shows new theoretical insights in the linear scaling rule and the necessary condition for the SDE approximation.