On the SDEs and Scaling Rules for Adaptive Gradient Algorithms
Authors: Sadhika Malladi
Kaifeng Lyu
Abhishek Panigrahi
Sanjeev Arora
Venue: NeurIPS 2022 Accept (2022)
Keywords: stochastic differential equations
Adam
RMSprop
scaling rule

Abstract: Approximating Stochastic Gradient Descent (SGD) as a Stochastic Differential Equation (SDE) has allowed researchers to enjoy the benefits of studying a continuous optimization trajectory while carefully preserving the stochasticity of SGD. Analogous study of adaptive gradient methods, such as RMSprop and Adam, has been challenging because there were no rigorously proven SDE approximations for these methods. This paper derives the SDE approximations for RMSprop and Adam, giving theoretical guarantees of their correctness as well as experimental validation of their applicability to common large-scaling vision and language settings. A key practical result is the derivation of a square root scaling rule to adjust the optimization hyperparameters of RMSprop and Adam when changing batch size, and its empirical validation in deep learning settings.

On the SDEs and Scaling Rules for Adaptive Gradient Algorithms
Sadhika Malladi*
Kaifeng Lyu"
Abhishek Panigrahi
Sanjeev Arora
Department of Computer Science Princeton University
{smalladi ,lyu, arora]scs. prooralacs..piimmttoo.... princeton.edu
Abstract
Approximating Stochastic Gradient Descent (SGD) Stochastic Differential Equation (SDE) has allowed researchers to enjoy the benefits of studying a contin- uous optimization trajectory while carefully preserving the stochasticity of SGD. Analogous study of adaptive gradient methods, such as R.Sprrpp and Adam, has been challenging because there were rigorously proven SDE approximations for these methods. This paper derives the SDE approximations for RMSprop and Adam, giving theoretical guarantees of their correctness well experimental validation of their applicability to common large-scaling vision and language set- tings. A key practical result is the derivation of square root scaling rule to adjust the optimization hyperparameters RMSprop and Adam when changing batch size, and its empirical validation in deep learning settings.
1 Introduction
Distributed synchronous optimization environments have enabled very rapid training of models (in terms of wall-clock time) by allowing a large batch size. Understanding large-batch stochastic optimization is crucial to enjoying the speed-up of this setting. In this context, Krizhevsky (2014); Goyal et al. (2017) empirically discovered the linear scaling rule (LSR) for Stochastic Gradient Descent (SGD). It calls for scaling learning rate proportionately to the batch size while fixing the number of epochs. It was recognized that the validity of this scaling rule stems from the benefits to generalization due to noise from mini-batch gradient estimation. But naive analysis, as done in Hoffer al. (2017), suggested that the scaling rule for SGD ought be square root instead of linear. Subsequently, Daggequuuutly__sttggbbbll Jastrzebski et al. (2017) pointed out that since the phenomenon involves varying the LR even down zero, the correct analysis should invoke continuous view, namely stochastic differential equation (SDE). The SDE view helps identify the correct scaling of the noise and leads to a derivation of the linear scaling rule (see Section 2.2
However, extending the SDE approach-ize i... continuous time approximations to popular adaptive optimization algorithms, like RMSprop (Tieleman and Hinton 2012) and Adam (Kingma and Ba 2015), has been challenging due to their use of coordinate wise normalization. By ignoring gradient noise, Ma et al. (2022) derived intuitive continuous approximations for full-batch RMSprop and Adam; however, this deterministic view precludes a scaling rule.
Recent papers have suggested square root scaling rule for adaptive algorithms: set the learning rate proportionately to the square root of the batch size while fixing the number of epochs. Based on perturbation theory, Granziol et al. (2022) proposed the square root scaling rule for RMSprop and Adam but could only reason about optimization behavior near convergence, not along the entire trajectory. A square root scaling rule was also empirically discovered for another adaptive gradient
Equual Contribution.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

method called LAMB Youu et al. 2020). which is an optimization method with layerwise adaptive learning rates, designed for better optimization and generalization in large-batch training. Instead of tuning learning rates while increasing batch size, LAMB used the square root scaling rule to automatically adjust the learning rate and achieve strong performance on vision and language tasks.
In this paper, we make the following contributions.
1. We propose new SDE approximations for two popular adaptive optimization algorithms, RM- Sprop and Adam (Definitions 4.1 and 4.4) in Theorems [.2 and 4.5 We prove that these approximations are !st-order weak approximations (Definition 2.4). providing a calculus-based guarantee of the approximation strength between the SDEs and their corresponding discrete processes as was done for SGD and its SDE Li al. (2019).
2. Our SDE approximations immediately yield square-root scaling rules (Definitions 5.1 and (..2) for adjusting the optimization hyperparameters of Adam and RMSprop when changing batch size to ensure that the resulting discrete trajectories are all Ist-order weak approximations of the same SDE (Theorem 5.3). Experiments (Figures | and and Appendix I validate the scaling rules in the vision and language modeling domains.
3. We provide efficient experimental verification of the validity of the SDE approximation for the adaptive algorithms in realistic deep nets (Definitions 5. and 5.2). Direct simulation of the SDE, e.g., Buler-Maruyamn is prohibitively expensive because it requires computing the full gradient and noise covariance at fine- geeiggrannd intervals. Hence we adapt (Definition 6.2) the new and efficient SVAG simulation for SGD (Li et al. 202T) for use with our proposed SDEs and rigorously prove its correctness (Theorem 6.3) Using SVAG, we provide evidence that the proposed SDE approximations track the analogous discrete trajectories in many common large-scale vision and language settings (Figure 3 and Appendix H).
Preliminaries
We O u, v², V to denote coordinate-wist operators for multiplication, taking squares, taking square roots. For ease of exposition we modify RMSprop and Adam to use th in the update for 0₂ We also assume that Un is nonzero if c is 0 to avoid division by zero
instead of using u 1
Definition 2.1. RMSprop (Tieleman and Hinton 2012) is an algorithm that updates O2 as follows,
Ok+1- On 199k O (vok + +e)
Uk+ f0, + (1 $)
where 02 is the parameter, gk is the stochastic gradient at step k. and Uy an estimate for the second moment of gk.
Definition 2.2. Adam Kingmma and Ba 2015) is an algorithm that updates 0g as follows,
m 31 m& + (1 - -ss
Uk+ 320%% + (1 - 32)g4
m m. ② (1 -
xx+1 1x+1 + (1 1/(11-11
0x+1 = - + ① (v UK +c)
where 0ş is the parameter, 9k is the stochastic gradient at step k my is the momentum, and iy is an estimate for the second moment of gk.
2.1 Noisy Gradient Oracle with Scale Parameter
We abstract the stochastic gradient as being provided by a noisy oracle for the full gradient. We formulate the oracle to highlight the phenomenon of interest: changing the batch size affects the scale of the noise.
Definition 2.3. A Noisy Gradient Oracle with Scale Parameter (NGOS) is characterized by tuple G2 = Z.). Given (noise) scale parameter a > 0, G. takes an input 0 and returns 9 () oz, where (f) the gradient of f at A, z the gradient noise drawn from the probability distribution (0( with mean zero and covariance matrix E(0). We use G- to denote the distribution of g given a and 0. The probability distribution Z, (+) can change with the scale a, but the covariance matrix Σ(0) is fixed across different noise scales.
Experiments in Appendix G.2 show that this change does not significantly impact performance.

---

For all in Definitions 2.1 and 2.2 we assume that 9: is drawn from noisy gradient oracle Go. In our setting, as is common when batches are sampled with replacement, primarily controlled through the batch size; in particular, a 1 1/VB (see Appendix F.1 for derivation). For sampling with replacement on finite dataset of size n, where f0)).... Fn are the loss functions for the n data points (and the average of these functions f())), this covariance matrix for given parameter 0 can be explicitly written as
1 Σ T1
0(
-0 -
2.2 SDE Approximation and Scaling Rules
Under appropriate conditions it becomes possible to approximate SGD via an Itô SDE, which uses Brownian motion to model the noise and has the following general form, where W, is Wiener process: dX;; = b(X,)dt )dt + o (-) The SGD update rule for a loss f is xx+1 = Fy - 796.-
where 17 is the learning rate and gk is given by the NGOS on input æg. The following is the canonical interpretation of SGD as an SDE:
dx, = - + (X,
(2)
Equation 2 hints relationship between learning and gradient noise specifically, the linear scaling rule-since scaling batch size by factor K scales the noise covariance by 1/K, which can be canceled by scaling 7) by K as well (Jastrzębski ensures the SDE approximation does not change when using a different batch size. With the same methodology. the current paper studies the SDE approximations for adaptive gradient algorithms to derive the square root scaling rule for them.
al. 2017). Therefore, the linear scaling rule
2.3 Quality of SDE Approximation and Theoretical Assumptions
The quality of the SDE approximation can be measured empirically (Section 6 and bounded theoreti- cally using a calculus-based guarantee, which was initiated in the context of deep learning in Li et al. (2019). In particular, the theoretical guarantee uses the following notion of approximation between discrete and continuous stochastic processes by comparing iteration k in the discrete process with continuous time Ay/e, where T/e > 0 the (effective) step size of the discrete process.
Definition 2.4 (Order-1 Weak Approximation, Li et al. (2019). Let {X7" t E [0, and (at) [T/neJ [T//n be families of continuous and discrete stochastic processes parametrized by 7e- We say k==
(X," and {a:"") } are order-1 I weak approximations of each other if for every test function g with at most polynomial growth (Definition B. there exists a constant C > independent of n) such that
max [//m]
[g(x"")) Eg(X kmn ) S Cne
A function Rd + R is said to have polynomial growth if there exist positive integers K1, K2 > 0 for all æ E R2d (Definition B. The above definition measures the strength of the approximation by the closeness of test function g computed on the iterates of both trajectories. The approximation becomes stronger in this sense as 7/ becomes smaller. In the SDE approximation of SGD, 7/ = 17 and k steps correspond to continuous time kn). A key difference between SGD and adaptive algorithms is that 1/- = 1)² for both RMSprop and Adam, which means k steps correspond to continuous time 2 We validate this time-scaling theoretically in Section 4
such that [g(x)) < Kn + (2x2) |||||==
Now we formalize the assumptions needed in the theory. Since our analysis framework is based upon calculus, it becomes necessary to assume differentiability conditions on the NGOS (Definition 2.5) Similar diffcrentiability conditions also appear in prior SDE works (Li et al., 2019 202T), and we note that lately it has become clear that restricting differentiable losses (via differentiable node activations such as Swish (Ramachandran et al. 2017)) does not hinder good performance.
Definition 2.5 (Well-behaved NGOS). The loss function f and covariance matrix function in a NGOS G, are well-behaved if they 5 (1) V (f( is Lipschitz and C** -moothh (2) The square root of covariance matrix Σ 112 is bounded, Lipschitz, and -moooh,, and (3) All partial
Note: C smoothness be relaxed using the mollification technique from Li et al. 2019
3

---

derivatives of V f0) and Σ (2) up to and including the 4-th order have polynomial growth. We also say that the NGOS Go is bell-behaved if f and Σ are rell-behaved.
Deriving an SDE approximation also requires conditions on the noise distribution in the NGOS. is allowed to be non non-Gausiian, but not heavy-tailed. We require an upper bound on the third moment of the noise so that the distribution is not very skewed. For other higher order moments, we require Ez ||| [/ namely the L--norm of random variable || 2. to grow at most linearly as a function of 0 (which is implied by ensuring an upper bound on all even order moments), We note that the following conditions are standard in prior work using Itō SDEs to study SGD.
Definition 2.6 (Low Skewness Condition). The NGOS go is said to satisfy the low skewness condition if there a function (0) of polynomial growth (independent of σ) such that [0)= [20] I
Kg(0)/ aa for all 0 E id and all noise scale parameters a.
Definition 2.7 (Bounded Moments Condition). The NGOS G- said to satisfy the bounded moments condition if for all integers m 2 1 and all noise scale parameters there exists a constant C2m (independent of σ)) such that Eg (0) ||| (l2rm] Im < C2m (1 + (00l2) for all 0 E Rd.
For well-behaved loss f()) and covariance Σ(0), the above two conditions are satisfied when 2, is the Gaussian distribution with mean zero and covariance [() That is, the Gaussian NGOS g ~ ((V o² (0)) satisfies the low skewness and bounded moments conditions. The low skewness condition holds because the odd moments of a Gaussian are all zeros, and the bounded moments
condition can be verified since Ez=N(o.E mm] Tm S Ew-N((, 2mm Tm |
and Σ (0) is Lipschitz.
The Gaussian NGOS with a a= 1 can be seen as an approximation of the gradient noise in
VB mini-batch training with batch size B, if (0) is set to match with (I. But this does not directly
imply that the gradient noise in mini-batch training satisfies the low skewness and bounded moments
conditions, as the noise is not exactly Gaussian. Assuming that the gradient of the loss function f.
at every data point is Lipschitz, these two conditions can indeed be verified for all batch sizes B 2 1.
2.4 Discussion on Heavy-Tailed Gradient Noise
We note that Definitions 2.6 and 2.7 allow some non- non-Gaussianity in the noise, but Ka (0 and O2m
could be large in practice. In this case, higher order moments of the gradient noise have non-negliginn
effects on training that the Itō SDE cannot capture. Zhang et al. (2020) and Simsskli et al. (2019)
presented experimental evidence that the noise is heavy-tailed. This motivated Zhou et al. (2020) to consider a Lévy SDE approximation (instead of Itô SDE) to study the (worse) generalization behavior of Adam. However, the quality of the Lévy SDE approximation was not formally guaranteed (e.g.,
in the sense of Definition 2.4. so finding a guaranteed approximation for adaptive optimization algorithms remains an open problem. Moreover, [i et al. (202I); Xie et al. (22TT) highlighted issues with the evidence, noting that the measurements used in Simsekli et al. (2019) are intended only for scalar values. When applied to vector valued distributions the measurement can (incorrectly) identify a multidimensional Gaussian distribution as heavy-tailed too (Li et al. 2021). It is in general difficult to estimate the moments of the noise distribution from samples, so the heavy-tailedness of real-world gradient noise remains an open question.
Our empirical results suggest that our assumptions in Definitions 2.6 and 2.7 are not too strong. In Section 6 we efficiently simulate the Itô SDE using an algorithm analogous to SVAG (Li et al.) (022). The simulation closely approximates the test accuracy achieved by the discrete trajectory, suggesting that even if heavy-tailed noise is present during training, it is not crucial for good generalization (Appendix H.. We remain interested in exploring the heavy-tailed analogs of our Itô SDEs. However, efficient simulation of such SDEs remains intractable and formal approximation guarantees are diffieult to prove, so are limited in assessing the utility of such approximations. We leave it for future work to investigate if and how heavy-tailed noise plays a role in adaptive optimization.
3 Related Work
We defer full discussion of empirical and theoretical works on adaptive gradient methods to Appendix A.] and only discuss works relevant to SDEs and scaling rules here.

---

Applications of SDE Approximations. There are applications of the SDE approximation beyond the derivation of a scaling rule. et al.] (2020) and Kunin al. (202T) assumed that the loss has some symmetry (i.e., scale invariance) and studied the resulting dynamics. Furthermore, [Li et al. (2020) used this property to explain the phenomenon of sudden rising error after LR decay in training. Xie et al. (202T) analyzed why SGD favors flat minima using an SDE-motivated diffusion model.
Past Work on Square Root Scaling Rule. As mentioned before, square root scaling was incor- rectly believed for few years to be theoretically justified for SGD. .ra.zizaz et al. (2022) decomposed the stochastic Hessian during batch training into a deterministic Hessian and stochastic sampling perturbation and assumes one of the components to be low rank to propose a square root scaling rule. (You et al. 2020) empirically discovered a square root scaling rule for language models trained by LAMB, layer-wiss variant of Adam. Xie et al. (2022) heuristically derived, but did not show an approximation bound for, a second-order SDE for approximating Adam, and they applied the SDE to study the time needed for Adam to escape sharp minima. Xie et al. (2022) did not discuss scaling rule, though their proposed SDE may admit one. Similarly, Zhou et (2020) derived a Lévy SDE for Adam, but no approximation bounds are given in the paper.
4 SDEs for Adaptive Algorithms
An SDE approximation operates in continuous time and thus implicitly considers the limit 1) - 0. In adaptive algorithms, the moment averaging parameters B, Br, B2 and 17 must be taken to limits such that the adaptivity and stochasticity can still be studied. For example, if 3, 82 remain fixed when 17 O, then the algorithm computes the moving averages in very small neighborhood, which averages out the effects of gradient noise and gradient history, causing the flow to turn into deterministic SignGD (Ma et al. 2022). We will need to assume (9 :1,32 - 1) - O, which implies the impact of the history grows as the learning rate decreases, and thus the adaptive features of these algorithms can still be studied in the continuous approximation (Ma et al. (2022). To keep the stochastic nature of the flow, we require the noise from mini-batching dominate the gradient updates.
4.1 Warm-Up: Linear loss
To build intuition for the SDE and the scaling rule, we first study simplified setting. In particular,
consider a linear loss f(0) = (0, g) and isotropic noise in the NGOS, namely gk ~ (g.g.00 o I2I). The R.Sprop update U = Bop + (1 ز) can be expanded as d _ Do + (1 - B) Σラ= g By linearity of expectation,
k-1
Elvu] = 30 + (1- B)
Σ j=0
(g + 021) = =3 vn + (1 3+2 + 0
This suggests that E[2x] is approximately g² + 0² after sufficient number of steps. Setting Vo = g 02 we see that the approximation [[wx] = g2 o becomes exact for all 2 0.
Using the linearity of variance (for independent variables), the standard deviation of each coordinate of Un is of scale O(11 3) Moreover, the expectation E(my) is of scale O(a²), so we know that es is nearly deterministic and concentrates around its expectation g2 + o²1 when ,3 is close to 1. Therefore, we take the approximation u% g2 + o²1 for all k 2 0. Ignoring e, the RMSprop update rule becomes:
θ 02 1 O (g" +0=+0011 00"1)-1/2
(3)
These dynamics depend on the relative magnitudes of g and a. We show that when o no SDE approximation can exist in Appendix [F22 Here, we explore the case where o
||g]]
which implies 0x+1 x 02 2 gr. Noting that gk ~ N o²I) gives 0 - 0₂ ~ N g. 1f "I)
approximately. Since Gaussian variables are additive, marginal distribution of 0g: 0 ~ N ((kn/a)g T approximately.
take telescoping
to obtain the
If an SDE approximation of R.Mpprop exists, then 0₂ can be closely approximated by a axxed random variable from the corresponding stochastic process at a fixed (continuous) time 4. Thus, to keep the
SDE unchanged upon changing the noise scale o, the hyperparameters must be adjusted to keep kn and ko7² unchanged, which implies 1) ~ and l f/- These observations intuitively yield
5

---

the square root scaling rule: noting that a changes with size B as ~ 1/VB suggests
r) v VB, and k L/B.
4.2 SDE Approximations for Adaptive Algorithms
Having established intuition in highly simplified setting for adaptive algorithms, now return to a more general and realistic setting. We derive the SDEs that are order-1 1 approximations of the discrete RMSprop and Adam algorithms under Definition both the parameters f and moment estimates. From the example of Section 4.1 we see that SDE approximations may exist if a 1/12 and > || (see Appendix G.1 for empirical validation of this assumption). In this case, we can prove that k ~ 17 is true not only for the simple setting above but also in general. This a key difference to the SDE for SGD each step RMSprop or Adam corresponds to a time interval of 1)² in SDEs, but each SGD step corresponds to time interval of r). In Section 4.1 v o² ~ 1/11² grows to infinity as 1l - O. This also happens in the general setting, so we track u - 0x/a² (instead of v4 directly) in the SDEs.
2.5 where the SDE states consist of
A Definition 4.1 (SDE for RMSprop). Let 0p 0T), e0 E7,, and C2 - (1 B)/n2. Define the state of the SDE as X3 (0,, uz) and the dynamics as
d0, = = P P. where P. == oodiag(u;))2 + Eg is preconditioner matrix, and We is the Wiener process.
(V)))) ob - (0,,)d))))))))
du, = cy(diag(<>())) - uz)dt
Theorem 4.2 (Informal version of Theorem C.2). Let uh & ux/0² and define the state of the discrete RMSprop trajectory with hyperparameters n, B,c (Definition 2.1) I (0x, uz). Then, for a well-behaved NGOS (Definition 2.3) satisfying the skewness and bounded moments conditions (Definitions) 2.6 and 2.7 the SDE in Definition satisfies
max
(xg()))
k=
where g and T are defined as in Definition 2.4 And the initial condition of the SDE Xo- Z0-
Remark 4.3. Section 4.1 suggested that the SDE approximation can only exist when a >> ||/||- This condition is reflected by keeping on = o1 a constant and C depends on og. When rj - 0, scales as 1/7, so a > ||/||2.
We need to find continuous approximations of the normalization constants Adam (Definition 2.2 As in the RMSprop case, we take (1 92) = C2. Then, we can estimate the normalization 3K in continuous time i. kxn²
1 exp(-c2!). We can do the analogous approximation for the other normalization constant 1 - , in Adam. Taking
constant
8 1 (1 c2n)2)//2²
(1 - 3, )/1² c₁ we can approximate it as 1 3 = 1 - cxp(-cit). This is a heuristic approach to deal with normalization constants, but we can indeed justify theory.
Definition 4.4 (Adam SDE). Let c & (1 B1)/n², 02 A (1 B2)/n² and define do. Co in Definition 4.1 Let (t) & exp(-cit) and (t) & exp(-c2t) Define the state of the SDE as Xx (O,, M4, Ut) and the dynamics as
d0: (t) (1)
P. p--1m,d m,,d/
dm, cı (VS(0,) m,,dd/ ope (0))(,)),,,
du; = c2 (diag( Σ(),)) uz)df,
where P. = oodiag(u.))1/ + Eov r is preconditioner matrix, W, is the Wiener process.
Our main theorem for Adam is given below. The initial steps of the discrete Adam trajectory can be discontinuous and noisy because of the normalization constants changing drastically. Hence, we introduce constant to and show that for any fo. can construct an SDE to be weak approximation for Adam starting from the [to/n? 11th step of Adam.
Theorem 4.5 (Informal version of Theorem D.2p. Define uı 0x/o² and let a. = (0z, My. ux) E
3d be the state of the discrete Adam trajectory with hyperparameters 17, B1, 12,6. Then, for well-behaved NGOS (Definition 2..3) satisfying the skewness and bounded moments conditions (Defiittins)) 2.6 and 2.7) and any to
0, the SDE in Definition 4.4 satisfies
max T=(to/>(T,...
[Eg(**) - Egg( xnp) S Ct²
(T///
where g and T are defined as in Definition 2.4 and the initial condition of the SDE is X.. 20/to/111

---

Proof Sketch. We provide a proof sketch for our SDE approximations here and defer the technical details Theorems C.2 and D.2. The proof follows the same steps as [i et al. (2019): first, we compute the approximation error of the continuous trajectory after one step in discrete time. Then, we use the single step error to bound the error over a finite interval of time. The proof extends standard SDE techniques in several ways. The given SDEs do not satisfy the Lipschitzness and smoothness conditions because the denominator can be unbounded. We thus construct an auxiliary SDE with an equivalent distribution to the desired SDE (Theorem CC5) but with better smoothness conditions, and we prove this SDE to be an order-1 I weak approximation of the discrete trajectory. Moreover, the SDE coefficients are time-dependent for Adam, unlike the ones for SGD, need to extend existing results to cover this case (see Appendix B.1
5 Square Root Scaling Rule
The SDE approximations in Definitions 4.1 and 4.4
motivate scaling rules for how to adjust the
optimization hyperparameters when changing the batch size. In order for do C1, C2. and co to remain constant, as required by the SDEs, needs change 17, 3, $1, S2, and accordingly.
Definition 5.1 (RMSprop Scaling Rule). When running RMSprop (Definition 2.1 with batch size B' = kB, use the hyperparameters Ty' T//K, 3 = - x(1 3), and é'
Definition 5.2 (Adam Scaling Rule). When rumning Adam (Definition 2.2 with batch size B' = xB, x(1
use the perraaaraeeters n/ //FF 31
B1). 8y x(1 82), and é' ___ Vx"
Theorem 5.3 (Validity of the Scaling Rules). Suppose have a well-behaved NGOS satissying the low skewness and bounded moments conditions.
Let o be the discrete RMSprop (Definition 2.1 trajectory with batch size and hyperparame- ters n, 8, and Let x (BB) be the trajectory with batch size k. and hyperparameters adjusted according to Definition
(B) and æ (kB) start frartfrrom the same initial point, then with g and T
5.1 If æ
defined as in Definition
( (Eg(a k//1)) S C(1 + xt)r?
max 2/
k=
2. Fix to > O. Let Ft (B) be the discrete Adam (Definition 2.2 trajectory with butch size B and hyper- parameters 17, Br, B2, and e. Let x (xB) be the trajectory with batch size uB and hyperparameters (RB)
adjusted according to Definition 5.2
If x (tt//>21
(B) and x [nto/pp] are equal, then with g and T
defined us Definition
max N== ===(t//11 [T/P]
Eg(x
(xB) (B)) - Eg(
[
[R/K]] )1 C(1 + +N)?
Proof. By the lincarity of covariance, scaling the batch size by к only modifies the NGOS by scaling by k/ Hence, both scaling rules ensure that oo. c1, c2. and eo (and thus, the SDEs) are unchanged when the batch size changes. The approximation bounds in Theorems 4.2oremm/422loenn 4.2 and 4.5 are in terms of n and since 7j scaled here by VK, the same method gives an upper bound Ckn Adding the approximation bounds for 7) and VKn together gives C(1 + x
Remark 5.4. The to condition on the Adam rule, holdover from the condition in Theorem 4.5 implies that our theory only directly applies when there is a warm-start phase of [to//7?1, where the marginal distribution of the trainable parameters at the end of the phase the same across different learning rates 17. Regardless, the scaling rules are shown to work in practice even without this phase.
The scaling rules depend on maintaining the same SDE approximation. so the bounded moments and low skewness conditions are sufficient (but not necessary) for this scaling rule work. [i et al.) (2021) provided an analogous discussion for SGD, and they show the scaling rule can hold even if there is heavy-tailed noise. We leave study of heavy-tailed gradient noise in adaptive algorithms as future work.
7

---

Test Accuracy
Test Accuracy
Test Accuracy
Test Accuracy
100
100
100
80
80
80
B-256
60
B=512
60
60
60
B-1024
B-2048
40
B-4096
40
40
40
B=8192
0.000 0.025 0.050 0.075' 0,100
h0o00 0.050 0.075 6
bo00 0.02 0.050 0.075 0,100
".000 0.025
0.050 0.075 0.100 Continuous time = g²
Continuous time = k"
Continuous time =
Continuous time =
(a) VGG-16 with Adam
(b) ResNet-50 with Adam
(c) VGG-16 with RMSprop
(d) ResNet-50 with RMSprop
Figure 1: Square root scaling rule experiments on CIFAR-10 with VGG-16 and ResNet-50 (details in Appendix D). We plot the mean and variance of 3 random seeds. Same color legend has been used across the plots. The performance gap between = 256 and 8192 is at most 3% in all cases.
Top-1 Validation Accuracy
Validation Log Perplexity
Validation Log Perplexity
10
70
10
B-256
B-512
B-1024
B-1024
60
B=2048
B=2048
8
B-512
B-4096
8-9096
40
B-1024
B-8192
8-8192
B=2040
6
B-4096
20
0=0192
B-16384
10
B-32768
2
0
0.00
0.01
0.02
0.03
0.04
0.00
0.02
0.04
0.000
0.005
0.010
Continuous time
n7
Continuous time =
Continuous time =
(a) ResNet-50 on ImageNet
(b) RoBERTa on Wikipedia +
(c) GPT on WikiText-103
Bookcorpus
Figure 2:
Large scale square root scaling rule experiments (details in Appendix Dy. Small and large batch models differ by at most 1.5% test accuracy in vision and 0.5 perplexity in language.
Experiments. Figures and show the square root scaling rule applied to ResNet-50 (He al. 2016) and VGG-16 (Simonyan and Zisserman 2014) trained on CIFAR-10 (Krizhevsky et al.). RoBERTa-large (Liu et al.) 2019) trained on the Wiki+Books corpus ZZuu et al. 2015). 12-layer GPT (Brown et al. 2020) on WikiText- 103 (Merity et al. (017)) and ResNet-50 trained on Ima- geNet (Deng et al.] 2009). We use the efficient language model pre-training recipe outlined in Izsak et al. (021). Appendix I has many additional settings, including ablations against other scaling rules
(Appendix
SVAG for Adaptive Algorithms
Validating the approximation strength captured in Definition 2.4 involves comparing the discrete algorithms and their SDEs on of test functions. However, obtaining the SDE solution through traditional simulations, e.g., Euler-Maruyama, is computationally intractable
al. (202)) proposed efficient simulation, SVAG, of the SDE for SGD in the finite LR regime: scale the constant LR by 1/E and take the limit -> co. In practice the simulation converges for small value of f. We adapt SVAG technique to simulate our proposed SDEs, which requires additionally adjusting the moment averaging hyperparameters (i.e., 3, B1, B2) and e.
Definition 6.1 (SVAG Operator). Given an NGOS G. (f, Σ, Za) with scale a (Definition 2.3 and hyperparameter > 1, the SVAG operator transforms Go into an NGOS Geo = (f.. E, Zeo) with scale &o. The NGOS Geo takes an input 0 and returns 9 = (f)g1 + (2)gg) where 91 are two
'One can also simulate the SDE by constructing ost-rrder weak approximations while taking + 0 along the scaling rules, but the batch size cannot be smaller than and hence 1) cannot go arbitrarily close to the limit.

---

Train Accuracy
Test Accuracy
2
11? Mean (Log Scale)
|1 Mean (Log Scale)
22
100
95
90
211
85
2.0
2
=
75
2"
70
1
70
65
28
60 0.00
0.02
0.04
0.06
60
0.00
0.02
0.04
0.06
0.00
0.02
0.04
0.00
0.02
0.04
Continuous Time n
Figure 3: SVAG on the Adam trajectory when training ResNet-50 CIFAR-10 matches the discrete trajectory (& = 1) on various test functions (see Appendix J for details). The closeness of the trajectories with respect to various test functions for different values of / implies the SDE approximation (Definition 4.4) is Ist-order weak approximation of Adam (Theorem 4.5
stochastic gradients from (0) and r.(() = +1 (-1) for [1,2}. The probability distribution Zee is defined such that 9 has the same distribution as Vf(0) + fox when & (0).
Lemma E. verifies that G.D does indeed compute stochastic gradients for f with covariance E. Applying the SVAG operator mini-batch training amplifies the noise scale by f. We then apply the square root scaling rule to adjust the learning rates and other hyperparameters accordingly, which yields the SVAG algorithm.
Definition 6.2 (SVAG Algorithm). For a loss f. SVAG operator hyperparameter 0, and optimiza- tion hyperparameters 7j, B, 31, B2, and e, compute the stochastic gradient as و = 71 (1)g + +2 (2) where and are defined as in Definition 6.1 and scale the optimization hyperparameters:
1. For RMSprop, set 7/2, 1- (-11-10111)) 8)/2. and cl and apply updates in Definition 2.1
2. For Adam, set 17 1// 81 1 -1) 82 + 1 (11- and el and apply updates as in Definition 2.2
The SVAG algorithm describes discrete trajectory that is a 1st-order approximation of the cor- responding SDE (Definitions) 4.1 and 4.4), thereby providing an efficient simulation of the SDEs.
Theorem 6.3 (SVAG algorithm approximates SDE). Assume the NGOS well-behaved and satisfies the bounded moments condition (Definitions 2.5 and 2.7
1. Let X, be the state of the RMSprop SDE (Definition 4.1 with hyperparameters n,, 3, and e. Let be the state of the analogous discrete SVAG algorithm with hyperparameter é.
2. Let X; be the state of the Adam SDE (Definition 4.4 with hyperparameters n, 31, 32, and e. Let T& be the state of the analogous discrete SVAG algorithu with hyperparameter f.
In both and 2, following holds for and T as in Definition
2.4
max [T]]
[B((x))
(E(( X (1) S CT21B
Proof. The main idea of the proof show that the SVAG operator transforms the noise distribution of well-behaved NGOS satisfying the bounded moments condition into one that is well-behaved and satisfies the bounded moments and the low skewness conditions (Lemma E.2). With these three conditions satisfied, we can directly apply Theorems 4.2 and 4.5 to complete the proof.
Li
Because the SDE scales time as 1/12², we must run SVAG for l² steps to match a single step of the discrete trajectories. Nevertheless, we note that in our setting and in Li et al. (2021), the approximation guarantee strong enough for small so this simulation still more efficient than Euler-Men Mally-Maruyama.

---

Experiments. Figure 3 compares the Adam SVAG trajectories (Definition 6.2) to 8 to the discrete one (f 1) on CIFAR-I0 (Krizhevsky al.) with ResNet-50 (He et al. 2015b). We use (+x)) and | as mathematically well-hehaved test functions to test the approximation strength (see Definition 2.4). We also measure the train and test accuracies, even though they are not differentiable (and hence, not covered by the theory). The converged SVAG trajectories are close to the discrete ones under these test functions, suggesting the SDE approximations are applicable to realistic deep learning settings, Additional details and settings, including large language models, are in Appendix
7 Conclusion
We derive SDEs that are provable Ist-order approximations of the RMSprop and Adam trajectories, immediately yielding formal derivations of square root scaling rules: increase the learning rate by vK and adjust the adaptive hyperparameters when increasing batch size by K. Experiments in the vision and language domains verify that applying these rules ensures that the values of several test functions, including test accuracy, are preserved. We furthermore design an efficient simulation for the SDEs, allowing us to directly validate the applicability of these SDEs to common vision and language settings. These SDEs can lead to deeper understanding of how adaptivity and stochasticity impact optimization and generalization, and we hope to extend our results formal identification of necessary and sufficient conditions for the approximation and its consequences to hold.
Acknowledgements
We thank Zhiyuan Li and Chao Ma for helpful discussion. We also thank Tianyu Gao and Alexander Wettig for helping us run language modeling experiments. This work is funded by NSF, ONR, Simons Foundation, DARPA and SRC.
References
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33:1877-1901, 202000 2020.
Jinghui Chen, Dongruo Zhou, Yiqi Tang, Ziyan Yang, Yuan Cao, and Quanquan Gu. Closing the generalization gap of adaptive gradient methods in training deep neural networks. In Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI'20, 2021. ISBN 9780999241 165.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: large-scale hier- archical image database. In 2009 IEEE Conference Computer Vision and Pattern Recognition, pages 248-255. 2009. doi: 10.1 2..1109/VPR..0099.50
Timothy Dozat. Incorporating Nesterov momentum into Adam. In Workshop, 2016.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(ul):2121-2159, 2011.
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 hour. arXiv preprint arXiv: (ooe:170.0000000 02677, 2017.
Diego Granziol, Stefan Zohren, and Stephen Roberts. Learning rates function of batch size: A random matrix theory approach to neural network training. Journal of Machine Learning Research, 23(173):1-65 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification, In The IEFE International Conference on Computer Vision (ICCV), December 2015a.
10

---

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification. In The IEEE International Conference on Computer Vision (ICCV), December 2015b.
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. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gamett, editors, Advances in Neural Information Processing Systems 30, pages 1731-1741. Curran Associates, Inc., 2017.
Peter Izsak, Moshe Berchansky, and Omer Levy. How to train BERT with an academic budget. arXiv preprint arXiv:2104 X04:22044007755, 2021.
Stanislaw Jastrzębski. Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in SGD, arXiv preprint 2017.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Alex Krizhevsky, One weird trick for parallelizing convolutional neural networks. arXiv preprint arXiv:1404.5500 1404.5997 2014.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. CIFAR-IO (Canadian Institute for Advanced Research). Uttp:t/://www cs. toronto.oo..........o... edw/"kriz/ torootoo..duuuukkii........ ccffrriffa.////// html
Daniel Kunin, Javier Sagastuy-Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka.
Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. In International Conference Learning Representations, 2021.
Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and dynamics of stochastic gradient algorithms I: Mathematical foundations. Journal of Machine Learning Research, 20(40); 1-47, 2019.
Zhiyuan Li and Sanjeev Arora. An exponential learning rate schedule for deep learning. In Interna- tional Conference on Learning Representations, 2020.
Zhiyuan Li, Kaifeng Lyu. and Sanjeev Arora. Reconciling modern deep learning with traditional optimization analyses: The intrinsic learning rate. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages Curran Associates, Inc., 2020.
14544-1455 14555.
Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling sgd with stochastic differential equations (sdes). In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 12712- Curran Associates, Inc., 2021.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. RoBERTa: A robustly optimized BERT pretraining approach. arXiv preprini arXiv: 1907. 2019.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019,
Liangehen Luo, Yuanhao Xiong, and Yan Liu. Adaptive gradient methods with dynamic bound of learning rate. In International Conference on Learning Representations, 2019.
Chao Ma, Lei Wu, and Weinan E. A qualitative study of the dynamic behavior for adaptive gradient algorithms. In Joan Bruna, Jan Hesthaven, and Lenka Zdeborova, editors, Proceedings of the 2nd Mathematical and Scientific Machine Learning Conference, volume 145 of Proceedings of Machine Learning Research, pages 671-692. PMLR, 16-19 Aug 2022.
11

---

Jerry Ma and Denis Yarats. Quasi-hyperbolie momentum and Adam for deep learning, In International Conference on Learning Representations, 2019.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations, 2017.
Prajit Ramachandran, Barret Zoph, and Quoc V Le. Searching for activation functions. arXiv preprint arXiv:1710.05941/ 2017.
Sashank J. Reddi, Satyen Kale, and Sanjiy Kumar. On the convergence of Adam and beyond. In International Conference Learning Representations, 2018.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. International Conference on Learning Representations, 2014.
Umut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. A tail-index analysis of stochastic gradient noise in deep neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Pro- ceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 5827-5837. PMLR, 09-15 Jun 2019.
Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for modern deep learning research. Proceedings of the AAAI Conference on Artificial Intelligence, 34 (19):13693-1369, Apr 2020.
Tijmen Ticleman and Geoffrey Hinton. Lecture 6.5 6.5-rmsprop: Divide the gradient by running aver- age of its recent magnitude.. 2012. URL https://pww: cs. toronto..... cerrontto.duuu/./. wwweeeduiiiiiiee csc321/ sliiess/ecture__ll seaetturrr_ssssse__eeeeeeesseeee 1ec6. pdf
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In International Conference Learning Representations, 2019.
Alexander Wettig, Tianyu Gao, Zexuan Zhong, and Dangi Chen. Should you mask 15% in masked language modeling? arXiv preprint arXiv:2202.08005. 2022.
Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In International Conference on Learning Representations, 2021.
Zeke Xie, Xinrui Wang, Huishuai Zhang, Issei Sato, and Masashi Sugiyama. Adaptive inertia: Disentangling the effects of adaptive learning rate and momentum. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 24430-24459, PMLR, 17-23 Jul 2022.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training BERT in 76 minutes. In International Conference on Learning Representations, 2020.
Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. In S. Bengio. H. Wallach, H. Larochelle. K. Grauman, N. Cesa- Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? In H. Larochelle. M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 15383-15393. Curran Associates, Inc., 2020.
Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, and Weinan E. To- wards theoretically understanding why SGD generalizes better than Adam in deep learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1288-21226. Curran Associates, Inc., 2020.
12

---

Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In 2015 IEEE International Conference on Computer Vision (ICCV), pages 19-27, 2015.
Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvornek, Xenophon Pa- pademetris, and James Duncan. AdaBelief optimizer: Adapting stepsizes by the belief in observed gradients. H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 18795-18888 18806. Curran Associates, Inc., 2020.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] Conditions for our results are outlined in Definitions 2.5 to 2.7 In particular, we discuss the limited applicability of our work to cases in which the gradient noise is heavy-tailed, though our empirical success suggests that the gradient noise satisfies our assumptions in many realistic vision and language settings, (c) Did you discuss any potential negative societal impacts of your work? [Yes] Our work is primarily theoretical in nature, but we discuss the broader impacts of our work in Appendix [A.2
(d) Have you read the ethics review guidelines and ensured that your paper conforms them?
[Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] Section has discussion of the assumptions on the gradient noise and the general requirements needed be able to construct a continuous approximation of the discrete algorithms. Moreover, Section 4.1 motivates the assumption that the gradient noise dominates the gradient. We validate this assumption experimentally in Appendix G.I Our assumptions do not stray
from prior works on SDEs (Li et al.] 2019 2021).
(b) Did you include complete proof' of all theoretical results? [Yes] Our proof's are written clearly in the appendix.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (cither in the supplemental matcrial or as a URL)? [Yes] We include the code for the vision experiments in the supplementary material. For the NLP experiments, use the code of Wettig et al. (2022).
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] Appendix contains the training details of all the experiments.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [No] Due to computational constraints, we didn 't include the error bars for all the experiments. However, we have made sure that the random seed for different runs in each experimental setting is the same.
(d) Did you include the total amount of compute and the type of resources used (e.g.. type of GPUs, internal cluster, or cloud provider)? [Yes] We ran our experiments on a cluster of 34 GPUs, where 24 are RTX 2080 GPUs and 10 are A5000 GPUs. Each experiment on CIFAR- 10 required a single RTX 2080 GPU, each experiment on ImageNet required single A5000 GPU, each pretraining experiment GPUs, each pretraining experiment on RoBERTa required set of RTX 2080 GPUs, and each finetuning experiment on RoBERTa required single RTX 2080 GPU.
GPT required set of 4 RTX 2080
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] We cite Wettig et al. (2022) for our language model experiments. We also cite the papers of all the deep learning architectures and datasets that we use for our experiments.
(b) Did you mention the license of the assets? [Yes] The github repository 5 of Wettig et al. (2022) has MIT license.
(c) Did you include any new assets either in the supplemental material or URL? [Yes] We include our code for the vision experiments in the supplementary material.
(d) Did you discuss whether and how consent using/curating? [N/A] The data we used for our experiments are publicly available.
obtained from people whose data you're 're
(e) Did you discuss whether the data you are using/curating contains personally identifiable information offensive content? [N/A] The data we used for our experiments publicly available.
14

---

5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applica- ble? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board
(IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent
on participant compensation? [N/A]
15


Reviews:
[{"review_id": "zUy67XYhnZb", "reviewer": "Reviewer_QQca", "summary": "The presented paper presents the SDE approximations for Adam and RMSProp optimizers and provides theoretical guarantees of their correctness. The derived SDE approximations revealed the theoretical explanations of the square-root scaling rule of learning rate with batch size increasing. To verify numerically the correctness of the proposed SDEs, authors extend the SVAG technique to the considered classes of SDEs. Numerical experiments consider computer vision and natural language processing models and confirm the correctness of the SDEs and square-root scaling rule.", "questions": "Please comment the following remarks\n1) what is $\\mathcal{G}$ in line 70?\n2) please explain why the test functions given in line 319 are satisfied definition 2.4.?\n3) how does the presented analysis correlate with results in https://arxiv.org/pdf/1706.02677.pdf%5B3%5D%20ImageNet and https://arxiv.org/pdf/1904.00962.pdf where the special warm up schedules were used?\n", "limitations": "The authors mention limitations in the assumptions before deriving the result SDEs (section 2.3).", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n1) authors present solid theoretical explanation of the empirical rule to manage learning rate with increasing of batch size\n2) the expected approximation quality of SDEs is verified in large scale problems from CV and NLP\n3) SVAG technique is modified for the considered SDEs to speed up simulations\n\nWeaknesses:\n1) the presented analysis assumes the not heavy-tailed noise in the gradient estimates, though some works demonstrate that noise in deep networks is indeed heavy-tailed, see e.g. https://tel.archives-ouvertes.fr/tel-03206456/document and https://arxiv.org/pdf/2106.05958.pdf -  that is left to future work\n2) no new practically important consequences of the presented analysis. Researchers use the square root scaling rule previously, now we know why this rule works for adaptive method and why the difference with SGD exists.   \n\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "p2oYn5WCgVv", "reviewer": "Reviewer_fuyA", "summary": "This paper derives SDE approximations of RMSProp and Adam when their hyperparameters changes according to a scaling rule. Experiments on VGG and ResNets were performed on CIFAR-10.", "questions": "Please see the \"Strengths and Weaknesses\" section.", "limitations": "As there are existing work studying similar problems, especially the Lévy process based SDEs, and some existing work suggests the gradient noises are heavy-tailed. It might be instructive to study whether SDEs in this paper makes more sense than existing work, and to compare and contrast their use cases in practice.", "rating": 5, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\nPros:\n- The scaling rules in Theorem 5.3 is interesting.\n- The SDE approximations of RMSProp and Adam are derived, and they appear to be correct.\n\nCons:\n\n- Experiment results (e.g., Fig 1-3) do not have confidence bands. Furthermore, in Figure 1 (a) and (b) all curves are very close but they are quite far away in (c) and (d). Does this mean the scaling rule is more respected by Adam but not RMSProp? I found it a bit hard to interpret the results. Perhaps the authors can add some interpretations in the figure caption and in the main texts.\n- Why in Figure 2(b) the curves do not start at the same point?\n- The proof technique seems not new, the theorems look more like extensions to the stochastic modified equation literature. Perhaps I missed something. Can the authors discuss more in the main text how the new analyses differ in technique from those literatures?\n- As some literature suggested, the gradient noise might be heavy-tailed, it is thus open to discuss which of the model (Gaussian vs. Lévy) makes more sense in practice — when would Gaussian noises occur in training practical neural nets?", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "XPLaJVWQRv_", "reviewer": "Reviewer_6FQK", "summary": "This paper derives SDEs corresponding to the dynamics of RMSProp and Adam. Based on the SDEs, the paper proposes to scale the learning rate by $\\kappa^{-1/2}$ when the batch size is scaled by $\\kappa$. Finally, the paper proposes a SVAG-type algorithm to simulate the training dynamics of the adaptive algorithms and ", "questions": "1. In addition to the learning rate $eta$, parameteters $beta$, $beta_1$ and $\\beta_2$ should also be changed if the batch size is scaled by a positive constant. How can this be interpreted?\n\n2. Does the SVAG algorithm for adaptive algorithms have any practical applications other than validating the training dynamics?\n\n3. It is mentioned in subsection 4.1 and inn the first paragraph of subsection 4.2 that the condition $sigma\\gg\\|g\\|_2$ is crutial in deriving the scaling scheme. How is this condition related to the conditions of Theorem 4.2?", "limitations": "The authors properly address this issue in appendix A.2.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The paper derives the scaling scheme for the learning rate of adaptive algorithms, which is very impressive.\n2. In addition to the SDEs, an approximation bound is given.\n\nWeaknesses:\n1. Readers, including myself, may be unfamiliar with the SVAG algorithm if they do not have adequate knowledge of stochastic simulation. More introduction might be required in this part.\n2. The SDEs and scaling schemes have been proposed in the literature, though these results have never been presented in such a rigorous way as this paper. The results of this paper might be a little incremental.", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Ia0BLShl_Dj", "reviewer": "Reviewer_AnTC", "summary": "This paper attains the SDE approximations for two popular optimization algorithms RMSprop and Adam, and validates their applicability empirically. The SDE approximations generate square-root scaling rules, which are validated in the vision and language modeling domains.", "questions": "1.  [1] proposes an exponential learning rate schedule. How is the square root schedule proposed in the paper compared to the exponential schedule?   \n2. [2] employs Levy SDE approximation. Can the authors explain why \"the=e quality of the Lévy SDE approximation was not formally guaranteed\" and the advantage of using Ito SDE in the paper?\n\nReferences\n\n[1] Zhiyuan Li and Sanjeev Arora. An exponential learning rate schedule for deep learning. arXiv preprint arXiv:1910.07454, 2019.   \n[2] Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, Weinan E. Towards Theoretically Understanding Why Sgd Generalizes Better Than Adam in Deep Learning. NeurIPS 2020.   ", "limitations": "Yes.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:   \n**1)** The paper is generally well-written and the presentation is clear.   \n**2)** The theoretical results are new and interesting. The empirical results are provided to validate the derived theories.   \n**3)** I did not check every detail of the proof but generally it is technically sound. The assumptions are also supported by empirical results.       \n\nWeaknesses    \nI did not find any obvious weakness of the paper.", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept