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

Abstract:
It is generally recognized that finite learning rate (LR), in contrast to infinitesimal LR, is important for good generalization in real-life deep nets. Most attempted explanations propose approximating finite-LR SGD with Itô Stochastic Differential Equations (SDEs), but formal justification for this approximation (e.g., (Li et al., 2019a)) only applies to SGD with tiny LR. Experimental verification of the approximation appears computationally infeasible. The current paper clarifies the picture with the following contributions: (a) An efficient simulation algorithm SVAG that provably converges to the conventionally used Itô SDE approximation. (b) A theoretically motivated testable necessary condition for the SDE approximation and its most famous implication, the linear scaling rule (Goyal et al., 2017), to hold. (c) Experiments using this simulation to demonstrate that the previously proposed SDE approximation can meaningfully capture the training and generalization properties of common deep nets.

heading: 1 Introduction
text: 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.

heading: 2 Preliminaries and Overview
text: 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.

heading: 2.1 Understanding the Role of Non-Gaussian Noise
text: 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.

heading: 3 Related Work
text: Applications of the SDE approximation in deep learning. One component of the SDE approximation is the gradient noise distribution. When the noise is an isotropic Gaussian distribution (i.e., ⌃(Xt) ⌘ I), then the equilibrium of the SDE is the Gibbs distribution. Shi et al. (2020) used an isotropic Gaussian noise assumption to derive a convergence rate on SGD that clarifies the role of the LR during training. Several works have relaxed the isotropic assumption but assume the noise is constant. Mandt et al. (2017) assumed the covariance ⌃(X) is locally constant to show that SGD can be used to perform Bayesian posterior inference. Zhu et al. (2019) argued that when constant but anisotropic SGD noise aligns with the Hessian of the loss, SGD is able to more effectively escape sharp minima. When noise covariance ⌃(Xt) is uniformly positive definite, Hu et al. (2019) showed that SDE approximation Equation (2) escapes strict saddle points in O(ln ⌘ 1) time, which matches the O(⌘ 1 ln ⌘ 1) escaping rate for SGD (Fang et al., 2019; Jin et al., 2017).
Recently, many works have used the most common form of the SDE approximation (2) with parameterdependent noise covariance. Li et al. (2020) and Kunin et al. (2020) used the symmetry of loss (scale invariance) to derive properties of dynamics (i.e., ⌃(Xt)Xt = 0). Li et al. (2020) further used this property to explain the phenomenon of sudden rising error after LR decay in training. Smith et al. (2020) used the SDE to derive the linear scaling rule (Goyal et al. (2017) and Definition 2.1) for infinitesimally small LR. Xie et al. (2021) constructed a SDE-motivated diffusion model to propose why SGD favors flat minima during optimization. Cheng et al. (2020) analyzed MCMC-like continuous dynamics and construct an algorithm that provably converges to this limit, although their dynamics do not model SGD.
Theoretical Foundations of the SDE approximation for SGD. Despite the popularity of using SDEs to study SGD, theoretical justification for this approximation has generally relied upon tiny LR (Li et al., 2019a; Hu et al., 2019). Cheng et al. (2020) proved a strong approximation result for an SDE and MCMC-like dynamics, but not SGD. Wu et al. (2020) argued that gradient descent with Gaussian noise can generalize as well as SGD, but their convergence proof also relied on an infinitesimally small LR.
LR and Batch Size. It is well known that using large batch size or small LR will lead to worse generalization (Bengio, 2012; LeCun et al., 2012). According to (Keskar et al., 2017), generalization is harmed by the tendency for large-batch training to converge to sharp minima, but Dinh et al. (2017) argued that the invariance in ReLU networks can permit sharp minima to generalize well too. Li et al. (2019b) argued that the LR can change the order in which patterns are learned in a non-homogeneous
synthetic dataset. Several works (Hoffer et al., 2017; Smith and Le, 2018; Chaudhari and Soatto, 2018; Smith et al., 2018) have had success using a larger LR to preserve the scale of the gradient noise and hence maintain the generalization properties of small-batch training. The relationship between LR and generalization remains hazy, as (Shallue et al., 2019) empirically demonstrated that the generalization error can depend on many other training hyperparameters.

