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

Introduction

Training with Stochastic Gradient Gescent (SGD) (1) and finite learning rate (LR) is largely considered essential for getting best performance out of deep nets: using infinitesimal LR (which turns the process into Gradient Flow (GF)) or finite LR with full gradients results in noticeably worse test error despite sometimes giving better training error [Wu et al., 2020, Smith et al., 2020, Bjorck et al., 2018].

Mathematical explorations of the implicit bias of finite-LR SGD toward good generalization have focused on the noise arising from gradients being estimated from small batches. This has motivated modeling SGD as a stochastic process and, in particular, studying Stochastic Differential Equations (SDEs) to understand the evolution of net parameters.

Early attempts to analyze the effect of noise try to model it as as a fixed Gaussian [Jastrzebski et al., 2017, Mandt et al., 2017]. Current approaches approximate SGD using a parameter-dependent noise distribution that match the first and second order moments of of the SGD (Equation 2). It is important to realize that this approximation is heuristic for finite LR, meaning it is not known whether the two trajectories actually track each other closely. Experimental verification seems difficult because simulating the (continuous) SDE requires full gradient/noise computation over suitably fine time intervals. Recently, Li et al. [2017, 2019a] provided a rigorous proof 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.

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 [Smith et al., 2020, Goyal et al., 2017].)

This paper casts new light on the SDE approximation via the following contributions:

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)

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.

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 gradient norm must be smaller than its variance.

Preliminaries and Overview

where each $\gamma_{k}$ is an i.i.d. random variable with the same distribution as $\gamma$. Taking learning rate (LR) $\eta$ 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 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 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_{t}$:

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.

[Krizhevsky, 2014, Goyal et al., 2017] When multiplying the minibatch size by $\kappa>0$, multiply the learning rate (LR) also by $\kappa$.

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 each iteration, e.g. the Euler-Maruyama method (3), which is called Noisy Gradient Descent in the rest of the paper. We are not aware of any empirical verification using conventional techniques, which we discuss in more detail in Section 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.

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 functions The discriminator net in GANs is an example of test function in machine learning.; 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 $|x_{t}|$, gradient norm $|\nabla\mathcal{L}(x_{t})|$ and trace of noise covariance $\operatorname{Tr}[\Sigma(x_{t})]$ 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

where $z_{k}\overset{\text{i.i.d.}}{\sim}N(0,I_{d})$. A basic step in analysis will be the following Error Decomposition:

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.

1 Understanding the Role of Non-Gaussian Noise

Some works have challenged the traditional assumption that SGD noise is Gaussian. Simsekli et al. , Nguyen et al. suggested that SGD noise is heavy-tailed, which Zhou et al. claimed causes adaptive gradient methods to generalize better than SGD. Xie et al. 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 Section 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 3 and Section F.3 that SGD (1) and NGD (3) with matching covariances achieve similar test performance on CIFAR10 ( $\sim 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 Itô 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 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-independent).

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 Section 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.

Related Work

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., $\Sigma(X_{t})\equiv I$), then the equilibrium of the SDE is the Gibbs distribution. Shi et al. 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. assumed the covariance $\Sigma(X)$ is locally constant to show that SGD can be used to perform Bayesian posterior inference. Zhu et al. argued that when constant but anisotropic SGD noise aligns with the Hessian of the loss, SGD is able to more effectively escape sharp minima.

Recently, many works have used the most common form of the SDE approximation (2) with parameter-dependent noise covariance. Li et al. and Kunin et al. used the symmetry of loss (scale invariance) to derive properties of dynamics (i.e., $\Sigma(X_{t})X_{t}=0$). Li et al. further used this property to explain the phenomenon of sudden rising error after LR decay in training. Smith et al. used the SDE to derive the linear scaling rule (Goyal et al. and Definition 2.1) for infinitesimally small LR. Xie et al. constructed a SDE-motivated diffusion model to propose why SGD favors flat minima during optimization. Cheng et al. 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. proved a strong approximation result for an SDE and MCMC-like dynamics, but not SGD. Wu et al. 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. 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.

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

where ${\bar{\gamma}}_{k}=(\gamma_{k,1},\gamma_{k,2})$ with $\gamma_{k,1},\gamma_{k,2}$ sampled independently and

2 SVAG Approximates the SDE

Let $\{X^{\eta}_{t}:t\in[0,T]\}$ and $\{x^{\eta}_{k}\}_{k=0}^{\lfloor\frac{T}{\eta}\rfloor}$ be families of continuous and discrete stochastic processes parametrized by $\eta$. We say $\{X^{\eta}_{t}\}$ and $\{x^{\eta}_{k}\}$ are order-$\alpha$ weak approximations of each other if for every $g\in G^{2(\alpha+1)}$, there is a constant $C>0$ independent of $\eta$ such that

We now show that SVAG converges weakly to the Itô SDE approximation in (2) when $l\to\infty$, i.e., $x_{lk}$ and $X_{k\eta}$ have the roughly same distribution. Figure 1 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.

Suppose the following conditionsThe $\mathcal{C}^{\infty}$ smoothness assumptions can be relaxed by using the mollification technique in Li et al. [2019a]. are met:

$\Sigma^{\frac{1}{2}}(X)$ is $\mathcal{C}^{\infty}$-smooth in $X$.

Let $T>0$ be a constant and $l$ be the SVAG hyperparameter (5). Define $\{X_{t}:t\in[0,T]\}$ as the stochastic process (independent of $\eta$) satisfying the Itô SDE (2) and $\{x_{k}^{\nicefrac{{\eta}}{{l}}}:1\leq k\leq\lfloor lT/\eta\rfloor\}$ as the trajectory of SVAG (5) where $x_{0}=X_{0}$. Then, SVAG $\{x_{k}^{\nicefrac{{\eta}}{{l}}}\}$ is an order-$1$ weak approximation of the SDE $\{X_{t}\}$, i.e. for each $g\in G^{4}$, there exists a constant $C>0$ independent of $l$ such that

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_{t}\}$ stays bounded away from the origin and infinity, then (ii) holds.

3 Proof Overview

Let $\{X^{x,s}_{t}:t\geq s\}$ denote the stochastic process obeying the Itô SDE (2) starting from time $s$ and with the initial condition $X^{x,s}_{s}=x$ and $\{x^{x,j}_{k}:k\geq j\}$ denote the stochastic process (depending on $l$) satisfying SVAG (5) with initial condition $x^{x,j}_{j}=x$. For convenience, we define $\widetilde{X}_{k}:=X_{\frac{k\eta}{l}}$and write $\widetilde{X}^{x,j}_{k}:=X^{x,\frac{j\eta}{l}}_{\frac{k\eta}{l}}$. Alternatively, we write $\widetilde{X}_{k}(x,j):=\widetilde{X}_{k}^{x,j}$ and $x_{k}(x,j):=x_{k}^{x,j}$.

Now for any $1\leq k\leq\lfloor\frac{lT}{\eta}\rfloor$, we interpolate between a SVAG solution $x_{k}$ and SDE solution $\widetilde{X}_{k}$ through a series of hybrid trajectories $\widetilde{X}_{k}(x_{j},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 $\widetilde{X}_{k}(x_{k},k)=x_{k}$ (i.e., SVAG solution after $k$ steps) and $\widetilde{X}_{k}(x_{0},0)=\widetilde{X}_{k}$ (i.e., SDE solution after $k$ time). This yields the following error decomposition for a test function $g\in G$ (see Definition 4.1).

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.5 and 4.6). 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(\frac{\eta^{2}}{l^{2}})$. Hence, the total error is $O(k\frac{\eta^{2}}{l^{2}})=O(\frac{\eta}{l})$.

Define the one-step increment of the Itô SDE as $\widetilde{\Delta}(x)=\widetilde{X}^{x,0}_{1}-x$. Then we have

Define the one-step increment of SVAG as $\Delta(x)=x^{x,0}_{1}-x$. Then we have

