PaperlayerView exact revision on arXiv

Source-linked AI summary

On the SDEs and Scaling Rules for Adaptive Gradient Algorithms

Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, Sanjeev Arora

arXiv:2205.10287v3cs.LG

TL;DR

Stochastic SDE approximations for adaptive optimizers had lacked rigorous guarantees, despite the value of continuous stochastic views of optimization. This paper derives and validates approximations for RMSprop and Adam, including square-root batch-size scaling rules.

  • Problem

    Extending stochastic SDE approximations to RMSprop and Adam has been challenging because their updates use coordinate-wise normalization.

    Source: §1 · p. 1

  • Method

    The paper derives first-order weak SDE approximations for RMSprop and Adam and adapts SVAG to simulate them efficiently.

    Source: §1 · p. 2, §4.2 · p. 6

  • Results

    Vision and language experiments show the scaling rules preserve several test-function values, including test accuracy, and efficient simulations validate SDE applicability in common settings.

    Source: §7 · p. 10

  • Takeaways & Limitations

    When increasing batch size by kappa, the rules increase learning rate by sqrt(kappa) and adjust adaptive hyperparameters.

    Source: §7 · p. 10

    The theory directly applies to Adam only after a warm-start phase of ceil(t0/eta^2) with matching parameter marginals across learning rates, though the rules work empirically without it.

    Source: §2 · p. 7

Abstract