heading: 4 Stochastic Variance Amplified Gradient (SVAG)
text: Experimental verification of the SDE approximation appears computationally intractable by traditional methods. We provide an algorithm, Stochastic Variance Amplified Gradient (SVAG), that efficiently simulates and provably converges to the Itô SDE (2) for a given training setting (Theorem 4.3). Moreover, we use SVAG to experimentally verify that the SDE approximation closely tracks SGD for many common settings (Figure 4; additional settings in Appendix F).

heading: 4.1 The SVAG Algorithm
text: For a chosen hyperparameter l 2 N+, we define
xk+1 = xk ⌘
l rL
l ̄k(xk), (5)
where ̄k = ( k,1, k,2) with k,1, k,2 sampled independently and
L l ̄k(·) :=
1 + p 2l 1
2 L k,1(·) +
1 p 2l 1
2 L k,2(·).
SVAG is equivalent to performing SGD on a new distribution of loss functions constructed from the original distribution: the new loss function is a linear combination of two independently sampled losses L k,1 and L k,2 , usually corresponding to the losses on two independent batches. This ensures that the expected gradient is preserved while amplifying the gradient covariance by a factor of l, i.e.,p⌘
l ⌃ l(x) = p ⌘⌃1(x), where ⌃l(x) := E[(rLl̄(x) rLl(x))(Ll̄(x) rLl(x)) > ]. Therefore, the Itô SDE that matches the first and second order moments is always (2). We note that SVAG is equivalent to SGD when l = 1, and both the expectation and covariance of the one-step update (xk+1 xk) are proportional to 1/l, meaning the direction of the update is noisier when l increases.

heading: 4.2 SVAG Approximates the SDE
text: Definition 4.1 (Test Functions). Class G of continuous functions Rd ! R has polynomial growth if 8g 2 G there exist positive integers 1,2 > 0 such that for all x 2 Rd, |g(x)|  1(1 + |x|22). For ↵ 2 N+, we denote by G↵ the set of ↵-times continuously differentiable functions g where all partial derivatives of form @
↵g @x
↵1 1 ···@x ↵d d
s.t. Pd
i=1 ↵i = ↵  ↵, are also in G.
Definition 4.2 (Order-↵ weak approximation). Let {X⌘t : t 2 [0, T ]} and {x ⌘ k} bT⌘ c k=0 be families of continuous and discrete stochastic processes parametrized by ⌘. We say {X⌘t } and {x ⌘ k} are order-↵
weak approximations of each other if for every g 2 G2(↵+1), there is a constant C > 0 independent of l such that
max k=0,...,bT⌘ c
Eg(x⌘k) Eg(X ⌘ k⌘)  C⌘↵.
When applicable, we drop the superscript ⌘, and say {Xt} and {xk} are order-↵ (or ↵ order) approximations of each other.
We now show that SVAG converges weakly to the Itô SDE approximation in (2) when l ! 1, i.e., xlk and Xk⌘ have the roughly same distribution. Figure 2 highlights the differences between our result and (Li et al., 2019a). Figure 4 provide verification of the below theorem, and additional settings are studied in Appendix F. Theorem 4.3. Suppose the following conditions2 are met:
(i) L ⌘ EL is C1-smooth, and L 2 G4. (ii) |rL (x) rL (y)|  L |x y|, for all x, y 2 Rd, where L > 0 is a random variable
with finite moments, i.e., ELk is bounded for k 2 N+. (iii) ⌃ 1 2 (X) is C1-smooth in X .
Let T > 0 be a constant and l be the SVAG hyperparameter (5). Define {Xt : t 2 [0, T ]} as the stochastic process (independent of ⌘) satisfying the Itô SDE (2) and {x
⌘/l k : 1  k  blT/⌘c} as the
trajectory of SVAG (5) where x0 = X0. Then, SVAG {x ⌘/l k } is an order-1 weak approximation of the SDE {Xt}, i.e. for each g 2 G4, there exists a constant C > 0 independent of l such that
max k=0,...,blT/⌘c |Eg(x⌘/lk ) Eg(X k⌘l )|  Cl 1 .
Remark 4.4. Lipschitz conditions like (ii) are often not met by deep learning objectives. For instance using normalization schemes can make derivatives unbounded, but if the trajectory {xt} stays bounded away from the origin and infinity, then (ii) holds. Remark 4.5. Though technically the weak approximation result (Theorem 4.3) only applies when the stochastic gradient is sampled independently at each step, experimentally we found the difference between performance of SGD and SVAG is negligible among different sampling methods, including random shuffling, sampling with and without replacement. (See detailed discussions in Appendix F.1 and Figure 7) This experimental evidence suggests the validity of SDE approximation doesn’t change with sampling methods used in practice.

