On the Validity of Modeling SGD with Stochastic Differential Equations (SDEs)
Authors: Zhiyuan Li
Sadhika Malladi
Sanjeev Arora
Venue: NeurIPS 2021 Poster (2021)
Keywords: stochastic gradient descent
stochastic differential equations
weak approximation
linear scaling rule

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., 2019) 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. 

On the Validity of Modeling SGD with Stochastic Differential Equations (SDEs)
Zhiyuan Li Sadhika Malladi Sanjeev Arora
Princeton University
fzhiyuanli, sma smalladi, arora) ars www.iestsspinnnnoo.ooommoo edu
Abstract
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 approx- imation 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 theo- retically 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 demonstrate that the previously proposed SDE approximation can meaningfully capture the training and generalization properties of common deep nets.
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 al., 2020; Smith al., 2020; Bjorek et al., 2018).
Mathematical explorations of the implicit bias of -R SGD toward good generalization have focused on the noise arising from gradients being estimated from small batches. This has motivated modeling SGD as stochastic process and, in particular, studying Stochastic Differential Equations (SDEs) to understand the evolution of net parameters.
Farly attempts to analyze the effect of noise try to model it as as fixed Gaussian (Jastrzebski al., 7017; Mandt al., 2017). Current approaches approximate SGD using a parameter-dependen noise distribution that match the first and second order moments of of the SGD (Equation (2)). It is important to realize that this approximation heuristic for finite LR, meaning it 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 the approximation remains heuristic. In the worst case, the LR needs to be exponentially small, i.c., -Q(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- GD. However, for common practical LR choice 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 (NcurIPS 2021).

---

Setting aside the issue of correctness of the SDE approximation, there is no doubt it has yielded impor- tant insights of practical importance, especially the linear sealing rule (LSR; see Definition 2.1) relat- ing batch size and optimal IR, 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 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 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 acale-invarican popular architectures. We give necessary condition for the SDE approximation to hold: at equilibrium, the squared gradient norm must be smaller than its variance.
nets in Arora et al. (2019b)), which includes most
2 Preliminaries and Overview
We use | denote the & norm of vector and @ denote the tensor product. Stochastic Gradient Descent (SGD) is often used to solve optimization problems of the form min [(x)) == E,C, (a) where {L. y E Т} is a family of functions from IR to IR and y is a I-valued variable, e.g., denoting random batch of training data. We consider the general case of an expectation over arbitrary index sets and distributions.
Tk+1 = - 1 (Ix),,
(SGD)
where each 7n is an i.i.d. random variable with the same distribution as 7, Taking learning rate (LR) 7) toward 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 propertics in practice. Stochastic processes give way (heuristicallyy model SGD as continuous-time evolution (i.c., 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 X₁:
dX,, -VLC(X) + (T))(x)) /2w,
(SDE approximation)
where W; is Wiener Process, and Σ(X) == E[(VL,( (X( V is the covariance of the gradient noise. When the gradient noise is modeled by white noise as above, is called an Itō SDE. Replacing W, with a more general distribution with stationary and independent increments (i.e., Lévy process, described in Definition A. yields Lévy SDE.
The SDE view swwwssscciiicall....mmmmmm the belief in key role played by noise covariance motivated the famous Linear Scaling Rule, 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 multiply- ing the minibatch size by 0, multiply the learning rate (LR) also by k.
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

---

Train Accuracy
Test Accuracy
1.00
0.75
0.75
GD
GD
0.50
0.50
SGD
SGD
0.25
NGD
0.25
NGD
0
10000 20000
30000
10000
20000
30000
Step
Step
Figure 1: Non-Gaussian noise is not essential to SGD performance. SGD with batch size 125 and NGD with matching covariance have close train and test curves when training on CIFAR-10, 7 = 0.8 for all three settings and is decayed by 0.1 step 24000. GD achieves 75.5% test accuracy, and SGD and NGD achieve 89.4% and 89.3%, respectively. We smooth the training curve by dividing into intervals of 100 steps and recording the average. For efficient sampling Gaussian noise, we use Group.orm instead of BatchNorm and turn off data augmentation. The sudden drop of accuracy when using GD is not coincidence, but a consequence of interplay between normalization and Weight Decay, See more discussion and implementation details in Appendix F.3.
valid is computationally challenging, because it requires repeatedly computing the full gradient and the noise covariance very fine time intervals, e-g. the Fuler-Marayama method Equation (16). We are not aware of any empirical verification using conventional techniques, which we discuss in more detail in Appendix I. Section 4 gives new, tractable simulation algorithm, SVAG, and presents theory and experiments suggesting it is reasonably good approximation to both the SDE and SGD,
Formalizing closeness of two stochastic processes. Two stochastic processes (c.g., SGD and SDE) track each other closely if they lead to similar distributions on outcomes (c.g., trained nets). Mathematies formulates closeness of distributions in terms of expectations of suitable classes of test functions'; ; 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 | gradient norm V (xx)) and trace of noise covariance (x([( 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
Pu+1 = 24 ((xx) + E (#x)=
(Noisy Gradient Descent/NGD)
(3)
lid N(0, Ia). A basic step in analysis will be the following Error Decomposition:
where 2n
[g(X,k) Eg(xx) = (Eg(X,))) - Eg((x)) + (+x) Eg((*))
(4)
Discretization Error
Gap non-Gaussian
Understanding the failure caused by discretization error: In Section 5, testable condition of SDE approximation derived for inver-invariant nets (i.e. nets using normalization layers). This condition only involves the Noise-Signal- Ratiao but not the shape of the noise. We further extend this condition to LSR and develops method predicting the largest batch size which LSR succeeds, which only takes 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
Thhe discriminator in GANs is an example of test function in machine learning.
3

---

similar test performance on CIFARIO ( (~~ 89%), suggesting that even if the gradient noise in SGD non- -on-Gausiian, modeling it by a Gaussian estimation 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 Itô SDE approximation (2) holds, but in fact the validity of the SDE approximation is a sufficient but not necessary condition for LSR work. In Section B.I, we provide a concrete example where LSR holds for all LRs and batch sizes, but the dynamics are constantly away from the Itō 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-independeee
SVAG does not require Gaussian gradient noise. In Section 4, we present an efficient algorithm SVAG to simulate the Itō 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.
3 Related Work
Applications of the SDE approximation in deep learning. One component of the SDE approxi- mation is the gradient noise distribution. When the noise is an isotropic Gaussian distribution (i.e., E(X,) = 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 assumne the noise is constant. Mandt et al. (2017) assumed the covariance (XX 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 Σ(X;) is uniformly positive definite, Hu et al. (2019) showed that SDE approximation Equation (2) escapes strict saddle points in 0(n (1) time, which matches the On 1 In escaping rate for SGD (Fang et al., 2019; Jin al., 2017).
Recently, many works have used the most common form of the SDE approximation (2) with parameter- dependent noise covariance. Li et al. (2020) and Kunin et al. (2020) used the symmetry of loss (scale invariance) to derive properties of dynamics (i.c., (X,)))x = 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) for infinitesimally small LR. Xie et al. (2021) constructed SDE-motivated diffusion model to propose why SGD favors flat minima during optimization. Cheng et al. (2020) analyzed MCMC-like continuous dynamics and construct algorithm that provably converges 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 ReL.U 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 hon-hooogeecous

---

SDE (LR=n) T time)
SVAG (LR=nII, IT/n steps)
SGD (LR=n, T/n steps)
1000
600
400
200
-200
-400
0.8 1.0 0.0
1.0
0.25
(a)
(b)
(c)
Figure 2: Itó SDE (2). SVAG (5), and SGD (1) trajectories (blue) sampled from distribution (green). Li et al. (2019a) show that VT, 3ry such that SDE (a) and SGD (c) are order-1 weak approximations (Definition 4.2) each other. Our result (Thcorem 4.3) shows that VT, n, 3/ such that SDE (a) and SVAG (b) are order-1 weak approximations of cach other. In particular, Li ct al. (2019a) requires an infinitesimal 1y and our result holds for finite 1).
synthetic dataset. Several works (Hoffer et al., 2017; Smith and Le, 2018; Chaudhari and Soatto, 2018; Smith et al., 2018) have had success using larger LR to preserve the scale of the gradient noise and hence maintain the generalization properties small-batch training. The relationship between LR and generalization remains hazy, (Shallue al., 2019) empirically demonstrated that the generalization етrоr can depend on many other training hyperparameters.
4 Stochastic Variance Amplified Gradient (SVAG)
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 Itô 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).
4.1 The SVAG Algorithm
For chosen hyperparameter N+, we define
Fk+1 T2 C (mx),
(5)
where "x (Nk,1, (e,2) with 7k. 72 sampled independently and
27 C.p,
1 /21 (.).
2
SVAG is equivalent to performing SGD on a new distribution of loss functions constructed from the original distribution: the new loss function linear combination of two independently sampled losses Cr... and L.,,,, 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 I, i.e.
VL' (c)) (z)( J]. Therefore, the Itō SDE that matches the first and second order moments is always (2). We note that SVAG is equivalent to SGD when 1 = 1, and both the expectation and covariance of the one-step update
/ (x) = V (x), where Σ (xx) := [[(VC'
(Ik+1* - 2x) are proportional to 1/l, meaning the direction of the update is noisier when | increases.
4.2 SVAG Approximates the SDE
Definition 4.1 (Test Functions). Class G' of continuous functions Ed - IR has polynomial growth if Ai E G there exist positive integers K1, K2 > 0 such that for all x E Rd (g(x)) & 61 + 22)2x)).
For a E N+, we denote by Go the set of a-times continuously differentiable functions where all partial derivatives of form
8" "1.1..00
s.t. d O; = a S a, are also in
1
Definition 4.2 (Order-a weak approximation). Let {Xp 1 T]} and (= k=0 be families of continuous and discrete stochastic processes parametrized by 7). We say {XP) and (2T) are order-a
5

---

weak approximations of each other if for every g E G2(a=1), there is constant C > 0 independent of such that
max
Eg(=)) Eg(X) knp S CCy
When applicable, we drop the superscript v, and say {Xe} and { are order-a (or a order) approximations of each other.
We now show that SVAG converges weakly to the Itō SDE approximation (2) when + oo, i.e., and Xg, 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 conditions are met:
(i) L = Ec, C -smooth, and E G
(ii) (2,(2) VC, (v)l S L I - y). for all æ, E Rd, where Ly > is a random variable with finite moments, i.e,, EL is bounded for k E N+.
(iii) Σi (X)
is - -moothh in X.
Let > be a constant and be the SVAG hyperparameter (5). Define {X, E [0,T]} as the
stochastic process (independent of q)) satisfying the 11 SDE (2) and {at 1 k: S [TT//]} as the trajectory of SVAG (5) where Tn = Xn. Then, SVAG {=k is an order-1 weak approximation of the independent of such that
SDE {Xe), for each E G there exists a constant
max [Bg(x>)) Eg 5
k= LIT/m)
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 {x²} 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.
4.3 Proof Overview
Let (X,"* 2 s} denote the stochastic process obeying the Itô SDE (2) starting from time s and with the initial condition X and {-
2 j} denote the stochastic process (depending on 1) satisfying SVAG (5) with initial condition æn = r. For convenience, we define X = X kn and
write N-. = x in
X
Alternatively, we write xx = N.. and Ix (x. == z..
Now for any 1 < 1 we interpolate between SVAG solution 2g and SDE solution X through a series of hybrid trajectories ()), i.e., the weight achieved by running SVAG for first j steps and then SDE from time j to k. The two limits of the interpolation are Xx( k:) (i.e., SVAG solution after steps) and ( 0) X (i.e., SDE solution after k time). This yields the following error decomposition for test function E G (see Definition 4.1).
[Eg(xx) - Eg( X )l [Eg((x) - Ex(Xx)) S Σ
Eg ( + 1)) - (x
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 we can show that the single-step approximation error from switching from SVAG to SDE is uniformly upper bounded by O 2 ). (See Figure 6 for demonstration) Hence, the total error is O n = O(7).
"The C smoothness assumptions be relaxed by using the mollification technique in Li et al. (2019a).

---

Lemma 4.6. Define the one-step increment of the lt SDE as (zx N=.0 2. Then we have
(i) (x(x) - 0((-))
(ii) (i) A (T) - - (z) + 0(1
(ii) ((1
(iv) V E 2412 0(1
Lemma 4.7. Define the one-step increment of SVAG A(x) = 21
x. Then we have
(i) (A(x) = -VV
(ii) (x(z)A(x)) = + VVC = [zxx) + 0(
(iii) [A(x) = 2 12 - to (3V((T) @ E(x) + VL = -1122)
(iv) V (EAA())44==010
where Mxx) B( (1) v((=)) and T denotes the symmetrization of tensor T, i.e. Tisk- = - g' Toyls where J' k sums over all permutation of i,
Though (i) and (ii) in Lemma 4.7 hold for any discrete update with LR = 7 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,
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.
Understanding the Failure of SDE Approximation and LSR
In this section, analyze how discretization error, caused by large LR, leads the failure of the SDE approximation (Section 5.1) and LSR (Section 5.2) for scale invariant networks (e.g., nets equipped with BatchNorm (loffe and Szegedy, 2015) and GroupNorm (Wu and He, 2018)). To get best generalization, practitioners often add Weight Decay (WD, a.k.a Ee regularization: see (7)). Intriguingly, unlike the traditional setting where l2 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 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 al., 2019a; Allen-Zhu et al., 2019a) is in general an open problem. Therefore, we rely on other natural test functions (Definition 5.1).
5.1 Failure of SDE Approximation
In Theorem 5.2, we show that the SDF-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 expectations of these natural test functions (Definition 5.
We consider the below SDE (6) with arbitrary expected loss L(x) and covariance Σ(x), and the moment-matchie matching SGD (7) satisfying (z) L(x) and E(x) ryE(x) where Σ(x) is the covari- ance of V (x)). In the entire Section 5, we will assume that for all 7, C, is scale invariant (Arora et al., 2019b; Li and Arora, 2020a), i.e., C. (z) = C, Ve and x E IRd \. {0)."
'The 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 ecale-invarian parameters.
7

---

(c(()) + xx dt + '¹² (X,)dW, (x,)dw,
(6)
1k+1 = - (3x) mx
(7)
We will measure the closeness of two distributions by three test functions: squared weight norm |
squared gradient norm [V<())) and trace of noise covariance Tr(S(x)). We say two equilibrium distributions are C-close if expectations of these test functions are within multiplicative constant C.
Definition 5.1 (C-closeness). Assuming the existence of the following limits, we use
R lim E|r,l 1-000
Rx = lim E|
Gx = lim I3 VL 1-000
Gx = lim V ((X.). t-
No lim E[Tr(
No == lim ((x)]]
and 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 each other iff
1 Rix Go
NN S.
(8)
C
R
G NO
We call No and N= the noise-to-signal ratio (NSR), and below we show that it plays an important &o Gx role. When the LR of SGD significantly exceeds the NSR of the corresponding SDE, the approxima- tion fails. Of course, we lack a practical way to calculate NSR of the SDE so this result existential rather than effective. Therefore we give condition in terms of NSR of the SGD that suffices to imply failure of the approximation. Experiments later in the paper show this condition effective at showing divergence from SDE behavior.
Theorem 5.2. If either (i). 1j G~ N=2 or (ii). N= < c 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 second order discretization error related to the norm of full-batch gradient, (((x)) 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 metries at the equilibrium of SGD Equation (11) 1) and SDE Equation (12), and thus conclude that a sufficiently large to-ssnall ratio (NSR) a necessary condition for C-closeness.
Here the role of scale invariance is to simplify the norm dynamics both SGD (10) and SDE (9) by removing the cross term. This is because of well-known property of scale-invariance, the orthogonality between gradient and the weight itself, i.e., (v), x) 0 for any 7, (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 7 v 2020), by Itô's lemma:
N N= (C² 1) and
1-- <
Noc Following the derivation in (Li et al., G
G
X, 1 - 22 X, 2 + ET(E)(X))).
Again by the orthogonality between gradient and weight, can be shown that for SGD (7),
E Ba+tl 2
E|ra| =(1 - m7 2 rx o V'E
E
(10)
= - 2 2 -
If both Ty and X, have reached their equilibriums, both LHS of (9) and (10) 0, and therefore
(2 7))AR = 10 + Noo,
(11)
2AR
N No.
(12)
Combining (11), (12), and (8), we have
G + NNo S 2NRx & 2ACR = CN...
Applying (8) again, we have o/ + No S Cn(Cxo Noo) S Co S C The proof completed by comparing the first and third, the second and the fourth terms respectively. D

---

PreNesNet32 (BN)
Preneswees2 (GN)
VGG19 (UN)
161
10"
Gell (tog scate)
20
10
1.01
10"1
10
10
10:2
101
10
10
10
10r
10a
10'
10'
10"
0.6
Final Train
Final Train
0.8
Final Train
Final
1st Phase Train
1st Phase Train
1st Phase Train
Ereer
0.4
0.2
0.0
0.0
llog scale)
Batch (log scale)
(log scale)
Figure 3: Experimental verification for theory predicting the failure of Linear Scaling Rule. We modify PreResNet-32 and VGG-19 19 be scale-invariant (according Appendix of (Li et al., 2020)). All three settings use the same LR schedule, LR= 0.8 initially and is decayed by 0.1 at epoch 250 with 300 epochs total budget. Here, G. and N₂ are the empirical estimations of Goo and No taken after reaching equilibrium in the first phase (before LR decay). Per the approximated version of Theorem 5.6, ie.. 1 KB S C BN BN /G we baseline runs with different batch sizes B report the maximal and minimal predicted critical batch size, defined as intersection of the threshold (C//N, C9) with the green and blue lines. respectively. We choose a threshold C 2, and consider LSR to fail if the final test error exeeeds lowest achieved test error by more than 20% of its value, marked by the red region on the plot. Further settings and discussion are in Appendix F.
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 partial answer, that the same gap happens already between order-1 and order-2 SDE approximation, when 1j Z Nx - This suggests G
failure of SDE approximation may be due to missing some second order term, and thus higher- order 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.
5.2 Failure of Linear Scaling Rule
In this section we derive similar necessary condition for LSSR to hold.
Similar to Definition 5.1, we will use R GB,,,, NB." 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 / / (.NNB... xx7) because the covariance scales inversely to batch size, Σ (x) = kΣ*B(r) (x).
Definition 5.4 ((C, Rt)-Linear Scaling Invariance). We say SGD (7) with batch size B and LR exhibits (C, k)-LSI if, for a constant C such that 0 < C VR.
1
RB,
NB.
GB.r
S C.
(13)
C
R
B
KN
B.
(.
We show below that (C, LSS fails if the NSR No. is too small, thereby giving a certificate for Go
failure of (C, R)-LSI even without a baseline run.
Theorem 5.5. For any f, C, and &. such that
NeB. kn oc
1
< (1
(14)
B.,
K c
K
SGD with batch size B and LR does not exhibit (C, k) -SI.
We now present a simple and efficient procedure to find the largest & which (C, K)-LSI will hold, providing useful guidance make hyper-paramett tuning more efficient. Before doing so, one must

---

Validation Accuracy
B = 128.18811 17 0.8
B = 1024, 6.4
0.8
0.8
0.6
/=
/ 8
0.6
1= 1
8
1= 2
1 = 16
= 2
/ 16
0.4
0.4
1 = 4 1 = 32
4
= 32
0.2
0.2
0
10000
20000
30000 Step
1000
2000
3000
4000
Effective Step Step
Effective Step
Figure 4: SVAG converges quickly and matches SGD (left) or shows failure the SDE approximation when LSR breaks (right). We train PreResNet32 with BN on CIFAR-10 for 300 epochs, decaying by 0.1 at epoch 250. SVAG takes smaller steps simulate the continuous dynamies in time, against "effective steps." and note that SVAG with = 1 is equivalent to SGD. We predict in Figure that LSR (and thus, the SDE approxination) breaks converges to a limiting trajectory different from SGD, suggesting that the SDE approximation did indeed break.
plot the accuracy
1024 for this training setting, and here we observe SVAG
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 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' 2 works well. One can estimate BB and & from baseline run. Then, one can straightforwardly compute the value for the k threshold given the theorem below. We conduct this process in Figure and Appendix F to test our theory.
Theorem 5.6. For any B, tp. C, and
N.. K > (1+
Bin
BB, 20
(~ 0
when
> 1),
(15)
G
GB.
B.
B.n
SGD with batch size B and LR r) does not exhibit (C, KLSL
Experiments
We provide our code at https:: //tp:::///tttbb/ com/ com/sadhii kamakaaaaaaam.aa lammmaaaad//aaaaa 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 and Appendix F. test SVAG on common architectures and datasets and report the results. Theorem 4.3 shows that SVAG converges to the SDE as 1 - oc, but we note that SVAG needs times as many steps as SGD to match the SDE. Therefore, in order for SVAG to be computationally efficient simulation of the SDE, we hope to observe convergence for small values of This is confirmed in Figure and Appendix F. The success of SVAG in matching SGD in many cases indicates that studying the Itô 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 I are interesting learning algorithms 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.
7 Conclusion
We present computationally efficient simulation SVAG (Section 4) that provably converges to the canonical order-1 SDE (2), which we use to verify that the SDE meaningful approximation for SGD in common deep learning settings (Section 6). We relate the discretization error to LSR (Definition 2.1): in Section we derive testable necessary condition for the SDE approximation and LSR to hold, and in Figure 3 demonstrate its applicability to standard settings.
10

---

Acknowledgement
The authors acknowledge support from NSF, ONR, Simons Foundation, DARPA and SRC. 7L. is also supported by Microsoft Research PhD Fellowship.
References
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019a. URL httpt:/ http::sss///ppeeeeepeeeeeee htt.://ooe.nntttteeereeeeeennnoo
neur seuurips.. cc/ -apee.
paper/2019 9/ file/ (tt.e//22dddde233332222233223 wtt.3/22222222222222222222222222
Zeyuan Allen-Zhut Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning. volume 97 of Proceedings of Machine Learning Research, pages 242 252, Long Beach, California, USA, 09-15 Jun 2019b. PMLR.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pages 322-332. PMLR, 2019a.
Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu. Theoretical analysis of auto rate-tuning by batch normalization. In International Conference on Learning Representationss 2019b.
Yoshua Bengio. Practical recommendations for gradient-based training of deep architectures. In Neural networks: Tricks of the trade, pages 437-478. Springer, 2012.
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL https:: h/op:://www wandb. com/. Software available from wandb.com.
Nils Bjorck, Carla P Gomes, Bart Selman, and Kilian Q Weinberger. Understanding batch normal- ization. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances Neural Information Processing Systems 31, pages 7694-7705. Curran Asso- ciates, Inc., 2018.
Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational inference, con- verges to limit cycles for deep networks, In International Conference on Learning Representations, 2018.
Xiang Cheng, Dong Yin, Peter Bartlett, and Michael Jordan. Stochastic gradient and langevin processes. In International Conference Machine Learning, pages 1810-1819. PMLR, 2020.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Proceedings of the 34th International Conference on Machine Learning-Voii Volume 70, pages 1019-1028. JMLR.org, 2017.
Simon S. Du, Xiyu Zhai, Barnabas Poezos, and Aarti Singh. Gradient descent provably optimizes prerpparameterized neural networks. In International Conference on Learning Representations, 2019.
Cong Fang, Zhouchen Lin, and Tong Zhang. Sharp analysis for nonconvex sgd escaping from saddle points. In Conference on Learning Theory, pages 1192-1234. PMLR, 2019.
Yuanyuan Feng, Lei Li, and Jian-Guo Liu. Semi-groups of stochastic gradient descent and on- line principal component analysis: properties and diffusion approximations. arXiv preprint 1712.07712 2017.
Priya Goyal, Piotr Dollár, Ross Girshick. Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint 1rXiv:1706.02677, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.

---

Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the general- ization gap in large batch training of neural networks. In I. Guyon, U. Luxburg, S. Bengio. H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems 30, pages 1731-1741. Curran Associates, Inc., 2017.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Fix your classifier: the marginal value of training the last weight layer. In International Conference on Learning Representations, 2018.
Wenqing Hu, Chris Junchi Li, Lei Li, and Jian-Guo Liu. On the diffusion approximation of nonconvex stochastic gradient descent. Annals of Mathematical Sciences and Applications, 4(1), 2019.
Sergey Ioffe and Christian Szegedy. Batch normalization: accelerating deep network training by reducing internal covariate shift. In Proceedings of the 32nd International Conference on International Conference on Machine Learning-Volde Volume 37, pages 448 456. JMLR. org, 2015.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31, pages 8571-8580. Curran Associates, Inc., 2018.
Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint arXiv: 171 17.6623,, 2017.
Chi Jin, Rong Ge, Praneeth Netrapalli, Sham M Kakade, and Michael Jordan. How to escape saddle points efficiently. In International Conference on Machine Learning, pages 1724 1732. PMLR, 2017.
Sato Ken-lti. Lévy processes and infinitely divisible distributions. Cambridge university press, 1999.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017.
P.E. Kloeden and E. Platen. Numerical Solution of Stochastic Differential Equations. Stochastic Modelling and Applied Probability. Springer Berlin Heidelberg, 2011. ISBN 9783540540625. URL https::tt..... / /books. google. com/ www/books books ? -iddBCvtssoom CMC.
Alex Krizhevsky, One weird trick for parallelizing convolutional neural networks. arXiv preprint 4rXiv:1404.5997, 2014.
Daniel Kunin, Javier Ergastuy-Brenal Surya Ganguli, Daniel L. K. Yamins, and Hidenori Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics, 2020.
Yann A. LeCun, Léon Bottou, Genevieve B. Orr, and Klaus-Robert Müller. Efficient BackProp. pages 9-48. Springer Berlin Heidelberg, Berlin, Heidelberg. 2012. ISBN 978-3-642-35289.. doi:
Qianxiao Li, Cheng Tai, and E. Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In Proceedings of the 34th International Conference on Machine Learning-Ver Volume 70, pages 2101-2110. JMLR. org, 2017.
Qianxiao Li, Cheng Tai, and Weinan. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. J. Mach. Learn. Res., 20:40-1, 2019a.
Yuanzhi Li. Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché-Bus Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 11674 1685. Curran Associates, Inc., 2019h.
Zhiyuan Li and Sanjeev Arora. An exponential learning rate schedule for deep learning, In Interna- tional Conference on Learning Representations, 2020a. URI. https: /ttp:r/e/s/////ee...e... net/ forum? i SaaammidddUU@TTSSF....
12

---

Zhiyuan Li and Sanjeev Arora. An exponential learning rate schedule for deep learning. In Interna- tional Conference on Learning Representations, 2020b.
Zhiyuan Li, Kaifeng Lyu, and Sanjeev Arora. Reconciling modern deep learning with traditional optimization analyses: The intrinsic learning rate. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Pro- cessing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6- 62, 2020, virtual, 2020. URL htttps://proceeeinne neurinng....nunnnnnnneee newr ips cc/ Abstract html.
paper/2020/ http////////aaa///555555/52/222
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:/608.0349 2016.
Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent as approximate bayesian inference. The Journal of Machine Learning Research, 18(1):4873 4907, 2017.
Mohammad Mohammadi, Adel Mohammadpour, and Hiroaki Ogata. On estimating the tail index and the spectral measure of multivariate a- a-stable distributions. Metrika, 78(5):549 561, 2015.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.
Thanh Huy Nguyen, Umut Simsekli, Mert Gürbüzbalaban, and Gaël Richard. First exit time anal- ysis of stochastic gradient descent under heavy-tailed gradient noise. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer. Florence d' Alché-Buc, Emily B. Fox, and Roman Gar- nett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 273 283, 2019, URL https::/procedddinn neurips. httppp.cc.pp...........eee..eeee cc/paper/ 9b0 www.:99885222e42e22 2019 httphahhhhaaaaaaa a 00w009000000022eeeeeeeeee33333 -bassetttaaaaaaaatt......aa..a.e html.
Philip Protter, Denis Talay, al. The euler scheme for lévy driven stochastic differential equations. The Annals of Probability, 25(1):393 423,, 1997.
Christopher J. Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl- -ick-Dicks,,i,, Roy Frostig, and George E. Dahl. Measuring the effects of data parallelism on neural network training. Journal of Machine Learning Research, 20(1 12):1 1-4,, 2019. URL http:// jmlr. org papers/v20.
18-789. html.
Bin Shi, Weijie Su, and Michael Jordan. On learning rates and schrödinger operators. arXiv preprint ar arXiv:2004.0699 2020.
Umut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. A tail-index analysis of stochastic gra- dient noise in deep neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, ed- itors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 1827-5837, PMLR, 09-15 Jun 2019. URL http:://pstt http:://pprcceddings
.mlr. wwwss/esssseeee v s wissseeekk 111 9a html.
L. N. Smith. Cyclical learning rates for training neural networks. In 2017 IEEE Winter Conference on Applications of Computer Vision (WACV), pages 474,472, 2017. doi: 10. 1.11099//AC..201..38.
Samuel L. Smith and Quoe V. Le. A bayesian perspective on generalization and stochastic gradient descent. International Conference on Learning Representations, 2018
Samuel L. Smith, Pieter-Jan Kindermans, and Quoe V. Le. Don't decay the learning rate, increase the batch size. In International Conference on Learning Representations, 2018
Samuel L. Smith, Erich Elsen, and Soham De. On the generalization benefit of noise in stochastic gradient descent, 2020.
Samuel I.. Smith, Benoit Dherin, David G. T. Barrett, and Soham De. On the origin of implicit regularization in stochastic gradient descent, 2021.
13

---

Jingfeng Wu, Wenqing Hu, Haoyi Xiong, Jun Huan, Vladimir Braverman. and Zhanxing Zhu. On the noisy gradient descent that generalizes as SGD. Hal Daumé III and Aarti Singh. editors, Proceedings of the 37th International Conference on Machine Learning. volume 119 of Proceedings of Machine Learning Research, pages 10367-10376, 13-18 Jul 2020.
Yuxin Wu and Kaiming He. Group normalization. arXiv preprint arXiv: 2018.
Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima, 2021.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why {adam} beats {sgd} for attention models, 2020. URL https:
//openreview..../m/m/.e net/ forum?
Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven HOI, and Weinan Towards theoretically understanding why sgd generalizes better than adam in deep learning, 2020.
Zhamxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior escaping from sharp minima and regularization effects. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 1654-7663 PMLR, 09-15 Jun 2019.
14


Reviews:
[{"review_id": "hezDM0UhUVJ", "reviewer": "Reviewer_Jauw", "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.", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "-\tIt is claimed in this paper that the discrete nature of SGD is not an essential ingredient to the generalization (Line 130—133), as evidenced by Fig. 4. However, recently Smith et al. (2021, https://openreview.net/pdf?id=rq_Qr0c1Hyo) show a contradictory result that discretization influences generalization significantly by proving that discrete time SGD actually minimizes a “modified loss”. Admittedly, the difference in theory may originate from the different minibatch sampling mechanism considered in either paper. The mechanism considered in the current paper is random with replacement, while that in Smith et al. (2021) considers permutation of minibatches in each epoch (sample without replacement in each epoch), which is perhaps more practically relevant. However, I assume the empirical analysis in the current paper is performed with the usual minibatch sampling mechanism that is considered by Smith et al. (2021), but why do the empirical results in the current paper suggest that discretization is not essential? Suppose one considers the permutated minibatches like Smith et al. (2021) for SVAG, would the process that SVAG approximate becomes another continuous-time stochastic process but with the modified loss of Smith et al. (2021)? How can we reconcile the results in the current paper and Smith et al. (2021)?\n-\tIn Fig. 4, SVAG generalizes better and is more stable than SGD as l increases at each effective step. Because SVAG is getting close to SDE as l increases, does it imply that the SDE would provide a good generalization performance given that we can exactly compute it? Would SDE yield a generalization upper bound for SGD?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations of the SVAG algorithm are adequately addressed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "bIAKfqbYeVy", "reviewer": "Reviewer_3hdD", "summary": "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. ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well written and easy to read. The motivations underlying their work is clear and are pertinent to me. The related work is concise but relevant. All the results are clear and are discussed. My principle concern relies in the fact that the authors show that SVAG converges weakly to the classical SDE approximation (2), however (except for the experiments Figure 4)  it is not clear how SVAG and SGD compare as soon as $l$ is not equal to $1$. However I still believe that their paper gives some valuable insight on how and when the Itô SDE approximates well SGD for macroscopic step sizes.\n\n\nHere are some minor concerns: \n- l207: \"(independent of $\\eta$): I assume you mean \"independent of $l$\"\n- maybe it is just me but it took me some time to understand that SVAG has the same first and second order moments as the Itô SDE (2) **when using a discretisation $\\mathrm{d} t = \\eta / l$** (and not $\\mathrm{d} t = \\eta$ as for SGD). This is clear when looking at the theorem but making it more explicit and explaining in the main text would make it easier to understand.\n- Figure 1: what is the value of $l$ here ?\n- Figure 3: how come there is this accuracy drop for gradient descent at step ~15000 ? \n- Section 5.1: I would have enjoyed a a more in depth discussion around theorem 5.2: do you think this condition or a similar condition holds for other losses than scale invariant losses ? what is particular for scale invariant losses for this results to hold ? why consider a regularising term $\\lambda \\vert x \\vert^2$, does the result still hold without regularisation ?\n- Section 5: I am not familiar with scale invariant losses, my overall concern is how relevant are your results outside this class which (to me) seem very specific. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations are discussed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "DzBWm63-nqB", "reviewer": "Reviewer_aTSh", "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.", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n* The SVAG algorithm is new and interesting. It provides a practical methodology for weakly approximating the SDE corresponding to SGD, $\\mathrm{d} X_t = -\\nabla \\mathcal{L}(X_t)\\mathrm{d}t + \\big(\\eta \\Sigma(X_t)\\big)^\\frac{1}{2}\\mathrm{d}W_t$. To the best of my knowledge, it is the first algorithm that does so using stochastic gradients and with no restriction on $\\eta$. The authors demonstrate how it can be applied as a diagnostic tool for testing if SGD and its SDE approximation have similar behaviour. This is useful as, for example, the linear scaling rule occurs when the SDE approximation holds (Smith et al., 2020).\n\n* The paper gives new theoretical insights into the failure of the SDE approximation and linear scaling rule – which is certainly of practical importance. Their theory fits nicely with a wide array of experiments.\n\n* The authors also discuss the role of non-Gaussian noise in their findings. In particular, they give experimental evidence that SGD and NGD can give similar test performances. This indicates that it is reasonable to model the gradient noise in SGD as Gaussian. Whilst these results are presented as secondary to the main results, they may be very relevant to the wider literature.\n\n* The experiments are thorough.\n\n* The paper is clearly written.\n\nWeaknesses:\n\n* The authors show that SVAG converges to the SDE in a weak sense, however in their experiment (figure 4), it appears they are making conclusions based on a single trajectory (which would only be justified by if SVAG converged to the SDE in a pathwise sense). If this is the case, then one or two sentences clarifying the limitations of the experiments would be helpful.\n\n* Similarly, the authors claim that the convergence of SVAG to the SDE for \"small values of $l$\" is \"confirmed in Figure 4\". I find this to be too strong a claim and would prefer the statement to be weaker, for example \"consistent with\" instead of \"confirmed\".\n\nConclusion:\n\nOverall, I believe this is a strong paper. In particular, it provides interesting theory and practical methodologies for testing the effectiveness of SDE approximations (which are already popular tools to studying SGD). The paper supports this with a wide variety of experiments.\n\nMinor comments:\n\n* In the introduction, the authors write \"but the proof needs the LR of SGD to be an unrealistically small (unspecified) constant\". How do we know if the LR is unrealistic if it is unspecified?\n\n* The authors write \"at equilibrium, the gradient norm must be smaller than its variance\". Is it not the squared gradient norm?\n\n* I'm not sure if plotting the train and test accuracies of standard gradient descent in Figure 3 is helpful.\n\n* It is not clear (from the main body of text) why the scale invariance of the loss function $\\mathcal{L}_\\gamma$ is essential to the theory in Section 5.\n\nTypos:\n\n* Stochastic Gradient Gescent (page 1)\n* $E\\big[\\big(\\nabla \\mathcal{L}_\\gamma (X) -  \\nabla \\mathcal{L}(X))\\big(\\mathcal{L}_\\gamma (X) -  \\nabla \\mathcal{L}(X)\\big)^{\\mathsf{T}}\\big]$ is missing a $\\nabla$ (pages 2 and 5)\n* $\\sqrt{\\frac{\\eta}{l}} \\Sigma^l(x) = \\sqrt{\\eta}\\Sigma^{1}(x)$ (page 5). Surely it is $\\frac{\\eta}{l} \\Sigma^l(x) = \\eta\\Sigma^{1}(x)$.\n* (ii) is used instead of (iii) in Lemma 4.5 (page 6)\n* The left quotation mark in ”equilibrium” is the wrong way round (page 7)", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "* I believe the authors could better address the potential limitations of their experiments (see \"Weaknesses\" part of main review).\n\n* Theorem 4.3 only applies on a finite time horizon $[0,T]$, whereas the optimization/MCMC literature is usually concerned with quantities obtained as $t\\rightarrow\\infty$. It could be worth highlighting this difference.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3.5 hours"}, "extra_text": {}}, {"review_id": "-mgoGX9ZQZD", "reviewer": "Reviewer_GfkN", "summary": "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:\n  - Proposing \"Stochastic Variance Amplified Gradient\" (SVAG), an efficient numerical method to test whether the trajectories of SGD and the corresponding SDE are close.\n  - 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.\n  - Providing necessary conditions for SDE and LSR to track SGD.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Significance**: this work studies the importance question of the validity of approximating SGD with SDE. However, this paper does not provide a formal answer. Rather, the paper provides an indirect argument via the proposed SVAG: SVAG provably follows a certain SDE and empirically tracks SGD well, hence such SDE should also be close to SGD.\n\nOn the positive side, since SVAG can be made to track SGD well empirically, this suggests we can potentially study SVAG (or its approximations) instead of SGD. Moreover, the paper provides conditions for the applicability of LSR (i.e. NSR not too small, batch size not too large) which are valuable contributions.\n\nHowever, the fact that SVAG interpolates between SGD (when taking $l=1$) and SDE (when taking $l \\rightarrow \\infty$) seems to be a negative result for using SDE to approximate SGD, since this means SDE and SGD sit on opposite sides of the spectrum and are _not_ close to each other. This suggests that SVAG is not the proper tool to study the SDE approximation to SGD, and I would like to have the authors' clarification on this.\n\n**Writing**: the paper states the messages of this work are clear overall, but some \n- Typo: end of line 65: \"the benefit of SGD\"\n- Typo: end of line 266: \"the equilibrium distributions\"\n- The first paragraph of Sec 5: what's the purpose of the discussion on weight decay?\n- Fig 4: the figures are a bit too crowded to read; perhaps also consider changing the color scheme since currently it's hard to tell between $l=1$ vs $l=32$.\n- Typo: end of line 828 (Thm E.1) should be a stop not a colon.\n\n========\nPost-rebuttal update: I thank the authors for their responses and have raised my score.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors discuss implications of their theoretical results and potential insights for practitioners. Immediate societal impacts beyond the research community do not apply.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4.5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a practical algorithm called SVAG, which is shown to converge in the weak sense
to the SDE approximation of SGD, for even moderately large step sizes.
Authors demonstrate its performance via experiments on neural networks, and show that SVAG can diagnose if SGD and its SDE approximation have similar behaviors. Authors further investigate the conditions in which the SDE approximation fails. 

The motivation of this paper is clear, and very timely. All the assumptions and the results are presented in a clear way. 

All reviewers agree that this paper is well-written and its results are interesting. Authors should address reviewers' suggestions in the final version of their work. Specifically, authors should add a discussion on the limitations of the experiments and amend the statement about SVAG converging for small values of l, clarify the bound in Li et al., 2019, and add more intuition about the importance of scale invariance.
Decision: Accept (Poster)