(Accelerated) Noise-adaptive Stochastic Heavy-Ball Momentum
Authors: Anh Quang Dang
Reza Babanezhad Harikandeh
Sharan Vaswani
Venue: Accepted by TMLR ()

Abstract: Stochastic heavy ball momentum (SHB) is commonly used to train machine learning models, and often provides empirical improvements over stochastic gradient descent. By primarily focusing on strongly-convex quadratics, we aim to better understand the theoretical advantage of SHB and subsequently improve the method. For strongly-convex quadratics, Kidambi et al. (2018) show that SHB (with a mini-batch of size $1$) cannot attain accelerated convergence, and hence has no theoretical benefit over SGD. They conjecture that the practical gain of SHB is a by-product of using larger mini-batches. We first substantiate this claim by showing that SHB can attain an accelerated rate when the mini-batch size is larger than a threshold $b^*$ that depends on the condition number $\kappa$. Specifically, we prove that with the same step-size and momentum parameters as in the deterministic setting, SHB with a sufficiently large mini-batch size results in an $O\left(\exp(-\frac{T}{\sqrt{\kappa}}) + \sigma \right)$ convergence when measuring the distance to the optimal solution in the $\ell_2$ norm, where $T$ is the number of iterations and $\sigma^2$ is the variance in the stochastic gradients. We prove a lower-bound which demonstrates that a $\kappa$ dependence in $b^*$ is necessary. To ensure convergence to the minimizer, we design a noise-adaptive multi-stage algorithm that results in an $O\left(\exp\left(-\frac{T}{\sqrt{\kappa}}\right) + \frac{\sigma}{\sqrt{T}}\right)$ rate when measuring the distance to the optimal solution in the $\ell_2$ norm. We also consider the general smooth, strongly-convex setting and propose the first noise-adaptive SHB variant that converges to the minimizer at an $O(\exp(-\frac{T}{\kappa}) + \frac{\sigma^2}{T})$ rate when measuring the distance to the optimal solution in the squared $\ell_2$ norm. We empirically demonstrate the effectiveness of the proposed algorithms.

