Compute Efficiency and Serial Runtime Tradeoffs for Stochastic Momentum Methods Depen Morwani*,1,2,† , Alexandru Meterez*,1,2 , Pranav Ajit Nair*,1,2 , and Sham Kakade1,2
arXiv:2606.19179v1 [cs.LG] 17 Jun 2026
2
1 Harvard University Kempner Institute at Harvard University
Abstract Stochastic momentum methods such as heavy ball (HB), Nesterov momentum, and variants of Accelerated SGD (ASGD) (Kidambi et al., 2018) are widely used in modern training, but their stochastic benefits depend on two distinct quantities: serial runtime, the number of iterations needed to reach a target accuracy, and compute efficiency (CE), the inverse total gradient-query or FLOP cost. Larger batches reduce serial runtime without hurting CE only when the contraction gap grows linearly with batch size. We study stochastic HB and ASGD for consistent linear regression with Gaussian covariates and prove finite-dimensional, discrete-time lower bounds on their batch-size tradeoffs. Our first result shows that HB does not improve the CE frontier over SGD for arbitrary spectra; rather, it preserves SGD-level CE over a larger batch-size window, allowing larger batches to reduce serial runtime until HB reaches its deterministic √ accelerated scale. This window can be a factor κ larger than the SGD critical batch size. For ASGD, the picture is more spectrum-dependent: for rapidly decaying power-law spectra, ASGD improves small-batch CE over HB/SGD, but as batch size grows it trades this CE advantage for improved serial runtime. Synthetic linear-regression experiments verify these qualitative regimes, including near-overlap of ASGD and HB for slowly decaying spectra and the predicted CE–serial tradeoff for rapidly decaying spectra.
1
Introduction
Training deep neural networks is typically done via a first order optimization method, such as gradient descent (GD), or in the case of large scale datasets and models that are memory prohibitive through minibatch Stochastic Gradient Descent (SGD) (Robbins and Monro, 1951), with the latter method taking gradient steps on a random subsample of the population. While most open source large scale models (Grattafiori et al., 2024; Guo et al., 2025; Team et al., 2026; Yang et al., 2025) are trained with preconditioned SGD, typically Adam variants (Kingma and Ba, 2014; Loshchilov and Hutter, 2017; Vyas et al., 2024), or the more recent approximate second order methods (Gupta et al., 2018; Jordan et al., 2024; Martens and Grosse, 2015), they all generally bake in a version of momentum. Assuming a noiseless quadratic model, under exact gradients classical momentum methods such as heavy ball (Polyak, 1964) and √ Nesterov’s accelerated gradient (Nesterov, 1983, 2013) improve over the GD rate of O(κ) to O( κ). However, this acceleration does not extend to the stochastic regime (Kidambi et al., 2018), limiting their practicality for current training regimes. Empirically, several works have found that momentum primarily accelerates in the large batch size regime, while in the small batch case its effects vanish (Fu et al., 2023; Kunstner et al., 2023; Shallue et al., 2019; Wang et al., 2023; Zhang et al., 2019). Most recently, Marek et al. (2025) conducted a large ∗
: Equal contribution. †: Work done while at Harvard.
1
Steps to target loss
SGD
105
Heavy Ball
Accelerated SGD
a = 0.9, target loss 2.16
1/B (iso-CE)
a = 2, target loss 0.0071 104
104 103
103 102
20 21 22 23 24 25 26 27 28 29 210 211 212
20 21 22 23 24 25 26 27 28 29 210
Batch size
Batch size
Figure 1: Batch-size tradeoffs for stochastic momentum methods in synthetic linear regression with power-law spectra. Plots show the number of serial steps needed to reach a fixed target loss as the batch size varies, after tuning hyperparameters for each method. (Left) For slow decaying power law spectra, HB and ASGD do not yield better CE, but they decrease the serial runtime. (Right) For fast decaying spectra, ASGD has better CE than HB, but has a shorter linear scaling window before it begins trading off CE for serial runtime. Experimental details are provided in Section 5.
scale empirical study in large language model (LLM) pretraining, and have found that at small batch sizes momentum does not bring any benefits over SGD. These works indicate a subtle interplay between compute efficiency and serial runtime in stochastic acceleration. Theoretically, there has been partial progress in understanding stochastic momentum at batch size 1, through the lower bounds established by Kidambi et al. (2018) and the asymptotic characterization of Ferbach et al. (2025); Lee et al. (2022). We refer the reader to Section 2 where we discuss these works further. Based on the the theoretical understanding, new momentum schemes have been proposed combining multiple momentum buffers that provably achieve acceleration at batch size 1 (Ferbach et al., 2025; Gupta et al., 2024; Jain et al., 2018; Liu and Belkin, 2018; Pagliardini et al., 2024; Vaswani et al., 2019), a family of methods we will refer to as Accelerated SGD (ASGD) (Ferbach et al., 2025; Jain et al., 2018; Kidambi et al., 2018; Lee et al., 2022). While these methods achieve accelerated rates and thus better compute efficiency at batch size 1, this regime is far from practical from a serial runtime point of view, especially in the current pretraining setups where the number of tokens can scale to several trillions. Thus, in large-scale pretraining, the key question is not only whether an optimizer improves sample complexity, but also whether it can maintain those gains at the large batch sizes needed to reduce serial training time. While keeping the data size fixed, one can linearly increase the batch size until the critical batch size (CBS) (Merrill et al., 2025; Meterez et al., 2025b; Shallue et al., 2019; Zhang et al., 2024) while achieving a linear decrease in the serial cost. While the CBS for SGD has been extensively studied empirically and theoretically in linear regression, to the best of our knowledge there are no results on the CBS of momentum algorithms. In this paper we provide a nearly full characterization of the serial runtime tradeoffs by providing lower bounds for arbitrary spectra for heavy ball, and power law spectra for ASGD. 2
1.1
Contributions
In the deterministic, full-batch regime, the classical story is clear: on a κ-conditioned quadratic, GD −1/2 has contraction gap of order κ−1 , while HB achieves the accelerated . Equivalently, √ gap of order κ HB reduces the number of serial iterations by a factor of order κ. Our goal is to understand what remains of this acceleration in the stochastic mini-batch regime, where we care about both the serial runtime and the compute efficiency. A useful way to compare stochastic momentum methods is to separate these two quantities. For a given algorithm, let s∗ (H, B) denote the best-achievable spectral radius, or contraction factor, at batch size B and with covariance spectrum H (suppressing dependence on the algorithm). We define the corresponding contraction gap by α(B) := 1 − s∗ (H, B). Algorithm
Batch-size regime
Serial runtime
Compute efficiency
SGD SGD
1 ≤ B ≲ Tr(H) λmax B ≳ Tr(H) λmax
Tr(H) Bλmin
λmin Tr(H) 1 Bκ
HB HB
κ
√ 1 ≤ B ≲ Tr(H) κ λmax √ B ≳ Tr(H) κ λmax
Tr(H) Bλmin
√ κ
λmin Tr(H) 1 √ B κ
Table 1: Serial-runtime and compute-efficiency implications of our lower bounds for SGD and HB under arbitrary spectra. Here κ = λmax /λmin , serial runtime scales as α(B)−1 , and compute efficiency scales as α(B)/B, where α(B) = 1 − s∗ (H, B) is the best-achievable spectral gap. SGD and HB have the same compute-efficiency frontier λ√ min / Tr(H), but HB preserves this frontier over √ a larger batch-size window, up to B ≍ (Tr(H)/λmax ) κ, thereby reducing serial runtime from κ to κ before saturating. In the full-batch limit, this gap satisfies αGD (∞) ≍ κ−1 and αHB (∞) ≍ κ−1/2 , which is the classical deterministic acceleration of HB over GD. The number of serial steps needed to reach a fixed target error scales as 1/α(B), while the total gradient-query cost, or total FLOPs up to constants, scales as B/α(B). Equivalently, the compute efficiency (CE) scales as α(B) . B Thus, higher CE means fewer total FLOPs to reach a target accuracy. If α(B) grows linearly with B, increasing the batch size reduces serial runtime without hurting CE. If α(B) grows sublinearly with B, larger batches still improve serial runtime, but only by sacrificing CE. Our first main result characterizes how this deterministic acceleration appears as the batch size is varied, where we compare HB to SGD, and is summarized in Table 1. The result holds for arbitrary spectra. HB improves the serial runtime over SGD, but it does not improve the CE frontier. More precisely, we prove a lower bound showing that HB cannot improve over the optimal CE scaling of SGD at any batch size. At small batch sizes, HB has the same linear-scaling behavior as SGD. Above the SGD critical batch size, however, HB can continue to convert larger batches into fewer serial steps while preserving the same CE scale. This continues until HB reaches its deterministic accelerated scale. Thus, HB’s benefit is not a better best-achievable CE frontier, but 3
rather a larger batch-size window over which SGD-level CE can be used to buy improved serial runtime. In practice, several newer stochastic momentum methods are also being considered. The majority of these schemes are paramterically equivalent to ASGD, which is known to improve CE over SGD at small batch sizes. Our second main result characterizes how this CE improvement changes with batch size. ASGD improves small-batch CE, but its CE-preserving linear-scaling window is shorter. We show that ASGD’s CE gains are concentrated in the small-batch regime. Unlike the case for HB, the serial-runtime picture for ASGD is more dependent on the full spectra. For rapidly decaying spectra, ASGD has an initial CE-preserving linear-scaling regime. After this regime, its contraction gap continues to improve with batch size, but only sublinearly. Thus, increasing the batch size still reduces ASGD’s serial runtime, but it does so by spending some of ASGD’s small-batch CE advantage. More concretely, for rapidly decaying spectra, the resulting picture is as follows. At small batch sizes, ASGD has better CE than HB, consistent with known upper bounds at B = 1. As B increases, ASGD moves along a CE–serial tradeoff curve: serial runtime improves, but CE decreases relative to the small-batch optimum. Eventually, ASGD reaches the deterministic accelerated serial-runtime scale. Beyond this point, further increasing the batch size no longer improves serial runtime and only decreases CE at the usual 1/B rate. Algorithm
Batch-size regime
Serial runtime
Compute efficiency
SGD
B≳1
a
d
B −1 d−a
HB HB
1 ≤ B ≲ da/2 B ≳ da/2
da B a/2
d
d−a B −1 d−a/2
ASGD
B≍1
d 2a−1
ASGD ASGD
1 ≲ B ≲ d1/2 B ≳ d1/2
d 2a−1 B − 2a−1 da/2
a2
a2
a2
d− 2a−1 a
a2
a−1
d− 2a−1 B − 2a−1 B −1 d−a/2
Table 2: Serial-runtime and compute-efficiency implications of our lower bounds under power-law spectra λi ≍ i−a with a > 1. Serial runtime scales as α(B)−1 and compute efficiency scales as α(B)/B, where α(B) = 1 − s∗ (H, B) is the best-achievable spectral gap. For these spectra, SGD has only a constant-size linear-scaling window. HB preserves SGD-level CE over the larger window B ≲ da/2 , while ASGD improves small-batch CE but loses CE as B grows through the intermediate regime. Thus, ASGD improves the CE–serial frontier most strongly at small batch sizes; for larger batch sizes the advantage is more modest, because part of the small-batch CE gain has already been converted into serial-time improvement. For more slowly decaying spectra, ASGD and HB are comparable. Figure 1 illustrates these CE–serial tradeoffs in synthetic linear regression experiments. The left panel shows the slowly decaying-spectrum regime, where ASGD and HB are nearly overlapping while HB improves serial runtime over SGD, at essentially the same CE as SGD. The right panel shows the rapidly decaying-spectrum regime predicted by our theory: ASGD improves small-batch CE over HB/SGD, and then spends this CE advantage to reduce serial runtime as the batch size grows. Table 2 summarizes the corresponding power-law scalings. 4
In summary, HB preserves SGD-level CE over a larger batch-size range and uses this to improve serial runtime, but it does not improve the CE frontier itself. ASGD improves the small-batch CE frontier, but its linear-scaling window is shorter; beyond this window, larger batches buy serial-time improvements by spending some of ASGD’s CE advantage.
2
Related Work
Acceleration with Noisy Gradients. The suboptimality of GD (Cauchy et al., 1847) for quadratic models has been well studied in literature (Nesterov, 2013). In their seminal works, Polyak (1964) and Nesterov (1983) have proposed different variants of momentum, namely heavy ball and Nesterov’s √ accelerated gradient method, which improve the deterministic rate to O( κ). Several works have shown that for SGD with batch size 1, both HB and NAG do not improve its compute efficiency (Jain et al., 2018; Kidambi et al., 2018). Extensions to these algorithms have been studied by introducing an extra momentum buffer (Gupta et al., 2024; Jain et al., 2018; Kidambi et al., 2018), a family of algorithms commonly refered to as accelerated SGD (Ferbach et al., 2025; Morwani et al., 2025), which provably improve the contraction at batch size 1. Closely related to our work is the work of Lee et al. (2022), who analyze heavy-ball momentum for high-dimensional random least squares in a proportional asymptotic regime. They consider mini-batch sizes β satisfying β/n → ζ > 0 as the sample size n and feature dimension d grow with d/n fixed, and show that acceleration appears only once the batch fraction crosses a spectrumdependent implicit conditioning ratio (ICR), a notion analogous to a critical batch fraction. For upper bounds, Liu and Belkin (2018) established upper bounds at arbitrary batch size for an algorithm named MaSS, which is schematically similar to ASGD. ASGD Variants in Practice. Several momentum schemes based on ASGD have been used in practice, in particular for large language model (LLM) pretraining. Morwani et al. (2025) have shown that the Schedule-Free algorithm (Defazio et al., 2024) can be rewritten as an ASGD equivalent. AdEMAMix (Pagliardini et al., 2024) uses a similar scheme, based on a fast and slow momentum buffer updated with different EMA parameters. Interestingly, Lion (Chen et al., 2023), an algorithm discovered via genetic algorithms, also interpolates between the current gradient and momentum, before applying sign on the update. Several other methods have been proposed in the literature that aggregate over more than 1 momentum buffer (Lucas et al., 2018; Ma and Yarats, 2018). DANA (Bordelon and Mori, 2026; Ferbach et al., 2025) is parametrically related to ASGD: it uses a single momentum buffer together with a direct gradient path in the parameter update, but, notably, chooses the parameters as a function of the model size and training time, rather than keeping them constant. Nature of the bounds. Our results are finite-dimensional, discrete-time lower bounds. For HB, the bound holds for arbitrary finite spectra; for ASGD, the stated comparison holds under the two-sided power-law spectral assumption. In both cases, the bounds hold uniformly over all stable parameter choices, including aggressive choices near the edge of stability, and only hide universal constants. Although we summarize the power-law consequences using large-d scaling notation, these scalings come from finite-d inequalities rather than asymptotic limits. For heavy ball, Lee et al. (2022) provide a lower bound in a proportional limit, showing that there is no acceleration in the small batch size regime. Divergence thresholds. The techniques used in this work to derive the lower bounds also, implicitly, characterize the edge of stability (Cohen et al., 2021, 2022) conditions of the discussed 5
algorithms. Recently, Andreyev et al. (2026) have derived stability thresholds for heavy ball and Nesterov momentum, which we also implicitly recover through the proof of Theorem 1.
3
Background and Preliminaries
In this section, we first introduce the setup necessary for our main theoretical results. Definitions. We denote as f (x) ≲ g(x) if there exists a constant c > 0 such that f (x) ≤ cg(x) for any x in the domain. Moreover, we denote as f (x) ≂ g(x) if g(x) ≲ f (x) ≲ g(x). For 2 matrices A and B we say that A ⪯ B if and only if B − A ⪰ 0, where a matrix B ⪰ 0 denotes that B is positive semidefinite (PSD). We study stochastic optimization for a noiseless linear regression problem with Gaussian covariates i.e. a consistent linear system. We define our notion of compute to be the total number of gradient queries. For the covariance dynamics of the iterate, let T denote the linear operator such that Σt+1 = T (Σt ). Following prior work on consistent linear systems (Meterez et al., 2025a; Wu et al., 2022a; Zou et al., 2021), the error contracts geometrically at a rate governed by the spectral radius of T . We train online, without repeating over the data. At each iteration, we sample from the population a minibatch of size B of covariates {(xi , yi )}B i=1 such that: yi = ⟨w∗ , xi ⟩,
xi ∼ N (0, H),
where w∗ ∈ Rd is the minimizer and H ∈ Rd×d , H ≻ 0 is the population covariance matrix of the covariates. We denote the eigendecomposition of H by H = QΛQ⊤ ,
Λ = diag(λ1 , . . . , λd ),
λmax = λ1 ≥ · · · ≥ λd = λmin > 0
We will use κ = λλmax to denote the conditioning number of H. Where stated, we will additionally min specialize to a power-law spectrum with exponent a defined as: λi ≂ i−a
λmin ≂ d−a
λmax ≂ 1
Moreover, we denote the population risk as: R(w) :=
1 E (⟨w, x⟩ − y)2 2
where the expectation is taken over the randomness induced by the stochastic gradient updates. We will recall a basic identity used repeatedly in the stochastic covariance analysis, namely the Gaussian fourth-moment formula. Proposition 1. For any deterministic matrix Σ and x ∼ N (0, H) we have the following equation: E[xx⊤ Σxx⊤ ] = 2HΣH + Tr(HΣ)H ⪯ 3 Tr(HΣ)H. For a minibatch of size B of covariantes with empirical covariance X̄ = B1 tion 1 generalizes as:
E[X̄ΣX̄] = 1 +
(1) PB
1 1 2 HΣH + Tr(HΣ)H ⪯ 1 + Tr(HΣ)H B B B
6
⊤ i=1 xi xi , Proposi-
(2)
Background. We analyze the ASGD algorithm (Jain et al., 2018; Kidambi et al., 2018; Morwani et al., 2025), with the following update rule: µt = βµt−1 + (1 − β)gt wt+1 = wt − η(µt + ζgt )
(3)
We denote as µt the momentum buffer at time t, the current iterate wt and the current minibatch gradient gt . The hyperparameters of this algorithm are the learning rate η, the momentum EMA parameter β and the weight placed on the current gradient ζ. Note that by setting ζ = 0, we recover the HB update: µt = βµt−1 + (1 − β)gt wt+1 = wt − ηµt
(4)
Therefore, intuitively, since ASGD has an extra hyperparameter we can optimize over, it is possible that it can improve over the rates achieved by HB.
4
Compute Efficiency Lower Bounds
We analyze these algorithms through the covariance dynamics of an augmented state vector zt dependent on the distance to the optimizer wt − w⋆ , and includes all variables needed to make the dynamics first-order linear in the covariance Σt = E[zt zt⊤ ]. For each method, the error covariance satisfies a linear recursion Σt+1 = T (Σt ) where T is a positive operator defined on the space of PSD matrices that depends on the data covariance H, batch size B and the algorithm hyperparameters. Since the excess risk is a linear function of Σt (Meterez et al., 2025a; Wu et al., 2022a,b; Zou et al., 2021), the risk contraction is mainly governed by the spectral radius ρ(T ) = maxi |ti | where ti are the eigenvalues of T . Thus, establishing lower bounds on ρ(T ), or equivalently upper bounds on the spectral gap s = 1 − ρ(T ) would directly imply lower bounds on the error rate for a consistent linear system. We leave the formal extension of our spectral bounds to risk bounds to future work. We first state a general HB lower bound valid for arbitrary spectra, and then specialize it to power-law spectra to identify the critical batch-size transition. We next give an analogous bound for ASGD, showing that for power-law spectra the ASGD lower bound has a strictly better scaling than the HB lower bound in the small-batch regime and reaches the accelerated scale at a smaller batch size.
4.1
Heavy Ball
We can rewrite the HB update from equation equation 4 in the following form: wt+1 = wt − η(1 − β) ĝt + β (wt − wt−1 ),
η > 0, β ∈ [0, 1).
We can rewrite the update as a linear recursion, where the vector form is the augmented state: "
"
#
"
wt+1 − w∗ (1 + β)I − η(1 − β)X̄t −βI = wt − w ∗ I 0 #
#"
wt − w∗ wt−1 − w∗
#
(5)
wt − w∗ Denote by zt = with covariance Σt = E[zt zt⊤ ], and denote the transition operator wt−1 − w∗ from equation equation 5 as Ât . We can decompose Ât = Mt + Ã, where à will be the deterministic 7
component and Mt the random component as: #
"
"
(1 + β)I − η(1 − β)H −βI Ã = I 0
−η(1 − β)(X̄t − H) 0 Mt = 0 0
#
Computing the covariance of the augmented state vector gives: Σt+1 = ÃΣt Ã⊤ + E[Mt zt zt⊤ Mt⊤ ] For the second expectation, we need to compute a 4th moment Gaussian term coming from E[X̄t zt zt⊤ X̄t ], only on the 11 block (since that is where X̄t is). We can compute this term by applying Proposition 1 and upper bound HΣH ⪯ Tr(HΣ)H. With an abuse of notation, we will write the recursion of Σt+1 with equality after applying this bound, since we only lose a small constant factor, thus obtaining: " ⊤
Σt+1 = ÃΣt à +
2η 2 (1−β)2 Tr HΣ11 t H B
0 0
0
#
Henceforth, we will express all matrices in the eigenbasis of H. Thus, after rotation H reduces to Λ with St = Q⊤ Σt Q and A = Q⊤ ÃQ. Thus, the recursion becomes: "1
St+1 = A St A
⊤
2
2
+ η (1 − β)
11 B Λ Tr(ΛSt )
0
0
0
#
(6)
Note that A has a block diagonal structure as A = blkdiag(A1 , . . . , Ad ), where each per–coordinate 2 × 2 block is: "
#
ai −β Ai = , 1 0
ai := (1 + β) − η(1 − β)λi .
(7)
Let TH,η,β,B denote the corresponding linear covariance update operator on symmetric 2d × 2d matrices as defined in equation 6, such that: St+1 = TH,η,β,B (St ) Let s(H, η, β, B) denote the spectral radius of TH,η,β,B
s(H, η, β, B) := ρ TH,η,β,B := max{ |s| : s eigenvalue of TH,η,β,B }. As we see later on (in Lemma 3), this spectral radius is attained by a real, nonnegative eigenvalue with a PSD eigenmatrix. We define the optimal spectral radius attainable at batch size B as s∗ (H, B) :=
inf
η>0, β∈[0,1)
s(H, η, β, B),
This brings us to our first main result. Theorem 1 (HB Compute Efficiency Lower Bound). For any data covariance matrix H and mini-batch size B ≥ 1, β, η > 0, the optimal spectral radius of HB satisfies: (
B λmin s (H, B) ≳ 1 − min , tr(H) ∗
s
λmin λmax
)
.
where ≳ absorbs universal constants. In particular, the transition to the accelerated regime occurs at batch √ Tr(H) κ crit size BHB = λmax . 8
Theorem 1 establishes a lower bound on the optimal spectral radius of HB, which in the case of consistent linear systems governs the error rate. From the theorem, we can see that there are, crit in effect, 2 regimes for HB: for batch size B < BHB , the algorithm cannot improve over the SGD crit scaling, whereas for batch size B ≥ BHB the best-achievable spectral gap is bounded above by the asymptotic contraction rate. Therefore, one can conclude that in order to minimize the number crit of serial steps without wasting compute, the optimal batch size to train at is BHB . In Corollary 1, we specialize the HB lower bound to a power law spectrum, in order to compare the rate with the ASGD bound in Section 4.2. Corollary 1 (HB on power-law spectra). Assume the eigenvalues of H satisfy λi ≂ i−a for some a > 1. Then, the optimal spectral radius of HB satisfies: n
o
s∗ (H, B) ≳ 1 − min Bd−a , d−a/2 , where ≳ absorbs universal constants. In particular, the transition to the accelerated regime occurs at batch crit size BHB ≂ da/2 We defer the full proofs to Appendix A. In the following section, we establish compute efficiency lower bounds for ASGD.
4.2
Accelerated SGD
The analysis for ASGD follows a very similar pattern as HB in Section 4.1. Note that we assume ζ > 0 and 0 < β < 1. The ASGD update rule from Equation 3 can be written as a linear recursion in the following augmented state: "
#
"
wt+1 − w∗ (1 + β)I − η(ζ + 1 − β)X̄t −βI = ∗ wt+1 − w + ηµt I − ηζ X̄t 0
#"
wt − w ∗ wt − w∗ + ηµt−1
#
(8)
We again decompose the transition matrix into a deterministic and a stochastic component: "
#
"
−η(ζ + 1 − β)(X̄t − H) 0 Mt = −ηζ(X̄t − H) 0
(1 + β)I − η(ζ + 1 − β) H −βI Ã = I − ηζH 0 "
#
#
wt − w ∗ Note that all the randomness is in Mt . Denoting by zt = and computing its wt − w∗ + ηµt−1 covariance we get: "
#
11 η 2 (ζ + (1 − β))2 Tr HΣ11 t H ζ(ζ + 1 − β) Tr HΣ t H Σt+1 = AΣt AT + 11 2 11 ζ Tr HΣt H B ζ(ζ + 1 − β) Tr HΣt H
Similarly, we will express all matrices in the eigenbasis of H. Reusing the same notation as in Section 4.1 after rotation in the eigenbasis of H we get Λ and St = Q⊤ Σt Q and A = Q⊤ ÃQ. After computing the 4th moment term using Proposition 1, the recursion becomes: 2 11 11 η 2 (ζ + (1 − β)) Tr ΛSt Λ ζ(ζ + 1 − β) Tr ΛSt Λ St+1 = ASt A⊤ + B ζ(ζ + 1 − β) Tr ΛSt11 Λ ζ 2 Tr ΛSt11 Λ
"
#
(9)
Let TH,η,β,ζ,B denote the corresponding linear covariance update operator for this update rule. Then, we have the following statement. 9
Theorem 2 (ASGD Lower Bound under power-law spectra). Assume the eigenvalues of H satisfy λi ≂ i−a for some a > 1. Then the optimal spectral radius of ASGD satisfies:
s∗ (H, B) ≳ 1 −
a2 − 2a−1 B d , 2 a
B 2a−1 d −a/2 d
a − 2a−1
B ≲ 1, ,
,
1 ≲ B ≲ d1/2 , B ≳ d1/2 ,
where ≳ absorbs universal constants. In particular, the lower bound saturates at the accelerated scale at batch crit size BASGD ≂ d1/2 . Theorem 2 distinguishes 3 scaling regimes as a function of batch size. In the first linear regime, linearly increasing the batch size allows for a linear decrease (up to constants) in the serial steps necessary to achieve a fixed target error. After this, there is a diminishing returns regime indicating that linearly increasing the batch size gives us a sublinear decrease in the serial runtime, since for √ a > 1 the batch exponent a/(2a − 1) < 1 (note that for large a this term behaves roughly like B). Finally, for large batch size we reach a saturation regime, saturating the lower bound at the accelerated scale. Increasing the batch size any further would not yield any improvements in the serial runtime, and will be, instead, wasting compute. However, note that in the power law setting, the ASGD lower bound reaches the accelerated scale at a smaller batch size than the HB lower bound. We expand upon this finding in Corollary 2. Corollary 2 (ASGD Accelerated Regime.). Assume the eigenvalues of H satisfy λi ≂ i−a for some a > 1. crit crit Let BHB and BASGD denote the smallest batch sizes at which the HB lower bound and the ASGD lower bound, respectively, reach the optimal accelerated spectral gap Θ(d−a/2 ). Then crit BHB ≂ da/2 ,
crit BASGD ≂ d1/2 .
Hence
crit BHB ≂ d(a−1)/2 , crit BASGD so for every a > 1, the ASGD lower bound reaches the accelerated scale at a strictly smaller batch size than the HB lower bound. Equivalently, throughout the interval
d1/2 ≲ B ≲ da/2 , the ASGD lower bound has already saturated at the accelerated scale. We defer the full proofs to Appendix B. Compute-serial runtime tradeoff. We now summarize the implications of the lower bounds above. Theorem 1 shows that, for arbitrary spectra, HB cannot improve over the SGD CE frontier. However, note that HB can use larger batches to reduce serial runtime, but this does not correspond to an improved compute-efficiency frontier over SGD. Beyond the HB critical batch size, the lower bound saturates at the deterministic accelerated scale. Under power-law spectra with a > 1, Theorem 2 shows that ASGD has a better small-batch spectral-gap scaling than HB and that its lower bound reaches the accelerated scale at the smaller crit batch size BASGD ≍ d1/2 , as stated in Corollary 2. Thus ASGD offers a sharper compute–serial tradeoff: small batches preserve its compute-efficiency advantage, while increasing the batch size up crit to BASGD converts part of this advantage into reduced serial runtime. We empirically confirm these regimes in synthetic linear-regression experiments in Section 5. 10
5
Experiments
We run synthetic experiments in linear regression on quadratics with power law data with a = 2.0, showing that the observed scaling is consistent with the lower bounds. We set the problem size to be D = 50000 and train for N = 500000, at batch sizes B ∈ {1, 2, 4, 8, 16, 32, 64, 128, 256, 1024}, averaged over 50 seeds. For both HB and ASGD we sweep over learning rate η ∈ {10−5 , 3 · 10−5 , 10−4 , 3 · 10−4 , 10−3 , 3 · 10−3 , 10−2 , 3 · 10−2 , 10−1 , 3 · 10−1 , 1.0, 2.0, 3.0, 5.0, 10.0}, momentum EMA parameter β ∈ {0.8, 0.9, 0.95, 0.99, 0.999, 0.9999} and ASGD hyperparameter ζ ∈ {0.05, 0.1, 0.2, 0.3, 0.5, 0.7, 0.9, 0.95, 0.99}, with ζ = 0 for HB, and we plot each curve at the best set of hyperparameters. We plot the number of steps required to reach a target loss as a function of the batch size.
6
Discussion and Conclusions
In this work, we have established compute efficiency lower bounds for heavy ball momentum and accelerated SGD, as a function of the problem instance and the batch size. Specializing the problem instances to power law spectra, we have directly compared the 2 algorithms showing that there is a performance to serial runtime tradeoff: one can train with ASGD at a smaller batch size for a longer serial runtime, achieving better final loss. From a theoretical perspective, our lower bound for ASGD improves over that of (Kidambi et al., 2018), most notably due to the bound applying for any power law spectrum and not a specifically constructed problem. We believe that the techniques used in this work to derive the lower bounds can also be applied to deriving upper bounds, and we leave this derivation to future work.
Acknowledgements The authors would like to thank Alex Damian and Jingfeng Wu, for helpful discussions. The authors would also like to thank Max Shad and Bala Desinghu for their help with the cluster. AM, DM, PN acknowledge the support of a Kempner Institute Graduate Research Fellowship. AM, SK, DM and PN acknowledge that this work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. SK and DM acknowledge support from the Office of Naval Research under award N0001422-1-2377 and the National Science Foundation Grant under award #IIS 2229881. DM is also supported by a Simons Investigator Fellowship, NSF grant DMS-2134157, DARPA grant W911NF2010021,and DOE grant DE-SC0022199.
References Arseniy Andreyev, Advikar Ananthkumar, Marc Walden, Tomaso Poggio, and Pierfrancesco Beneventano. Momentum further constrains sharpness at the edge of stochastic stability. arXiv preprint arXiv:2604.14108, 2026. Blake Bordelon and Francesco Mori. Theory of optimal learning rate schedules and scaling laws for a random feature model. arXiv preprint arXiv:2602.04774, 2026. Augustin Cauchy et al. Méthode générale pour la résolution des systemes d’équations simultanées. Comp. Rend. Sci. Paris, 25(1847):536–538, 1847. 11
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems, 36:49205–49233, 2023. Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.00065, 2021. Jeremy M Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, Zachary Nado, George E Dahl, et al. Adaptive gradient methods at the edge of stability. arXiv preprint arXiv:2207.14484, 2022. Aaron Defazio, Xingyu Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, and Ashok Cutkosky. The road less scheduled. Advances in Neural Information Processing Systems, 37:9974– 10007, 2024. Damien Ferbach, Katie Everett, Gauthier Gidel, Elliot Paquette, and Courtney Paquette. Dimensionadapted momentum outscales sgd. arXiv preprint arXiv:2505.16098, 2025. Jingwen Fu, Bohan Wang, Huishuai Zhang, Zhizheng Zhang, Wei Chen, and Nanning Zheng. When and why momentum accelerates sgd: An empirical study. arXiv preprint arXiv:2306.09000, 2023. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024. Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Peiyi Wang, Qihao Zhu, Runxin Xu, Ruoyu Zhang, Shirong Ma, Xiao Bi, et al. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025. Kanan Gupta, Jonathan W Siegel, and Stephan Wojtowytsch. Nesterov acceleration despite very noisy gradients. Advances in Neural Information Processing Systems, 37:20694–20744, 2024. Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pages 1842–1850. PMLR, 2018. Prateek Jain, Sham M Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Accelerating stochastic gradient descent for least squares regression. In Conference On Learning Theory, pages 545–604. PMLR, 2018. Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https:// kellerjordan.github.io/posts/muon/. Rahul Kidambi, Praneeth Netrapalli, Prateek Jain, and Sham Kakade. On the insufficiency of existing momentum schemes for stochastic optimization. In 2018 Information Theory and Applications Workshop (ITA), pages 1–9. IEEE, 2018. Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. arXiv preprint arXiv:2304.13960, 2023. 12
Kiwon Lee, Andrew Cheng, Elliot Paquette, and Courtney Paquette. Trajectory of mini-batch momentum: batch size saturation and convergence in high dimensions. Advances in Neural Information Processing Systems, 35:36944–36957, 2022. Chaoyue Liu and Mikhail Belkin. Accelerating sgd with momentum for over-parameterized learning. arXiv preprint arXiv:1810.13395, 2018. Ilya Loshchilov and Frank Hutter. arXiv:1711.05101, 2017.
Decoupled weight decay regularization.
arXiv preprint
James Lucas, Shengyang Sun, Richard Zemel, and Roger Grosse. Aggregated momentum: Stability through passive damping. arXiv preprint arXiv:1804.00325, 2018. Jerry Ma and Denis Yarats. Quasi-hyperbolic momentum and adam for deep learning. arXiv preprint arXiv:1810.06801, 2018. Martin Marek, Sanae Lotfi, Aditya Somasundaram, Andrew Gordon Wilson, and Micah Goldblum. Small batch size training for language models: When vanilla sgd works, and why gradient accumulation is wasteful. arXiv preprint arXiv:2507.07101, 2025. James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning, pages 2408–2417. PMLR, 2015. William Merrill, Shane Arora, Dirk Groeneveld, and Hannaneh Hajishirzi. Critical batch size revisited: A simple empirical approach to large-batch language model training. arXiv preprint arXiv:2505.23971, 2025. Alexandru Meterez, Depen Morwani, Costin-Andrei Oncescu, Jingfeng Wu, Cengiz Pehlevan, and Sham Kakade. A simplified analysis of sgd for linear regression with weight averaging. arXiv preprint arXiv:2506.15535, 2025a. Alexandru Meterez, Depen Morwani, Jingfeng Wu, Costin-Andrei Oncescu, Cengiz Pehlevan, and Sham Kakade. Seesaw: Accelerating training by balancing learning rate and batch size scheduling. arXiv preprint arXiv:2510.14717, 2025b. Depen Morwani, Nikhil Vyas, Hanlin Zhang, and Sham Kakade. Connections between schedule-free optimizers, ademamix, and accelerated sgd variants. arXiv preprint arXiv:2502.02431, 2025. Yurii Nesterov. A method for unconstrained convex minimization problem with the rate of convergence o (1/k2). In Dokl. Akad. Nauk. SSSR, volume 269, page 543, 1983. Yurii Nesterov. Introductory lectures on convex optimization: A basic course, volume 87. Springer Science & Business Media, 2013. Matteo Pagliardini, Pierre Ablin, and David Grangier. The ademamix optimizer: Better, faster, older. arXiv preprint arXiv:2409.03137, 2024. Boris T Polyak. Some methods of speeding up the convergence of iteration methods. Ussr computational mathematics and mathematical physics, 4(5):1–17, 1964. Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pages 400–407, 1951. 13
Christopher J Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. Journal of Machine Learning Research, 20(112):1–49, 2019. Kimi Team, Tongtong Bai, Yifan Bai, Yiping Bao, SH Cai, Yuan Cao, Y Charles, HS Che, Cheng Chen, Guanduo Chen, et al. Kimi k2. 5: Visual agentic intelligence. arXiv preprint arXiv:2602.02276, 2026. Sharan Vaswani, Francis Bach, and Mark Schmidt. Fast and faster convergence of sgd for overparameterized models and an accelerated perceptron. In The 22nd international conference on artificial intelligence and statistics, pages 1195–1204. PMLR, 2019. Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409.11321, 2024. Runzhe Wang, Sadhika Malladi, Tianhao Wang, Kaifeng Lyu, and Zhiyuan Li. The marginal value of momentum for small learning rate sgd. arXiv preprint arXiv:2307.15196, 2023. Jingfeng Wu, Difan Zou, Vladimir Braverman, Quanquan Gu, and Sham Kakade. Last iterate risk bounds of sgd with decaying stepsize for overparameterized linear regression. In International conference on machine learning, pages 24280–24314. PMLR, 2022a. Jingfeng Wu, Difan Zou, Vladimir Braverman, Quanquan Gu, and Sham Kakade. The power and limitation of pretraining-finetuning for linear regression under covariate shift. Advances in Neural Information Processing Systems, 35:33041–33053, 2022b. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025. Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter at which batch sizes? insights from a noisy quadratic model. Advances in neural information processing systems, 32, 2019. Hanlin Zhang, Depen Morwani, Nikhil Vyas, Jingfeng Wu, Difan Zou, Udaya Ghai, Dean Foster, and Sham Kakade. How does critical batch size scale in pre-training? arXiv preprint arXiv:2410.21676, 2024. Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, and Sham Kakade. Benign overfitting of constant-stepsize sgd for linear regression. In Conference on learning theory, pages 4633–4635. PMLR, 2021.
14
A
Heavy Ball Analysis
A.1
Proof of Proposition 1
Proof. The proof is a simple application of Isserlis’s theorem. Elementwise, we have that: E[xx⊤ Σxx⊤ ]ij =
X
Σkl E[xi xk xl xj ] =
kl
X
Σkl (Hik Hlj + Hil Hkj + Hij Hkl )
kl
Assembling the result in matrix form gives the first part. For v ∈ Rd , define u := H 1/2 v. Then v ⊤ HΣHv = v ⊤ H 1/2 (H 1/2 ΣH 1/2 )H 1/2 v = u⊤ H 1/2 ΣH 1/2 u. Also,
Tr(HΣ) v ⊤ Hv = Tr H 1/2 ΣH 1/2 u⊤ u. Since H ⪰ 0 and Σ ⪰ 0, we have H 1/2 ΣH 1/2 ⪰ 0. Hence there exists an orthogonal matrix P and a diagonal matrix D ⪰ 0 such that H 1/2 ΣH 1/2 = P DP ⊤ . Since D ⪰ 0, we have D ⪯ Tr(D) I. Therefore, u⊤ H 1/2 ΣH 1/2 u = u⊤ P DP ⊤ u ≤ u⊤ P (Tr(D)I)P ⊤ u = Tr(D) u⊤ u.
Using Tr(D) = Tr H 1/2 ΣH 1/2 = Tr(HΣ), we obtain v ⊤ HΣHv ≤ Tr(HΣ) v ⊤ Hv. Since this holds for all v ∈ Rd , it follows that HΣH ⪯ Tr(HΣ)H.
A.2
Deriving the Secular Equation
Recall that:
"1
St+1 = A St A
⊤
2
2
+ η (1 − β)
11 B Λ Tr(ΛSt )
0
0
0
#
Let γ = Tr(ΛSt11 ). Pushing a vec through this equation coupled with the fact that vec(ASt A⊤ ) = (A ⊗ A)vec(St ) gives us:
vec(Λ) 2 2 2η (1 − β) γ 0 svec(St ) = (A ⊗ A)vec(St ) + 0 B 0 vec(Λ) 2η 2 (1 − β)2 γ 0 =⇒ (sI − A ⊗ A)vec(St ) = 0 B 0 We can break the above equation up into d, 2 × 2 equations each using a block of A: 15
λi
2
2
2η (1 − β) γ 0 (sI − Ai ⊗ Ai )vec(St,i ) = 0 B
0
λi 0 2η 2 (1 − β)2 γ =⇒ vec(St,i ) = (sI − Ai ⊗ Ai )−1 0 B 0
λi
0 To simplify γ, we multiply the above Equation by and sum over i ∈ {1, . . . , d} to get: 0
0
d h 2 X
2
1=
2η (1 − β) B
i=1
λi −1 0
i
λi 0 0 0 (sI − Ai ⊗ Ai ) 0 0
(10)
Inverting the Kronecker. Suppose Mi is an eigenbasis for Ai , and Di are the eigenvalues of Ai . Then, Mi ⊗ Mi is an eigenbasis for Ai ⊗ Ai and Di ⊗ Di are the eigenvalues of Ai . Thus, we have: (sI − Ai ⊗ Ai )−1 = Mi ⊗ Mi (sI − Di ⊗ Di )−1 Mi⊤ ⊗ Mi⊤ with "
#
"
#
r 0 Di = i,+ 0 ri,− and
r r Mi = i,+ i,− 1 1
(11)
Equations (10) and (11) on further simlification give us:
1=
λ2i s(s + β) 2η 2 (1 − β)2 X 2 )(s − r 2 )(s − β) B (s − ri,+ i,− i
Let ϕi (s) =
s(s + β) 2 ) (s − r 2 ) . (s − β) (s − ri,+ i,−
The final secular Equation can be writen as 1 =
d λ2i ϕi (s) c X , B i=1 1 − cλ2i ϕi (s)
16
c = η 2 (1 − β)2 .
(12)
Validity of vectorizing Equation 9 The vectorized equation is acting on the 4d dimensional space defined by the d 2 × 2 blocks, while the symmetric matrices occupy a 3d space here. So the secular equation has extra roots. But we will argue that the maximum eigenvalue of the secular equation is still associated with a per-diagonal block PSD matrix. First, within the space of per-diagonal symmetric matrices, per-diagonal PSD matrix has the maximum eigenvalue. This follows from KT theorem, as per-diagonal PSD matrix form a total cone. Then, let’s consider the solutions of secular equation outside the symmetric space. Let M be an eigen vector of TH,η,β,B which is neither symmetric nor anti-symmetric. By linearity of TH,η,β,B , TH,η,β,B (M ) = λM =⇒ TH,η,β,B (M T ) = λM ⊤ Now, let’s decompose M into its symmetric and anti-symmetric parts: Ms :=
M + M⊤ , 2
Ma :=
M − M⊤ . 2
Then, by linearity, TH,η,β,B (Ms ) =
TH,η,β,B (M ) + TH,η,β,B (M ⊤ ) λM + λM ⊤ = = λXs , 2 2
and similarly TH,η,β,B (M ) − TH,η,β,B (M ⊤ ) λM − λM ⊤ = = λMa . 2 2 Therefore, any eigenvalue carried by an eigenmatrix M is also carried by a symmetric eigenmatrix and an anti-symmetric eigenmatrix. Also, if M is anti-symmetric, Tr(ΛS11 ) = 0 =⇒ TH,η,β,B (M ) = T∞ (M ). Thus, the eigenvalue of an anti-symmetric matrix for the stochastic operator coincide with the eigenvalue for the deterministic operator. In addition, note the following: (i) A is block diagonal, (ii) give the stochastic opertor TH,η,β,B , gauging at the block diagonal eigenmatrices of the determinsitc operator T∞ suffices. For a block diagonal eigenmatrix, the eigenvalues of T∞ are that of Ai ⊗ Ai , with the eigenmatrix corresponding to the largest eigenvalue being a symmetric eigenmatrix, with rank-1 block diagonal entries of the form ui uTi , where ui represents the eigenvector of Ai corresponding to its largest eigenvalue. Thus, on the determinisitc operator T∞ , the eigenvalue of a symmetric eigenmatrix is always greater than or equal to that of an anti-symmetric matrix. By Conjecture 3, the spectral radius for TH,η,β,B is obtained by a real eigenvalue and from Lemma 2, the spectral radius of TH,η,β,B is lower bounded by the spectral radius for T∞ , thus, the dominating eigenvalue of the secular equation is given by a symmetric PSD matrix. TH,η,β,B (Ma ) =
A.3
The Zero-Noise Heavy Ball Analysis
In the deterministic setting, the Heavy-Ball iteration in Equation 6 reduces to St+1 = ASt A⊤
(13)
The following technical lemma governs the maximal learning rate of the deterministic Heavy– Ball algorithm. Lemma 1 (Stability band of deterministic HB). Assume β, η > 0. If the Heavy–Ball iteration in Equation 13 is stable (i.e. all eigenvalues of every Ai lie in the open unit disk), if and only if β < 1,
0<η< 17
2(1 + β) . (1 − β)λmax
Proof. Fix β ∈ [0, 1) and first consider a generic 2 × 2 matrix "
#
a −β A= , 1 0
a ∈ R.
Let r± be the eigenvalues of A, i.e. the roots of z 2 − az + β = 0. We recall the standard discrete-time stability criterion for a degree–2 polynomial. Writing this polynomial in the form z 2 + pz + q = 0 with p = −a and q = β, the Jury/Schur stability test states that the roots lie strictly inside the unit disk, |z| < 1, if and only if |q| < 1, 1 + p + q > 0, 1 − p + q > 0, 1 − q > 0. (These are the degree–2 Jury conditions.) In our case, substituting p = −a and q = β into the above gives |q| < 1 ⇐⇒ |β| < 1, 1 + p + q > 0 ⇐⇒ 1 − a + β > 0 ⇐⇒ a < 1 + β, 1 − p + q > 0 ⇐⇒ 1 + a + β > 0 ⇐⇒ a > −(1 + β), 1 − q > 0 ⇐⇒ 1 − β > 0. Thus, for the matrix A, the eigenvalues satisfy |r± | < 1 if and only if |β| < 1 and −(1 + β) < a < 1 + β. We now apply this to the deterministic HB iteration on the quadratic. In particular, we apply these conditions to each state matrix Ai , where ai = (1 + β) − η(1 − β)λi . Stability for all i requires −(1 + β) < ai < 1 + β for all λi ∈ [0, λmax ]. The upper bound ai < 1 + β is satisfied, since η > 0 by assumption. The lower bound must hold in particular at the largest curvature λmax , where ai is smallest: (1 + β) − η(1 − β)λmax > −(1 + β). Rearranging, 2(1 + β) . (1 + β) − η(1 − β)λmax > −(1 + β) ⇐⇒ 2(1 + β) > η(1 − β)λmax ⇐⇒ η < (1 − β)λmax Finally, stability also requires |β| < 1 and 1−β > 0, which in our nonnegative-β setting is exactly 0 ≤ β < 1. Combining these conditions gives the claimed deterministic HB stability band.
A.4
Compute Efficiency (CE) lower bounds
Let s(H, η, β, B) denote the spectral radius of TH,η,β,B : s(H, η, β, B) := ρ TH,η,β,B := max{ |s| : s eigenvalue of TH,η,β,B }. As we see later on (in Lemma 3), this spectral radius is attained by a real, nonnegative eigenvalue with a PSD eigenmatrix. We define the optimal spectral radius attainable at batch size B as s∗ (H, B) := inf s(H, η, β, B), η>0, β∈[0,1)
∗
and the corresponding spectral gap as 1 − s (H, B). 18
Theorem 3 (HB–SGD compute efficiency lower bound). For any covariance matrix H and mini-batch size B ≥ 1, β, η > 0, the optimal spectral gap satisfies (
B λmin , s∗ (H, B) ≥ 1 − 8 min tr(H)
A.5
s
λmin λmax
)
.
Helper Lemmas
We will refer to any eigenvalue s solving the secular equation equation 12 as an observable eigenvalue, i.e., an eigenvalue whose eigenmode has nonzero coupling to the scalar observable Sw . Recall the secular equation for the observable eigenvalue s of TH,η,β,B : 1 = F (s) :=
d c X λ2i ϕi (s) , B i=1 1 − c λ2i ϕi (s)
c = η 2 (1 − β)2 ,
(14)
with s(s + β) ϕi (s) = 2 ) (s − r 2 ) , (s − β) (s − ri,+ i,−
ri,± =
ai ±
q
a2i − 4β
2
,
ai = (1 + β) − η(1 − β)λi .
(15)
2 The poles of ϕi are at s = β and s = ri,± ; for a deterministically stable HB choice, all these poles lie in (0, 1). We now formalize the intuition that adding stochastic gradient noise can only slow convergence. In particular, it cannot yield a larger spectral gap than the deterministic Heavy–Ball dynamics. In the zero-noise case (full batch, B = ∞), the Heavy–Ball iteration on the [wt , wt−1 ] state induces the deterministic covariance recursion St+1 = ASt A⊤ , where A = blkdiag(A1 , . . . , Ad ) is the block-diagonal state matrix in the H-basis (with Ai as defined in Equation 7). Let T (∞) (S) := ASA⊤ (16) denote this deterministic (zero-noise) covariance operator.
Lemma 2 (Stochastic spectral radius dominates deterministic). For every (η, β, B) we have ρ TH,η,β,B
≥ ρ T (∞)
= ρ(A)2 .
Equivalently, the stochastic spectral gap is no larger than the deterministic gap: 1 − s(H, η, β, B) ≤ 1 − ρ(A)2 . Before proving this, we record a Perron–Frobenius–type fact for the covariance operators, which follows from the Krein–Rutman theorem (Perron–Frobenius for positive operators on cones). Corollary 3 (Perron–Frobenius for covariance maps). Let S 2d denote the space of real symmetric 2d × 2d matrices and let K := { S ∈ S 2d : S ⪰ 0 } be the cone of positive semidefinite (PSD) matrices. Consider either of the covariance operators T ∈ TH,η,β,B , T (∞) . Then: • T is a positive operator on the cone K, i.e. if S ⪰ 0, then T (S) ⪰ 0. • The spectral radius ρ(T ) := max{ |s| : s eigenvalue of T } is attained by a real, nonnegative eigenvalue. That is, there exists S⋆ ∈ K \ {0} and a real λ⋆ ≥ 0 such that T (S⋆ ) = λ⋆ S⋆ , λ⋆ = ρ(T ). 19
In particular, each covariance operator admits a principal eigenpair (λ⋆ , S⋆ ) with λ⋆ = ρ(T ) and S⋆ ⪰ 0. The corollary shows that the leading asymptotic mode of the covariance dynamics can always be represented by a PSD eigenmatrix. We now prove Lemma 2. Proof of Lemma 2. For brevity, write T (B) := TH,η,β,B . From the exact recursion and the definition of T (∞) in equation 16, we can write T (B) (S) = T (∞) (S) + NB (S), where the linear “noise map” NB is "1 2
2
NB (S) := η (1 − β)
#
B Λ Tr(ΛΣ11 )
0
0
0
"
#
Σ11 Σ12 S= . Σ⊤ 12 Σ22
,
If S ⪰ 0 then Σ11 ⪰ 0, and both Λ Σ11 Λ and Λ tr(ΛΣ11 ) are PSD. Hence NB (S) ⪰ 0, and therefore T (∞) (S) ⪯ T (B) (S)
for all S ⪰ 0.
Apply Corollary 3 to the deterministic operator T (∞) acting on K. There exists an eigenpair (λ∞ , S∞ ) with S∞ ⪰ 0,
T (∞) (S∞ ) = λ∞ S∞ ,
S∞ ̸= 0,
λ∞ = ρ T (∞) .
Since S∞ ⪰ 0 and NB is PSD-valued on K, the domination relation implies T (B) (S∞ ) = T (∞) (S∞ ) + NB (S∞ ) ⪰ λ∞ S∞ . By induction,
n
T (B) (S∞ ) ⪰ λn∞ S∞
for all n ≥ 1.
Fix a matrix norm ∥ · ∥ on S 2d that is monotone on the PSD cone: 0 ⪯ Y ⪯ X implies ∥Y ∥ ≤ ∥X∥. The spectral norm ∥ · ∥2 has this property, since ∥X∥2 = λmax (X) for X ⪰ 0 and Y ⪯ X implies λmax (Y ) ≤ λmax (X). The induced operator norm of (T (B) )n is T
(B) n
n
:= sup S̸=0
T (B) (S) . ∥S∥
Evaluating this supremum at S = S∞ and using monotonicity on the cone, T
n
(B) n
≥
T (B) (S∞ ) ∥S∞ ∥
≥
λn∞ S∞ = λn∞ . ∥S∞ ∥
By Gelfand’s formula, ρ T (B) = lim
n→∞
T (B)
n 1/n
≥ lim λ∞ = λ∞ = ρ T (∞) .
n→∞
Finally, T (∞) acts as A ⊗ A on the vectorized state, so its eigenvalues are products λi λj of eigenvalues of A. Therefore 2
ρ T (∞) = max |λi λj | = max |λi |
i,j
i
Combining the two inequalities yields the claimed bound. 20
= ρ(A)2 .
A.6
Proof
We start with a step size bound for HB–SGD. Lemma 3 (Stability step-size cap). Assume HB–SGD with parameters (η, β, B) is stable, i.e., all eigenvalues of TH,η,β,B satisfy |s| ≤ 1. Then the step size must satisfy 2(1 + β) 2B , η ≤ min tr(H) (1 − β)λmax
.
Proof. First, let us show that stability forces F (1) ≤ 1. Any real s with F (s) = 1 corresponds to an eigenvalue s of TH,η,β,B . In the stable regime, all eigenvalues of TH,η,β,B satisfy |s| ≤ 1, so in particular there can be no real eigenvalue s > 1. ϕi (s) = O(1/s) implies F (s) → 0 as s → ∞. Therefore, if F (1) > 1, by the Intermediate Value Theorem, there exists s > 1 with F (s) = 1, which would correspond to an eigenvalue s > 1 and hence contradict stability. This proves the claim. At s = 1, we can evaluate ϕi (1) explicitly. Using 2 2 (1 − ri,+ )(1 − ri,− ) = (1 + β)2 − a2i = 2(1 + β) η(1 − β)λi − η 2 (1 − β)2 λ2i ,
we obtain ϕi (1) =
1+β 1+β 1 1 2 )(1 − r 2 ) = (1 − β) · η(1 − β)λ · 2(1 + β) − η(1 − β)λ . (1 − β) (1 − ri,+ i i i,−
In the stable regime, each denominator 1 − cλ2i ϕi (1) is positive (otherwise F would blow up at or before s = 1 and there would be an eigenvalue with s ≥ 1). Using x/(1 − x) ≥ x for x ∈ [0, 1), we have d d c X λ2i ϕi (1) c X ≥ λ2 ϕi (1). 1 ≥ F (1) = B i=1 1 − cλ2i ϕi (1) B i=1 i Substituting c = η 2 (1 − β)2 and the expression for ϕi (1) gives 1 ≥
d η(1 + β) X λi . B 2(1 + β) − η(1 − β)λi i=1
(17)
We first derive the bound η ≤ 2B/tr(H). Dropping the negative term in each denominator, 2(1 + β) − η(1 − β)λi ≤ 2(1 + β), we obtain 1 ≥
d d η(1 + β) X λi η(1 + β) X λi η tr(H) ≥ = , B 2(1 + β) − η(1 − β)λi B 2(1 + β) 2B i=1 i=1
which completes the proof of this case. We now show the second bound of η ≤ 2(1 + β)/((1 − β)λmax ) must hold. Lemma 2 implies that if the deterministic HB iteration is unstable, i.e. ρ(A) > 1, then ρ(TH,η,β,B ) > 1. Now if η > 2(1 + β)/((1 − β)λmax ), then Lemma 1 implies ρ(A) > 1, contradicting our assumed stability. This proves our second claim. Lemma 4 (Gap in terms of η and β). For any stable HB–SGD parameters (η, β, B), the spectral gap satisfies 1 − s(H, η, β, B) ≤ min 4ηλmin , 1 − β . 21
Proof. By Lemma 2, the stochastic covariance operator TH,η,β,B has spectral radius at least that of the deterministic (full-batch) HB operator: s(H, η, β, B) ≥ ρ T (∞) = ρ(A)2 ,
where A = blkdiag(A1 , . . . , Ad ) is the deterministic HB state matrix in the [wt , wt−1 ] state. Hence 1 − s(H, η, β, B) ≤ 1 − ρ(A)2 . We first bound the gap by 1 − β. For each coordinate i, the per–coordinate block Ai has eigenvalues ri,± satisfying ri,+ ri,− = β. Thus max{|ri,+ |, |ri,− |}2 ≥ |ri,+ ri,− | = β. Taking a maximum over i yields ρ(A)2 ≥ β, so 1 − ρ(A)2 ≤ 1 − β. Therefore 1 − s(H, η, β, B) ≤ 1 − β. We now show 1 − s(H, η, β, B) ≤ 4ηλmin . Let Amin be the 2 × 2 HB block corresponding to the smallest eigenvalue λmin , with eigenvalues rmin,± solving z 2 − amin z + β = 0,
amin = (1 + β) − η(1 − β)λmin .
As Amin is a block of A, we have ρ(A) ≥ max{|rmin,+ |, |rmin,− |}. Hence 1 − s(H, η, β, B) ≤ 1 − ρ(A)2 ≤ 1 − max{|rmin,+ |, |rmin,− |}2 . We distinguish the real and complex cases for rmin,± . For case of real roots (the overdamped case), Assume rmin,± ∈ R, which corresponds to a2min > 4β. 2 We have 1 − ρ(A)2 ≤ (1 − rmin,+ ) ≤ 2(1 − rmin,+ ). The roots satisfy (rmin,+ − 1)(rmin,− − 1) = 1 − (rmin,+ + rmin,− ) + rmin,+ rmin,− = 1 − amin + β = η(1 − β)λmin . √ Since the roots are real and satisfy r+ r− = β with r− ≤ r+ , we must have r− ≤ β. Using that β ≤ 1 (stability), 1 − ρ(A)2 ≤ 2(1 − rmin,+ ) = 2
p η(1 − β)λmin η(1 − β)λmin √ = 2η(1 + β)λmin ≤ 4ηλmin , ≤2 1 − r− 1− β
which completes the proof of this case. For case of complex roots (underdamped case), assume rmin,± are complex conjugates, where √ 2 amin < 4β. Then both have modulus β, so 1 − ρ(A)2 ≤ 1 − |rmin,+ |2 = 1 − β. For the underdamped case (complex roots), due to that amin = (1 + β) − η(1 − β)λmin , the condition a2min < 4β is equivalent to: η(1 − β)λmin > (1 −
p
β)2
=⇒
η(1 +
p
β)λmin > 1 −
p
Therefore, 1 − β = (1 −
p
β)(1 +
p
β) ≤ (1 +
which completes the proof of the complex case. 22
p
β)2 ηλmin ≤ 4ηλmin ,
β.
Now we are equipped to complete the proof of Theorem 3. Proof of Theorem 3. Fix H and any stable HB–SGD parameters (η, β, B). Combining Lemma 4 (the step size cap) with Lemma 3 (the bound on 1 − s(H, η, β, B)), we obtain three simultaneous upper bounds on the same quantity 1 − s(H, η, β, B): 8Bλmin 8(1 + β)λmin , , 1−β . 1 − s(H, η, β, B) ≤ min{4ηλmin , 1 − β} ≤ min tr(H) (1 − β)λmax
Using the definition of s∗ (H, B), it remains to bound: 8Bλmin 8(1 + β)λmin , 1 − s (H, B) ≤ sup min , 1−β . tr(H) (1 − β)λmax β∈[0,1)
∗
The first term inside the minimum does not depend on β), so the proof consists in bounding:
sup min β∈[0,1)
8(1 + β)λmin , 1−β . (1 − β)λmax
The first function in the min is increasing in β ∈ [0, 1), while the second is decreasing. Consequently, the sup is achived at the at the crossing point β ⋆ , where 8(1 + β ⋆ )λmin = 1 − β⋆ (1 − β ⋆ )λmax
(1 − β ⋆ )2 = 8(1 + β ⋆ )
⇐⇒
Hence
sup min β∈[0,1)
8(1 + β)λmin , 1−β (1 − β)λmax
s
⋆
=1−β =
8(1 + β ⋆ )
λmin . λmax
λmin =4 λmax
s
λmin , λmax
⋆
using 1 + β ≤ 2. This completes the proof. Corollary 4 (HB on power-law spectra). Assume the eigenvalues of H satisfy λi ≂ i−a for some a > 1. Then, the optimal spectral radius of HB satisfies: n
o
s∗ (H, B) ≳ 1 − min Bd−a , d−a/2 , where ≳ absorbs universal constants. In particular, the transition to the accelerated regime occurs at batch crit size BHB ≂ da/2 Proof. Under the power-law spectrum λi ≂ i−a ,
λmin ≂ d−a ,
λmax ≂ 1,
a > 1,
Theorem 3 gives (
Bλmin s (H, B) ≥ 1 − 8 min , tr(H) ∗
Now,
s
λmin λmax
Bλmin B d−a ≂ Pd . −a tr(H) i=1 i
Since a > 1, we have d X i=1
i−a ≂ 1 =⇒
Bλmin ≂ B d−a . tr(H)
23
)
.
Also, s
Therefore
(
λmin ≂ λmax
s
s
)
Bλmin , min tr(H)
λmin λmax
d−a = d−a/2 . 1 n
o
n
o
≂ min B d−a , d−a/2 ,
=⇒ s∗ (H, B) ≳ 1 − min B d−a , d−a/2 .
B
Accelerated SGD Analysis
We now turn our attention to Accelerated SGD (ASGD) algorithm and we establish rates for it in a similar technical way. We begin by establishing the rates for the deterministic operator T∞ , followed by deterministic and stochastic conditions for the learning rate. Recall that: # " 2 11 11 η 2 (ζ + (1 − β)) Tr ΛSt Λ ζ(ζ + 1 − β) Tr ΛSt Λ ⊤ . (18) St+1 = ASt A + B ζ(ζ + 1 − β) Tr ΛSt11 Λ ζ 2 Tr ΛSt11 Λ Let TH,η,β,ζ,B denote the corresponding linear covariance update operator on symmetric 2d × 2d matrices, defined by St+1 = TH,η,β,ζ,B (St ) (19) Denoting p = [(ζ + (1 − β))2 , ζ(ζ + 1 − β), ζ(ζ + 1 − β), ζ 2 ]. Let γ = Tr(HSt11 ) and A = blkdiag(A1 , . . . , Ad ), where each per–coordinate 2 × 2 block is "
(1 + β) − η(ζ + 1 − β)λi −β Ai = 1 − ηζλi 0
#
Pushing a vec through Equation 9 coupled with the fact that vec(ASt A⊤ ) = (A ⊗ A)vec(St ) gives us: svec(St ) = (A ⊗ A)vec(St ) +
η2γ (p ⊗ vec(Λ)) B
=⇒ (sI − A ⊗ A)vec(St ) =
η2γ (p ⊗ vec(Λ)) B
We can break the above equation up into d, 2 × 2 equations each using a block of A: (sI − Ai ⊗ Ai )vec(St,i ) = =⇒ vec(St,i ) =
η 2 γλi p B
η 2 γλi (sI − Ai ⊗ Ai )−1 p B
(20)
Let, q := ζ+1−β,
a := (1+β)−η(ζ+1−β)λi ,
b := −β, 24
c := 1−ηζλi ,
M (s)i := (sI−Ai ⊗Ai )−1 .
Let {e1 , e2 , . . . , ed } represent the standard basis vector. To simplify γ, we multiply Equation 20 by λi e1 and sum over i ∈ {1, . . . , d} to get: 1 = F (s) := with
η2 X 2 ⊤ λ e M (s)i p B i i 1
(21)
sq 2 + b2 ζ 2 + 2ab s ζ q (s + bc) ((s − bc)2 − a2 s)
(s − bc) e⊤ 1 M (s)i p =
(s + β) s(ζ + 1 − β)2 + β 2 ζ 2 − 2β(1 + β)sζ(ζ + 1 − β) + βηζλi s(ζ + 1 − β)2 − β 2 ζ 2
=
h
(s − β + βηζλi ) (s − 1)(s − β 2 ) + 2ηλi s(1 + ζ − β 2 ) − β 2 ζ + η 2 λ2i β 2 ζ 2 − s(ζ + 1 − β)2
i
As argued for the Heavy-Ball case, one can show that the maximum eigenvalue of TH,η,β,ζ,B will occur at a symmetric PSD matrix, thus, justifying the vectorization of Equation 9. Bounding T∞ . By the same argument as in the Heavy–Ball case, the spectral radius of the stochastic covariance operator dominates that of the deterministic operator: ρ(TH,η,β,ζ,B ) ≥ ρ(T∞ ),
T∞ (M ) = AM A⊤ .
Since A = blkdiag(A1 , . . . , Ad ), restricting to the space of symmetric matrices gives ρ(T∞ ) = ρ(A ⊗ A) = ρ(A)2 = max ρ(Ai )2 . i
Thus it suffices to understand the eigenvalues of each block "
#
(1 + β) − η(ζ + 1 − β)λi −β Ai = . 1 − ηζλi 0 We now record two lower bounds on ρ(A)2 . First bound (product bound / complex-root regime). Let ri,+ , ri,− be the two roots of the characteristic polynomial of Ai . Their product is ri,+ ri,− = β(1 − ηζλi ). Hence ρ(Ai )2 = max{|ri,+ |, |ri,− |}2 ≥ |ri,+ ri,− | = β(1 − ηζλi ). Taking i corresponding to λmin , we obtain ρ(A)2 ≥ β(1 − ηζλmin ), and therefore 1 − ρ(A)2 ≤ (1 − β) + ηβζλmin . Second bound (real-root regime). be its two real roots. Then
(22)
Let Amin denote the block corresponding to λmin , and let r+ ≥ r− 2 ρ(A)2 ≥ ρ(Amin )2 ≥ r+ ,
so 2 1 − ρ(A)2 ≤ 1 − r+ ≤ 2(1 − r+ ). From the characteristic polynomial at λmin ,
(1 − r+ )(1 − r− ) = ηλmin (1 − β)(1 + ζ). 25
Also, since r− + r+ = amin ≤ 1 + β, we have 2r− ≤ 1 + β Therefore 1 − r+ =
=⇒
1 − r− ≥
1−β . 2
ηλmin (1 − β)(1 + ζ) ≤ 2ηλmin (1 + ζ), 1 − r−
and hence 1 − ρ(A)2 ≤ 4ηλmin (1 + ζ).
(23)
Combining equation 22 and equation 23, we obtain n
o
1 − ρ(A)2 ≲ min ηλmin (1 + ζ), (1 − β) + ηβζλmin .
B.1
Deterministic learning-rate stability
Lemma 5 (Deterministic ASGD stability band). Assume η > 0. The deterministic ASGD iteration St+1 = ASt A⊤ is stable, i.e. all eigenvalues of every Ai lie in the open unit disk, if and only if 2(1 + β) . 0<η< 1 − β + ζ(1 + β) λmax Proof. For each coordinate i, write "
#
a −β Ai = i , ci 0 Its characteristic polynomial is
ai = (1 + β) − η(ζ + 1 − β)λi ,
ci = 1 − ηζλi .
r2 − ai r + βci = 0. Applying the degree–2 Jury criterion to r2 + pr + q = 0 with p = −ai , q = βci , the roots lie in the open unit disk if and only if |q| < 1, 1 + p + q > 0, 1 − p + q > 0. The condition 1 + p + q > 0 becomes 1 − ai + βci > 0, which simplifies to ηλi (1 − β)(1 + ζ) > 0. This holds automatically for η > 0, λi > 0, 0 < β < 1, and ζ > 0. The condition 1 − p + q > 0 becomes 1 + ai + βci > 0, that is, 2(1 + β) − ηλi (ζ + 1 − β) + βζ > 0. Since (ζ + 1 − β) + βζ = 1 − β + ζ(1 + β), this is equivalent to 2(1 + β) . 1 − β + ζ(1 + β) λi Imposing this for every i yields the claimed condition at λmax . η<
26
(24)
B.2
A stochastic step-size cap
Let
s(H, η, β, ζ, B) := ρ TH,η,β,ζ,B ,
α(H, η, β, ζ, B) := 1 − s(H, η, β, ζ, B).
Recall from the secular equation equation 21 that any real solution of F (s) = 1 is an eigenvalue s of TH,η,β,ζ,B . Since F (s) → 0 as s → ∞, stability implies F (1) ≤ 1. Evaluating the secular equation at s = 1 gives d η X (1 − β 2 )(1 + ζ) + βηζλi (ζ + 1 − β + βζ) . F (1) = λi B i=1 (1 − β) + βηζλi 2(1 + β) − ηλi (1 − β + ζ(1 + β))
(25)
Define S c := {i : βηζλi > 1 − β}.
S := {i : βηζλi ≤ 1 − β},
(26)
Let xi := 1 − β(1 − ηζλi ) = (1 − β) + βηζλi ,
Ni := (1 − β 2 )(1 + ζ) + βηζλi (ζ + 1 − β + βζ).
Following the definition of xi and Ni , Equation 25 can be written as, η X Ni , F (1) = λi B i xi 2(1 + β) − ηλi 1 − β + ζ(1 + β) Since,
2(1 + β) − ηλi 1 − β + ζ(1 + β) ≤ 2(1 + β), we have,
1 1 ≥ . 2(1 + β)xi xi 2(1 + β) − ηλi 1 − β + ζ(1 + β) X Ni η =⇒ F (1) ≥ λi . 2B(1 + β) i xi =1
X (1 − β 2 )(1 + ζ) + βηζλi (ζ + 1 − β + βζ) η λi 2B(1 + β) 1 − β(1 − ηζλi )
Splitting in two cases: Small eigenvalues: βηζλi ≤ 1 − β xi = (1 − β) + βηζλi ≤ 2(1 − β) =⇒ Hence F (1) ≥
1 1 ≥ . xi 2(1 − β)
X η Ni λi . 4B(1 + β) i∈S 1 − β
Now Ni = (1 − β 2 )(1 + ζ) + βηζλi (ζ + 1 − β + βζ) ≥ (1 − β 2 )(1 + ζ), Ni (1 − β 2 )(1 + ζ) =⇒ ≥ = (1 + β)(1 + ζ). 1−β 1−β Substituting back we get, η(1 + ζ) X F (1) ≥ λi . 4B i∈S 27
(27)
Higher eigenvalues: βηζλi ≥ 1 − β xi = (1 − β) + βηζλi ≤ 2βηζλi =⇒ Therefore
X Ni η λi . 4B(1 + β) i∈S c βηζλi
F (1) ≥ =⇒ F (1) ≥
1 1 ≥ . xi 2βηζλi
i Xh 1 (1 − β 2 )(1 + ζ) + βηζλi (ζ + 1 − β + βζ) . 4B(1 + β)βζ i∈S c
=
η(ζ + 1 − β + βζ) X (1 − β 2 )(1 + ζ) c |S | + λi . 4B(1 + β)βζ 4B(1 + β) i∈S c η(ζ + 1 − β + βζ) X (1 − β)(1 + ζ) c |S | + λi . 4Bβζ 4B(1 + β) i∈S c
=⇒ F (1) ≥
(28)
Summing up Equation 27, 28 we get: F (1) ≥
η(1 + ζ) X (1 − β)(1 + ζ) c η(ζ + 1 − β + βζ) X λi . λi + |S | + 4B i∈S 4Bβζ 4B(1 + β) i∈S c
Since,
ζ + 1 − β + βζ 1−β =ζ+ , 1+β 1+β
we have, F (1) ≥
η(1 − β) X (1 − β)(1 + ζ) c η(1 + ζ) X ηζ X λi + λi + λi + |S |. 4B i∈S 4B i∈S c 4B(1 + β) i∈S c 4Bβζ
Since
η(1 + ζ) X η X ηζ X λi = λi + λi , 4B i∈S 4B i∈S 4B i∈S
we get F (1) ≥
η X ηζ η(1 − β) X (1 − β)(1 + ζ) c λi + Tr(H) + λi + |S |. 4B i∈S 4B 4B(1 + β) i∈S c 4Bβζ
(29)
We now assume the power-law spectrum λi ≂ i−a ,
λmin ≂ d−a ,
λmax ≂ 1,
a > 1.
Let k := |S c |. Since the spectrum is monotone, S c consists of the top k eigenvalues up to constants, so X i∈S
d X
λi ≂
i−a ,
Tr(H) ≂ 1.
i=k+1
For a > 1 we have:
d X
i−a ≂ k 1−a .
i=k+1
Since all terms in equation 29 are nonnegative, we may drop the last two terms and obtain F (1) ≥
η X ηζ λi + Tr(H). 4B i∈S 4B 28
Using the power-law estimates above, this gives F (1) ≳
η 1−a ηζ k + . B B
Since stability implies F (1) ≤ 1, we obtain 1≳
η 1−a ηζ k + , B B
hence the stochastic step-size cap η≲
B.3
B . k 1−a + ζ
(30)
Proof of the ASGD rate bound
We define the optimal ASGD spectral radius and gap at batch size B by s∗ (H, B) :=
inf
η>0, β∈[0,1), ζ>0
s(H, η, β, ζ, B),
α∗ (H, B) := 1 − s∗ (H, B).
For every stable choice of (η, β, ζ), the deterministic dominance argument gives α(H, η, β, ζ, B) ≤ 1 − ρ(A)2 . Combining this with equation 24, we obtain n
o
α(H, η, β, ζ, B) ≲ min ηλmin (1 + ζ), (1 − β) + ηβζλmin .
(31)
We split into two cases. Case 1: ζ ≥ 1.
In this regime, using equation 31 and equation 30, α(H, η, β, ζ, B) ≲ ηλmin (1 + ζ) ≲
1+ζ λmin ≂ d−a . ζ
Thus the branch ζ ≥ 1 is at best of order d−a , so it cannot yield acceleration. Case 2: 0 < ζ < 1.
In this regime, 1 + ζ ≂ 1, so the first term in equation 31 is ηλmin (1 + ζ) ≂ ηd−a .
For the second term, using equation 26, (1 − β) + ηβζλmin ≤ (1 − β) + ηζd−a ≲ ηζk −a + ηζd−a ≲ ηζk −a , since k ≤ d implies k −a ≥ d−a . Therefore α(H, η, β, ζ, B) ≲ η min{d−a , ζk −a }.
(32)
Now combine equation 32 with the two step-size caps equation 30 and Lemma 5. For every stable choice of (η, β, ζ) with 0 < ζ < 1, (
α(H, η, β, ζ, B) ≲ min
)
B d−a d−a B ζ k −a , , , k −a . k 1−a + ζ ζ k 1−a + ζ
(33)
At this point, η and β have been eliminated: the right-hand side is an upper bound valid for every stable (η, β, ζ), and it depends only on the induced threshold k and on ζ. Thus ( ∗
α (H, B) ≲
sup 1≤k≤d, 0<ζ<1
min
)
B d−a d−a B ζ k −a , , , k −a . k 1−a + ζ ζ k 1−a + ζ 29
(34)
We now optimize over ζ for fixed k. Set a
k d
ζk := This is the crossover point where
.
d−a = ζk −a .
For fixed k, define B d−a d−a B ζ k −a Ak (ζ) := 1−a , Bk (ζ) := , Ck (ζ) := 1−a , Dk := k −a . k +ζ ζ k +ζ Then Ak and Bk are decreasing in ζ, Ck is increasing in ζ, and Dk is constant. If 0 < ζ ≤ ζk , then ζk −a ≤ d−a , so Ak (ζ) ≥ Ck (ζ),
Bk (ζ) ≥ Dk .
Hence in this region min{Ak (ζ), Bk (ζ), Ck (ζ), Dk } = min{Ck (ζ), Dk }, which is nondecreasing in ζ. If ζk ≤ ζ < 1, then d−a ≤ ζk −a , so Ak (ζ) ≤ Ck (ζ),
Bk (ζ) ≤ Dk .
Hence in this region min{Ak (ζ), Bk (ζ), Ck (ζ), Dk } = min{Ak (ζ), Bk (ζ)}, which is nonincreasing in ζ. Therefore, for each fixed k, the right-hand side of equation 34 is maximized at the crossover point ζ = ζk . Substituting ζk = (k/d)a , we obtain α∗ (H, B) ≲ sup min 1≤k≤d
B d−a
k d
k 1−a +
a , k −a
We now optimize over k. Regime I: B ≲ 1.
Let f (k) := k 1−a +
a
k d
.
Its minimum is obtained by balancing the two terms: a a k 1−a k ≂ =⇒ k ≂ d 2a−1 . d At this value, a(a−1) f (k) ≂ d− 2a−1 , and therefore a2 B d−a ≂ B d− 2a−1 . f (k) Also, a2
k −a ≂ d− 2a−1 . Hence, when B ≲ 1, the first term is the smaller one, and we obtain a2
α∗ (H, B) ≲ B d− 2a−1 . 30
.
(35)
Regime II: 1 ≲ B ≲ d1/2 . In this regime, the optimum is obtained by balancing the two terms in equation 35: B d−a a ≂ k −a . 1−a k + kd Equivalently, a −a
B≂d k
k
1−a
a
k d
+
= da k 1−2a + 1.
For B ≳ 1, this gives a
d B
k≂
1 2a−1
.
Substituting back, a2
a
α∗ (H, B) ≲ k −a ≂ B 2a−1 d− 2a−1 . Regime III: deterministic ceiling.
Finally, stochastic noise can only slow convergence, so ∗ α∗ (H, B) ≤ αdet (H),
∗ where αdet (H) is the optimal full-batch ASGD gap. The standard deterministic ASGD optimization gives ∗ αdet (H) ≲ d−a/2 .
Therefore
α∗ (H, B) ≲ d−a/2
for all B.
The crossover expression a2
a
B 2a−1 d− 2a−1 matches the deterministic ceiling d−a/2 exactly when B ≂ d1/2 . Combining the three regimes, we obtain a2 − 2a−1 , B ≲ 1, B d ∗
α (H, B) ≲
2
a
B 2a−1 d −a/2 d
a − 2a−1
,
31
,
1 ≲ B ≲ d1/2 , B ≳ d1/2 .
(36)