from arXiv · show

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 $\textit{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.

1 Introduction

Adaptive optimizers lacked stochastic SDE approximations with theoretical guarantees, leaving their scaling behavior difficult to establish. This paper derives approximations for RMSprop and Adam, obtains square-root scaling rules, and validates them experimentally.

1 Introduction

  • For SGD, the SDE view explains the linear scaling rule by matching the noise change from larger batches with a proportional learning-rate increase.Earlier naive analysis instead suggested square-root scaling for SGD.Source: §1 · p. 1
  • Prior square-root proposals for adaptive methods were limited: perturbation theory addressed behavior near convergence, while LAMB provided empirical evidence on vision and language tasks.LAMB uses layerwise adaptive learning rates and applied square-root scaling rather than tuning learning rates as batch size increased.Source: §1 · pp. 1–2
  • The paper derives RMSprop and Adam SDEs with first-order weak approximation guarantees, yielding square-root scaling rules for batch-size changes.The rules adjust the optimization hyperparameters so discrete trajectories at different batch sizes approximate the same SDE.Source: §1 · p. 2
  • Experiments validate the rules in vision and language settings, and adapted SVAG simulations show the SDEs track discrete trajectories.Direct SDE simulation is prohibitively expensive because it requires frequent full-gradient and noise-covariance computations.Source: §1 · p. 2

2 Preliminaries

The preliminaries define RMSprop and Adam, model batch-dependent gradient noise with a noisy oracle, and set out the SDE approximation framework and its conditions. They also discuss how heavy-tailed noise bounds the scope of the Itô-SDE analysis.

2 Preliminaries

  • RMSprop updates parameters using a second-moment estimate, while Adam also uses momentum; the paper adopts the current estimate v_k in parameter updates.The paper assumes v_0 is nonzero when epsilon is zero to avoid division by zero.Source: §2 · p. 2

2.1 Noisy Gradient Oracle with Scale Parameter

  • The NGOS models a stochastic gradient as the full gradient plus mean-zero noise scaled by sigma, with covariance fixed across noise scales.In the paper's batch-sampling setting, batch size primarily controls the noise scale.Source: §2.1 · pp. 2–3

2.2 SDE Approximation and Scaling Rules

  • The SDE framework links learning rate and gradient noise: SGD's linear scaling preserves its SDE, while the paper uses the same approach to derive adaptive methods' square-root rule.The paper measures approximation by comparing test functions on discrete iterates and continuous trajectories.Source: §2.2 · p. 3, §2.3 · p. 3

2.3 Quality of SDE Approximation and Theoretical Assumptions

  • The approximation theory assumes smooth, well-behaved loss and covariance functions plus bounded moments and low noise skewness; RMSprop and Adam use effective step size eta^2.Gaussian noise satisfies the stated noise conditions under well-behaved loss and covariance; mini-batch noise conditions are verified under Lipschitz per-example gradients.Source: §2.3 · pp. 3–4

2.4 Discussion on Heavy-Tailed Gradient Noise

  • The Itô-SDE analysis permits some non-Gaussian noise but does not capture consequential higher moments; experiments suggest its assumptions are not too strong and heavy-tailed noise may not be crucial for good generalization.The paper leaves heavy-tailed adaptive-SDE analysis for future work because efficient simulation and formal approximation guarantees remain difficult.Source: §2.4 · p. 4

3 Related Work

Related work uses SDEs to study optimization dynamics and scaling, but prior adaptive-method analyses did not provide the same formal approximation guarantees. The paper situates its contribution against those applications and proposals.

3 Related Work

  • SDE-based work has studied dynamics under loss symmetry, explained sudden error rises after learning-rate decay, and analyzed why SGD favors flat minima.These studies illustrate applications of SDEs beyond deriving scaling rules.Source: §3 · p. 5
  • Earlier square-root scaling work included a near-convergence perturbation analysis and LAMB's empirical language-model result; the cited LAMB work is You et al. (2020).The perturbation proposal assumes a low-rank component of the stochastic Hessian.Source: §3 · p. 5
  • Prior Adam SDE proposals studied escape from sharp minima or used Lévy noise, but the cited works did not provide approximation bounds; Xie et al. (2022) did not discuss a scaling rule.Source: §3 · p. 5

4 SDEs for Adaptive Algorithms

The adaptive SDE analysis requires moment parameters to approach 1 as the learning rate shrinks and batch noise to dominate gradient updates. It develops intuition in a linear-loss setting, then gives RMSprop and Adam approximations with conditions that include Adam’s delayed start.

4.1 Warm-Up: Linear loss

  • When noise dominates the gradient in the linear-loss warm-up, RMSprop’s parameter iterates are approximately Gaussian, motivating hyperparameter changes that preserve the SDE as noise scale changes.The warm-up also notes that no SDE approximation exists when σ is much smaller than the gradient norm.Source: §4.1 · pp. 5–6

4 SDEs for Adaptive Algorithms

  • RMSprop and Adam SDEs use time steps of η² per algorithm step, track normalized second moments, and require β parameters to approach 1 while mini-batch noise dominates updates.In the general setting, the state includes parameters and moment estimates; the noise scale is taken proportional to 1/η.Source: §4 · p. 5, §4.2 · p. 6

4.2 SDE Approximations for Adaptive Algorithms

  • For well-behaved noisy gradient oracles satisfying skewness and bounded-moment conditions, the RMSprop SDE is an order-1 approximation of its discrete trajectory.The RMSprop state tracks u_k = v_k/σ²; the scaling regime keeps σ0 = ση constant as η shrinks.Source: §4.2 · p. 6
  • Adam’s initial updates can be discontinuous and noisy, so its order-1 weak approximation is established from any chosen time t0 > 0, initialized at the corresponding discrete iterate.The proof sketch uses a smoother auxiliary SDE with equivalent distribution and extends existing techniques to time-dependent coefficients.Source: §4.2 · pp. 6–7

5 Square Root Scaling Rule

The scaling rules adjust adaptive-optimizer hyperparameters as batch size changes, with the learning rate scaled by the square root of the batch-size factor. Their validity is stated for noisy gradient oracles satisfying low-skewness and bounded-moment conditions.

5 Square Root Scaling Rule

  • When batch size changes from B to κB, both RMSprop and Adam scale the learning rate as η′ = η√κ.The RMSprop and Adam rules also adjust their averaging parameters; the Adam rule scales ϵ′ = ϵ√κ.Source: §5 · p. 7
  • For Adam at batch size κB, the rule sets β′1 = 1 − κ(1 − β1) and β′2 = 1 − κ(1 − β2).The rule also sets ϵ′ = ϵ√κ.Source: §5 · p. 7
  • The scaling-rule validity theorem assumes a well-behaved noisy gradient oracle satisfying low-skewness and bounded-moment conditions.The supplied theorem passage states these assumptions but does not include its conclusion.Source: §5 · p. 7

1. Let x(B)

The supplied theorem fragments specify trajectories with batch sizes B and κB, adjusted hyperparameters, and a shared starting point.

1. Let x(B)

  • The compared discrete RMSprop trajectories use the notation x(B) and x(κB), with g and T defined as in Definition 2.4.Source: §1 · p. 7
  • The setup compares RMSprop trajectories at batch sizes B and κB, with the larger-batch trajectory using the prescribed scaling rule.The trajectories are stipulated to start from the same initial point.Source: §1 · p. 7

2. Fix t0 > 0. Let x(B)

The scaling rules preserve the SDE parameters when batch size changes, with approximation bounds that grow with the learning-rate adjustment. Experiments report small performance gaps between small- and large-batch models.

2. Fix t0 > 0. Let x(B)

  • Scaling the batch size by κ leaves σ0, c1, c2, and ϵ0—and therefore the SDEs—unchanged under the rules.The approximation bound for the two learning rates is C(1 + κ)η2.Source: §2 · p. 7
  • The theory directly applies to Adam only after a warm-start phase of ⌈t0/η2⌉ with matching marginal parameter distributions across learning rates.The rules are nevertheless shown to work in practice without this phase.Source: §2 · p. 7
  • Bounded moments and low skewness are sufficient, but not necessary, conditions for the scaling rules; heavy-tailed noise in adaptive algorithms remains future work..Source: §2 · p. 7
  • Small- and large-batch models differ by at most 1.5% test accuracy in vision and 0.5 perplexity in language.On CIFAR-10, the performance gap between B = 256 and B = 8192 is at most 3% in all cases.Source: Figure 2 · p. 8, Figure 1 · p. 8

6 SVAG for Adaptive Algorithms

The section introduces SVAG, an efficient discrete simulation of the adaptive-algorithm SDEs, and evaluates its approximation against discrete trajectories. Results on Adam show close trajectories on test functions, while the theory applies under stated noise conditions.

6 SVAG for Adaptive Algorithms

  • SVAG adapts the noise scale by ℓ and applies the square-root scaling rule to adjust learning rates and other hyperparameters.For RMSprop and Adam, η is set to η/ℓ, ϵ to ϵℓ, and the relevant β parameters are adjusted.Source: §6 · p. 9
  • SVAG provides an efficient discrete simulation that is a 1st-order approximation of the RMSprop and Adam SDEs.Traditional Euler-Maruyama simulation is computationally intractable in this setting.Source: §6 · pp. 8–9
  • The SVAG approximation guarantee assumes a well-behaved NGOS satisfying the bounded-moments condition.The proof also uses the SVAG-transformed noise's bounded-moments and low-skewness properties.Source: §6 · p. 9
  • Because SDE time scales as k = t/η2, SVAG requires ℓ2 steps to match one discrete-trajectory step.For small ℓ, the approximation guarantee remains sufficient and the simulation is more efficient than Euler-Maruyama.Source: §6 · p. 9
  • On CIFAR-10 with ResNet-50, Adam SVAG trajectories up to ℓ = 8 are close to the discrete trajectory on the measured test functions.The tests include gradient-noise covariance trace, gradient norm, and train and test accuracy; the accuracy measures are outside the stated theory.Source: Figure 3 · p. 9, §6 · p. 10

7 Conclusion

The paper derives provable 1st-order SDE approximations for RMSprop and Adam, yielding square-root scaling rules for batch-size changes. Vision and language experiments report preservation of several test-function values, including test accuracy.

7 Conclusion

  • The RMSprop and Adam SDEs are provable 1st-order approximations and formally derive square-root scaling rules.When batch size increases by κ, the learning rate increases by √κ and adaptive hyperparameters are adjusted.Source: §7 · p. 10
  • Vision and language experiments find that the scaling rules preserve values of several test functions, including test accuracy.The paper also designs an efficient simulation to validate SDE applicability in common vision and language settings.Source: §7 · p. 10
  • The authors identify understanding adaptivity and stochasticity's effects on optimization and generalization, and establishing necessary and sufficient conditions, as directions for future work..Source: §7 · p. 10

Checklist

The checklist notes scope limits tied to gradient-noise assumptions and reports experiment-resource details. Related work covers Adam variants, while the broader-impact discussion connects formal scaling rules to training costs and incentives.

Checklist

  • The authors identify limited applicability under heavy-tailed gradient noise, while empirical success suggests the assumptions hold in many realistic vision and language settings.The checklist also reports that error bars were omitted for some experiments due to computational constraints.Source: Checklist · p. 14

A.1 Additional Recent Works

  • The related-work discussion catalogs Adam variants addressing convergence, generalization, stability, and large-language-model pretraining.It describes methods including AMSGrad, AdaGrad refinements, Nadam, AdamW, AdaBound, AdaBelief, QHAdam, and LAMB.Source: §A.1 · p. 16

A.2 Broader Impact

  • Formal scaling rules remove the need for hyperparameter grid search and may reduce training runs, but may also encourage training very large models..Source: §A.2 · p. 16

B SDE Approximation Theorem

The theorem gives conditions under which a stochastic gradient algorithm is a first-order weak approximation of an SDE over a finite time interval. Its proof controls the difference between their one-step changes and extends that control across iterations.

B SDE Approximation Theorem

  • Under the theorem’s stated conditions, the stochastic gradient algorithm is a first-order weak approximation of the SDE when both start from the same point.The result applies over a finite interval and to test functions with at most polynomial growth.Source: B SDE Approximation Theorem · p. 17

B.1 Proof for Theorem B.2

  • The proof bounds the trajectory error by summing one-step discrepancies, using moment bounds and regularity of the SDE’s expected test functions.The argument also establishes uniform moment bounds for the stochastic gradient iterates.Source: B SDE Approximation Theorem · p. 17, §B.1 · pp. 18, 20

B.2 Proof for Lemma B.3

  • Itô–Taylor expansion is used to compare the SDE’s one-step moments with those of the stochastic gradient algorithm.The comparison applies the expansion to products of coordinate increments to compute moments.Source: §B.2 · p. 21

C RMSProp SDE Proof

This section develops theoretical and empirical support for adaptive-method SDE approximations, proving RMSprop and Adam approximation results under stated noise and starting-time conditions and examining the applicability conditions for the SVAG operator. Experiments test those conditions and the scaling behavior of an update-rule modification.

C RMSProp SDE Proof

  • Under bounded-moment, low-skewness noise assumptions, RMSprop’s SDE is an order-1 weak approximation of its discrete iterations.The theorem relates parameters through σ0 = ση, ϵ0 = ϵη, and c2 = (1 − β)/η^2; the proof uses an auxiliary SDE with umin = 2^−c2T min_i u0,i.Source: C RMSProp SDE Proof · p. 22

C.1 Verifying Condition (1)

  • The drift and diffusion functions for the auxiliary RMSprop SDE are Lipschitz and belong to G4.The proof uses bounded, Lipschitz Σ^1/2 and denominators bounded below by a constant.Source: §C.1 · p. 23

C.2 Verifying Condition (2)

  • The excerpts report discrete RMSprop’s second moment as η²Σij(θ)(√ui + ϵ0/σ0)(√uj + ϵ0/σ0) + O(η⁴), and auxiliary-SDE first, second, and third moments as η²b(x) + O(η⁴), η²σ(x)σ(x)⊤ + O(η⁴), and O(η⁴).The discrete-moment computation is stated for well-behaved NGOS and Zσ satisfying bounded-moments and low-skewness conditions.Source: §C.2 · pp. 23–24

C.3 Verifying Condition (3)

  • For each coordinate i, |∆i| = νiη²|∂if(θ) + σzi| ≤ η²νi(1 + |∂if(θ)|)(1 + σ|zi|) ≤ ηνi(1 + |∂if(θ)|)(1 + σ0|zi|).By the Lipschitzness of ∇f(θ) and the bounded moments condition for Zσ, Item 2 can be proved by taking powers and expectations on both sides of the inequalities.Source: §C.3 · p. 25

D Adam SDE Proof

  • For any t0 > 0, Adam’s SDE is an order-1 weak approximation of its discrete iterates from k0 = ⌈t0/η^2⌉ through T.The theorem assumes well-behaved noise with bounded moments and low skewness, and initializes the SDE at Xt0 = xk0.Source: D Adam SDE Proof · pp. 25–26

D.1 Verifying Condition (1)

  • Adam’s drift and diffusion functions are Lipschitz and belong to G4.The proof applies the same argument as the RMSprop Lipschitz verification.Source: §D.1 · p. 26

D.2 Verifying Condition (2)

  • Adam’s discrete-step and auxiliary-SDE moments are compared to verify Condition (2), with matching second moments up to O(η^4).For the discrete moments, the second-moment term is σ0^-2η^2Σij for p = d + i and q = d + j, and O(η^4) otherwise; the auxiliary-SDE moments are η^2b(x) + O(η^4), η^2σ(x)σ(x)^⊤ + O(η^4), and O(η^4) for the third moment.Source: §D.2 · pp. 27–28

D.3 Verifying Condition (3)

  • Adam’s proof also verifies Condition (3) with a constant bound independent of the step size.The stated bound applies for all k ≤ N and x ∈ Xk.Source: §D.3 · p. 28
  • The proof verifies Condition (3) using moment bounds with constants independent of η_e for all k ≤ N and x ∈ X_k.Source: §D.3 · p. 28

E Analysis of SVAG Operator

  • Applying SVAG preserves the noise process’s mean-zero and covariance properties, and for scale increases it preserves well-behavedness and bounded moments while satisfying low skewness.For ℓ ≥ 1, the resulting noise scale is ℓσ; the analysis verifies the low-skewness condition using a third-moment bound.Source: E Analysis of SVAG Operator · pp. 29–30

F.1 How does the noise scale change with batch size?

  • The noise scale varies inversely with the square root of batch size: σ ∼ 1/√B.A batch gradient averages independent, mean-zero noise vectors, reducing covariance by the batch size and yielding the stated noise-scale relation.Source: §F.1 · p. 30

F.2 What happens when the noise does not dominate the gradient?

  • When gradient noise is much smaller than the average gradient, no SDE approximation exists; experiments show noise dominates for small batches and tends to do so late in training for larger batches.The experiment reports this pattern for RMSprop on ResNet-50 and VGG-16, including tests with small and large ϵ.Source: §F.2 · p. 31, Figure 4 · p. 31

G Experimental Verification of Assumptions

G.1 Noise Dominates the Gradient

  • The general SDE applicability condition requires gradient variance to dominate the norm of the average gradient at every step; measurements support it for small batches and for large batches near training’s end.The analysis preceding the experiments uses a fixed-gradient linear-loss approximation, while the more general condition is stated in terms of sampling-seed expectations.Source: §G.1 · p. 31, Figure 4 · p. 31

G.2 Using vk instead of vk+1 in the update rule

  • ≈1% test-accuracy drop, while the scaling rule retains at most a 3% accuracy gap between batch sizes 256 and 8192.With Adam on CIFAR-10, trajectories and test functions match across batch sizes despite replacing v_{k+1} with v_k in the update rule.Source: §G.2 · p. 32

H SVAG Experiments

SVAG is a computationally efficient simulation of the RMSprop and Adam SDEs, with an order-1 weak approximation whose error scales as 1/ℓ. Across the reported settings, its trajectories converge and closely track the corresponding discrete optimizers.

H SVAG Experiments

  • SVAG is an order-1 weak approximation of the SDE, with approximation error scaling as 1/ℓ and apparent convergence for computationally tractable small ℓ.The corresponding discrete trajectory is matched after kℓ^2 SVAG steps for k optimizer steps.Source: H SVAG Experiments · p. 32
  • On CIFAR-10, SVAG converges and closely tracks RMSprop at ϵ = 10^-30 and ϵ = 10^-8, and Adam at both values.The RMSprop experiments use Figures 7 and 8; Adam uses Figures 3 and 9.Source: H SVAG Experiments · p. 32
  • With Adam’s update modified to use v_k instead of v_{k+1}, test accuracy drops by approximately 1%, while trajectories across batch sizes remain similar.The experiment covers ResNet-50 and VGG-16 on CIFAR-10.Source: Figure 6 · p. 33
  • On Wikipedia + Books and WikiText-103, SVAG converges and closely tracks Adam.The GPT experiments use sequences of length 64 for computational reasons.Source: H SVAG Experiments · p. 33

I Square Root Scaling Experiments

Experiments test the square root scaling rule across vision and language tasks and compare it with partial and linear scaling variants. The reported results show similar behavior across batch sizes in the evaluated settings, while the rule’s relative advantage over linear scaling varies by dataset.

I Square Root Scaling Experiments

  • Across CIFAR-10 settings, test accuracy varies by at most 3% across batch sizes; ImageNet validation-accuracy gaps are at most 3% or 1.5%, depending on ϵ.On Books+Wikipedia, log perplexity and gradient norms remain close across batch sizes, with similar downstream-task performance; on WikiText-103, log perplexity matches and gradient norms remain close.Source: I Square Root Scaling Experiments · pp. 35–36, Figure 20 · p. 39, Figure 22 · p. 40

I.1 Ablation study on the proposed scaling rule

  • Scaling η, ϵ, β1, and β2 together consistently beats the tested partial-scaling variants on CIFAR-10; scaling only η, ϵ, and β1 is the closest alternative.The comparison uses ResNet-50 with Adam across batch sizes.Source: §I.1 · p. 39

I.2 Ablation against linear scaling rules

  • On CIFAR-10, linear scaling can match square root scaling in validation accuracy, but square root scaling tracks the smaller-batch training trajectory better; on ImageNet, linear scaling brings a large validation-accuracy drop.The CIFAR-10 linear trajectories catch up after learning-rate decay; ImageNet experiments were not extensive because of computational issues.Source: §I.2 · pp. 40–41

J.1 A note on learning rate schedule and warm-up

  • The theorems do not directly cover time-varying learning rates or batch sizes, although experiments found the scaling rules held with learning-rate schedulers using a special warm-up.The authors describe extensions to time-varying settings as a possible proof direction.Source: §J.1 · p. 41

J.2 CIFAR-10

  • RMSprop and Adam baselines used batch size 256 at η=10−3; runs at 1024, 4096, and 16384 scaled hyperparameters while preserving total continuous time and the warmup and decay schedules.RMSprop decayed its learning rate by 0.1 at 400 and 450 epochs; Adam decayed by 0.1 at 300 epochs.Source: §J.2 · p. 41

J.3 ImageNet

  • The ImageNet setup uses an Adam ResNet-50 baseline with batch size 256, trains for 90 epochs, and compares small and larger ϵ settings.The reported batch-size experiments use the proposed scaling rule.Source: §J.3 · p. 44

J.4 Books and Wikipedia (Academic BERT)

  • For Books and Wikipedia, the study pretrains a 24-layer RoBERTa on length-128 sequences and evaluates pretrained models on GLUE tasks using a hyperparameter grid search.The corpus combines Bookcorpus and English Wikipedia.Source: §J.4 · p. 44

J.5 WikiText-103 (GPT)

  • The WikiText-103 setup uses a 12-layer GPT pretrained with Adam; its dataset contains 103 million tokens extracted from Wikipedia.The data are split into training and validation sets at a 9:1 ratio.Source: §J.5 · p. 46

Related papers

Loading 2205.10287v3…