Published in Transactions on Machine Learning Research (03/2025)
(Accelerated) Noise-adaptive Stochastic Heavy-Ball Momentum
Anh Dang
onh. fl__dangsss..o
Simon Fraser University
Reza Babanczhad Samsung AI, Montreal
beberczzhad@gmaiilcco com
Sharan Vaswani
wwsuueri..horamn@gmaaa.coo
Simon Fraser University
Reviewed on OpenReview: https: // openreview. net/ forum? id- Okep 10
Abstract
Stochastic heavy ball momentum (SHB) is commonly used to train machine learning models, and often provides empirical improvements over stochastic gradient descent. By primarily focusing on strongly-convex quadratics, we aim to better understand the theoretical ad advantage of SHB and subsequently improve the method. For strongly-convex quadratics, Kidambi et al. (2018) show that SHB (with a mini-batch of size 1) cannot attain accelerated convergence, and hence has no thcorctical benefit over SGD. They conjecture that the practical gain of SHB is by-product of using larger mini-batches. We first substantiate this claim by showing that SHB can attain accelerated rate when the mini-batch size is larger than threshold b" that depends on the condition number τ. Specifically, we prove that with the same step-size and momentum parameters as in the deterministic setting, SHB with a sufficiently large mini-batch size results in an O (exp(-T/vR) + o)) convergence when measuring the distance to the optimal solution in the l norm, where T is the number of iterations and o² is the variance the stochastic gradients. We prove lover-bound which demonstrates that a н dependence in b" is necessary. To ensure convergence to the minimizer, design noise-adaptive multi-stage algorithm that results in an O exp (-T/vR) (T) rate when measuring the distance to the optimal solution in the f norm. We also consider the general smooth, strongly-conver setting and propose the first noise-adaptive SHB variant that converges to the minimizer at an O(exp(-T/x> + +) rate when measuring the distance to the optimal solution in the squared &2 norm. We empirically demonstrate the effectiveness of the proposed algorithms.
1 Introduction
Heavy ball (HB) or Polyak momentum (Polyak, 1964) has been extensively studied for minimizing smooth, strongly-conver quadratics in the deterministic setting. In this setting, HB converges to the minimizer at an accelerated linear rate (Polyak, 1964; Wang et al., 2021) meaning that for a problem with condition number R (sce definition in Scction 2), iterations of HB results in the optimal (cxp(-T/vR)) convergence. For general smooth, strongly-convex functions, Ghadimi al. (2015) prove that HB converges to the minimizer al. linear but non-accelerated rate. In this setting, Wang al. (2022) prove accelerated linear rate for HB, but under very restrictive assumptions (c.g. one-dimensional problems or problems with a diagonal hessian). Recently, Goujaud et al. (2023) showed that HB (with any fixed step-size or momentum parameter) cannot achieve accelerated convergence on general (non-quadratic) strongly-convex problems, and consequently has no theoretical benefit over gradient descent (GD).
While there is good theoretical understanding of HB in the deterministic setting, the current understanding of stochastic heavy ball momentum (SHB) is rather unsatisfactory. SHB is commonly used to train machine

---

Published in Transactions on Machine Learning Research (03/2025)
learning models and often provides empirical improvements over stochastic gradient descent (SGD). Further- more, it forms the basis of modern adaptive gradient methods such as Adam (Kingma & Ba, 2014). As such, it is important to better understand the theoretical advantage of SHB over SGD. Previous works (Defazio, 2020; You et al., 2019) have conjectured that the use of momentum for non-convex minimization can help reduce the variance resulting in faster convergence. Recently, Wang el al. (2023) analyze stochastic moment in the regime where the gradient noise dominates, and demonstrate that in this regime, momentum has limited benefits with respect to both optimization and generalization. However, it is unclear whether momentum can provably help improve the convergence in other settings. In this paper, we primarily focus on the simple setting of minimizing strongly-convex quadratics, with the aim of better understanding the theoretical benefit. of SHB and subsequently improving the method.
We first consider the general smooth, strongly- convex setting and aim to design an SHB variant that matches the theoretical convergence of SGD. In this setting, Sebbouh et al. (2020); Liu et al. (2020) use SHB with a constant step-size and momentum parameter, obtaining linear convergence minimizer. In order to attain convergence to the solution, Sebbouh el al. (2020) use a sequence of constant- then-decreasing step-sizes to achieve an O
the neighborhood of the
rate, where o² is the variance in the stochastic gradients and the sub-optimality is mcasured in the squared & norm. In contrast, in the same setting, SGD can attain an O (exp (-T/x) + o²/z) convergence to the minimizer. To the best of our knowledge, in this setting, there is no variant of SHB that can converge to the minimizer at a rate matching SGD.
Contribution 1: Noise-adaptive, non-accelerated convergence to the minimizer for smooth, strongly-convex functions. In Section 3, we propose an SHB method that combines the averaging interpretation of SHB (Sebbouh et al., 2020) and the exponentially decreasing step-sizes (Li et al., 2021; Vaswani et al., 2022) to achieve an O (exp (-T/x) + convergence rate that matches the SGD rate. Importantly, the proposed algorithm is noise-adaptive meaning that it does not require the knowledge of o², but recovers the non-accelerated linear convergence rate (matching Ghadimi et al. (2015)) when a = 0. Morcover, the algorithm provides an adaptive way to set the momentum paramcter, alloviating the need to tune this additional hyper-parameter.
Next, we focus on minimizing strongly-convex quadratics, and aim to analyze the conditions under which SHB is provably better than SGD. A number of works (Kidambi et al., 2018; Paquette & Paquette, 2021; Loizou & Richtárik, 2020; Bollapragada et al., 2023; Lee et al., 2022) have studied SHB for minimizing quadratics. In this setting, Kidambi et al. (2018) show that SHB (with batch-size 1 and any choice of step-size and momentum parameters) cannot attain an accelerated rate. They conjecture that the practical gain of SHB is a by-product of using larger mini batches. Similarly, Paquette & Paquette (2021) demonstrate that SHB with small batch-sizes cannot obtain a faster rate than SGD. While Loizou & Richtárík (2020) prove an accelerated rate for SHB (for any batch-size) in the "LI sense", this does not imply acceleration according to the standard sub-optimality metrics. Recently, Bollapragada et al. (2023): Lee et al. (2022) use results from random matrix theory to prove that SHB with constant step-size and momentum can achieve an accelerated rate when the mini-batch size is sufficiently large. Compared to these works, we use the non-asymptotic analysis standard in the optimization literature, and prove stronger worst-case results.
Contribution 2: Accelerated convergence to the neighborhood for quadratics. Our result in Section 4.1 substantiates the claim by Kidambi al. (2018). Specifically, for strongly-convex quadratics, we prove that SHB with a mini-batch size larger than a cerlain threshold b* (that depends on k) and constant step-size and momentum parameters can achieve O (exp(- + a) non-asymptotic convergence up to neighborhood of the solution where the sub-optimality is measured in the la norm. For problems such as non- non-paraeetric regression (Belkin et al., 2019; Liang & Rakhlin, 2020) or feasible linear systems, where the interpolation property (Ma et al., 2018; Vaswani et al., 2019) is satisfied, a = 0 and SHB with a large batch-size results in accelerated convergence to the minimizer.
Contribution 3: Lower Bound for SHB. Our result in Section 4.2 shows that there exist quadratics for which SHB (with constant step-size and momentum) diverges when the mini-batch size is below certain threshold. Morcover, the lower-bound demonstrates that dependence in b' is necessary.
The result in Section 4.1 only demonstrates convergence to the neighbourhood of the solution. Next, we aim to design an SHB algorithm that can achieve acceleraeed convergence to the minimizer.
2

---

Published in Transactions on Machine Learning Research (03/2025)
Contribution 4: Noise-adaptive, accelerated convergence to the minimizer for quadratics. In Section 4.3, we design a multi-stage SHB method (Algorithm 1) and prove that for strongly-convex quadratics, Algorithm 1 (with sufficiently large batch-size) converges to the minimizer at an accelerated O (exp(-T/va] + (F)) rate where the sub-optimality is measured in the l2 norm. Algorithm is noise- adaptive and has a similar structure as the algorithm proposed for incorporating Nesterov acceleration in the stochastic setting (Aybat et al., 2019). In comparison, both Bollapragada et al. (2023); Lee et al. (2022) only consider accelerated convergence to a neighbourhood of the minimizer. In concurrent work, Pan et al. (2024) make a stronger bounded variance assumption in order to analyze SHB for minimizing strongly-conee quadratics. They propose a similar multi-stage algorithm and under the bounded variance assumption, prove that it can converge to the minimizer at an accelerated rate for any mini-batch size. In Section 4.3, we argue that the bounded variance assumption is problematic even for simple quadratics and the algorithm in Pan et al. (2024) can diverge for small mini-batches (see Fig. 2).
In settings where T > n, the batch-size required by the multi-stage approach in Algorithm can be quite large, affecting the practicality of the algorithm. In order to alleviate this issue, we design two phase algorithm that combines the algorithmic ideas in Sections 3 and 4.1.
Contribution 5: Partially accelerated convergence to the minimizer for quadratics. In Section 4.4, we propose a two-phase algorithm (Algorithm 2) that uses constant step-size and momentum in Phase 1, followed by an exponentially decreasing step-size and corresponding momentum in Phase 2. By adjusting the relative length of the two phases, we demonstrate that Algorithm 2 (with sufficiently large batch-size) can obtain partially accelerated rate.
Contribution 6: Experimental Evaluation. In Section 5, we empirically validate the effectiveness of the proposed algorithms on synthetic benchmarks. In particular, for strongly-convex quadratics, we demonstrate that SHB and its variants can attain an accelerated rate when the mini- bani-batch size is larger than a threshold. While SHB with a constant step-size and momentum converges to a neighbourhood of the solution, Algorithms and are able to counteract the noise resulting in smaller sub-optimality.
2 Problem Formulation
We consider the unconstrained minimization of a finile-sum objective / Rd + == Ei For supervised learning, n represents the number of training examples and fa is the loss of example i. Throughout, we assume that f and each f/ are differentiable and lower-bounded by f and fū, respectively. We also assume that each function fa is L,-smooth, implying that f is L-smooth with L == max; L,, Furthermore, f is considered to be rongly convex while each Si is convext. We define x # as the condition number of the problem, and denote w" to be the unique minimizer of the above problem. 'We primarily focus strongly- convex quadratic objectives where fi(w)= == jwu" A - (di, w) and (w) = 1 ET fi(w) = w" Aw (d, w), where A, are symmetric positive semi-definite matrices. Here, L = Amax [A] and p Amin [A] > 0, where Amax and Amin refer to the maximum and minimum eigenvalues.
In each iteration & E [T] == {0, (..,,..,7) T}, SHB samples a mini-batch B& (b: : [Bel) of examples and uses it compute the stochastic gradient of the loss function. The mini-batch is formed by sampling without replacement. We denote (ix(w))) to be the average stochastic gradient for the mini-batch BN, meaning that V ffa(ws) the direction of V (wx) together with a momentum term computed using the previous ilerate. Specifically, the SHB update is given as:
EiEB, f((w))) and E|V fax (wx)) [wy] = V (wz.) At iteration k SHB takes a descent step in
Wk++1 = W0 Og (ix(ww) + 8x (wx - wk--)
(1)
where Wk+1, Wk, and Wx-1 the SHB iterates and w-1 wa; (ox) and {Bx)T=) is the sequence of step-sizes and momentum parameters respectively. In the next section, we analyze the convergence of SHB for general smooth, strongly-convex functions.
We include definitions of these properties in App.
3

---

Published in Transactions on Machine Learning Research (03/2025)
3 Non-accelerated linear convergence for strongly-convex functions
We first consider the non-accelerated convergence of SHB in the general smooth, strongly-convex setting. Following Loizou et al. (2021); Vaswani et al. (2022), we define x² := E. as the measure of stochasticity. We develop an SHB method that (i) converges to the minimizer at the O (exp (-T/<) + - rate, (ii) is noise-adaptive in that it does not require the knowledge of o² and (iii does not require manual tuning the momentumm parameter. In order to do so, we use an alternative form of the update (Sebbouh et al., 2020) that interprets SHB as a moving average of the iterates 24 computed by stochastic gradient descent. Specifically, for 20 Wh,
Ak+1
Wk+1
Wk +
2k 26 == 2k-1 Th (wxk)
(2)
Ak+1 +1
Ap+1 +
where MMK } are parameters to be determined theoretically. For any {me, } sequence, if a ++ d 3. = T then the update in Eq. (2) is equivalent to the SHB update in Eq. (1) (Sebbouh et al., 2020, Theorem 2.1). The proposed SHB method combines the above averaging interpretation of SHB and exponentially decreasing step-sizes (Li et al., 2021; Vaswani et al., 2022) to achieve noise- adaptive non- accelerated convergence rate. Specifically, following Li et al. (2021); Vaswani et al. (2022), we set a = U where v is the problem-dependent scaling term that captures the smoothness of the function and * is the problem-independenn term that controls the decay of the step-size. By setting {Л,, Th appropriately, the following theorem (proved in App. B) shows that the proposed method converges to the minimizer at an O (exp (-T/x) + a //)r) rate. In contrast, Sebbouh et al. (2020) use constant- sestaantthen-doeaying step-sizes to obtain a sub-optimal (x2/12 2 f rate.
m
Theorem 1. For L-smooth, strongly-convex functions, SHB (Eq. (2)) with 2 1, == +. 7 (7) (1 (1 -ma))*)) converges Tik k+1, Th unk, 7/0 and d := 1-2nL
T ? ( (<Am(7)))
CC d²
E w w01 s C. I|w w01l exp
T
where (= V (n-1)6 and C4, Cs are polynomial in x and poly-logarithmic in T.
This rate matches that of SGD with exponentially decreasing slep-size (Li et al., 2021; Vaswani et al., 2022). In the deterministic setting, when b = n, then by Lemma 7, ç = O, and SHB matchos the non-accelerated linear rate of GD and HB (Ghadimi et al., 2015). Non-parametric regression (Belkin
-1.5
1.0
-2.0
alpha_
alpha_000 _aaa
0.8
-2.5
alpha_0/0 2/k
-3.0
beta_k
log(alpha
0.6
-3.5
beta
et al., 2019; Liang & Rakhlin, 2020) or feasible lin- car systems (Loizou & Richtárik, 2020) satisfy the interpolation (Ma et al., 2018; Vaswani et al., 2019) property. For these problems, the model is able to completely interpolate the data meaning that the noise at the optimum vanishes and hence a = For this case, SHB matches the convergence rate of constant step-size SGD (Vaswani et al., 2019). Com- pared to Sebbouh et al. (2020), the rate in Theorem the
-4.0
0.4
-4.5
0.2
-5.0
-5.5
0.0
0
20
40
80
100
iterations
has the same optimal O(1/T)) dependence variance term, but results in a worse dependence on the constants. On the other hand, our algorithm results in a better dependence on the bias term: for = 100, == 10, -1=1
Figure 1: Variation in a% == 1+
and B8
A
1+
where 7 = (3/2)
d 1-2 (1 (1-
)
O(cxp(-T)) Theorem versus the O(1/T2) rate obtained by Scbbouh et al. (2020). Consequently, when using the full dataset or under the interpolation setting, the rate in Theorem recovers that of deterministic HB (Ghadimi et al., 2015), while that in Sebbouh et al. (2020) does not. For general strongly-convex
4

---

Published in Transactions on Machine Learning Research (03/2025)
functions, Goujand et al. (2023) prove that HB (with any step-size or momentum parameter) cannot achieve accelerated convergence rate on general (non-quadratic and with dimension greater than 1) smooth, strongly-convex problems. Furthermore, we know that the variance term (depending on o²) cannot be decreased at a faster rate than 2(4r) (Nguyen et al., 2019). Hence, the above rate is the best-achievable for SHB in the general strongly-convex setting.
We reiterate that the method does not require knowledge of o² and is hence noise-adaptive. Furthermore, all algorithm parameters are completely determined by the μ,, L and 7 sequence. Hence, the resulting algorithm does not require manual tuning of the momentum. In Fig. 1, show the variation of the (ak, 3k)) paramcters, and observe that the method results in more aggressive decrease in the step-size (compared to the standard O(4))) rate). This compensates for the increasing momentum parameter. The above theorem requires knowledge of L and μ which can be difficult to obtain in practice. Hence in App. C, we consider the effect of misestimating and i on the convergence rate of SHB. These are the first results that consider the effect of parameter misspocification for SHB.
Next, focus eerongly-conver quadratics where SHB can obtain an accelerated convergence rate.
4 Accelerated linear convergence for strongly-convex quadratics
In this section, we focus strongly-convex quadratics and in Section 4.1, we prove that SHB with large batch-size attains accelerated lincar convergence to neighbourhood determined by the noise. In Section 4.2, we prove a corresponding lower-bound that demonstrales the necessity of a large batch-size to atlain acceleration. The exponentially decreasing step-sizes in Section 3 are too conservative to obtain an accelerated rate to the minimizer. Consequently, in Section 4.3, we design a multi-stage SHB algorithm that achieves accelerated convergence to the minimizer. Finally, in Section 4.4, we design a two-phase SHB algorithm that has simpler implementation, but can only attain partially accelerated rates.
4.1 Upper Bound for SHB
In the following theorem (proved in App. D), we show that for strongly-convex quadratics, SHB with a batch- size b larger than certain problem-dependent threshold b", constant step-size and momentur parameter converges to neighbourhood of the solution an accelerated linear rate. We note that the measure of suboptimality in this section is expressed as the norm, whercas in the previous section, it was represented as the squared norm. By Jensen's 's inequality, an upper-bound on E w - w" implies an upper-bound on w*|l
E |wr
Theorem 2. For L-smooth, μ strongly-convex quadratics, SHB (Eq. (1)) with as = 4. for S 1.
B& 8 (1 ?aui). 2 batch-size b s.t. b2 b" n max
converges as:
1+ 1-- 1+ 0.2
6/2VR Va
VaT (3 2/K
(4' 1 2/Rv()) llw0 -w"l + 12vax
min { 1,
S
E UT - w"ll S
exp
max
Vol
where =
E [V[((>)]]
|= / (n- and C' : 3526.
n-h
The first term in the convergence rate represents the bias. Since 1 1-2VEV > 3 when b > b", the initial sub-optimality ||wo - =#|l is forgotten at an accelerated linear rate proportional to exp(-T/vR) Moreover, since the bias term depends on 5, using a larger batch-size (above 6*) leads to a smaller ( resulting in faster convergence. We note that a is a constant independent of T to avoid the dependence of b* on T. In the deterministic case, when b == n and ( = O, we recover the non-asymptotic accelerated convergence for HB (Wang et al., 2021). Similar to the deterministic case, the accelerated convergence requires "warmup" number of iterations meaning that T needs to be sufficiently large to ensure that 6v2x
exp (- T Ve vR 2 max f -2/Rv/)) < size of the neighbourhood. The above theorem uses X2 = ([] as the measure of stochasticity, where x² 2 because of the hnnssstheess of the problem. Compared to constant. step-size SGD that achieves
Va The second term represents the variance, and determines the
5

---

Published in Transactions on Machine Learning Research (03/2025)
O(cxp(-T/kk) + x), SHB with a sufficiently large batch-size results in an accelerated O(cxp(-T/vVR + x) rate. We observe that if k is large, larger batch-size is required to attain acceleration. Tikewise, using a smaller step-size requires a proportionally larger batch-fraction to guarantee an accelerated rate. On the other hand, as n increases, the relative batch-fraction (oqual to b/n) required for acceleration is smaller. The proof of the above theorem relies on the non-asymplolic result for HB in the deterministic setting (Wang et al., 2021), coupled with an inductive argument over the iterations.
The above result substantiates the claim that the practical gain of SHB is by-product of using larger mini-batches. In comparison to the above result, Loizou & Richtárik (2020) also prove an accelerated rate for SHB, but mcasure the sub-optimality in terms of |Bjrr - problem's 's stochasticity and only shows convergence in expectation, a weaker form of convergence compared to our mean-square result. In contrast to Bollapragada et al. (2023, Theorem 3.1) which results in an (T exp(-7//R) a log(d)) rate where d is the problem dimension, we obtain a faster convergence rate without an additional T dependence in the bias term, nor an additional log(d) dependence in the variance term. In order to achieve accelerated rate, threshold b' scales as O (y,.4.2) When >> O(K²), our result implies that SHB with nearly constant (independent of n) mini-batch size can attain accelerated convergence to a neighbourhood of the minimizer. In contrast, (Bollapragada et al., 2023, Theorem 4) require a batch-size of ((dx3/2) to attain an accelerated rate in the worst-case. This condition is vacuous in the over-pan parameterized regime when d > n. Hence, compared to our result, Bollapragada et al. (2023) require a more stringent condition on the batch-size when > VR. On the other hand, Lee et al. (2022) provide an average-cast analysis of SHB as d, - oc, and prove an accelerated rate when b 2 n vR where is the average condition number. In the worst-case (for example, when all data points are the same and K == K 1), Lee et al. (2022) require b = n in order to altain an accelerated rate.
-="==""ll w*lll This does not effectively model the
&
In the interpolation setting described in Section 3, the noisc at the optimum vanishes and = implying that ==0. O. In this setting. we prove the following Corollary 1 in App. D. Hence, under interpolation, SHB with a sufficiently large batch-size results in accelerated convergence to the minimizer, matching the corresponding result for SGD with Nesterov acceleration (Vaswani et al., 2022, Theorem 6) and ASGD (Jain et al., 2018).
Corollary 1. For L-smooth, le strongly-convex quadratics, under interpolation, SHB (Eq. (1)) with the same parameters as in Theorem 2 and batch-size b s.t. b 2 b' := 11t (where C is defined in Theorem 2 converges as:
6 V exp
(- T Va V 2
E WT w0l
max
4 111-2
1120 - -w"ll
Va
When the noise X * 0 but is assumed to be known, Corollary 4 (proved in App. D) shows that the step-size and momentum parameter of SHB can be adjusted to achieve sub-optimality (for some desired >> 0) at an accelerated linear rate. In the above results, the batch-size threshold depends on k. In the following section, we prove a lower-bound showing that a dependence on ks is necessary.
4.2 Lower Bound for SHB
For SHB with the same step-size and momentum as Corollary 1, we show that there exists quadratics for which SHB with batch-size lower than a certain threshold diverges.
6

---

Published in Transactions on Machine Learning Research (03/2025)
Theorem 3. For I-smooth, j strongly-convex quadratic problem f(w) == a Ei= 2 samples and dimension d = n 100 such that w" = 0 and each A; is an n-by-n matrix of all zeros except at the (i, position, we run SHB (1) with a₁ a t,, 3x= = (1 \vaa)) If b
1 w A, with
and
A :- wh for a c > 1, after 6T iterations, we have that: wk-1
E || [2] > c 1
The above lower-bound demonstrates that the dependence on k is necessary in the threshold b* for the batch-size. We note that the designed problem with n d corresponds to feasible lincar system and therefore satisfies interpolation. Intuitively, Theorem 3 shows that in order to atlain an accelerated rate for SHB, it is necessary to have a large batch-size to effectively control the error between the empirical Hessian b EiEB, A, iteration and the true Hessian. When the batch-size not large enough, the aggressive updates for accelerated SHB increase this error resulting in divergence. Importantly, the above lower-bound also holds for the step-size and momentum parameters used in Bollapragada et al. (2023). We note that our lower-bound result still leaves open the possibility that there are other (less aggressive) choices of the step-size and momentum that can result in an (accelerated) convergence rate with smaller batch-size. The proof of the above theorem in App. E takes advantage of symbolic mathematics programming (Meurer et al., 2017), and maybe of independent interest. In contrast to the above result, the lower bound in Kidambi et al. (2018) shows that there exist strongly-convex quadratics where SHB with a batch-size of 1 and any choice of step-size and momentum cannot result in an accelerated rate.
;
We have shown that for strongly-convex quadratics (not necessarily satisfying interpolation), SHB (with large batch-size) can result in accelerated convergence to the neighbourhood of the solution. Next, design a multi-stage algorithm that ensures accelerated convergence to the minimizer.
Algorithm 1: Multi-stage SHB
Algorithm 2: Two-phase SHB
Input: T (iteration budget), (batch-size) Initializatinn: Wn, W-1 Wo, k 0
Input:
T (iteration budget), (batch-size), e E (0, 1) (relative phase lengths)
I == [n(v))
W
T In(v2) (W() the Lambert W function?)
384 VK
Initialization: Wo; w-1 = Wo, k =
To
Set To = cT
WE [1, 11, T. 112//
(2-v2) (1/++)) (n() + In( In(vR))
for k = 0; k < To; h = k do 1, set a, B according to
Choose
for i= 0; i < i+1;i=i+1 do
Theorem 2
Set = 2-1,
- B = (1 3 (o.
Use Update 1
Tp Wi
end
for k= 0; k < Tú; k = k + do
for
k= To T Sct m) ) 1,k
do according to Theorem 1 Use Update 2
Sample batch B& and calculate V (2x)
T++1 = Is - a, [x(( + B, (ax xz-1)
end
end
Wi+1 = TT;
return WT
end
return W1+1
4.3
Multi-stage SHB
In this section, we propose to use multi-stage SHB algorithm (Algorithm 1) and analyze its convergence rate. The structure of our multi-stage algorithm is similar to Aybat et al. (2019) who studied Nesterov acceleaation in the stochastic setting. For a fixed iteration budget T, Algorithm allocates T T/2 iterations to stage zero and divides the remaining T/2 iterations into / stages. The length for each of these stages
2The Lambert W function is defined as: for x, R. y = W(x) yexp(y)
7

---

Published in Transactions on Machine Learning Research (03/2025)
increases exponentially, while the step-size used in each stage decreases exponentially. This decrease in the step-size helps counter the variance and ensures convergence to the minimizer. Theorem 4 (proved in in App. F) shows that Algorithm converges to the minimizer at an accelerated linear rate.
Theorem 4. For L-smooth, μ strongly-convex quadratics with k 1, for T 2 T 3:2% max (4K, } In(2)
Algorithm with b 2 b" = max
1+ G+ in
converges as:
T
11000 -w"l Cs X
E |loT w01 S S exp
8/R
vT
where C == 3526 and C7, Css are polynomial in and poly-logarithmic in T.
From Theorem 4, we see that Algorithm achioves a convergence rate of O (xp minimizer. It is important to note that in comparison Theorem 1, the sub-optimality above is terms of E f|wy - w'll (instead of E |aa wer--="l'" [l²). Hence, the above rate is optimal for strongly-convex quadratics since the bias term decreases at an accelerated linear rate while the variance term goes down as 1/VF. Unlike
(-
T
+
x
to the
vR
vT)
in Corollary 4, Algorithm 1 does not require the knowledge of x and is hence noise-adaptive. When X = 0, Algorithm 1 matches the rate of SHB in Corollary 1.
In concurrent work, Pan et al. (2024) design a similar multi-stage SHB algorithm. However, the algorithm's analysis requires a bounded variance assumption which implics that for all E [T], there exists a ā < oc such that E f(ws) - V fux(w))l( < 0². For strongly-conver quadratics, this assumption implies that the algorithm iterates lie a compact set (Jain et al., 2018). Note that this assumption is much stronger than that in Theorem 4 which only requires that the variance at the optimum be bounded. With this bounded variance assumption, Pan et al. (2024) prove that their multi-stage SHB algorithm converges to the minimizer at an accelerated rate without any condition on the mini-batch size. This is in contrast with our result in Thcorem 4 which requires the mini-batch size to be large enough. This discrepancy is because of the different assumptions on the noise. In Fig. 2a, we use the same feasible linear system as in Theorem 3 and demonstrate that with a batch-size 1, the algorithm in Pan et al. (2024) can diverge. This is because the iterates do not lie on a compact set and ā can grow in an unbounded fashion for O(T) iterations (see Fig. 2b), demonstrating that the bounded variance assumption is problematic even for simple examples.
With this assumption, Pan et al. (2024) prove that their multi-stage algorithm converges at a rate of Ô (Tk exp(-T//<) + dā (for similar definition of suboptimality as in Theorem 4). The above upper- bound implics that their algorithm can only achieve a sublinear rate even when solving feasible linear systems with a large batch-size (Jain et al., 2018). In comparison, Algorithm 1 with a large batch-size can achieve an accelerated linear rate when solving feasible linear systems. From a theoretical perspective, the ò (K1/4 exp(-T/vZ) + bound in Theorem 4 is better in the bias term (by a factor of T) and hence requires fewer "warmup" iterations. It is also better in the variance term in that it does not incur a dimension dependence. While the bound established by Pan et al. (2024) holds for T> R(Vk), our analysis requires T 2 ((K//R) to achieve the convergence guarantce. This additional dependence on K is an artifact of our simplified proof that analyzes each stage independently. Specifically, we use the result from Wang el al. (2021) that introduces additional vK "warm-up" iterations. These additional vK iterations in each stage introduce the additional и dependence. To conclude, compared to Pan et al. (2024), we achieve better convergence guarantees with a simpler analysis under more realistic assumptions for larger iterations.
Finally, we note that if the variance is guaranteed to be bounded for some problems, the proposed algorithms can exploit this additional assumption and achieve rates comparable to Pan ct al. (2024) without a large batch-size requirement. Please refer to App. G for detailed analysis.
In Theorem 4, we observe that the batch-size threshold b* depends on a₁ 2- 0(/r). In order to understand the implications of this requirement, consider the case when T = i/n (for some ve > 0). In this
case, b" = n max
For practical problems, n is of the order of millions compared to T which
8

---

Published in Transactions on Machine Learning Research (03/2025)
10440
1052
10
200
norm
1944
logl011
1070
1653
130
1937
SGD
20000
40000
60000
80000
300000
6000
20000
terations
kerations
(a) Multi-stage SHB (Pan et al., 2024) with d 1 diverges exponentially fast in the first few thousands iterations.
(b) The variance 8² is increasing.
Figure 2: Divergence of Multi-stage SHB (Pan et al., 2024) with 6 b= on the synthetic example in Theorem with Æ 5000. We set wa = 100 and rum the algorithm in (Pan et al., 2024) with C = 2. We consider 5 independent runs, and plot. the average gradient norm [[V(( against the number of iterations. In Fig. 2b, plot the (log) variance log [E f(W)) f(u(w)))))) against the number of iterations. We observe that multi-stage SHB diverges and the variance 8² increases, showing that the bounded variance assumption in Pan et al. (2024) is problematic.
is in the thousands and hence v << 1. Furthermore, when n >> O(x2), b* predominantly determined by the condition number.
An alternative way to reason about the above result is to consider fixed batch-size b as input. In this case, the following corollary presents the accelerated convergence of multi-stage SHB but only for a range of feasible T.
Corollary 2. For L-smooth, strongly-convex quadratics with > 1, Algorithm with batch-size such
that b > 10" = nl a-T attains the same rate as in Theorem 4 for T E 33.2 max (4x, Ci where C = 25 and C₁ is defined in the proof of Theorem in App. F.
[n- Db
In(2)
3(m
Cx²
We have seen that a complicated algorithm can result in the optimal accelerated rate for range of T. Next, we design a simple-to-implemen algorithm that attains partially accelerated rates for all T.
4.4 Two-phase SHB
We design a two-phase SHB algorithm (Algorithm 2) that has a convergence guarantee for all T, but can only obtain partially accelerated rate with a dependence on x² for q E [2, 1]. Here q = corresponds to the accelerated rate of Section 4.1, while q corresponds to the non-accelerated rate of Section 3. Algorithm 2 consists of two phases in phase consisting of To iterations, it uses Eq. (1) with constant step-size and momentum parameter (according to Theorem 2); in phase consisting of Ti T - To iterations, it uses Eq. (2) with an exponentially decreasing πe sequence and corresponding A (according to Theorem 1). The relative length of the two phases is governed by c == To/r. In App. H, we analyze the convergence of Algorithm 2 with general and prove Theorem 11. For a specific setting when c = 1, we prove the following corollary.
Corollary 3. For I-mmooth, μ strongly-convex quadratics with x 4, Algorithu with batch-size such that b 2 b'
14 I and 2 results in rate of O (exp (- 2 } for all T.
We observe that Algorithm with a sub-optimal convergence rate of O (exp SGD and the non-accelerated SHB algorithm in Section 3. Compared to the accelerated SHB in Section 4.1,
(exx(-T7/< (-T/x²7) + o/VT), is faster than
9

---

Published in Transactions on Machine Learning Research (03/2025)
the two-phase algorithm converges to the minimizer (instead of the neighbourhood). However, even in the interpolation setting when a 0, the two-phase algorithm (without any knowledge of σ) can only attain a partially accelerated rate.
5 Experimental Evaluation
10P
10"4
horm
10-
My
10
M uy
10.7
10
10
1000
2000 3003 4000
5000
5000
7000
1000 2000 3000 4000
5003
6300
7000
10
1000 2000 3000 4003 5000 6000 7000 Beration
feration
heratton
(a)
1000 and
10
(b) 500 and
10-2
(c) K 200 and r 10-2
Canctens 1egg
Gradeen Nom
10
10-
10-3
My W/
30
10-
Taration
Raration
(d) K 1000 and r 10"
(e) K = 500 and r 10"
(1) A. = 200 and r = 10
Norm
M/
10
May 10
10"11
1000
2020
1003
4200
5050
1000
7000
1005 2000 3030 0000
1000
6200
7000
10
1400
2020
1000 4000
5000
6020
7000
(g) 1000 and r 10-6
(h) 500 and r 10-6
(i) K 200 and r 10-6
SHB
2P-SHB
Multi-SHB-CNST
Multi-SHB
SGD
Nesterov-EXP
Figure Comparing SHB, SHBti-SHB, Multi- CNST, 2P-SHB, SGD, Nesterov-EXP, for the squared loss on synthetic datasets with different x and noise r. Both SGD and SHB converge the neighborhood, but SHB attains accelerated rate. MuBti-SHB, Multi-SHB-CrSS and 2P-SHB result in smaller gradient norms and have similar convergence as Nesterov-EXP.
For our experimental evaluation", we consider minimizing strongly-convex quadratics. In particular, we generate random synthetic regression datasets with n 10000 and d = 20. For this, we generate a random w" vector and a random feature matrix X Ra We control the maximum and minimum eigenvalues of the resulting X matrix, thus controlling the L-smoothness and p-strong-convexity of the resulting quadratic problem. The measurements E R are generated according to the model: y = Xw" + 8 where ~ (O, rla) corresponds to Gaussian noise. We vary λ E {1000, 500, 200} and the magnitude of the These choices are motivated by Aybat et al. (2019). By controlling r, we can
noise r E {10-2, control the variance in the stochastic gradients. Using these synthetic datasets, we consider minimizing the unregularized linear regression loss: f(w) = 2
1 ||x w yl In this case, A = XTX, d 2 x and
A, = X,, d, = 2y X,
We compare the following methods: SHB with a constant step-size and momentum (set according to Theorem 2) with a = 1 (SHB), Multi-stage SHB (Algorithm 1) (Multi-SHB), Two-phase SHB (Algorithm 2) with e = 0.5 (2P-SHB), and use the following basclines SGD (SGD), SGD with Nesterov acceleration and exponentially decreasing step-sizes (Vaswani et al., 2022) (Nesterov-Ex0
Additionally, we consider a heuristic we refer The heuristic has the same
to as Multi-stage SHB with constant momentum parameter
structure as Algorithm 1, but the momentum parameter in each stage is fixed i.e. ,,, = (1 - 1/2xx)² We will
TThe code available at https://github.... www.isi.coo//ohhaeeeeeee httaaaaaaaaaaaaceaeeeetttee noise adaptive_snt
10

---

Published in Transactions on Machine Learning Research (03/2025)
see that this heuristic can result in better convergence than Multi-SHB. However, analyzing it theoretically is nontrivial. For each compared method, we use mini-batch size b = 0.9n to ensure that it is sufficiently large for SHB to achieve an accelerated rate for our choices of K. We note that using b = 0.9n on noisy regression problem has enough stochasticity to meaningfully compare optimization methods. We fix the total number of iterations 7 7000 and initialization ion wi = For each experiment, we consider 3 independent runs, and plot the average result. We will use the full gradient norm as the sub-optimality mcasure and plot against the number of iterations.
From Fig. 3, observe that: (i) both SGD and SHB converge to the neighborhood of the minimizer which depends on the noise r. However, SHB attains an accelerated rate, thus converging to the neighborhood faster. (ii) Multi-SHB, Multi-SHB-CNSS and 2P-SHB can better counteract the noise, and result in smaller gradient norm after reaching the neighborhood at an accelerated rate. (iii) The Multi-SHB-CNSS heuristic results in slightly better empirical performance than Multi-SHB when x. is relatively small. (iv) 2P-SHB results consistently better performance compared Multi-SHi (v) Across problems, the SHB variants have similar convergence as Nesterov-EXx
10a
Grotent Norm Dogi
10a 101
rl0g)
10'
10'
30-1
10"1
10-1
30
10-1
10-
10
10-40
10-11
10
250
500
750
1000
1250
1500
1750
2000
10-12
250
500
750
1000
1250
1500
1750 2000
10-13
250
500
750
1000
1250
1500
1750 2000
(a) K 2048
(b) K
= 1024
(c) k= 512
10
30*
Godden Norm (0g)
10-1
30-1
10"4
20-1
20-1
10-1
10-44
10-40
10-14
10-10
250
500
750
1900
1250 1500 3750 2000
10-10
250
200 750
1000
1250 1500 1750 2000
10-1
0
250
500
750
1000
1250
1500 1750 2000
Bleaaion
(d) / 256
(e) r
= 128
(1) K 64
10a
10a
30*
Ondert Norm bogi
10-1
30-1
10-2
10-4
10-4
10-4
10-3
20-7
10-
10
500
1000
1250 1500 2000
10
300
1000
1250
1500
1750 2000
20-11
500
1000
1250 1500 1750 2000
Beraiion
Benation
Ianation
(g)
K 32
(b) K
16
(i) 8
SHB-1.0
SHB-0.7
SHB-0.4
SHB-0.1
KAP
SHB-0.9
SHB-0.5
SHB-0.3
NON-ACC-SHB
SQRT-KAP
SHB-0.8
SHB-0.6
SHB-0.2
SGD
Comparison of SHB-C, NON- SCD and basclines KAP, SQRT-KAP for the squared loss on synthetic datasets with different k. For large K, SHB can converge in an accelerated rate if the batch-size is larger than the threshold b* The performances of SGD and NON-ACC- SHB are similar and significantly slower than SHB when & is large.
Figure 4:
Next, we consider solving synthetic feasible linear systems with different values of K, and examine the convergence of SHB with different batch-sizes. The data generation procedure is similar above, however, there is no Gaussian noise (s = 0) and hence interpolation is satisfied. In particular, the measurements y E R" are now generated according to the model: and batch-size = En for g (0.1,0.2, 0.3, 0.4, 0.5, 0.7, 0.8, 0.9, 1.0]. We fix the total number of iterations T = 2000. For each experiment, we consider 5 independent runs, and plot the average result. We will use the
Xw". We vary K E {8, 16, 32, 64, 128, 256, 512, 1024, 2048]
11

---

Published in Transactions on Machine Learning Research (03/2025)
full gradient norm the performance measure and plot it against the number of iterations. We compare the following methods: accelerated SHB with constant step-size and momentum (set according to Theorem 2) with a = 1 and varying batch-size & (SHB-E), non-accelerated SHB with a constant step-size and momentum (set according to Thcorem 1) and a fixed batch-size = 0.3n ANON-ACC-SHB), SGD with constant step-size and a fixed batch-size b = 0.3n (SGD). We also add the following baselines to understand the dependence of x: line proportional exp (=x² (KAP) and line proportional to exp (SQRT-KAP). The baselines are calculated by multiplying the initial gradient norm with the corresponding exponential term calculated at each iteration.
From Fig. 4, we observe that (i) when i large, using SHB with smaller batch-sizes can result in divergence, (ii) SHB can only attain acceleration when the batch-size is larger than some x-dependent threshold, and the extent of acceleration depends on the batch-size, (iit) across problems, the performance of SGD and NON-ACC- SHB is similar and slower than SHB when is large, (iv) the larger batch-size, SHB converges at rate similar to the SQRT-KAP bascline, (v) across problems, SGD converges at rate similar to the KAP baseline. This verifies our theoretical results in Sections 4.1 and 4.2.
Finally, in App. 1.2, we consider the algorithm proposed in Pan et al. (2024). We observe that with a sufficiently large batch-size, the method converges and has similar performance to the proposed SHB variants.
6
Conclusion
For the general smooth, strongly-conver sctting, developed novel variant of SHB that uses exponentially decreasing step-sizes and achieves noise-adaptive non-accelerated linear convergence for any mini-bate size. This rate matches that of SGD and is the best achievable rate for SHB in this setting (given the negative results in Goujaud et al. (2023)). For strongly-convex quadratics, we demonstrated that SHB can achieve accelerated linear convergence if its mini-batch size is above a certain results imply that for strongly-convex quadratics where n >> O(*2), SHB (and its multi-stage and two-phase variants) with a nearly constant (independent of n) mini-batch size can be provably better than SGD, thus quantifying the theoretical benefit. of SHB. In the future, we aim to close the gap between the upper and lower-bounds on the mini-batch size required for SHB to attain an accelerated rate. Furthermore, we aim to improve our lower-bound and characterize the behaviour of SHB with any step-size and momentum parameter. On the more practical side, we hope to develop SHB variants that can attain an accelerated rate when the batch-size is large, and automatically default to non-accelerated rates for smaller batch-sizes.
threshold. Our
Acknowledgements
We would like to thank Saurabh Mishra and Duy Anh Nguyen for helpful discussions and feedback the paper. This research partially supported by the Natural Sciences and Engineering Research Council of Canada (NSERC) Discovery Grant RGPIN-2022-048816
References
Necdet Serhat Aybat, Alireza Fallah, Mert Gurbuzbalaban, and Asuman Ozdaglar. universally optimal multistage accelerated stochastic gradient method. Advances in neural information processing systems, 32, 2019. (cited 7, 10)
Mikhail Belkin, Alexander Rakhlin, and Alexandre B Tsybakov. Does data interpolation contradict statistical optimality? In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 1611-18 1619. PMLR, 2019. (cited on 2, 4)
Raghu Bollapragada, Tyler Chen, and Rachel Ward. On the fast convergence minibatch heavy ball momentum. Xixio preprint 220iv:2066.00 07553, 2023. (cited on 2, 3, 6, 7)
Chih- Chih-Chung Chang and Chih-Jen Lin. Libsvm: a library for support vector machines. ACM transactions on intelligent systems and technology (TIST), 2(3):1 27, 2011. (cited 48)
12

---

Published in Transactions on Machine Learning Research (03/2025)
Aaron Defazio. Momentum via primal averaging: theoretical insights and learning rate schedules for non-convex optimization. arXiv preprinl orXin:2010.000 00000, 2020. (cited on 2)
Euhanna Ghadimi, Hamid Reza Reyzmahdavian, and Mikael Johansson. Global convergence of the heavy-ball method for convex optimization. In 2015 Europeon control conference (ECC), pp. 310 315. IEEE, 2015. (cited on 1, 2, 4)
Baptiste Goujaud, Adrien Taylor, and Aymeric Dieuleveut. Provable non-accelerations of the heavy-ball method, 2023. (cited on 1, 5, 12)
Robert Mansel Gower, Nicolas Loizou, Xun Qian, Alibek Sailanbayev, Egor Shulgin, and Peter Richtárik. Sgd: General analysis and improved rates. In Internalional conference on machine learning, pp. 5200-5209 PMLR, 2019. (cited on 29)
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, pp. 545 604. PMLR, 2018. (cited on 6, 8)
Rahul Kidambi, Pranceth Netrapalli, Prateck Jain, and Sham Kakade. On the insufficiency existing momentum schemes for stochastic optimization. In 2018 Information Theory and Applications Workshop (ITA), pp. 9. IEEE, 2018. (cited on 1, 2,
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6880 69800, 2014. (cited on 2)
Kiwon Lee, Andrew N. Cheng, Courtney Paquette, and Elliot Paquette. Trajectory of mini-batch momentum: Batch size saturation and convergence in high dimensions, 2022. URL
https://arxiny httpp::///aiivoorr/abb//////
2220///b./22....//// 010099. (cited on 2, 3, 6)
Xiaoyu Li, Zhenxun Zhuang, and Francesco Orabona. A second look at exponential and cosine step sizes: Simplicity, adaptivity, and performance. In International Conference on Machine Learning, pp. 6553 6564. PMLR, 2021. (cited on 2, 4)
Tengyuan Liang and Alexander Rakhlin. Just interpolate: Kernel "ridgeless" regression can generalize. 2020. (cited on 2, 4)
Yanli Liu, Yuan Gao, and Wotao Yin. An improved analysis of stochastic gradient descent with momentum. Advances in Neural Information Processing Systems, 03:18261-18271 2020. (cited 2)
Sharon Lohr. Sampling: design and analysis. CRC press, 2021. (cited on 25, 43)
Nicolas Loizon and Peter Richtárik. Momentum and stochastic momentum for stochastic gradient, newton, proximal point and subspace descent methods. Computational Optimization and Applications, 77(3): 653-710, 2020. (cited on 2, 4, 6)
Nicolas Loizon, Sharan Vaswani, Issam Laradji, and Simon Lacoste-Julien. Stochastic polyak step-size for sgd: An adaptive learning rate for fast convergence, 2021. (cited on 4)
Siyuan Ma, Raef Bassily, and Mikhail Belkin. The power of interpolation: Understanding the effectiveness of SGD in modern pee-maaraetrrized learning. In Proceedings of the 35th International Conference on Machine Learning, ICML, 2018. (cited on 2, 4)
Aaron Meurer, Christopher P. Smith, Mateusz Paprocki, Ondfej Čertík, Sergey B. Kirpichev, Matthew Rocklin, AMIT Kumar, Sergiu Ivanov, Jason K. Moore, Sartaj Singh, Thilina Rathnayake, Scan Vig, Brian E. Granger, Richard P. Muller, Francesco Bonazzi, Harsh Gupta, Shivam Vats, Fredrik Johansson, Fabian Pedregosa, Matthew Curry, Andy R. Terrel, Stěpán Roučka, Ashutosh Saboo, Isuru Fernando, Sumith Kulal, Robert Cimrman, and Anthony Scopatz. Sympy: symbolic computing in python. PeerJ Computer Science, 3:e103, January 2017. ISSN 2376-5992. doi: N0.7717/peerj-ca 10.3. URL https://dbi.org/// /tt///110.7777711 177117/peerrcca 103. (cited on 7, 33, 35)
13

---

Published in Transactions on Machine Learning Research (03/2025)
PHUONG HA Nguyen, Lam Nguyen, and Marten van Dijk. Tight dimension independent lower bound the expected convergence rate for diminishing step sizes in sgd. Advances in Neural Information Processing Systems, 32, 2019. (cited on 5)
Rui Pan, Yuxing Liu, Xiaoyu Wang, and Tong Zhang. Accelerated convergence of stochastic heavy ball method under anisotropic gradient noise, 2024. URL https://arxiv, www./bs///////////o http:s//aarrd..or//aa///221.1 14567. (cited on 3, 8, 9, 12, 48, 49)
Courtncy Paquette and Elliot Paquette. Dynamics of stochastic momentum methods on large-scale, quadratic models. Advances in Newral Information Processing Systems, 34:9229 00-9299-9240, 2021. (cited on 2)
Boris T Polyak. Some methods of speeding up the convergence of iteration methods. Ussr compulational
mathematics and mathematical physics, 4(5):1-17, 1964. (cited on 1)
Othmane Sebbouh, Robert Mansel Gower, and Aaron Defazio. On the convergence of the stochastic heavy
ball method. ArXiv, abs/2006.0786; 2020. URL
https://apis wwwanticscholeatee.......o.oeoee htapap/..pp...mmmm.ccccchc..... wwwp:/orooo/oooooaaaeeaoo 219688738. (cited on 2, 4, 16)
Sharan Vaswani, Francis Bach, and Mark Schmidt. Fast and faster convergence of sgd for wert-parameterized models and accelerated perceptron. In The 22nd International Conference on Arlificial Intelligence and Statistics, pp. 1195-1204 PMLR, 2019. (cited on 2, 4)
Sharan Vaswani, Issam Laradji, Frederik Kunstner, Si Yi Meng, Mark Schmidt, and Simon Lacoste-Julien, Adaptíve gradient methods converge faster with con-pparameterization (but you should do lino-search), 2021. (cited on 48)
Sharan Vaswani, Benjamin Dubois-Taine, and Reza Babanezhad. Towards noise-adaptive, problem-adaptivg (accelerated) stochastic gradient descent. In International Conference on Machine Learning, pp. 22015-22059. PMTRR, 2022. (cited 4, 6, 10, 19, 20, 22, 48)
Kun Wang, Chi-Heng Lin, and Jacob D Abernethy. modular analysis of provable acceleration via polyak's momentum: Training a wide relu network and a deep linear network. In International Conference on Machine Learning, pp. 10816-10827. PMLR, 2021. (cited on 1, 5, 6, 8, 29, 44)
Jun-Kun Wang, Chi-Heng Lin, Andre Wibisono, and Bin Hu. Provable acceleration of heavy ball beyond quadratics for a class of polyak-lojasiewtcz functions when the non-convexity is averaged-out, In International Conference on Machine Learning, pp. 22839 22864. PMLR, 2022. (cited on 1)
Runzhe Wang, Sadhika Malladi, Tianhao Wang, Kaileng Lyu, and Zhiyua Li. The marginal value of momentum for small learning rate sgd. arXiv preprint 1rX1::2307.15196 2023. (cited on 2)
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes. Xiv preprint ar arXic:194..00 00962, 2019. (cited on 2)
14

---

Published in Transactions on Machine Learning Research (03/2025)
Supplementary Material
Organization of the Appendix
A Definitions
Proofs for non-accelerated rates
C Proofs for non-accelerated rates with misestimation
Proofs for upper bound SHB
E Proofs for lower bound SHB
F Proofs for multi-stage SHB
Proofs for SHB with bounded noise assumption
H Proofs for two-phase SHB
T Additional experiments
A Definitions
Our main assumptions are that each individual function fi is differentiable, has a finite minimum fi, and is L-smooth, meaning that for all and w,
fuv f(w)) + (Vfi(w) w) + L
(Individual Smoothness) which also implics that f is L-smooth. A consequence of smoothness is the following bound the norm of the stochastic gradients,
|1 w||
[w)( < 2L(f(()) ff). We also assume that each fi is convex, meaning that for all and w,
fi()) fi(w) (Vf(w), v).
(Convexity)
We will also assume that f is μ strongly-convex, meaning that for all v and w,
f(v) 2 f(w) + (Vf(w), e- w) + J || w||?
(Strong Convexity)
15

---

Published in Transactions on Machine Learning Research (03/2025)
B Proofs for non-accelerated rates
We will require (Sebbouh et al., 2020, Theorem H.I). We include its proof for completeness.
Theorem 5. For L-smooth, 1 strongly-convex functions, suppose (Mx ). is decrcasing sequence such
that 1/o = r and < Th 22. Define AE == 1-2n1. (1 (1 -maH)))) A := || - =- + X (w, 1
E := A& + 2m (f - )), ax:= I+AE+1' Bx. == XE I-nel 0² == E.ff((w*) - ft 0. Then SHB 1+
Eq. (1) converges as
B[ S 51 - B[E] + 2T
(3)
where (= V (n-1)6
Proof. We will first expand and bound the term Ac+1,
Ak+1= =
||u+11 - w² + > - w²))l
1 w=" 17--ap -ax V fix(wu)) + Ba(wx= wk-1) + AR+1 a -axaaaafleee fiu(wx) + =x((ux - wk-1)ll?
(SHB step)
1w'x- w² ax + AR++ (w)) + Bx(1 + Ax+ () -x- || IWE w² T/E V (e)(w))) + (1 - mp) (uxk wk- 1)11
(definition of as and 3k)
|lwk - w² + A1 (w - w1-1) 1) - 7x lelduee - wk-1) + [aa((
As + n (wn Wk-1) 1) + Viu ((((w))) 112
-2TT w² + (wx - wk-1), (xx(wx wk-1) + V (wx))
Av + nh f (wx))1 + n2p Ar | - wk-21l 11?
Snan
+ 2n - wk-1, (ww)) - 27//112 (10 - ,,uu, wk-1)
2mx - -w", -w", ffu((w)) - 21 wk-1, f(t((w))) 27 ||u.- -wh-- -w 11? we + Ti V Fin (ww)l ue (w> - Wk-1, v.f((w))) - 2m WK-1, fin(w)))
< Ax mh (X2 ||wee wn-1ll 11| + 22x - -." -
-w--i) -2 w', V Jin (wx))
(by L-smoothness of fak)
Add BK-1 27 1 (f(wx)- - f*) on both sides,
Ak+1 + BE+1 AA - 7//14 (22 Ilwx - w/-1ll 111 + 2dx (w0 - wu, .uk - -k wk-1i) + LLnllfuwwr) - fill + 2n / (wi - Wk- v(u(wu))
27k(wk w", V.
2Tp - Wk-1, V.cik (w))) + 277 (f(wk) - -²)
S Ax - mmp (>ā Il - wu 1| 112 + 2dx (wr - w .ux - Wk- 13) - -211e - w², ((x()) + 2 - finl - 27 (1 - mu/1 1. faae(w)))
+ 2 A6+1 (f(w) - -
Taking expectation w.r.t bt. (uk) - fhh = [fsx(w)s) - fa((w")] + +Ifu(u")= - fii then
B[Ast1 + Bb+1] < [[AN] - E ma (X 1wwe - wx-1l 112 + 2Nx (wi - w",,, wh wk w )
2n Bl - n²,, ([(w)))) + 2 + 2T [ |
-
- (1 Wk-1, [(( + 2 Ef - f*}
(Using Lemma 2)
16

---

Published in Transactions on Machine Learning Research (03/2025)
Since f is strongly-convex, -2mx (wy w", vf(ww)) - |lw. w01 -2=12-2nnn ( - -f*), then
B E[Ex+1] S [[[As]] - mANlEllwu w011 + XR 1ww wk- 11 + 2dx (ur - w² WR. Wk-
Au
+ 2 1200² + 2f
fr] 27/ 2x 7 (Kux 1 v((w)))) (nnBls(wH) [ + 27x+ AR+ [E[[ (Wx) -
- (1 -1-mia [[A]) + 2L 12 + 2f (wx)) f*1
-2m (1 11,1) 1wxx w 1 [[((w))]]
22n,Blf(wx)) [4] + 27k +1 A++ 1 [l(((()) f
By convexity, - (Vf(w*)), wh wk-1) f(wk-1) - fw)) [f( - -f(())) ff*, then
B[E&+1] (1 - [Ac))B[Ad + 2Tx 2TT 14 2 ((x)) )
S [f
-2m (W) - 20 (ww) - - 27 [[[(xx - fr]
< - AA + Ao [ + 2 + Al 5"]
- 21 - (f f - 2f - f41 + 27 16 - f*1
$ (1 MEN) + 2Lxxç² o
+ +2E[f(wR)) - f*1 (2Ini mTAAe(1 M&p1) - -7p- + 7
(Theorem 5 first part)
We want show that 2T - T76: Ax m=)1s) - TJK + 1 5. which is equivalent to
Mh+1 16+1 S Th 2LT + > nap)).
RHS
77 (1 2LTR + Ax - 71H14)
mp (1 2ImA) 7)>x((1 (1 mEpL)
m1 2LTH) 1 1 1-2 (1 (1 m1a)))) (1 mal)
(definition of Ax.)
1
1 2nL
MkH + 1- 27L (1 (1 mat1))ttt
= m1 -2L1))
1-2nL 2n2
(1 (1- (1 (1mn)))ttt)) + 22 (1)
(since 2 m))
11
20
1-2nL
(1 - (1 m111))
2nL
(1 - (1 - 11111))6t1)
(since 7p 2 7k+1)
P
1 LHS.
Hence,
E[Ex+1] $ $(1 B(E)) + 2Txc2?n000
17

---

Published in Transactions on Machine Learning Research (03/2025)
Theorem 1. For L-smooth, strongly-convex functions, SHB (Eq. (2)) with 2 1, v= II -F}) 1/4
T k+1 7h u 1Ju and d 1-29L (1 (1 m1 converges mak
E ||T--- w¹1l C1 11w0- w01l exp (-m4m((77)) C. Co
7
o²
T
where <= v (n-1)6
and Ca, Cs are polynomial x and poly-logarithmic in T.
Proof. From the result of Theorem 5 we have
E < (1 E + n
Unrolling the recursion starting from wo and using the exponential step-sizes "
T
(1
p**
T
T
2k (1 (1-
E[Er] v E
II
+
2T
02
II
7
4L
4L
1=1
k=1+1
T
T
< 1
w w"|/eep
11
+
2.x 20²
2k
exp
4L
4L
#=1
k=1
k+++
=D
(No and 1 exp(-x:))
Using Lemma 3 to bound C then the first term can be bounded as
|| Ww w²1 exp
T -
< Iwo
w w0112
c2 exp (- (-4r T
4x In(T/+))
4L
where x = and c2 exp In))
Using Lemma to upper-bound D, we have D S
then the second term can be bounded as
GALo2c2522 (In(T/+))2 2T
2f
e²
Hence
E(Er] < |1100 =='l 02 exp (-
T
64Lo² c2 (In(T/+))2
4x
6²
22
By Lemma 1, then
E [||WT-I- - 201 < C2 Ilw0 - w"1l2exp exp
(- 1n((/s)
T γ
o²
(04L(210(11)) (In(T)))
02
T
6272
Cr.
Let Ca(ks T) = exp
2n
and Cs T) = ((L())"(( then
expl 2x
ex
E -luu-11 w²1 Ca
w011 exp (-x ( + T C.Cs
T
k
18

---

Published in Transactions on Machine Learning Research (03/2025)
B.1 Helper Lemmas
Lemma 1. For 8T: == | w² + (wT wr-1)l/
- (ww)), Er 2 cr 1|W--1 -w"l
where CL
4(1-2)
n22 [1 exp (-277)
Proof.
B[ = [[AT] + B[BT] 2 B[Br] 2 BI f*]
Hence, we want to lower-bound ÀT ar and we do this next
1 1 ? Ar 1T
117 T+1
(Using the definition of nip and A²)
1
2T.
1 Y
(-T'7" MY
exp
(Since 1 - S exp((-1)
2L
1 1-7
[ exp
(- []7]]
(Since y= (7)
Putting everything together, and using strong-convexity of f
B[Er] 2
4(1 7))
[
exp
117
E | w011
112
2L
:CL
We restate (Vaswani et al., 2022, Lemma 2, Lemma 5, and Lemma 6) that we used in our proof. Lemma 2. If
a² E[f(w") frl.
and each function fi is y strongly-conuer and L-smooth, then
oB = Es(sa(w") - -fāl S k
n2 0²
n
Lemma 3. For 7 =7) 1/T
(TT
T
TT
2r
A =
In (T/7)
(n(T/>)
t=1
Lemma 4. For 7 (7) 1/T
and any N > 0, with C2 = exp
(1 (
4k (n==2(1n(7/-))² e2pT
24 exp
?
(=++1
19

---

Published in Transactions on Machine Learning Research (03/2025)
C Proofs for non-accelerated rates with misestimation
A practical advantage of using Eq. (2) with exponential step-sizes robustiesss misspecification L and μ. Specifically, in App. C.1, we analyze the convergence of SHB (Eq. (2)) when using an estimate L (rather than the true smoothness constant). In App. C.2, analyze the convergence of SHB when using estimate á for the strong-convexity parameter.
C.1 L misestimation
Without loss of generality, we assume that the estimate L off by multiplicative factor i.c. î. -= for 22 VL O. We note that is deterministic estimate of L. Here ez quantifies the estimation error with corresponding to an exact estimation of L. In practice, it is typically possible to obtain lower-bounds on the smoothness constant. Hence, the VL > 1 regime is of practical interest.
Similar to the dependence of SGD on smoothness mis-estimation obtained by Vaswani et al. (2022): Theorem 6 shows that with any mis-estimation on L we can still recover the convergence rale of to the minimizer w". Specifically, Theorem 6 demonstrates convergence rate of max[ 1o +Ay max/
o² O exp (-²)+ T
The first two terms in Theorem 6 are similar to those in Theorem 1. For VL I L, the third term is zero and the rate matches that in Theorem 1 upto a constant that depends on Vt. For Vr > 1, SHB initially diverges for kn iterations, but the exponential step-size decay ensures that the algorithm eventually converges to the minimizer. The initial divergence and the resulting slowdown the rate is proportional to vz, Finally, we note that Vaswani et al. (2022) demonstrate similar robustness for SGD with exponential step-sizes, while also proving the necessity of the slowdown in the convergence.
0 exp min(22.1]7) x
Theorem 6. Under the same settings Theorem 1, SHB (Eq. (2)) with the estimated i = _ results in the following convergence,
27.
E |UT-1 w'
1w w¹12
2
min(vL, 13T
CL
exp
2x
In(T/+)
02 32L 22 (1((/T))
( maxx
{1, v
In In(2/+)0²)
x
Cr
e2T2T
4L
(n(xx(,,n(ur.))
(o2 2A, -
VL
where 02 = eex
(m(*/-)) ) hi = T In(ve)
#z In(" and A, max B[s(w,, - - and
Zx In(
4(1-2) CL
exp
4
Proof. Suppose we estimate L to be L. Now redefine
1
7/k
2 7k
2nL (1 - (11-(1 mx/1)*))
Mkl
Å. Wk w² + Sx(wEe w/- -/- 1|(
В, dn (f - f(w"))
& A Br
20

---

Published in Transactions on Machine Learning Research (03/2025)
Follow the proof of Theorem 5 until Theorem 5 first part step with the new definition,
Blex+1] 5 - + 2
21Il((w) fr] (2L - nedu(1 Th - mh + TR+12U+11 (4)
G can be bound as
G 2L nk (1 11(1)) 776 + T7 A
n/ (2L1j 1) - An (1 m1111) - 111 Ru+t
m2(22m - 1) + mh (1 - 2Lm) -1211 (1- (1 -mss))+t) m++1R++1
(definition of )x)
2m (Lape in) 1 - 2012 (1 (1- -1)+114) + netide+1
(m++1 S m)))
11
(TT Ln)
-M++1 + 114114+1
2nk(+ in)
Hence Eq. (4) can be written
E[ 11 S (1 Mn))E(Ea) + 2Lnc3 4E[f(w)) f'Ima(Ln Lm)
First case pz $ 1 then Lnk- in and will recover the proof of Theorem with slight difference including VL-
[ < l1w0 w01l =='1222epp exp (- VLT 7
32L oo2ck2 (In(T/+))
2n In(T/r)
6²
22
Second if VL > 1
Let Ko = T [n(ve] In(]/>)
then for k < kn regime, Libe 1n1>0 >
E S E( + 2n00211002 + [E[f(wx)= f ('ImaLLn im)
Let A, max Bsf f*] and observe that Laj- - in S Link v--1 then n
[( S m..)) + 2Tx 2000 + +4Tn/ A
vs. -
VL VI. -
= (1 Wul
[(EA] 2L(x 0² + 20,
2
mi
2L
VL
Since vL >> 1
E(E++1] (1
P
1)(Eat can
2L
Unrolling the recursion for the first ko iterations we get
1--1
ho
kg
E|
Eol
II
(1
P ++))
+ 65
2T. 75)
2T.
k=1
k=++1
Bounding the first term using Lemma
(_ 14
21, x
< exp
2T,
21

---

Published in Transactions on Machine Learning Research (03/2025)
Bounding the second term using Lemma 4 similar to (Vaswani et al., 2022, Section C3)
kg
ho-1
(1
M 70)
16x² Ko In(T/r)²
7k
2T. 2T,
S exp
2x(1 - 7)
2202
7
i=++1
Put everything together,
B[ÊAJ) 11000 =0'l1 exp (-H 11 7 - -0
gko
16k²
Ko In(7/>)2
-=xppexp
(22 2L -
+ c5 exp
2n(1 - 7)
22n7
T
Now consider the regime h 2 ko where LTpE - Ln S 0
E[E++1] 5
11
E + 2Txc 0² V2
2L
4L
02,02
<1 2L
+
2L T
Unrolling the recursion from k = ko to T
T
1 122220"
T Σ k=ko
T
Blêr|
< E [kol II
(1
1
7k) +
7x
II (1 T,)
2L
2L.
k=
1=k+1
Bounding the first term using Lemma 3,
II
T
(1 14 -² S exp
g - - T+11
2T.
1 1-9
k=
Bounding the second term using Lemma 4 similar to (Vaswani et al., 2022, Section C3)
Σ II (1
1 11 s exp 2L
T+1
16k² (T ku + (TT+))
2k((1
-1)
62g2
72
Hence, put everything together
EiEr| ElE,, exp (- 14 -
12nç202
T+1
16k²
(T- ko (T/72
(- 2L 1-
exp
2L
21r(1 - 7)
e2g2
T
Combining the bounds for two regimes
70 - T+1
T+1
(1wo 1o00-w111- exp
Je 7- A0
2xx1 -
16k² ko In(2/7)2
E S exp
w0
+ cs exp
2L
1-7
2L 1 1-7
7)
22g2
T2
12
22,
7T41
16x² (T' Ko + 1) In(T/²)2
21.
exp
2x(1 - 7)
e2ar2
T2
11 7 - 1 T+1
7 T+1
16x² c2g2
Ko In (T/++))
=1lw0 - 11 exp
+ C3 exp
2L 1-9-
21x(1 - 7)
T
16x²
(T'
Ko + 1) (T/T22
2T.
exp
2x(1 (1)
22n²
T2
Using Lemma 3 to bound the first term and noting that
1 S
T+1
27 let C2 =exp ( (2(27//))
In(T/>)'
E(Er] $ Ilw0 21 exp (- T 7
1602x3
Ko ln(T)/22
160gk² (T - ko + 1) (T/7)3
(-
2x In In(T/t)
+ Gs
e2n²
T2
21
e2g2
T
22

---

Published in Transactions on Machine Learning Research (03/2025)
Substitute the value of cs and ko we have
B[GT] < l1w0 =='12e exp
2k (-2x In(T/))
T
vi,n5200 802n2 In(T/+)2
+
LT
62G2
+ 32 2 +
-1
L c2n2 In(vr.) In(//)
VL
T
6222
Combining the statements from V2 S and UL > 1 gives us
min, 13T
[[GT] S 1ww0 w
02
exp
2x.
In(T/+)
32Tc2h (T/+r)
22
max
1,
4L
Vi
In( In(T/s)n(² (n, [K(20² 2 vI.
VL )
The next step is to remove the L from the LHS, and obtain a better measure of sub-optimality. By Lemma 1,
4(1-7)
E[ 2
11-88 exp
|10--1 - w=1
1
2L
Note that Cy >> O is constant w.r.t T. Hence,
E |urr-1 - 10 ww w=1 02 exp
min, 1)T 7
CL 32L IT(7/)
2k
In(T/+)
C2
1, In(T/>))<( In(z)) (nCO0² + +20, V 1
v2
max
CL
e2n2T
4L
VL
C.2 misestimation
Next, we analyze the effect misspecifying μi the strong-conveeiity parameter. We assume we have access to an estimate á μ where vn is the degree of misspacificatinn. We note that fi is a deterministic estimate of p. We only consider the case where we underestimate μ, and hence V/, < This the typical case in practice for example, while optimizing regularized convex loss functions in supervised learning, á is set to the regularization strength, and thus underestimates the true strong-convexity parameter.
Theorem 7 below demonstrates an O exp (-"a7)
convergence to the minimizer. Hence, SHB with T
underestimate of the strong-convexitt results in slower convergence to the minimizer, with the slowdown again depending the amount of misspecification.
Theorem 7. Under the same settings as Theorem 1, SHB (Eq. (2)) with the estimated Upt for Up results in the following convergence,
E f|uT-1 - < l/wb w0 11 C2
exp (-
T
2x 1n(//))
32T (In(T/>))2 2302720
T o²
Cμ
where 0g exp
2r [/{)/-])
and C₂
461 4(1-2)
- exp
2x
21
Proof. Suppose we estimate a be a. Now redefine
d
1-2nL (1 - (1 m)ff)))) Ag
| w² * Su wx-1)
nn
23

---

Published in Transactions on Machine Learning Research (03/2025)
B, = 2f f(w")((w)))))))))>)>>))))) E = Ar + Bx
Follow Theorem first part steps with the new definition, the only difference was at the step where we use strongly-conver on f for -2m (wa '², Vf((w)) S - ll w01 2m fnwl)(wx) f*].
[[ - (1 7k A.) [ 2 m= (1) l|:0x w0
= (1 - The P B[ + 2 + me (f - (1) 1l w¹11
< (1 Thk B[& 2 + MbHlU; -1) [f((wk)) f
(since f is strongly-convex)
= - ())))) 2Ix 1/,00 + 2M (up ()((( - -s"
Since 14 < then (p - (w)) - f - 0 s0
E (1 - 1 + 2 2 n2
Hence, following the same proof as Theorem
PTT })
32LÇ2 "2LC202c2n2 (In(T/-))
BiEr| S 1ww w w"|l²00 exp
exp
+
2x In(#//))
22T
By Lemma 1,
E(ET] 4(1 7)
1 - exp (- (-
Yull?
| - w0
122112
2L
Note that C₁ > 0 is constant w.r.t T. Hence,
(- UnT
32I c2x (In(T/+)) v2
|wr-1- - w²1 $ 1100 - w w011 C2
cxp
2x In(T/T)
+
T
D Proofs for upper bound SHB
Lemma 5. For L-smooth and μ strongly-conver quadrolics, SHB (Eq. (1)) with ag = a i and a S 1,
34 3 (1 1 /app) 2 batch-size size b satisfies the following recurrence relation,
aCo x ((b)
|[| pf Aoll 2aCo {(b)
Σ k=O
T---
B ||Al
T-i-k
k=0
where AІ: = Wk-1 - w²
w²
Co v 3 / K ((b) = v
n-b (n-
and
1
ya 2
Proof. With the definition of SHB (1), [k((w) is the mini-batch gradient at iteration k; then, for quadratics,
Wk+t1 w= W/k
(1 + I) aA
-BIa
Wk w
Vf(wk)
V 7fsk(wa) (wk)
+
la
0
Wik 1 w0*
Apt1
Ak+1 H.r + adgg
24

---

Published in Transactions on Machine Learning Research (03/2025)
Recursing from k 0 to T taking norm and expectation w.r.t to the randomness in all iterations.
E| Arll 2 |H Boll + al
Σ HT- 6x
Using Theorem and Corollary 6, for any vector v, || oll S Co ot ||a|| where p = vB. Hence,
E| E| S Co p² | + Coa
Σ Ee- ---1-= B ||ōx|| k=
(a +
L
In order to simplify 8u, we will the result from Lemma 7 and Lohr (2021),
- b (n 1)b E. [f(ww) (f(w)))
6x E (fww) [fa(w))]]]
(Sampling with replacement where b the batch-size and the total number of examples)
n b
E, ((xx) - Vf(w") - V.((wA)) - v
(Vf(w*)= 0)
1)b
< 3 n b
E. Vf(wk) vf(("))] E V - + }|(
(n b
((a b+ +)+ S 3/0 + b² + co1)
<3 3 n b
L lw. w w² L2 l - w E IV
(n 1)b
(Using the L smoothness of f and fi)
$ 3 n1
21.2 | - 2²11++x' x2]
+
(n b)b
(wh is independent of the randomness and by definition x2 E, (1()((l)
3 3
n -
2T.22 Il[wE w²112 122 + || -
-w"l""""" w²113] + x2]
(1 w01 2 0)
(n b)b
EM 19 < 3
n1
222 ||all + xx]
(Definition of Ax)
(n 1)
n b (n b)b
E <
222 +
-(()
(Taking square-roots, using Jensen's inequality on the LHS and V S Vo vi on the RHS)
= [|] & VbL | + <(b)x
Putting everything together,
aCo ((b)
T-1
B|||r/|| P A ((b)
Σ k=0
p T |
T
+
Σ
p
T 1-k
km
25

---

Published in Transactions on Machine Learning Research (03/2025)
Theorem 2. For I-smooth, l strongly-convex quadratios, SHB (Eq. (1)) with og = a = i for S
Bk = 3 = (1 Vaji 2 batch-size b s.t. bz b" == n max
1+ Cot 1+ (m
converges as:
E F|wT - w 6/2VK Va
V 2/K
(3
+ 12
min { 1,
S v
exp
max
(4'
l1w0 w0ll
1
where x=
E [V ; <= /
n-b (ল- and C 3526.
Proof. Using Lemma 5, we have that,
T-
aCo5x
T-1
E B| S Cop p" | v2aCo <
Σ p
T --
p
L
We use induction prove that for all T 2 1,
E Arll s 2Co [o+ v[va] 1 A + 20o 5ax
L(1 -
where p+ Vsva <1
Base case: By Theorem 8, Co 2 1 hence ||oم|| s 2C% | | + 2C
T((-p)
Inductive hypothesis: For all & 10. T 1}, || S 2Co 10+ VVvVa"/la/ |1 +
2C. as
I(1-M)
Inductive step: Using the above incquality,
aCo
T-
E |rTl SCop" p² | V20Co
T Axll
T-1-8
L
k==0
<oo 1 / v A + V20Co
1- |
+
(Since S, a > 0)
L
k==0
4c V5 v A + v2aCo5
Σ p 2Co [p+ 1p+v Vival
2Co 2C as X
L(1 p) ]
||oo|| +
aCo 1
1Co611-11 n²
(Sum of geometric series and using the inductive hypothesis)
L 1-p
2/20CH
p VLva A
Coo p+ + V Vaj" | +
|
P
2v2a208 <²x
T--
aCo CX
1-p" pT 1 1-p p
+
pL(1 p)
L
k=
First, need to prove that 2v2ach'
PT
Σt=
| S Co 1pt vSvāl" 1Aol-
p
P
(Pv/sva
1
2V2aCh5 OT
p p+ p+V5va)
2V2aC35 OT
A
|
P
P
ptviva
(Sum of geometric scries)
26

---

Published in Transactions on Machine Learning Research (03/2025)
<2V2V00 vs (o+ v(va)*
|1ooll
Hence, we require that,
2V2 V S Co = s 8C
Hence it suffices to choose s.t.
S6 322
(Since Co 3/
S 3223k
11 S 3525
(Using the definition of ()
n - 1)b
n +
3 K
Since the batch-size b satisfies the condition that: is satisfied, and Ç 3225K"
for C= 15552 3526, the above requirement
1+ n-l Cx
Next, we need to show D = 2V2a"
pL(1-p) 0²
#1 EE (+)²
+ aCo5x 1-p² 1-p <
25 I(1-p)
C/ S²x TT
T-1 Σ
aCu5x 1 PT
D =
pL(1 P)
L 1-p P
K=
(1) T
C2 52x
-1
aCo5x 1- 1-p" PT
PT
(Sum of geometric series)
PL(1 P)
(1) -1
1-p p
2v2a200 53
1- p" p
aCoSx
PL(1 p)
1- pT
L(1-p)
2/ (3x
aCoSx
L(1 p))
L(1-p)
Since we want D
2nCo SX we require that 1(1-p)
2 2/2a² 52x
aCo Sx
<
L(1 2)2
L(1 p)
2v2Cuaa <
1-p
Ensuring this imposes an additional constraint on Ç. We require ( such that,
S 1-p -
(Since p=
Va 2
2/2Coa
Vavk
Hence suffices to choose such that,
5S 12v2k
(Since Co s 3
Since the condition the batch-size ensures that 5
1_ this condition is satisfied. Hence, 3222
E Arll s 2C |p+ VCva A + 2C, 6xx
L(1 -11-p) p)
27

---

Published in Transactions on Machine Learning Research (03/2025)
This completes the induction.
In order to bound the noise term as 12 min 11, va batch-size that ensures Ç I va. Using the definition of Ç, we require that,
} we will require an additional constraint on the
b
n1 1)6 Sva
3
which is satisfied by the condition on the batch-size. From the result of the induction,
A + 20oa5x
E B| s 20%
p + v[va
L(1
Va
T ||bo|| + 2Co 10aasx2 2C
20%
+ V(va
Va (p- 1
2/K
L
Va
2
Va
(1 112/kv((
| + 20
aaaaxy 200
2/R Va
20
2/
L
(1
1 S (1 (1-2VFVVS) S because of the constraint on batch-size)
3v2
K
V@
(3
1 1-2//// // [|] I
2a 2
R
2/K
K
max
(4"
a Va
33v
Va
{3
2/1 vc) {3] 1
12
min {1, Val
5
VR
1
max
1 2/m
(SS Va)
Va
2/K
(4"
WT
w¹ll
5
6v2 2
V K
exp
T
Va
max
{3
(4' [ (.1-2VRV3)) llw0
w" + 12yax min { 1,
5
a
k 2
P (For all x, - x S exp((-1p (-r)
Va
Corollary 1. For L-smooth, strongly-convex quadraties, under interpolation, SHB (Eq. (1)) with the same parameters as in Theorem 2 and batch-size b s.t. b 2 b" := n 1+ (where C' is defined in Theorem 2) converges
6/2 Va
T Va
{3 ( (1122/((() llw00 - -w"ll
E WT - w"l S
V exp
max
Vk
Proof. Under interpolation X 0. This removes the additional constraint on 6* that depends on the constant. a, finishing the proof.
Corollary 4. Under the same conditions of Theorem 2, for target error > setting = min (xtx )p
and T
2
log
12v2 yr wo wo--""]] "I ensures that ||t w0 S Ve.
(-(1/
Vae
Proof. Using Theorem 2, we have that,
6v2 V exp
(LV -VVVV) T
12vax
min { 1, C
E fluT w w01ls
VK Ilwb w11
Va
p
Va
28

---

Published in Transactions on Machine Learning Research (03/2025)
Using the step-size similar to that for SGD in (Gower et al., 2019, Theorem 3.1), we that to get Ve accuracy first we consider 12yax < YE that implies S 1 2
va( vd(1-2//V/E VK
We also need Gv2 Va
V exp
|100
w0 S Le
Taking log on both sides,
va (1 (11-2//V) T
VE vā
S log
VK(
6v2 vK
1|w0 w
2//6
(12V2 v/ 11000 w"1
T
log
va 2/Kv()
VaE
D.1 Helper Lemmas
We restate (Wang et al., 2021, Theorem 5) that we used in our proof.
Theorem 8. Let H (1 3)Ia aA
BId
waexedwhere A E Rdaxd is positive definite matrix. Fix
la
a vector up E IRd. If 3 is chosen to satisty
1 2 max { (1- vadmin(A))
( / Vadmax (4) then
|| | s (VB) k C. l|vall
where the constant
v2(B 11
Cu min {h (B, admin (A)), h (B, admaa (A))]
and h == (1 / (3 (1 + )2)²).
Lemma 6. For posilive definile malrix A, denole = Amar (A) L. Set a Amaz 7 for and Amaa (A)
B = (1 1 admim (A)) = (1 va h (3,2) := = - (3 (1 2 (3 (1+ v
2/x Then, Co == / /miin(((BBaaami (B (A))}
(2
3 V and
Proof. Using the definition of (3,2) with the above setting for B and simplifying,
1
a() 3ap
1
2
aap
16
ap
Va
3a
(a= +
K
2
16k
3
3
(a
R
2/16
16k
15
(#2 1)
K
16
16
22 + 3)
2/2
8V2/K
< 3
(8
Vh(Baau)
15
V15a
16
Now we nced to bound
V2(++) / MBBaL)
Using the definition of h (3,2) and simplifying,
a(3.aL)=(2v)) Vap - aL + a + +2 + aL
4 ap)
29

---

Published in Transactions on Machine Learning Research (03/2025)
3/2
3/2
a
4a
2
a²
(setting a/L and expanding above)
Vk
2 3/2
2k
16x²
0
1 1/2
1
1
K
VR
2
k3/2
2k
16x²
1
1
Va
1
Vk
2n
2x
16x²
Since 6 2 1,
2
2x3/2
>
and
2x
1
16k²
0, hence
aL) 2
4
1
aa va 1
VK
2k3/2
2k
16
2
1
+
VK
3/3/2
2k
16
Both 2 VK 2x372 and +
1
16 are decreasing functions of for 2
2x
Hence, RHS(K) = 4 - (1-((
2x3/7 (1 + 2x
is increasing function of k. Since,
16k
h(3, aL) 2 RHS(K) 2 RHS(() for all 2 1,
1
1
15a
MB3 aT) 2 a
4 - 2
(8 1)
2
2
16
16
Using the above lower-bound for v2(1+8) VM(,at)
have
22(1 + 3)
8v2
3
<
V /M3,aL)
15a
Va
Putting everything together we get,
{ K 3 a
Co S max
(3Va
Co S 3
Va
Lemma 7. For batch sampling method where each batch is sampling without replacement from the dalaset.
n - b = (n 1)b
[VV V f(ww))
EE
[Vf((
where ff((w)) : Eieg Vf((w))
Proof. First, [[V](((*))] E [ Eien V[((wk)]
: E Eie Vf(w)) Vf(w)) Then
will calculate the variance f((ww)).
Var (V[((w))) Var
V(((k))
iEB
Var
(((x)
IEB
Var
Vi(ww))X
where X, is an indicator if sample is in the batch B
j.KEN
Var (V.((w)))
Var [V +2
Σ
Cov (or (((()))
fz
jfk
30

---

Published in Transactions on Machine Learning Research (03/2025)
Denote Vf((ww)== V, and Vu((w)) = V& for simplification, hence
Var V.X. = Var [X = V b b b
(a sample is in the batch with probability 5))
n2 n
Cov [V,X,, V = , V -V
V,
(E X X E[X
n 2
Since E|X,Xx, = Pr[both samples i are in the batch] =
b 6--22 2
n2
and [[X,] = B[Xx]] = 2
b
6(b - 1)
b²
Cov VAXi, = V
n(n 1)
n
b(b b(b-1) n)
n
Plug back to Var (V)) then,
J.KEN
b(n - b)
b(b n) 1)
Var (v)) =
2
Σ 1,V& jtk
b2
n²
n2
ViV
n b
n
2
Σ
(n- 1)b
n?
n
1=1
jzt
n b
Σ
v?
Σ
V2
Σ
(n 1)b
n?
34h
b
[E[V)] (E)
(n 1)b
n Var (V:)
(n 1)6
E (VVs(w)) - v((wx))
n b
E
(n 1)b
E
Proofs for lower bound SHB
Before looking at the general lower-bound for n samples, it is instructive to consider the lower-bound arguments for a 2-sample example. The arguments for the general n-sample example are similar.
Theorem 9. For L-smooth, ji strong-convex quadratics problem fww) = 1 - 3 wwwt A&w with 2 samples and dimension d = n = such that w" = 0 and each A, is a 2-by-2 matrix all zeros except at the (i, i) position, SHB (1) with aę 1. 8* = 3 (1 [aau) With batch-size 1, when > 6, after
37 itcrations, we have the following: if Au = Wkk for a c -c= 1.1
Wk
E|[ AST 12 c ||
Proof. By definition, the 2 samples are: A = (0 ( 1₂ - (o 9)) and hence 1 = 0
0 L
Calculating the smoothness and strong-convexity of the resulting problem, F. = 1/, = 5, K 4. By the
31

---

Published in Transactions on Machine Learning Research (03/2025)
definition of SHB (1), we have that,
Wk++1 w²
(1 + 3)1a aAA
wh w² -3la
Wk w
La
WK-1 w
Let wa " be the first and second coordinate of wy respectively, A is the element in (d.j))position of A. Since w" = 0, the above update can be written
(1) (2)
1 B a.l (1.1)
(1)
0
3
0
1 3 aAk ((.2)
(2.2)
W
W
(2)
Henco, we can separate the two coordinates and interpret the update as SHB in dimension for each coordinate.
Subsequently, we only focus on the second coordinate which corresponds L (A in matrix A.
Wk+11 Wkk -aAk2u + B(wk Wk-1)
[k++1
1 3
21
-B
Wk
(a 1/1 1/L= 2/1)
Wk-1
Denoting As =
Wk wk-1)
the above update is
AK+1 = H4
where H2 is either H1 : 1 B -8) (corresponding to A = 0) or H2 == (11+ + 3 -
1
(corresponding to A3 = L) with probability 0.5.
In order to prove divergence, will analyze three iterations of the update in expectation. We enumerate across 8 possible sequences (depending which sample is chosen): (1, 1), (1, 2,, For example, if the sequence is (1, 11,1,2) the corresponding update (across 3 iterations) is:
3 H(1.12) As where H(12)) == H,
We denote Hi to be the matrix corresponding to the i-th permutation. For example, Hi H(1,11)- Next, we analyze the suboptimalityy || in expectation.
A k43 1 =
|
(probability each of the sequences is 1/8)
/sin(Ox) |cos( (0x. )
Representing A polar coordinates, for a O2 E [0, 2m), As = reche where rkk E RR and Φu =coos
Blll A k
72 8 8
71 dx
32

---

Published in Transactions on Machine Learning Research (03/2025)
In order to analyze the divergence of SHB, we define the norm square increase factor
Ell Au- 11
|Axl
E| A+
V
| 2
Ei=1 | 11
(I = || = r 11
%h | I1?
Σ ==
1HibEll 1|
10x11
|H 1|
(l = 1)
==1
V depends on 0b and hence it is a function of 0g. Using symbolic mathematics programming (Meurer et al., 2017), we can calculate as an expression of 3,0,
1
8
sin (
|'
V
H4
8
cos
8 sin(20) + 3 + 3,5 sin(20) + 3 cos((0) - -333
5,8 s2n( - 2 cos(20) 6,3
2,33 sin(20) + 3 cos(29) 3,38
2,3² s2n - 3,3² cos(22 (0) + 5,3² cos(20) +
We first verify that (3300) is monotonically increasing w.r.t 8 E [0, 1] by laking derivative of (33 0) w.r.l B. We plot the derivative for e [0, 1] and 0 [0, 2m.. From Fig. 5a, we can see that the derivative (3, 0) is positive for 3 [0, and 0 E [0, 2T].
Choosing 3 0.63 (corresponding to 6). we plot v against 9 in Fig. 5b and minimize & w.r.t 0 finding the minimum to be 1.1. Since min(V)- 1.1 1, the sub-optimality increasing in expectation for any A when 8 = 0.63. Hence, since (3, is monotonically increasing with respect to 3 (Fig. 5a), when k > 6 (correspond to 3 > 0.63), for an arbitrary A..
|| > A |
where e 2 1.1 for all Unrolling the recursion starting from 3T,
El| 13 > ot ||
Since e 2 1.1 1, the second coordinate will diverge and SHB will diverge consequently (Fig. 5c).
33

---

Published in Transactions on Machine Learning Research (03/2025)
10'
10
300
04
SHB
3
20
200
630
100
(a) 3D plot of derivative of (33 0) with respect to for E [0, 1] and 0 E 2x). The whole plane is above hence (3,0) 0) is monotonically in- creasing for B [0, 1] for any
(b) Plot of v against 0 for B 0.63
(c) Plot of SHB SGD for the 2- sample case with 6 b= 1, == =6 6. SHB diverges while SGD converges
Figure 5: Figures for 2-sample SHB lower bound proofs
Theorem 3. For I-smooth, j strongly-convex quadratic problem f(w) == 1 Ei=1 samples and dimension d 100 such that w" 0 and each A, is an n-by-n matrix of all zeros except at the (i, position, we run SHB (1) with og = = + 3x = (1 \Vaī)) <
1 w A with n
n 2
n and
) for c 1, after 67 iterations, we have that: As wk (
A [ > ct ||ul?
E
Proof. Denote L = maxien A,. and minggn A(ii), For the strongly-convex quadratic objective function f(w) = 1, E 2 ",,, w" =
Since each A, is diagonal, similar to Theorem 9, we can separate the coordinates and consider SHB in dimension for each the coordinates. Subsequently, we only focus on coordinate u that corresponds to the largest A (f.4) ie. u = arg maxien (1.1) The update for this coordinate is given by:
wh+1 = = a fif + B(wk - wk-1),
where ful(w)) EieB. f((ww)) = (Eiers A(wr)) wk. Hence,
Σ (u.w) wk + -(wk wk--1)
Wk+1 = W
at b
Similar to Theorem 9, we calculate the smoothness of f(w) as Amax (V2f(w)) = Amax (2 A.)
Hence + 3 and the full update can be written as:
wk+1 = wk
Σ
wk 3(wk wk-1)
bL
Bk.
/Wk+1]
Wk Wk-1
WF
In each iteration, we randomly sample (without replacement) examples. Hence, the probability that Ag is in the batch is n- When A, is in the batch, EieB, Alu.a) T.. On the other hand, when A,, is not in the
34

---

Published in Transactions on Machine Learning Research (03/2025)
batch, Eien, Afr Similar to Theorem 9, we define As := rewritten as:
Wk
Hence, the update can be
wk-1
A = H4
(1+ B -B 1
where H5 is either H1 =
w-P == n-b
(corresponding to when A, is not in the batch) or
0
H2 (1-10- + B 3
b (corresponding to when Au is in the batch).
w.p P2 =
We will use the same technique as in Theorem and analyze six iterations of the update in expectation using symbolic mathematics programming (Meurer et al., 2017). For this, we denote Hz to be the matrix corresponding to the i-th permutation of 26 possible sequences and P to be the probability of that sequence. Therefore, P: is a product of Pl. corresponding to matrices H1, H2 in the i-th sequence. For example, when Hi (1111111) = H,H, H H H1, PI ps. Writing the suboptimality ||Ax]] in expectation,
2
E A el 2
Σ 1=1
P: |HA
Representing A in polar coordinates, for a 0g E [0, 2T), As := TROE where r& E R4 and Φx = (sin((())
|cos(
The norm square increase factor FliA+1o
given by:
||
E A K
V
|
r Σ i= Pi HHox
rI |
Σ Pe 11
Using symbolic mathematics programming, we write V as a function of b, 3, (sce Fig. 6 for the complete expression) and analyze (b, 0). Similar to Theorem 9, we first show that Vb, 3,0) is monotonically increasing w.r.t B. Using the expression of (b, 3, 0) = (b.33,)) for each b E [n - 1], we plot W V' (b, 3, B) for 3 E [0.25, 1), E [0, 2n] and observe that V/, is positive. In Fig. 7a, we show an example plot of V " (6, B 0) when b 70. Furthermore, we discretize B and 0 numerically verify that for any E [n 1], v (6, 3, 0) is greater than 0. In Table 1, we show an example for values of vi V (b, 0) when b = 70. Hence for every 6 [n - 1], (b, 3, 0) is monotonically increasing function in [3.
ag
Next, for each batch-size E [n- - 1], we minimize (b, 3,0) and find (b)) as the smallest 8 such that (b. 3, 0) 1. In Fig. 7b, when b 70, we plot minimum of (bb, (, 0) w.r.t 0 and show the corresponding (b).. Since (6,33,0 monotonically increasing w.r.t 8, we conclude that for a fixed batch-size b [n 1],
vO E [0, 2m), vB (8* (8),),,)) 1), (,0,) 1.
2
From the definition of B, we can calculate the corresponding & for any 3 E [0.25, 1) as k= 2(1 - Hence, for a fixed batch-size b, the coordinate (and hence SHB) will diverge n. > (b) (corresponding to (b))..
n- From Theorem 2, see that the batch factor equal to Iπ- must be sufficiently small to ensure convergence of SHB. In particular, SHB converges at an accelerated rate if (n-1)p n-b S ट.. Hence, in order to derive the
lounr, we plot log (ύ- against log(k* (b)) in Fig. 7c. We observe that for larger K (b), the batch factor is smaller. In other words, when is large, SHB requires a larger batch-size avoid divergence.
35

---

Published in Transactions on Machine Learning Research (03/2025)
-88 sin (20) + _ + 9 sin (20) + 8 cos (20)
9,3
25,3 sin (20)
7 cos (20) + 25,3 + 19.3 sin (20) + 12.3 cos (20)
1888 +
7,8 sin (20) + 2,5 cos (20) 12 88 sin (20) 7 cos (20) + 5,9"
4.8" sin (20) - 6,9 cos (20) + 8,9 2 sin (20) + 2 f cos (20)
500 sin (20) 50 cos (20)
550,8
2800,8 sin (20) 8003 cos (20) 3200,3
3500,3 sin (20)
s* + 8 cos (29)
b
6
b
2350,9 cos (20)
4350,9
1300,3
sin (20)
400,3 cos (20)
1600
1400,8 sin (20)
2350 cos (20)
2850,9
1300 sin (20)
800 cos (20)
2000g
100 sin (20)
550,9 cos (20)
350g
100 sin (20)
800 cos (20)
800
b
b
b
100 sin (20)
50,92 cos (20)
150,
50 cos (20)
50
1000000 sin (20)
25000 cos (20)
125000
300000, sin (20)
ba
180000,3 cos (20)
420000.3
30000
(20)
170000,3 cos (20)
110000
230000 sin (20)
240000, cos (20)
12
12
420000,3 7 sin (20)
45000, cos (20)
45000,3
30000,3 sin (20)
180000,3 cos (20)
180000,3
82
6a
8
62
6
82
6
40000 253
sin (20
+
8
25000 cos (20)
25000
10000000, sin (20)
5000000, cos (20)
15000000,
82
62
b²
8²
62
12000000, sin (20)
16000000, cos (20)
24000000,
11000000,$* sin (20)
5000000, cos (20)
17000000
12
3 sin (20)
16000000, cos (20)
16000000,
6000000,3 sin (20)
9000000,9 cos (20)
15000000
5000000 cos (20)
6
83
8
6
5000000
500000000
2.3 sin (20
1)
1000000000.4
100000000, sin (28)
500000000,8 cos (28)
500000000,3
8
84
B4
400000000, sin (20)
8 cos (20)
1200000000.3
500000000 cos (20)
500000000
6 10000000000 sin (29)
64
b4
B4
b
B
25000000000,
cos (20)
35000000000
2500000000 cos (20)
25000000000
500000000000 cos (29)
500000000000
6
b
Figure 6: L as function of b, 5,8,0
In order to quantify the relationship between '= (b) and b, we calculate the best-fit line for our plot in Fig. 7c using linear regression. The slope corresponding to the best fit line -0.6 and the y-intercept -3.8.
Hence, we conclude that,
n - b
log
< 10.6log( (llo(("(())) (6)) - -.3
(n 1)6
b
e 3.3
(n 1)b
(b))o..8
(n - 1)b
(x* (b))))).
(n b)e
Previously, we have shown that (b, N. 0) > 1 for all K > K² (b). Hence, Vb K. 0) > 1 when
0.6
(k* (b))
(n 1)b
(n b)e
n
(n 1)b
3.3 3,330... 0.6
b
n1
1 +
2.300
1-1
b
12 72 c3
Therefore, when the batch-size < 20000
0)). the norm square increase factor V (b H, 0) will be
greater than I which leads to divergence. For an arbitrary AR,
E || A | > cE |
E 6T
2 > of |
Since e <> 1, SHB will consequently diverge. In Fig. 7d, we plot the gradient norm of SHB for k = 10 and n, SHB diverges empirically verifying our lower-bound. ]
observe that when = 10 < 1t 1.,.05000000
36

---

Published in Transactions on Machine Learning Research (03/2025)
50
35
30
20
2.0
8
0.9
10-
0.8
0.6
2
0.4
0.3
0.3
0.4
0.5
0.6
0.7
0.8
0.9
(a) 3D plot of derivative of [b 8) with respect to 8 (b) Plot of minimum (b, 3, with respect to for
for a sample = 70, B E [0.25, 1) and 0 E 2n]. The sample b = 70, B E [0.25, 1) and 0 E [0,2m] 8* (b) is whole plane is above 0 hence (770,, 3, 0) is monotonically smallest such that 1. increasing for 3 E [0, 1] for any 0.
10a0
1850
10
Chedd
10a
SGD
10:4
SHB
10
250
300
750
1000
1230
1500
1750
2000
lg(kkappa))
Berations
(c) Plot of log batch factor log (m-178 against log(x*). (d) Plot SHB vs SCD for the n-sample case with & 10, 100, = 10 <
n. SHB diverges while SGD
Using linear regression, the slope of the best fit line is converges
0.6 and the y-intercept is 3.8
Figure Figures for n-sample SHB lower bound proofs
37

---

Published in Transactions on Machine Learning Research (03/2025)
0
/
0.250
0.333
0.416
0.500
0.583
0.666
0.749
0.833
0.915
0.999
0
0.120
0.299
0.667
1.364
2.608
4.719
8.168
13.643
22.137
35.083
0.2x
0.203
0.387
0.701
1.216
2.037
3.311
5.238
8.102
12.283
18.281
0.4w
0.490
0.891
1.532
2.515
3.970
6.055
8.962
12.911
18.151
24.968
0.6
0.584
1.114
2.012
3.466
5.734
9.160
14.193
21.423
31.632
45.902
0.8x
0.354
0.748
1.477
2.754
4.892
8.333
13.703
21.875
34.094
52.153
1.0x
0.120
0.299
0.667
1.364
2.608
4.719
8.168
13.643
22.137
35.082
1.2x
0.203
0.387
0.701
1.216
2.037
3.311
5.238
8.102
12.283
18,281
1.4x
0.490
0.891
1.532
2.515
3.970
6.055
8.962
12.911
18.151
24.968
1.6
0.584
1.114
2.012
3.466
5.734
9.160
14.193
21.423
31.632
45.902
1.8x
0.354
0.748
1.477
2.754
4.892
8.333
13.703
21.875
34.094
52.153
2x
0.120
0.299
0.667
1.364
2.608
4.719
8.168
13.643
22.137
35.083
Table 1: Values (b,
3, 0) when 70 for different E [0.25, 1) and E [m,
38

---

Published in Transactions on Machine Learning Research (03/2025)
F Proofs for multi-stage SHB
Theorem 4. For L-smooth, i strongly-convex quadratics with x > 1, for T T
3.2" Im(2)
max (4Ks }
Algorithm with b 2 b" max
converges as:
1+ Cx2
E |eer - w01 5 exp (-
T
1100 - 1won--w11 + C2 X
8/8 (~8VR)
VT
where C: :=:352 and C7, Cs are polynomial in x and poly-logarithmic in T.
Proof- Stage zero consists of To = T iterations with a = + and B = (1 2V Let T; be the last iteration
2
in stage i, 1-=-1 T. Using the result of Theorem 2 with a for To iterations in stage zero and defining A. Wt w",
6v2 Va B1 < Vv exp
Va T
12vax
{1
51 S1
}
E WT w0
VR
exp
|wo
w'll+
min
4 VR
Va
T
12x
min {
S
|/wo
w0ll || +
8VK
Va
T
12x
S <6V2VK exp
||wo w=
8/
We split the remaining I iterations into I stages. For stage i E [1, 11, set a; and choose a, 2 Using Theorem 2 for stage i,
Vai Ti Via
12
E I 11 S S622
exp a3
E E|AT; -1ll +
min
/
Vas
s 6/ Vv cxp
1
T;
E|
12x
42//2 vK
2222
S sexp (12+5) + ln(vF)
T.
E
+
12x
2:/24 VR
2
T. Now we want to find T; such that (4/ + 5) In(2) + In(vR)
2
20411/24
4
T,
>
21/2
(2))
(2
22)
4
T.
12 (1(
(2 2
T,
(11/2 5) (n( In(vR)) 21n + 2l((K)) 2 (1/2+ + 5) In(x)
21++1//24 K
22 -
T
1 <
- expl exp((((/2++
exp
2(+11)/24 kR
Define Pi = exp((-(ff -4/2+5)) If we unroll the above for stages we have:
E |||
I p Ar:
12x
Σ
2-1/2
2 1/2
II Pj jmi-1
I
exp Σ
(12 + 5) 1 In k)
E |AA
12x
Σ
2 1/2 exp
Σ
(i/2+
-ine
1=1
J=#11
39

---

Published in Transactions on Machine Learning Research (03/2025)
S eep (-12/4 lm BTA +
12x
2- exp
2 1/2
Σ
(j/2) -in
jmi+1
12x
++++++
Sexp (-pp--1/4--0 - knx) E 11
+
2-1/2 exp
-
4
(2-24100))
12x
I Sexp(-19) km)) E B| ||
2-1/2 2
1/2
exp (mp(-inn))
(since 2 2 c)
i=1
- exp exp (-10/4 1 1m E B|
12x Σ
2(-
exp (-ilnx)
(since 1 2 (2)
12 XM
1
S Sexp (-12/4 kmx) E E| ||
(Simplifying Σ )
1L(K 1)
2(1)
12 X
I exp (-1/44 E B
1)
2(1)
Putting together the convergence from stage and stages [1,
(OVEVR exp
w" + 12x
12 X
1
BT S exp (-12/4)
||w0
+
8VF
1) 2(4)
1
(6V2VR exp
T
12x
12
1
1lw00
w=| +
+
2(4)
8y
11
1u(n 1) 2(4)
1
(6V2VK exp
T
100) w001)
24 x4
1
1/200
(5)
24)
8V
u(16 1) 2(4)
Now we need to bound the number of iterations in Σ i= Ti
4 2:/2 Ln)) +n(vk)) In(v)) +1 (2 v2)
$ 88)
2/12 (T2( 5 (n(2) + In(vk)) +
1=1
21/2 <8/R K
di + In In(vK)
11
(For i 2 1)
1=1
8V1
41 In In(vk)
21/2
2(1+13/2
8T 41 + In
V2- +
s 16 57 + In (R))
2(1+1)/2
Assume that / In(vk) In this case,
Σ T. < 1
(6)
We need to sct I s.t. the upper-bound on the total number of iterations in the stages is smaller than the available budget on the iterations which is equal to T/2. Hence,
T
192 in(v2)
V exp (In(V2) 1) (1 (nIn(2))
>
40

---

Published in Transactions on Machine Learning Research (03/2025)
= exp (In(/2)) (1 (nn(v22)
T In( In(v2) [n((2) S W T (ln(22)
384 VK
384 VR
(where W is the Lambert function)
Hence, it suffices to set I
1 W
T In(v2) 384 vF
).
We know that,
In(v2)
1
W
T Lul v2)
1n(v2)
384 /R
2 2
2
1/(2ln(V2)
= exp (7)
T In (n(22)
N
vi/e
exp
w
2
384 VK
1/(2In(v2)
In(v2)
V16
384 VW
(since exp(W(a))= = )Wa))
W
T In(v2)
354 V'
For x e², ((x) $ V 1 + 2 (x). Assuming T
381v e 50 In(v2)
T TIn(22)
2 6²,
384 Vn
1/1n(2)
T Iul (2)
exp ()) 2 v1/e
384 v/x
1 + 2 logg
(
T (n(v2)
384 VR
Since 2 = (cxp(x))1n12)
In(2)/
T In( (n((2)
2(1) (exp(1/2))\u(2) 2 1/e))nn(3)
384 /R
1 + 2 log"
(n(2)) 38
T In( 384 VR
= (cxx((122))lu()) 2 (VI/e) 2m(2))
1 + log²
T (n(22) 384 VK
1 + 2og
T l2( 384 vR
384
21/2 2
Tln(v2)
29 3 K 1 + 2 log?
T lu(v2)
384 VK
1
ln(
T
2 / (1+2l
( T in(v2)
Define C =
384 vR
meaning that 27/2 S
G.
Using the overall convergence rate,
l2)
Elllut - w²ll <
1
6V2VR exp
T
|lwo
w'll
+ ((x-1) 1) 2(1)
24 1
2 )
8VK
v VT
T 8VK
24 24xx VC
<
6V22vk exp
||w0
w +
VT
m/x - 1) VT
T
24 VCI
= E - -w'
exp
|1100 - w0 || +
VT
8/K
iu( 11) VT
41

---

Published in Transactions on Machine Learning Research (03/2025)
We assumed that I 2 In(vk) meaning that we want s.t.
1 w In(v2)
T [n((2)
) 2 In(vR)
384 vn
Since W(x)= 2 log 4++ for >> 0 we need to have:
T In(v2)
W
-12
In InVE
In(v2)
384 VK
4 T In(v2)
+1+1
38 VK
log
2 In (V()) -1
2
384
(2lu(v -2)2 -1
T>
4 4ln 2
Hence, it suffices to choose
/ K
T
16 2l)())))
In 2)
Since e 2
T> 96 K
3.22ku R0
16 - In((s)
In 2)
In(v2)
Therefore to satisfy all the assumptions we need that
3 / 384 V
T' 2 max
In(v2) In(v2) e²
3 (3.2
3.2 2 e² VK
max
In(2)
(n(
Let C3- 3.2" max (oxx(4)))) then In(2)
T 2 VRC's
With this constraint then
1 S
1
Hence the final convergence rate can be presented as
VT WC.
T
24
VC1
E W|wr--w'l - w0ll 56v2 CIVE
exp
11100 - w01l +
KV/VC's
8VK
-(x-1)
VT
=6V2
C
1/4
exp
T
|lw0
w= 11
24Xx VT
C3
8V
1u(15 -1 1) VT
$
384 Vx ))
In(v2))
T
X 24h/C1
<6
+ 2log2
exp
11100 - w01l +
8/K
vT p(k - 1)
Let Cr T) == 61 + 2 log
[
)²
t and Cs T)
/ then x(x-1)
E ||wT - w' || 5
exp
T
|ww0 - w= + Cs x
80R
VT
C
42

---

Published in Transactions on Machine Learning Research (03/2025)
Corollary 2. For L-smooth, j strongly-convex quadratics with 1, Algorithm with batch-size such Chat b2 2 b" 1+ xx
attains the same rate in Theorem 4 for T E [.2 v max (4x, Ci V 3[m
[n 1)b
In(2)
where C = 26 and C₁ is defined in the proof of Theorem 4 in App. F.
Proof. By the batch-size constraint in Theorem 4, need
From Theorem 4, a= =- S 7
hence
b
n
n- 3(2/01)
71
1+
3T/C),)
3(T/2)² [1
"
$ -
(n 1)b
(T/c,)² )² $
3(n - b)
(n - 1)b
T
3(n b)
G Proofs for SHB with bounded noise assumption
Theorem 10. For L-smooth and strongly-convex quadratics, SHB (Eq. (1)) with ag = aa- i and a S 1.
8x = 3 = (1 /ap) bounded gradient noise E, [V](w)) Vfi(wx))l z o², and batch-size b has the following convergence rate,
K E | - w w01l 332 exp
Va T 2 VR
11100 - w=I+ 6ÇØ
Proof. With the definition of SHB (1), [u((w) is the mini-batch gradient at iteration k; then, for quadratics,
Wh+1
w=
(1 + I) aA
-BIa
w
Vf(w)) - V fisu (wx)
ta
Wh
w
Ia
0
1h-1 w
Arti
A
5%
Ak+1= HAs + adg
Recurring from 6: 0 T taking norm and expectation w.r.t to the randomness in all iterations.
AT < H A oE
Σ
- 8%
In order simplify sR, we will use the result from Lemma 7 and Lohr (2021),
[x 12] = E& V f(wx) - V Sae (wxx))
n E, [V((w)) - v(((w)]]
no
(Sampling with replacement where is the batch-size and n is the total number of examples)
n g2
n
43

---

Published in Transactions on Machine Learning Research (03/2025)
' b n b á
[x[1o=11] V
Using Theorem 8 and Corollary (Wang et al., 2021), for any voctor v, |70 Hence,
Copt ||aa| where p vB.
Coa
B||l Copp" 1Aoll
+
T- E |0ull
(a= 7.)
L
Rm0
Coa
=opp" |1 +
T---- cō
L
Co Coaço
T-1
<Cp'' | +
L
k=0
p' | + Coa
1 - pT
L
1 1-p P
Co Coa
CC p² A +
(1 - L)L
We have Co Co 3 p p=
2VR then,
V@
T'
K 22 agō
E
E1 3
1
| +3
a
2/K
a Va
K
Vu
5
53
a
1
2/Rr)
|| +
Va
T'
+ Co
E ||w - w0ll 5322
exp
2
Vir
|110
w01
Since the proof of Theorem 4 analyzes each stage independently and successively uses Theorem 2, we can repeat the proof of Theorem 4 using Theorem 10, yiclding a similar O (exp dependence on the batch-size.
T
+ rate without any a
/
44

---

Published in Transactions on Machine Learning Research (03/2025)
H Proofs for two-
two-phase phase SHB
Theorem 11. For L-smooth; j strongly-convex quadratics with K > 4, Algorithm 2 with batch-size such that b 2 b' 14 m-r results in the following convergence, 6x
WT-I- w0ll 5 1
202x
exp
k9 8181(((1 - c) c) (-
T
?
1w0 w'll
Cr.
12 Rl/6L22 Ço
(-11-0)11) - (~ Sx
8vLez In(T)
+
exp
veL
eVveL
v(I-c77
where 0 = 1 In(ov/AI-e) (
n-b
captures the dependence on the batch-size, C2 exp (xl([|-]))
CL 4(1-2) In(x)
(n-11j6
exp (-22) and C' == 35.6.
Proof. After To iterations, by Theorem 2 with a = 1, and x = [
E Ilwz - will s 6/2/n exp (-2 To
(A/R) ) 1l10 ww + 12v3cx
(since in Theorem 2, 5= = / n-b
After Ti iterations, by Theorem with y = (1/11)2/23 and T= 1.
Ti (- 64Lo2c2 (In(T)))²
Eler] < Wyo - w0ll c2 exp
(<4x
22g2
Ti
Taking square-root on both sides and using that VE[ET] 2 E((ET)
Ti7
8 LV222 (In(Ti))
E E <
WTo w"ll Ve2e
+
8x In
e7
VTi
Taking expectation over the randomness in iterations t = 0 to To 1,
(- (- Ti SN In
8VLCz (In(Ti)
ElVEr) S VC
E Iuw wi1l+
e7
VT.
Using the above incquality and using that x² S 2L
Ti7
(- To
12/3v2LC
= E[VET] Sveeex
6 6V2VR exp
14/h (-4v%)
1110 w=1+
/
8 0
(In (Ti))
e7
VTi
To
Tiy
- exp
1lw0 - -w"l
AVK
8k (n(Ti)
12 e02 Ço
Ti
WVLc2 (In(Ti)
11
exp
In
eq
/Ti
For Ti 2 e and since 7 1
To
Ti ))
E <bv 2 exp
(88n(TT)
Vk
+
K
11100 -w"ll
45

---

Published in Transactions on Machine Learning Research (03/2025)
12V6L vei
So
T
8V (In(T)))
exp
8x In(Ti) +
1
vTi
Consider term : = /22./22// exp
(~81m(75)
To
1t )) l1wo - -w We have,
VR
To
T
cT
(1 c)T
K
VK
cv K + 1 -
Suppose I Ra = T then
q= In(ev/ +1
In(x)
Since c E (0,1), E (0.5, 1), then
T 7
A <622 exp
(~
k2 8 In((1 - c) ccc
11200 w'll
Consider the noise B = 12v6Los Co
7i exp
+ BVLc (In(Ti))
8x In((T)
V/i
B 12 Ço
Tı
SVLC20CK3/2 (In(Ti))
exp
8k Snln(Ti)
e7
VT1
12v6Lo2 Co
(1-c)Ty c)T' (1-
LVLoo
(In ([((1))]])
P
exp
8N c)T²)
ey
V(I-c]T
12/6Le250 So
[1-c)T²
8 8VL02
In(T)
P
exp
8x
Eq
V(I-c)T
(Since E (0,
Putting everything together,
(-A T
7
BI ETl 66v 202k exp
1120-v - w=1l
|110
K# 81881((11 - 7))
12/6Lez Ço
exp (-11-c)1 c1)'((
8 8VLoz oCK
In(T)
(-8+
e7
(11
Using Lemma 1 with er 4(1-7) [1 exexpp (-47) then
(-T T
7
Ver Belw-- w0ll < 6662222 202 exp
|10b w''l
Kd 119811((11 - c)T)
12 12vble2 Ço
(1 (-11-c)77 c]]
8VIez In(T)
+
exp
81s (-8rln(TT)
= EB| [ - w0ll < 16
202k exp
Rd (-
T
7
| - w"ll
811((11 - T)
12V6Lo260 50
(-1--0711) cc)
8vLoz In(T)çn3/2
+
exp
VeLu
e7/eL
c)-
46

---

Published in Transactions on Machine Learning Research (03/2025)
Corollary 3. For L-smooth, μl strongly-convex quadratics with 4, Algorithm with batch-size such that b 2 b' 1+ m-T and 2 results in rate of O eep (- x
T
(T)
for all T.
6x
Proof. From Theorem 11, 9 1 In(o/k+t-o) In(x)
plug in 2 then
In /E+1
9 =1
In((s
Since q is monotonically decreasing with respect to and K >> 4 then
In
44+1 0.7
2
q 51
In(4)
Hence the bias term the rate Theorem 11 converges as (exp (- 2) 0 (exp (exp (- The noise term converges at the rate of
VT hence the convergence rate of Algorithm 2 when c = 0.5 is
0 exp + + VT)
47

---

Published in Transactions on Machine Learning Research (03/2025)
Additional experiments
1.1 Quadratics experiments on LIBSVM datasets
To conduct experiments for smooth, strongly-convex functions, we adopt the settings from Vaswani et al. (2022). Our experiment involves the SHB variant and other commonly used optimization methods. The comparison will be based on two common supervised learning losses, squared loss for regression tasks and logistic loss for classification. We will utilize linear model with 6g-regularization å ||w|| in which À = 0.01. To assess the performance of the optimization methods, use ijenn and revl data sets from LIBSVM (Chang & Lin, 2011). For each dataset, the training iterations will be fixed at T = 100n, where n is the number of samples in the training dataset, and we will use a batch-size 100. To ensure statistical significance, each experiment will be run 5 times independently, and the average result and standard deviation will be plotted. We will use the full gradient norm the performance measure and plot it against the number of gradient evaluations.
2
The methods for comparison are: SGD with constant step-sizes (K-CNST), SGD with exponentially decreasing step-sizes (Vaswani et al., 2022) (K (K-EXP), SGD with exponentially decrcasing step-sizes and SLS (Vaswani et al., 2022; 2021) (SLS-EXP), SHB with constant step-sizes (set according to Theorem 5) (SHB-CNST), SHB with exponentially decreasing step-sizes (set according to Theorem 1) (SHB-EXP), SHB with exponentially decrcasing step-sizes (set according to Theorem 1) and SLS (Vaswani et al., 2021) -SHB-SLS-EXP). revl
ijcnn
10P
10
Gradent Nerm (193)
Gredient (leg)
10-
10-1
10-2
10'
1003
10
10a
(Gradlent
evaluntionss 1000
(Gradient eush antless 1000
(a) ijenn
(b) rcvl
K-CNST
SHB-CNST
KEXP
SHB-EXP
SLS-EXP
SHB-SLS-EXP
Fignre 8: Sanared loss on iiemn and revl datasets
10a
licnn
10"
rev1
Dratiem Norm (ee))
1-0-2
20
Gradee korm
10-3
10-3
10-1
10
10"*
10a
10
2000
(Gratient avalustional 1000
(Gradert avel coulattion)) 1000
(a) ijenn
(b) rcvl
KCNNST
SHB-CNST
KXXP
SHB-EXP
SLS-EXP
SHB-SLS-EXP
Figure 9: Logistic loss on ijenn and rcv1 datasets
From Fig. 8 and Fig. 9, we observe that exponentially decreasing step-sizes for both SHB and SGD have close performance and they both outperform their constant step-sizes variants. We also note that using stochastic line-search by Vaswani et al. (2021), SHB-SLS-EXP matches the performance of the variant with known hnest.
1.2 Comparison to Pan et al. (2024) multi-stage SHB
In this section, we consider minimizing smooth, strongly-convex quadratics. The data generation procedure is similar to Section 5. We vary E {2000, 1000, 200, 100} and the magnitude of the noise E {10-2, 10-4, 10-6 10-$3. For cach datasot, we use a batch-size b = 0.9n to ensure that it is sufficiently large for SHB to achieve an accelerated rate with all of our chosen k. We fix the total number of iterations T 7000 and initialization wi = 0. For each experiment, consider 3 independent runs, and plot the
48

---

Published in Transactions on Machine Learning Research (03/2025)
101
30
1002 2033
1030
1000
#30
300
13
1003
2003
WD0 4530
2000
10
1900
2000
3900 ex3
3000
0030
1330
29.01
1200
2003
100 eOE 5030 6000 2001 Bens0r
(a) K 2000, -= 10
(b) 2000, r 10
(c) K 2000, = 10
(d) K 2000, = 10 8
Oet beeeak
-22
700
AND
(e) K 1000, r 10-2
(f) 1000, 10-4
(g) 1000, 10-6
(h) 1000, 5 10-8
Mix
18 2X0 203
203 X00 20W0 Tendian
16
390 9030
130
K = 500, r == 10
K = 500, r 10
(k) K = 500, == 10 -
(1) K = 500, r = 10
13"
-85 A82
43) 0220 10% 126 XX6
.82
A53
4.3 1230
190
136 A40
1990 196
ARS 981 400 1630 12.0
124. AN3 J08 600 1030 LOG AGI
(m) K 200, r 10
(n) 200,
10-
(o) R = 200, r= 10-6
(p) K 200, r 10
3505
3905
(4) K 100, r 10-2
(r) K 100, r 10-4
(s) K 100, r = 10 (1) K 100, r 10-8
Multi-SHB-CNST
Multi-SHB
Multi-SHB-PANI PAN-2
Multi-SHB-PAN-TKAPP
2P-SHB
Comparison of Multi-SHB, Multi-SHB-PP PAN-T-AAP With sufficiently large batch-size, the method by Pan et al. (2024) is able to avoid the divergence behaviour in Fig. 2. The performance of SHB variants is similar to the method in Pan et al. (2024), and it consistently lies in-between the two extremes.
Figure 10:
2P-SHB, Multi-SHB-PAN
and
average result. We will use the full gradient norm the performance measure and plot it against the number of iterations.
We compare the following methods: Multi-stage SHB (Algorithm 1) (Multi-SHB), our heuristic Multi-stage SHB (Algorithm 1) with constant. momentum parameter Mulli-SSB-CNSST Two-phase SHB (Algorithm 2) with c= 0.5 (2P-SHB), Multi-stage SHB (Pan et al., 2024) with = 2 (Multi-SHB-PARe and C = TVR (MuLt1--HH--PAN---AAP)
We observe that with sufficiently large batch-size the method by Pan et al. (2024) is able to avoid the divergence behaviour in Fig. 2. Furthermore, the performance of 2P-SHB, Multi-SHB Multi-SHB-CMS is similar to the method in Pan et al. (2024), and it consistently lies in-between the two extremes.
49


Reviews:
[{"review_id": "X2LnKn7Y7r", "reviewer": "Reviewer_B4mw", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "While the work shows potential, three major issues need to be addressed before considering the possibility of publication:\n\n- The proofs of Theorems 6 and 7 are currently incorrect due to the following reasons:\n\nThe step size $\\eta_k$ depends on $\\hat{L}$, making it random. Contrary to the claim made at the top of page 20, the argument used in the proof of Theorem 5 cannot be extended here because the step size cannot be factored out of the expectation. As a result, the subsequent derivations are invalid.\n\nOne potential way to address this issue is to assume that $\\eta_k$ is $\\mathcal{F}_k$-measurable. However, is this assumption realistic? Perhaps an online estimator of $L$ could be employed.\n\nEven if this problem is resolved, the authors cannot perform a case disjunction as done in the current proofs since $\\hat{L}$ remains random. Indicator functions must be used, complicating the proof substantially.\n\n\n- The second significant concern is whether the improvements claimed by the authors are valid. For instance, in Theorem 1:\n  The authors claim a convergence rate of \n    \\[\n    O\\left(\\exp(-T/\\kappa) + \\frac{\\kappa^2}{T}\\right),\n    \\]\n    where $\\kappa = L/\\mu$ is the condition number. They assert that this result is better than Sebbouh et al. (2020, Appendix H), which provides a convergence rate of\n    \\[\n    O\\left(\\frac{\\kappa^2}{T^2} + \\frac{\\kappa^2}{T}\\right).\n    \\]\n    While it is true that the residual term is improved, the primary term $O(1/T)$ appears to have a larger constant in the authors' results. This discrepancy must be clarified, as it undermines the claim of superiority.\n \n- The manuscript suffers from several instances of imprecision and lack of rigor. Examples include:\n\nIn Theorem 1:\n    What is $\\eta$ in the definition of $\\lambda_k$?\n          The dependencies on constants, especially the condition number $\\kappa$, are not fully exposed. For example, $c_2$ and $c_L$ implicitly depend on $\\kappa$ but this is not explicitly stated.\n\n\nIn Theorem 2:\n    Does $\\chi$ depend on $i$? If not, is it the supremum?\n        The definition of $\\Delta_k = \\| w_k - w \\|$ is inconsistent with subsequent sections (e.g., Theorem 3).\n\n\n\nIn Theorem 4:\n     The dependence on $\\kappa$ is again hidden in constants such as $C_1$ and $C_3$. This needs to be explicitly addressed.\n \n In the proof of Theorem 2:\n    The condition $\\rho + \\sqrt{\\zeta} \\sqrt{a} < 1$ is introduced without being stated in the assumptions.\n      Additionally, this inequality is only feasible if $\\zeta < \\frac{1}{4\\kappa}$. Is this assumption realistic?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors address the minimization of a $\\mu$ strongly convex function. They also suppose that the functional is $L$ smooth and they study the convergence results of Stochastic Heavy Ball-type algorithms.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The paper addresses an interesting and important topic, but significant corrections and clarifications are required. The authors need to:\n\n- Correct the proofs of Theorems 6 and 7.\n \n- Reevaluate and clarify their claims of improved results.\n\n- Ensure rigor and consistency throughout the manuscript, particularly in the handling of constants and notation.", "broader_impact_concerns": "Nothing to report.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "VEEp1TVAaP", "reviewer": "Reviewer_iV4k", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n- The paper has some very nice-looking results.\n- Although I did not have time to review the proofs, the results seem logical.\n- The definition of the different variance measures that are used are also interesting and more meaningful than what was done in previous works.\n\n\nWeaknesses:\n- The storytelling can be slightly improved: explaining why you need to study the multi-stage SHB (to deal with the problem of convergence to a neighborhood)? Why do you need the 2-stage version of it? why not use the exponentially decreasing step sizes that worked for Th1? For me, all of these things are missing and make it harder to connect the dots.\n\n- There are also some typos: like in the contribution section where the rate for multi-stage SHB is stated with $\\sigma /T$ instead of $\\sigma /\\sqrt{T}$. Also, Algorithm 2 does not specify how to choose $a$.\n\n- The authors also kind of claim that they propose the multistage SHB and at the same time say that previous works already considered it, so which is which? I understood that you analyzed it using a different (more realistic) noise assumption, but you did not invent it. You need to make this clear.\n\n- For quadratic functions the obtained stochastic term $\\sigma/\\sqrt{T}$ seems to be suboptimal, and the authors neglected to discuss why this is the case.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper studies the Stochastic Heavy Ball (SHB) method and proves an improved rate for general smooth and strongly convex functions compared to previous works. The rate is obtained using an exponentially decreasing step size; notably, the obtained rate does not need knowledge of the variance. \n\nFor smooth, strongly convex quadratic, the authors show that for sufficiently large batch size, SHB with a constant step size enjoys an accelerated rate (up to a neighborhood defined by the noise $\\sigma$ that can be made as small as possible through the choice of a parameter $a$), they also show that this need for a large batch size is necessary.\n\nTo solve this convergence up to a neighborhood, the authors analyze the multi-stage SHB algorithm and show that by using it they can keep the accelerated rate (always for quadratic) and improve the stochastic rate from $\\sigma$ to $\\sigma/\\sqrt{T}$. \n\nFinally, the authors show that by only using two stages of the previous algorithm, they still get a slightly accelerated rate ($\\kappa^{0.7}$ instead of $\\kappa^{0.5}$ ) for the bias term and keep the same $\\sigma/\\sqrt{T}$ stochastic term as before.\n\n**Update:** I have just realized that you have Theorems with the norm squared and Theorems with only the norm, I believe it would be best to stick to one choice.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I think this is an interesting work, but it needs some slight improvements in terms of writing: \n\nExplain what happens for quadratic functions when you use the same step size schedule as for Th1 and why it prevents having the accelerated bias term.\n\nExplain what is the problem with the multistage SHB that you would consider the two-stage version, is it a problem with the number of stages? I am also curious about what would happen with a three or four-stage version and how can such versions be defined.\n\nYou should also discuss the noise term you obtain with the multistage and two-stage versions and compare with your Th2 if you can choose the best $a$ (which seems to give the $\\sigma/T$ dependence at the expense of a $T$ factor multiplying the exponential.\n\n\nMaybe also drop unimportant constants in the main paper and only keep what is important (to keep the focus of the reader on what is really important).\n\n**Update:** You should use one convention (either norm or norm squared ) in your Theorems or mention somewhere that you use both (but it is better to only have one or the other).\n\nAnd now it seems that based on TH2, the best choice of $a$ can lead to a $\\sigma/T^2$ term, can you comment on this?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "BEI1Arlkuc", "reviewer": "Reviewer_Lpjn", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n1. This paper studies the convergence rate of SHB for smooth and strongly convex functions, and I believe both the problems and results are interesting. \n2. I appreciate the authors' effort to state the conditions and implications of each theorem precisely, making the paper more readable.\n3. This paper has a rather detailed comparison to related work, which helps the readers to understand the contribution of the paper.\n\n**Weaknesses and Questions**\n1. For the results in Theorem 1, it seems that there are two \"drawbacks\" listed below, could the authors explain more on these two points?  \n\n    (1) one need to fix the number of iteration $T,$ and the set the hyperparameters depends on $T.$ This seems to be not common for practical algorithms. \n\n    (2) In the convergence rate, there's a term $\\frac{\\gamma}{\\ln(T/\\tau)}$ that also depends on $T$. Moreover, $\\lim_{T \\rightarrow \\infty} \\frac{\\gamma}{\\ln(T/\\tau)} = 0,$ since $\\lim_{T \\rightarrow \\infty} \\gamma = 1.$ Thus, for large $T, $ it seems to me that the convergence rate is smaller than $\\frac{1}{\\kappa}.$\n\n2. In Theorem 2, there's a term that depends on the noise is not vanishing in $T.$ While I believe this type  of upper-bound appears in proofs of  stochastic optimization algorithms, I still wonders:\n\n   (1) what is the technical difficulty of obtaining a convergence noise term in $T$ as well as obtaining an accelerated convergence rate? In Theorem 1, you prove a vanishing noise term for general smooth and strongly convex function with non-accelaerated rate.\n\n   (2) The results, in corollary 2 regarding the interpolation case is a bit vague to me. It seems to me that, in practice, it is not common that one can have both strongly convex and interpolation. For example, for linear regression problems, if one considers overparameterized regime, then the problem is not strongly convex; if one considers underparameterized regime or ridge regression, then there's no interpolation in general. Thus could the authors give a practical example for this case?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the convergence rate of stochastic heavy ball methods (SHB), and provides various results under different settings. The main results are summarized as follows:\n1 For general smooth and strongly convex functions, the authors prove a $\\mathcal{O}(\\exp(- T/\\kappa) + \\sigma^2/T)$ rate.\n\n2 For smooth and strongly convex quadratic function, the authors prove a $\\mathcal{O}(\\exp(- T/\\sqrt{\\kappa}) + \\sigma^2)$ accelerated rate for batch size $b>b^*$. Specifically, for interpolation case where $\\sigma = 0,$ the Theorem implies a convergence to global optimal.\n\n3 For smooth and strongly convex quadratic function, the authors propose a multi-stage algorithm and provide a  $\\mathcal{O}(\\exp(- T/\\sqrt{\\kappa}) + \\sigma^2/\\sqrt{T})$ accelerated rate. The authors also provide a two-stage algorithm which provides a $\\mathcal{O}(\\exp(- T/\\kappa^{0.7}) + \\sigma^2/\\sqrt{T}).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. While I appreciate the authors' effort to state each theorem precisely, I think many theorems are too \"precise\" and contain too many notations. Also, sometimes the same parameters are separated into two terms. To name a few:\n\n    - In Theorem 1, $c_2, c_L$ and $\\exp( \\dots)$ terms are all depend on $T$, could you put all the terms depend on $T$ together, so one can see the dependence on $T$ more clearly?\n\n    - In Theorem 1, I didn't see the utility of $\\tau.$ Why not just pick a specific $\\tau$?\n\n    - In Theorem  2 and Corollary 2, $\\max \\{ 3/4, 1-2\\sqrt{kappa} \\sqrt{\\zeta} \\}  \\geq 3/4,$ so why not directly use $3/4$ in the convergence rate as an upper bound.\n   \n    - In Theorem 4, again $C_1, C_3, \\kappa^{1/4} \\exp(\\dots)$ terms both depends on $\\kappa$ and $T,$ could you make term depends on $\\kappa$ and $T$ together, so one could clearly see the order of $\\kappa$ and $T$ in the theorem?\n\n2. For the numerical experiment, \n    \n    - In the theorems, your emphasize the for $n \\gg \\kappa^2,$ only constant batch size guarantees the accelerated convergence. However, in the experiments in Figure 3, you only consider cases where $n \\ll \\kappa^2$ and pick $b = 0.9n$ as the batch size. I suggest having some experiments showing the convergence in the case of  $n \\gg \\kappa^2$ and constant $b$  if possible, to strengthen the theoretical results.\n  \n    - In both Figure 3 and Figure 4, it is not clear to me how the convergence rate depends on $\\kappa.$ Thus, is it possible to make a plot showing the relationship between the convergence rate and $\\kappa?$ In particular, under the conditions of Theorem 2 and Theorem 4, is it possible see a $\\mathcal(1/\\sqrt{\\kappa})$ dependence on the convergence rate w.r.t $\\kappa?$\n\n3. A few references are out-of-date, for example, Bollapragada et al. (2022) and Lee et al. (2022). Also, There's no Theorem 3.1 in the newest version of Bollapragada et al. (2022). I suggest the authors to double-check the references and cite the latest version if possible.", "broader_impact_concerns": "No concerns on the ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents new convergence analysis for stochastic momentum in the strongly convex and smooth setting, that scale gracefully as the batch size increases, or as the problems get closer to interpolation. Gracefully, because the rates converge to the faster rates established for the full batch momentum method as we approach interpolation. The authors also provide a lower bound for quadratics that explains why one their results forces the batch size to depend on the condition number. Because of these results, I support accepting this paper. This is also inline with reviewers, who also supported acceptance.
Decision: Accept (poster)