heading: 4.3 Proof Overview
text: Let {Xx,st : t s} denote the stochastic process obeying the Itô SDE (2) starting from time s and with the initial condition Xx,ss = x and {x x,j k : k j} denote the stochastic process (depending on l) satisfying SVAG (5) with initial condition xx,jj = x. For convenience, we define eXk := X k⌘ l and write eXx,jk := X x, j⌘l k⌘ l . Alternatively, we write eXk(x, j) := eXx,jk and xk(x, j) := x x,j k .
Now for any 1  k  b lT⌘ c, we interpolate between a SVAG solution xk and SDE solution eXk through a series of hybrid trajectories eXk(xj , j), i.e., the weight achieved by running SVAG for the first j steps and then SDE from time j to k. The two limits of the interpolation are eXk(xk, k) = xk (i.e., SVAG solution after k steps) and eXk(x0, 0) = eXk (i.e., SDE solution after k time). This yields the following error decomposition for a test function g 2 G (see Definition 4.1).
|Eg(xk) Eg(X k⌘ l )| = |Eg(xk) Eg( eXk)|  Xk 1 j=0 Eg( eXk(xj+1, j + 1)) Eg( eXk(xj , j))
Note that each pair of adjacent hybrid trajectories only differ by a single step of SVAG or SDE. We show that the one-step increments of SVAG and SDE are close in distribution along the entire trajectory by computing their moments (Lemmas 4.6 and 4.7). Then, using the Taylor expansion of g, we can show that the single-step approximation error from switching from SVAG to SDE is uniformly upper bounded by O(⌘ 2
l2 ). (See Figure 6 for demonstration) Hence, the total error is O(k ⌘2 l2 ) = O( ⌘ l ).
2The C1 smoothness assumptions can be relaxed by using the mollification technique in Li et al. (2019a).
Lemma 4.6. Define the one-step increment of the Itô SDE as e (x) = eXx,01 x. Then we have
(i) Ee (x) = ⌘l rL(x) +O(l 2), (ii) Ee (x)e (x)> = ⌘
2 l ⌃(x) +O(l 2),
(ii) Ee (x)⌦3 = O(l 2), (iv) q E|e (x)⌦4|2 = O(l 2).
Lemma 4.7. Define the one-step increment of SVAG as (x) = xx,01 x. Then we have (i) E (x) = ⌘l rL(x),
(ii) E (x) (x)> = ⌘ 2 l ⌃(x) + ⌘2 l2 rL(x)rL(x) > = ⌘ 2 l ⌃(x) +O(l 2),
(iii) E (x)⌦3 = ⌘ 3 l2 3 l 1 2 ⇤(x) + ⌘3 l3
⇣ 3rL(x)⌦ ⌃(x) +rL(x)⌦3 ⌘ = O(l 2)
(iv) p E| (x)⌦4|2 = O(l 2),
where ⇤(x) := E(rL 1(x) rL(x)) ⌦3, and T denotes the symmetrization of tensor T , i.e., T ijk = 1 6 P i0,j0,k0 Ti0j0k0 , where i 0 , j 0 , k 0 sums over all permutation of i, j, k.
Though (i) and (ii) in Lemma 4.7 hold for any discrete update with LR = ⌘l that matches the first and second order moments of SDE (2), (iii) and (iv) could fail. For example, when decreasing LR according to LSR (Definition 2.1), even if we can use a fractional batch size and sample an infinitely divisible noise distribution, we may arrive at a different continuous limit if (iii) and (iv) are not satisfied. (See a more detailed discussion in Appendix B.2) SVAG is not the unique way to ensure (iii) and (iv), and any other design (e.g. using three copies per step and with different weights) satisfying Lemma 4.7 are also first order approximations of SDE (2), by the same proof.

heading: 5 Understanding the Failure of SDE Approximation and LSR
text: In this section, we analyze how discretization error, caused by large LR, leads to the failure of the SDE approximation (Section 5.1) and LSR (Section 5.2) for scale invariant networks (e.g., nets equipped with BatchNorm (Ioffe and Szegedy, 2015) and GroupNorm (Wu and He, 2018)). To get best generalization, practitioners often add Weight Decay (WD, a.k.a `2 regularization; see (7)). Intriguingly, unlike the traditional setting where `2 regularization controls the capacity of function space, for scale invariant networks, each norm ball has the same expressiveness regardless of the radius, and thus WD only regularize the model implicitly via affecting the dynamics. Li et al. (2020) explained such phenomena by showing for training with Normalization, WD and constant LR, the parameter norm converges and WD affects ‘effective LR’ by controlling the limiting value of the parameter norm. That paper also gave experiments showing that the training loss will reach some plateau, and gave evidence of training reaching an ”equilibrium” distribution that it does not get out of unless if some hyperparameter is changed. Throughout this section we assume the existence of equilibrium for SGD and SDE.
To quantify differences in training algorithms, we would ideally work with statistics like the train/test loss and accuracy achieved, but characterizing optimization and generalization properties of deep networks beyond the NTK regime (Jacot et al., 2018; Allen-Zhu et al., 2019b; Du et al., 2019; Arora et al., 2019a; Allen-Zhu et al., 2019a) is in general an open problem. Therefore, we rely on other natural test functions (Definition 5.1).

heading: 5.1 Failure of SDE Approximation
text: In Theorem 5.2, we show that the SDE-approximation of SGD is bound to fail for these scale-invariant nets when LR gets too large. Specifically, using above-mentioned results we show that the equilibrium distributions of SGD and SDE are quite far from each other with respect to expectations of these natural test functions (Definition 5.1).
We consider the below SDE (6) with arbitrary expected loss L(x) and covariance ⌃(x), and the moment-matching SGD (7) satisfying EL (x) = L(x) and ⌃(x) = ⌘⌃(x) where ⌃(x) is the covariance of rL (x). In the entire Section 5, we will assume that for all , L is scale invariant (Arora et al., 2019b; Li and Arora, 2020a), i.e., L (x) = L (cx), 8c > 0 and x 2 Rd \ {0}. 3
3The results in this section can be extended straightforwardly to the case where the network is not entirely scale invariant, where the C-closeness is defined for the corresponding metrics of the scale-invariant parameters.
dXt= r L(Xt) +
2 |Xt|
2 dt+ ⌃ 1/2 (Xt)dWt (6)
xk+1 = xk ⌘r L k(xk) +
2 |xk|
2
(7)
We will measure the closeness of two distributions by three test functions: squared weight norm |x|2, squared gradient norm |rL(x)|2, and trace of noise covariance Tr[⌃(x)]. We say two equilibrium distributions are C-close if expectations of these test functions are within a multiplicative constant C. Definition 5.1 (C-closeness). Assuming the existence of the following limits, we use R1 := lim
t!1 E|xt|2, R1 := lim t!1 E|Xt|2,
G1 := lim t!1 E|rL(xt)|2, G1 := lim t!1 E|rL(Xt)|2, N1 := lim
t!1 E[Tr[⌃(xt)], N1 := lim t!1 E[Tr[⌃(Xt)]]
to denote the limiting squared weight norm, gradient norm and trace of covariance for SGD (7) and SDE (6). We say the two equilibriums are C-close to each other iff
1 C  R1 R1 , G1 G1 , ⌘N1 N1  C. (8)
We call N1G1 and N1 G1 the noise-to-signal ratio (NSR), and below we show that it plays an important role. When the LR of SGD significantly exceeds the NSR of the corresponding SDE, the approximation fails. Of course, we lack a practical way to calculate NSR of the SDE so this result is existential rather than effective. Therefore we give a condition in terms of NSR of the SGD that suffices to imply failure of the approximation. Experiments later in the paper show this condition is effective at showing divergence from SDE behavior.
Theorem 5.2. If either (i). ⌘ > N1 G1 (C2 1) or (ii).N1G1 < 1 C2 1 , then the equilibria of SDE (6) and SGD (7) are not C-close.
The high-level idea behind Theorem 5.2 is the observation that the norm dynamics of SGD (10) and SDE (9) differ by a second order discretization error related to the norm of full-batch gradient, ⌘ 2E|rL(xk)|2. Thus intuitively, the two dynamics can be close only when the difference is tiny and negligible. We can make the argument formal by comparing the relationships between the above defined three metrics at the equilibrium of SGD Equation (11) and SDE Equation (12), and thus conclude that a sufficiently large noise-to-signal ratio (NSR) is a necessary condition for C-closeness.
Here the role of scale invariance is to simplify the norm dynamics of both SGD (10) and SDE (9) by removing the cross term. This is because of a well-known property of scale-invariance, the orthogonality between gradient and the weight itself, i.e., hrL (x), xi = 0 for any x, . (Lemma E.7)
Proof of Theorem 5.2. We will prove the the contrapositive statement: if the equilibriums of (7) and (6) are C-close, then ⌘  N1
G1 (C2 1) and 1C2 1  N1 G1 . Following the derivation in (Li et al., 2020), by Itô’s lemma:
d dt E|Xt|2 = 2 E|Xt|2 + ETr[⌃(Xt)]. (9)
Again by the orthogonality between gradient and weight, it can be shown that for SGD (7), E|xk+1|2 E|xk|2 =(1 ⌘ )2E|xk|2 + ⌘2E|rL k(xk)|2 E|xk|2
=⌘ ( 2+⌘ )E|xk|2+⌘2E|rL(xk)|2 + ⌘2ETr[⌃(xk)] (10)
If both xk and Xt have reached their equilibriums, both LHS of (9) and (10) are 0, and therefore (2 ⌘ ) R1 = ⌘G1 + ⌘N1, (11)
2 R1 = N1. (12) Combining (11), (12), and (8), we have
⌘G1 + ⌘N1  2 R1  2 CR1 = CN1.
Applying (8) again, we have ⌘G1 +N1  C⌘(G1 +N1)  C2N1  C3⌘N1. The proof is completed by comparing the first and third, the second and the fourth terms respectively.
Remark 5.3. Since the order-1 approximation fails for large LR, it’s natural to ask if higher-order SDE approximation works. In Theorem E.4 we give a partial answer, that the same gap happens already between order-1 and order-2 SDE approximation, when ⌘ & N1
G1 (C2 1). This suggests
failure of SDE approximation may be due to missing some second order term, and thus higherorder approximation in principle could avoid such failure. On the other hand, when approximation fails in such ways, e.g., increasing batch size along LSR, the performance of SGD degrades while SDE remains good. This suggests the higher-order correction term may not be very helpful for generalization.

heading: 5.2 Failure of Linear Scaling Rule
text: In this section we derive a similar necessary condition for LSR to hold.
Similar to Definition 5.1, we will use RB,⌘1 , GB,⌘1 , NB,⌘1 as test functions for equilibrium achieved by SGD (7) when training with LR ⌘ and mini-batches of size B. We first introduce the concept of Linear Scaling Invariance (LSI). Note here we care about the scaled ratio NB,⌘1 /(NB,⌘1 ) because the covariance scales inversely to batch size, ⌃B(x) = ⌃B(x). Definition 5.4 ((C,)-Linear Scaling Invariance). We say SGD (7) with batch size B and LR ⌘ exhibits (C,)-LSI if, for a constant C such that 0 < C < p ,
1 C 
R B,⌘ 1
R B,⌘ 1
, N
B,⌘ 1
N B,⌘ 1
, G
B,⌘ 1
G B,⌘ 1
 C. (13)
We show below that (C,)-LSI fails if the NSR N1G1 is too small, thereby giving a certificate for failure of (C,)-LSI even without a baseline run. Theorem 5.5. For any B, ⌘, C, and  such that
N B,⌘ 1 G B,⌘ 1 < (1 1  ) 1 C2 1 1  , (14)
SGD with batch size B and LR ⌘ does not exhibit (C,)-LSI.
We now present a simple and efficient procedure to find the largest  for which (C,)-LSI will hold, providing useful guidance to make hyper-parameter tuning more efficient. Before doing so, one must
choose an appropriate value for C, which controls how close the test functions must be for us to consider LSR to have “worked.” It is an open question what value of C will ensure that the two settings achieve similar test performance, but throughout our experiments across various datasets and architectures in Figure 3 and Appendix F, we find that C = p 2 works well. One can estimate GB,⌘1 and NB,⌘1 from a baseline run. Then, one can straightforwardly compute the value for the  threshold given in the theorem below. We conduct this process in Figure 3 and Appendix F to test our theory.
Theorem 5.6. For any B, ⌘, C, and
 > C 2(1 +
N B,⌘ 1 G B,⌘ 1 ), (⇡ C2 N B,⌘ 1 G B,⌘ 1 when N B,⌘ 1 G B,⌘ 1
1), (15)
SGD with batch size B and LR ⌘ does not exhibit (C,)-LSI.

heading: 6 Experiments
text: We provide our code at https://github.com/sadhikamalladi/svag. Figure 3 provides experimental evidence that measurements from a single baseline run can be used to predict when LSR will break, thereby providing verification for Theorem 5.6. Surprisingly, it turns out the condition in Theorem 5.6 is not only sufficient but also close to necessary. Figure 4 and Appendix F.1 test SVAG on common architectures and datasets and report the results. Theorem 4.3 shows that SVAG converges to the SDE as l ! 1, but we note that SVAG needs l times as many steps as SGD to match the SDE. Therefore, in order for SVAG to be a computationally efficient simulation of the SDE, we hope to observe convergence for small values of l. This is confirmed in Figure 4 and Appendix F.1. The success of SVAG in matching SGD in many cases indicates that studying the Itô SDE can yield insights about the behavior of SGD. We note our experiments are limited in the sense that it only confirms the closeness of train/test accuracy curves, which doesn’t verify the weak convergence guaranteed in Theorem 4.3. But at least, the experiments suggest that SDE and SVAG with large l are interesting learning algorithms to study, with similar or even better generalization than SGD. Moreover, in the case where we expect the SDE approximation to fail (e.g., when LSR fails), SVAG does indeed converge to a different limiting trajectory from the SGD trajectory.

heading: 7 Conclusion
text: We present a computationally efficient simulation SVAG (Section 4) that provably converges to the canonical order-1 SDE (2), which we use to verify that the SDE is a meaningful approximation for SGD in common deep learning settings (Section 6). We relate the discretization error to LSR (Definition 2.1): in Section 5 we derive a testable necessary condition for the SDE approximation and LSR to hold, and in Figure 3 we demonstrate its applicability to standard settings.

heading: Acknowledgement
text: The authors acknowledge support from NSF, ONR, Simons Foundation, DARPA and SRC. ZL is also supported by Microsoft Research PhD Fellowship.

Review summaries:
SDE has been used as an approximation to help understand the implicit bias of SGD. This paper studies the validity of such approximations, which hasn't been formalized in prior work for realistic choices of learning rates. The contributions are three-folds:   - Proposing "Stochastic Variance Amplified Gradient" (SVAG), an efficient numerical method to test whether the trajectories of SGD and the corresponding SDE are close.   - Empirically verifying SVAG can be made to follow SGD. Since SVAG provably follows SDE under certain conditions, this means SDE also tracks SGD well by transitivity.   - Providing necessary conditions for SDE and LSR to track SGD. 

The authors consider a discrete time algorithm (SVAG) which is indexed by a parameter $l$ and which exactly corresponds to SGD when $l= 1$.  They show that this algorithm weakly converges to SGD's continuous SDE approximation as $l$ goes to $\infty$, and this without having to take the step-size infinitely small. They also provide a setting in which the SDE approximation provably fails to capture SGD's dynamics. Finally they give a testing rule for the SDE approximation to hold. 

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.

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.