Though (i) and (ii) in Lemma 4.6 hold for any discrete update with LR $=\frac{\eta}{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 Section 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.6 are also first order approximations of SDE (2), by the same proof.

Understanding the Failure of SDE Approximation and LSR

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

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 then 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 will measure the closeness of two distributions by three test functions: squared weight norm $|x|^{2}$, squared gradient norm $|\nabla\mathcal{L}(x)|^{2}$, and trace of noise covariance $\operatorname{Tr}[\Sigma(x)]$. We say two equilibrium distributions are close to each other if expectations of these test functions are within a multiplicative constant.

We call $\frac{N_{\infty}}{G_{\infty}}$ and $\frac{\overline{N}_{\infty}}{\overline{G}_{\infty}}$ 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.

If either (i). $\eta>\frac{\overline{N}_{\infty}}{\overline{G}_{\infty}}(C^{2}-1)$ or (ii).$\frac{N_{\infty}}{G_{\infty}}<\frac{1}{C^{2}-1}$, then the equilibria of SDE (6) and SGD (7) are not $C$-close.

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 $\eta\gtrsim\frac{\overline{N}_{\infty}}{\overline{G}_{\infty}}(C^{2}-1)$. This suggests 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.

2 Failure of Linear Scaling Rule

In this section we derive a similar necessary condition for LSR to hold.

Similar to Definition 5.1, we will use $R_{\infty}^{B,\eta},G_{\infty}^{B,\eta},N_{\infty}^{B,\eta}$ as test functions for equilibrium achieved by SGD (7) when training with LR $\eta$ and mini-batches of size $B$. We first introduce the concept of Linear Scaling Invariance (LSI). Note here we care about the scaled ratio $N_{\infty}^{B,\eta}/(\kappa N_{\infty}^{\kappa B,\kappa\eta})$ because the covariance scales inversely to batch size, $\Sigma^{B}(x)=\kappa\Sigma^{\kappa B}(x)$.

We say SGD (7) with batch size $B$ and LR $\eta$ exhibits $(C,\kappa)$-LSI if, for a constant $C$ such that $0<C<\sqrt{\kappa}$,

We show below that $(C,\kappa)$-LSI fails if the NSR $\frac{N_{\infty}}{G_{\infty}}$ is too small, thereby giving a certificate for failure of $(C,\kappa)$-LSI even without a baseline run.

For any $B$, $\eta$, $C$, and $\kappa$ such that

SGD with batch size $B$ and LR $\eta$ does not exhibit $(C,\kappa)$-LSI.

We now present a simple and efficient procedure to find the largest $\kappa$ for which $(C,\kappa)$-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 2 and Appendix F, we find that $C=\sqrt{2}$ works well. One can estimate $G_{\infty}^{B,\eta}$ and $N_{\infty}^{B,\eta}$ from a baseline run. Then, one can straightforwardly compute the value for the $\kappa$ threshold given in the theorem below. We conduct this process in Figure 2 and Appendix F to test our theory.

SGD with batch size $B$ and LR $\eta$ does not exhibit $(C,\kappa)$-LSI.

Experiments

Figure 2 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 Section 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\to\infty$, 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 Section F.1. The success of SVAG in matching SGD in many cases indicates that studying the Itô SDE can yield insights about the behavior of 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.

Conclusion

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 2 we demonstrate its applicability to standard settings.

Acknowledgement

The authors acknowledge support from NSF, ONR, Simons Foundation, Schmidt Foundation, Mozilla Research, Amazon Research, DARPA and SRC. ZL is also supported by Microsoft Research PhD Fellowship.

References

Appendix A Preliminaries on SDE

Here, we review the common approximation schemes for SDEs and discuss why they are not efficient enough to be applied to the Itô SDE approximation for SGD. We adapt the information in Chapters 13 and 14 of . In general, an Itô SDE can be written as

where $\mu$ and $\sigma$ are called the drift and diffusion coefficients respectively. The standard Itô SDE (2) used to approximate SGD sets $\mu(X_{t},t)=-\nabla\mathcal{L}(X_{t})$ and $\sigma(X_{t},t)=(\eta\Sigma(X_{t}))^{1/2}$.

Suppose we want to solve the SDE on a time interval $[0,T]$. First, we discretize the time interval into $N$ equal steps $\tau_{1},...,\tau_{N}$ of size $\Delta t$. We will construct a Markov chain $Y$ that is a weak approximation in $\Delta t$ (Definition 4.2) to the true solution, and let $Y_{0}=x_{0}$ where $x_{0}$ is the initialization for the SGD trajectory.

where $\Delta W_{n}\overset{i.i.d.}{\sim}{\mathcal{N}}(0,\Delta t)$. In the ML setting, computing a single step in this Markov chain requires computing the full gradient (for $\mu(Y_{n},\tau_{n})$) and the covariance of the gradient (for $\sigma(Y_{n},\tau_{n})$). As such, modeling a single step in the recurrence requires making one pass over the entire dataset. The error of the approximation scheme scales with $\Delta t$, so making $N$ larger (thereby requiring more recurrence steps) will improve the quality of the approximate solution. We furthermore note that storing the gradient covariance matrix requires a large amount of memory. Each weight parameter in the network must be modeled by its own recurrence equation, so for modern day deep networks, this approximation seems computationally intractable.

The Euler-Maruyama scheme is considered the simplest approximation scheme for an Itô SDE. A variety of other schemes, such as the Milstein and stochastic Runge-Kutta schemes, have been derived by adding a higher order corrective term, taken from the stochastic Taylor expansion, to the recurrence computation. In particular, these schemes all still require the computation of $\mu$ and $\sigma$ at each step of the recurrence, so they remain computationally intractable for the Itô SDE used to approximate SGD.

A.2 Preliminary on Stochastic Process

We call a $m$-dimensional stochastic process $X=\{X_{t}:t\geq 0\}$ a Lévy process if it satisfies the following properties:

Independence of increments: For any $0\leq t_{1}<t_{2}<\cdots<t_{n}<\infty,X_{t_{2}}-X_{t_{1}},X_{t_{3}}-X_{t_{2}},\dots,X_{t_{n}}-X_{t_{n-1}}$ are independent;

Stationary increments: For any $s<t,\ X_{t}-X_{s}$ is equal in distribution to $X_{t-s};$

Continuity in probability: For any $\varepsilon>0$ and $t\geq 0$ it holds that $\lim_{h\rightarrow 0}P(|X_{t+h}-X_{t}|>\varepsilon)=0.$

We call a counting process $\{N(t):t\geq 0\}$ a Poisson process with rate $\lambda>0$ if it satisfies the following properties:

the number of events (or points) in any interval of length $t$ is a Poisson random variable with parameter (or mean) $\lambda t$.

Appendix B Discussion on Non-Gaussian Noise

In Section B.1 we give an example where LSR holds while SDE approximation breaks. In Section B.2, we show this example to a more general setting – infinitely divisible noise. We also explain why decreasing LR along LSR will not get a better approximation for SDE, while decreasing LR along SVAG will, since both operation preserves the same SDE approximation. In Section B.3, we discuss the possibility where the noise is heavy-tailed and with unbounded covariance.

Let $Z(t)$ be a $1$-dimensional Poisson process (Definition Definition A.2), where $Z(t)$ follows Poisson distribution with parameter $t$. We assume the distribution of the gradient on single sampled data $\gamma$, $\nabla L_{\gamma}(x)$ is the same as $Z(1)$ for any parameter $x$. For a batch ${\bm{B}}$ of size B (with replacement), since Poisson process has independent increments, $\nabla\mathcal{L}_{\bm{B}}(x):=\frac{1}{B}\sum_{\gamma\in{\bm{B}}}\nabla\mathcal{L}_{\gamma}(x)\overset{d}{=}\frac{Z(B)}{B}$.

Thus for any constant $T$ and initialization $x_{0}=0$, performing SGD starting from $x_{0}$ for $\frac{T}{B}$ steps with LR $B\eta$ and batch size $B$, the distribution of $x_{\frac{T}{B}}$ is independent of $B$, i.e.,

Thus LSR holds for all batch size $B$. Below we consider the corresponding NGD (3), $\{\hat{x}_{k}\}$, where

and $\{z_{i}\}_{i=0}^{\frac{T}{B}-1}\overset{i.i.d.}{\sim}N(0,1)$.

Thus it holds that $\hat{x}_{\frac{T}{B}}=-\eta T+\eta\sum_{k=0}^{\frac{T}{B}-1}z_{k-1}\overset{d}{=}-\eta(T+W_{T})$, where $W_{T}$ is a Wiener process with $W_{0}=0$, meaning the NGD final iterate is also independent of $B$, and constant away form the final iterate $x_{\frac{T}{B}}\overset{d}{=}-\eta Z(T)$. Indeed we can show the same result for Itô SDE (2), $\textup{{d}}X_{t}=-\textup{{d}}t+\sqrt{\eta}\textup{{d}}W_{t}$:

Thus we conclude that LSR holds but SDE approximation fails. Since NGD achieves the same distribution as Itô SDE, the gap is solely caused by non-gaussian noise.

However, the reader might still wonder, since batch size is always at least $1$, there’s always a lower bound for LR $\eta$ when going down along the ladder of LSR, and thus a discrete process with a finite step size of course cannot be approximated by a continuous one arbitrarily well. So isn’t this example trivial? In Section B.2, we will see even if we are allowed to use fractional batch size, and thus allow $\eta\to 0$, LSR can still hold without Itô SDE approximation.

B.2 Infinitely Divisible Noise and Lévy SDE

To understand why decreasing LR along LSR will not get a better approximation for SDE, and how LSR can hold without Itô SDE approximation when $\eta\to 0$, we assume the noise is infinitely divisible below for simplicity, which allows us to define SGD with fractional batch sizes and thus we can take the limit of $\eta\to 0$ along the ladder of LSR.

For SGD with batch size $B$, such a random loss function can be found when $m$ is a factor of $B$, where it suffices to define $\mathcal{L}^{m}$ as $m$ times the same loss with a smaller batch size $\frac{B}{m}$.Batch loss of nets with BatchNorm is not necessarily divisble, because (13) doesn’t hold, as the individual loss depends on the entire batch of data with the presence of BN. Still, it holds for ghost BatchNorm with $B$ equal to the number of mini-ghost batches. In other words, we can phrase LSR in a more general form, which only involves the distribution of the noise, but not the generating process of the noise (e.g. noise from sampling a batch with replacement).

Keep LR the same. Replace $\nabla\mathcal{L}_{\gamma}$ by $\nabla\mathcal{L}_{\gamma}^{m}$ and multiply the total number of steps by $m$.

Formally, $X_{t}$ is the solution of the following SDE driven by a Lévy process.

Why decreasing LR along LSR will not get a better approximation for SDE: The Lévy SDE is equal to the Itô SDE only when the noise is strictly Gaussian. Thus the gap induced by non-Gaussian noise will not vanish even if both SGD and NGD decrease the LR along LSR, as it will converge to the gap between Itô SDE and Lévy SDE. See Figure 5 for a summary of the relationships among SGD, NGD, Itô SDE, and Lévy SDE.

Since decreasing LR along LSR converges to a different limit than SVAG does, it’s natural to ask which part of the approximation in Lemma 4.6 fails for the former. By scrutinizing the proof of Lemma 4.6, we can see (i) and (ii) still hold for any stochastic discrete process with LR $\frac{\eta}{l}$ and matching first and second order moments, while the term $\frac{\eta^{3}(3-l^{-1})}{2l^{2}}\Lambda(x)$ now becomes $\frac{\eta^{3}}{l}\Lambda(x)$ for SGD along LSR, which is larger by an order of $l$. Therefore, the single-step approximation error becomes $O(l^{-1})$ and the total error after $\lfloor\nicefrac{{Tl}}{{\eta}}\rfloor$ steps remains constant.Such error does not only occur in the third order moment. It also appears in the higher moments. Therefore simply assuming the noise distribution is symmetric (thus $\Lambda=0$) won’t fix this gap.

SDE approximation is not necessary for LSR, even for LR $\eta\to 0$: We also note that though derives LSR by assuming the Itô SDE approximation holds, this is only a sufficient but not necessary condition for LSR. 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 Itô SDE limit. The loss landscape and noise distribution are constant, i.e., parameter-independent. This is also an example where the gap between SGD and Itô SDE is solely caused by non-gaussian noise, but not the discretization error.

B.3 Heavy-tailed Noise and Unbounded Covariance

experimentally found that the distribution of the SGD noise appears to be heavy-tailed and proposed to model it with an $\alpha$-stable process. In detail, in Figure 1 of , they show that the histogram of the gradient noise computed with AlexNet on CIFAR-10 is more close to that of $\alpha$-stable random variables, instead of that of Gaussian random variables. However, a more recent paper pointed out a fundamental limitation of methodology in : made a hidden but very restrictive assumption that the noise of each parameter in the model is distributed identically. Moreover, their test (Theorem B.3) of the tail-index $\alpha$ works only under this assumption. Thus the empirical measurement in ($\widehat{\alpha}<2$) doesn’t exclude the possibility that that stochastic gradient noise follows a joint multivariate Gaussian.

Let $\{X_{i}\}_{i=1}^{K}$ be a collection of i.i.d. random variables with $X_{1}\sim{\mathcal{S}}\alpha{\mathcal{S}}(\sigma)$ and $K=K_{1}\times K_{2}$. Define $Y_{i}:=\sum_{j=1}^{K_{1}}X_{j+(i-1)K_{1}}$ for $i\in\{1,\ldots,K_{2}\}$. Then the estimator

We provide the following theoretical and experimental evidence on vision tasks to support the argument in that it is reasonable to model the stochastic gradient noise by joint Gaussian random variables instead of $\alpha$-stable random variables even for finite learning rate. (Note SVAG (e.g., Figure 4) only shows that when LR becomes infinitesimally small, replacing the noise by Gaussian noise gets similar performance.)

In Figure 2, we find that the trace of covariance of noise is bounded and the empirical average doesn’t grow with the number of samples/batches (this is not plotted in the current paper). However, an $\alpha$-stable random variable has unbounded variance for $\alpha<2$.

In Figures 3, 19, 18, and 20, we show directly that replacing the stochastic gradient noise by Gaussian noise with the same covariance gets almost the train/test curve and the final performance.

Applying the test in Theorem B.3 on joint multivariate Gaussian random variables can yield an estimate ranged from $1$ to $2$ for the tail-index $\alpha$, but for Gaussian variables, $\alpha=2$. (Theorem B.4)

Another recent work also confirmed that the noise in stochastic gradient in ResNet50 on vision tasks is finite. However, they also found the noise for BERT on Wikipedia+Books dataset could be heavy-tailed: the empirical variance is not converging even with $10^{7}$ samples. We left it as a future work to investigate how does SDE approximate SGD on those tasks or models with heavy-tailed noise.

Specifically, when $d=K_{1}$ and $m=1$, taking $\Sigma=\beta\bm{1}\bm{1}^{\top}+(1-\beta)I$, we have

Appendix C Omitted Derivation in Section 4

We prove Theorem 4.3 in this section. The derivation is based on the following two-step process, following the agenda of :

Showing that the approximation error on a finite interval ($N=\lfloor\frac{Tl}{\eta}\rfloor$ steps) can be upper bounded by the sum of expected one-step errors. (Theorem C.1, which is Theorem 3 in )

Showing the one-step approximation error of SVAG is of order $2$, and so the approximation on a finite interval is of order $1$. (Lemmas 4.5 and 4.6)

Let us consider generally the question of the relationship between one-step approximations and approximations on a finite interval. Let $T>0$ and $N=\lfloor lT/\eta\rfloor$. Let us also denote for convenience $\widetilde{X}_{k}:=X_{\frac{k\eta}{l}}$. Further, let $\{X^{x,s}_{t}:t\geq s\}$ denote the stochastic process obeying the same Equation 2, but with the initial condition $X^{x,s}_{s}=x$. We similarly write $\widetilde{X}^{x,j}_{k}:=X^{x,\frac{j\eta}{l}}_{\frac{k\eta}{l}}$ and denote by $\{x^{x,j}_{k}:k\geq j\}$ the stochastic process (depending on $l$) satisfying Equation 5 but with $x_{j}=x$.

The following result is adapted from to our setting, which relates one-step approximations with approximations on a finite time interval. To prove it, we will construct hybrid trajectories interpolating between SVAG (5) and the SDE (2), as shown in Figure 6.

Then, for each $g\in G^{4}$, there exists a constant $C>0$, independent of $l$, such that

Let $T,l>0$, $N=\lfloor lT/\eta\rfloor$ and for convenience we also define $\widetilde{X}_{k}:=X_{\frac{k\eta}{l}}$. Further, let $\{X^{x,s}_{t}:t\geq s\}$ denote the stochastic process obeying the same Equation 2, but with the initial condition $X^{x,s}_{s}=x$. We similarly write $\widetilde{X}^{x,j}_{k}:=X^{x,\frac{j\eta}{l}}_{\frac{k\eta}{l}}$ and denote by $\{x^{x,j}_{k}:k\geq j\}$ the stochastic process (depending on $l$) satisfying Equation 5 but with $x_{j}=x$. Alternatively, we write $\widetilde{X}_{k}(x,j):=\widetilde{X}_{k}^{x,j}$ and $x_{k}(x,j):=x_{k}^{x,j}$. By definition, $\widetilde{X}_{k}(x_{k},k)=x_{k}$ and $\widetilde{X}_{k}(x_{0},0)=\widetilde{X}_{k}$.

Thus we have for any $1\leq k\leq\lfloor\frac{lT}{\eta}\rfloor$, we can decompose the error as illustrated in Figure 6,

By assumption (ii), we know the there is some $K^{\prime}\in G$,

Suppose $u^{1},\ldots,u^{k}\in G^{4}$ uniformly, that is, $u^{1},\ldots,u^{k}\in G$ and there’s a single $K_{0}\in G$ such that $\left|\tfrac{\partial^{s}u}{\partial x_{(i_{1})},\dots x_{(i_{j})}}(x)\right|\leq K_{0}(x)$, for $s=1,2,3,4$ and $i_{j}\in\{1,2,\ldots,d\},j\in\{1,\ldots,s\}$. Let assumption (i),(ii) in Thm. C.1 hold and $K_{1}(x),K_{2}(x)$ be such functions. Then, there exists some $K\in G$, independent of $l,r$, such that

Using Taylor’s theorem with the Lagrange form of the remainder, we have for any $j\in\{1,\ldots,k\}$,

Taking expectations over the first term, using assumption (i) of Thm. C.1, we get

Taking expectations over the second term, using assumption (i) of Section C.1 and Lemma D.1, we get

Note that by assumption (ii) of Section C.1, we have

We can deal with the third term similarly to the second term and thus we conclude

C.2 One-step approximation

To obtain (i)-(iii), we simply apply Lem. D.2 with $\psi(z)=\prod_{j=1}^{s}(z_{(i_{j})}-x_{(i_{j})})$ for $s=1,2,3$ and $i_{j}\in\{1,\ldots,d\}$ respectively. (iv) is due to Lemma D.1. ∎

Next, we estimate the moments of the SVAG iterations below. See 4.6

Recall $\Delta(x)=-\frac{\eta}{l}\nabla\mathcal{L}_{{\bar{\gamma}}}(x)$, where $\mathcal{L}_{{\bar{\gamma}}}(x)=\frac{1+\sqrt{2l-1}}{2}\mathcal{L}_{\gamma_{1}}(x)+\frac{1-\sqrt{2l-1}}{2}\mathcal{L}_{\gamma_{2}}(x)$. Taking expectations, (i) and (ii) are immediate. Note $|\Delta(x)|=O(l^{-0.5})$, (iv) also holds.

Below we show (iii). For convenience, we denote $\sqrt{2l-1}$ by $c$, $\nabla=\nabla\mathcal{L}(x)$, $\widetilde{\nabla}_{i}=\nabla\mathcal{L}_{\gamma_{i}}(x)-\nabla\mathcal{L}(x)$, for $i=1,2$ and $\widetilde{\nabla}=\frac{1+c}{2}\widetilde{\nabla}_{1}+\frac{1-c}{2}\widetilde{\nabla}_{2}=\nabla\mathcal{L}_{\gamma}(x)-\nabla\mathcal{L}(x)$. We have

Appendix D Auxiliary results for the proof of Thm. 4.3

Let $\alpha\geq 1$, there exists a $K\in G$, independent of $l$, such that

where $i_{j}\in\{1,\dots,d\}$ and $C>0$ is independent of $l$.

Using Cauchy-Schwarz inequality, Itô’s isometry, we get

where $C$ depends only on $\alpha$. Now, using the linear growth condition (4.3 (ii)) and the moment estimates in Theorem 19 in , we obtain the result. ∎

We prove the following Itô-Taylor expansion, which is slightly different from Lemma 28 in .

Suppose that $b,\sigma\in G^{3}$, and $X^{x,0}_{t}$ is the solution of the following SDE, with $X^{x,0}_{0}=x$.

That is, there exists some function $K\in G$ such that

We also prove a general moment estimate for the SVAG iterations Equation 5.

Let $\{x_{k}:k\geq 0\}$ be the generalized SVAG iterations defined in Equation 5. Suppose

where $C>0$ are independent of $l$ and $k$, which immediately implies, for all $k=0,\ldots,\lfloor\frac{lT}{\eta}\rfloor$,

Appendix E Omitted proofs in Section 5

In this section, we provide the missing proofs in Section 5, including Theorem 5.5, Theorem 5.6 and the counterpart of Theorem E.4 between $1$st order SDE (2) and $2$nd order SDE (32), which is Theorem E.1. We also provide the derivation of properties for scale invariant functions in Section E.5.

We will prove the theorem by showing the contrapositive statement: if the equilibriums of (7) and (6) are $C$-close, then $\eta\leq\frac{\overline{N}_{\infty}}{\overline{G}_{\infty}}(C^{2}-1)$ and $\frac{1}{C^{2}-1}\leq\frac{N_{\infty}}{G_{\infty}}$. Following the derivation in , by Itô’s lemma and scale invariance of $\mathcal{L}_{\gamma}$:

It can be shown that for SGD (7), it holds that

If both $x_{k}$ and $X_{t}$ have reached their equilibriums, both LHS of (23) and (24) are , and therefore

Applying (8) again, we have $\eta\overline{G}_{\infty}+\overline{N}_{\infty}\leq C\eta(G_{\infty}+N_{\infty})\leq C^{2}\overline{N}_{\infty}\leq C^{3}\eta N_{\infty}.$ ∎

E.2 Proof of Theorem 5.6

Suppose $(C,\kappa)$-LSI hold, similar to Equation 25, we have

Thus combining (27) ,(28) and (8), we have

Therefore we conclude that $\kappa\leq C^{2}(1+\frac{N_{\infty}^{B,\eta}}{G_{\infty}^{B,\eta}})$.

E.3 Proof of Theorem 5.5

Suppose $(C,\kappa)$-LSI hold, by (8), we have

E.4 Necessary condition for C𝐶C-closeness between 111st order and 222nd order SDE approximation

In this section we will present a necessary condition for $C$-closeness between $1$st order and $2$nd order SDE approximation, similar to that betweeen $1$st order approximation and SGD. The key observation is that the missing second order term $\eta G_{\infty}$ in $1$st order SDE, also appears in the $2$nd order SDE, as it does for SGD. Thus we can basically apply the same analysis and show the similar conclusion (Theorem E.4).

Below we recap the $1$st and $2$nd order SDE approximation:

$1$st order SDE approximation (with $\overline{\Sigma}=\eta\Sigma$):

(32) is an order-2 weak approximation of SGD (1):

The following theorem is an analog of Theorem 5.2.

If the equilibriums of (31) and (32) exist and are $C$-close for some $C>0$, then

where $\lambda$ is usually of scale $10^{-4}$ in practice and thus can be omitted when calculating upper bound.

Since $\mathcal{L}$ is scale-invariant, so $\nabla\mathcal{L}(X)^{\top}X=0$, which implies $|\nabla\mathcal{L}^{\prime}(X)|^{2}=|\nabla\mathcal{L}(X)|^{2}+\lambda^{2}|X|^{2}$. Plug in $\mathcal{L}^{\prime}$, we have

If we compare Equation 37 to Equations 25 and 26 (we recap them below), it’s quite clear $2$nd order is much closer to SGD in terms of the relationship between $R_{\infty}$, $G_{\infty}$ and $N_{\infty}$. Thus $1$st and $2$nd order SDE approximation won’t be $C$-close if $\frac{\overline{N}_{\infty}}{\overline{G}_{\infty}}$ is larger than some constant for the exact same reason that $1$st SDE is not $C$-close to SGD.

In detail, by combining (37), (26) and (33), we have

which imples \eta\leq\big{(}C^{2}(1+\frac{\eta\lambda}{2})-1\big{)}\min\{\frac{\overline{N}_{\infty}}{\overline{G}_{\infty}},\frac{\overline{N}_{\infty}^{\prime}}{\overline{G}_{\infty}^{\prime}}\}. ∎

derived a similar equation to Equation 35 in Appendix F of their paper.

E.5 Properties of Scale Invariance Function

These properties are proved in . We include them here for self-containedness.

$\left\langle x,\nabla\mathcal{L}(x)\right\rangle=0$.

$\forall c>0$, $c\nabla\mathcal{L}(cx)=\nabla\mathcal{L}(x)$.

For (1), by chain rule, we have $\left\langle x,\nabla\mathcal{L}(x)\right\rangle=\lim_{t\to 0}\frac{\mathcal{L}((1+t)x)-\mathcal{L}(x)}{t}=0$.

$\left\langle x,\nabla\mathcal{L}(x)\right\rangle=0$, $x^{\top}\Sigma(x)x=0$.

Appendix F Experiments

We use the models from Github Repository: https://github.com/bearpaw/pytorch-classification. For VGG and PreResNet, unless noted otherwise, we modified the model following Appendix C of so that the network is scale invariant, e.g., fixing the last layer. Such modification doesn’t lead to change in performance, as shown in . We use Weights & Biases to manage our experiments .

We verify that SVAG converges for different architectures (including ones without normalization), learning rate schedules, and datasets. We further conclude that for most the standard settings we consider (excluding the use of large batch size in Figures 4 and 14 and GroupNorm on CIFAR-100 in Figure 13), SVAG with large $l$ achieves similar performance to SGD, i.e. SVAG with $l=1$.

Theorem 4.3 only holds if each step of SGD is a Markov process, which is in part determined by how each example is sampled from the dataset. We describe three common ways that examples can be sampled from the dataset during training.

Random shuffling (RS): RS is standard practice in experiments and is the default implementation in PyTorch. A random shuffled order of datapoints is fixed at the start of each epoch, and each sample is drawn in this order. SGD with this sampling scheme can be viewed as a Markov process per epoch, although not per step. We use this method in our experiments, but our theory for SVAG (Section 4) does not cover this sampling method.

Without replacement (WOR): WOR requires drawing each sample i.i.d. from the dataset without replacing previously drawn ones. SGD with this sampling scheme can be viewed as a Markov process per step, so our theory for SVAG (Section 4) does cover this case.

With replacement (WR): WR requires drawing each sample i.i.d. from the dataset with replacement. SGD with this sampling scheme can be viewed as a Markov process per step, so our theory for SVAG (Section 4) does cover this case.

In Figure 7, we observe that SVAG (including SGD) behaves similarly when using all three of these sampling methods. Therefore, although our theory does not directly apply to the commonly used RS scheme, we can heuristically apply Theorem 4.3 to understand its behavior.

We furthermore note that our findings do not match the conclusion in that SGD with RS has a different implicit bias compared to WOR and WR. We suggest two possible reasons for this discrepancy: (1) Their result holds when $\eta$ is smaller than an unmeasurable constant, so it may be the case that their results do not apply to the constant LR regime we use SVAG in. (2) Their result concerns behavior after a single epoch and our experiments run for hundreds of epochs.

F.1.2 Further Verification of SVAG on more architectures on CIFAR-10

In Figure 8, we demonstrate that SVAG converges and closely follows SGD for PreResNet32 with BatchNorm (left), PreResNet32 (4x) with BatchNorm (middle) and PreResNet32 with GroupNorm (right).

In Figure 9, we demonstrate that SVAG converges and closely follows SGD for VGG16 without Normalization (left), VGG16 with BatchNorm (middle) and VGG16 with GroupNorm (right).

F.1.3 Further Verification of SVAG on more complex LR schedules

We verify that SVAG converges and closely follows the SGD trajectory for networks trained with more complex learning rate schedules. In Figure 11, we use the triangle (i.e., cyclical) learning rate schedule proposed in , visualized in Figure 10. We implement the schedule over 320 epochs of training: we increase the initial learning rate $0.001$ linearly to $0.8$ over $80$ epochs, decay the LR to $0.001$ over the next $80$ epochs, increase the LR to $0.4$ over $80$ epochs, and decay the LR to $0.001$ over the remaining 80 epochs. As seen in Figure 11, SVAG converges to the SGD trajectory in this setting.

We further test SVAG on the cosine learning rate schedule proposed in with $\eta_{\text{max}}=0.8$ and $\eta_{\text{min}}=0.001$ with total training budgets of $160$ epochs. We visualize the schedule in Figure 10. In Figure 12, we see that SVAG converges and closely follows the SGD trajectory, suggesting the SDE (2) can model SGD trajectories with complex learning rate schedules as well.

F.1.4 Further Verification of SVAG on more datasets (CIFAR-100 and SVHN)

F.2 Further Verification of Necessary Condition for LSR

We further verify the necessary condition for LSR (Theorem 5.6) using different architectures and datasets. Figure 15 tests the condition for ResNet-32 and wider PreResNets trained on CIFAR-10. Although our theory requires strict scale-invariance, we find the condition to still be applicable to the standard ResNet architecture , ResNet32, likely because most of the network parameters are scale-invariant. Figure 16 tests the condition for wider PreResNets and VGG-19 trained on CIFAR-100. We require the wider PreResNet to achieve reasonable test error, but we note that the larger model made it difficult to straightforwardly train with a larger batch size.

In Figure 15 and Figure 2, $G_{t}$ and $N_{t}$ are the empirical estimations of $G_{\infty}$ and $N_{\infty}$ taken after reaching equilibrium in the second to last phase (before the final LR decay), where the number of samples (batches) is equal to $\max(200,50000/B)$, and $B$ is the batch size.

Per the approximated version of Theorem 5.6, i.e., $B^{*}=\kappa B\lesssim C^{2}B{N_{\infty}^{B}}/{G_{\infty}^{B}}$, we use baseline runs with different batch sizes $B$ to report the maximal and minimal predicted critical batch size, defined as the x-coordinate of the intersection of the threshold ($\nicefrac{{G_{t}}}{{N_{t}}}=C^{2}$) with the green and blue lines, respectively. Both the green and blue line have slope $1$, and thus the x-coordinate of intersection, $B^{*}$, is the solution of the following equation,

For all settings, we choose a threshold of $C^{2}=2$, and consider LSR to fail if the final test error exceeds the lowest achieved test error by more than 20% of its value, marked by the red region on the plot. Surprisingly, it turns out the condition in Theorem 5.6 is not only necessary, but also close to sufficient.

F.3 Additional Experiments for NGD (Noisy Gradient Descent)

We provide further evidence that SGD (1) and noisy gradient descent (NGD) (3) have similar train and test curves in Figures 18, 19, and 20. To perform NGD, we replace the SGD noise by Gaussian noise with the same covariance as was done in . In , the authors trained a network using BatchNorm, which prevents the covariance of NGD from being exactly equal to that of SGD. Hence, we use GroupNorm in our experiments, which improves NGD accuracy. We note that each step of NGD requires computing the full-batch gradient over the entire dataset (in this case, done through gradient accumulation), which is much more costly than a single SGD step. Each figure took roughly $7$ days on a single RTX 2080 GPU.