Accelerated Convergence of Stochastic Heavy Ball Method under Anisotropic Gradient Noise
Authors: Rui Pan
Yuxing Liu
Xiaoyu Wang
Tong Zhang
Venue: ICLR 2024 poster (2024)
Keywords: optimization theory
convergence analysis
heavy-ball momentum
learning rate schedule
near-optimal convergence rate
Primary area: optimization

Abstract: Heavy-ball momentum with decaying learning rates is widely used with SGD for optimizing deep learning models. In contrast to its empirical popularity, the understanding of its theoretical property is still quite limited, especially under the standard anisotropic gradient noise condition for quadratic regression problems. Although it is widely conjectured that heavy-ball momentum method can provide accelerated convergence and should work well in large batch settings, there is no rigorous theoretical analysis. In this paper, we fill this theoretical gap by establishing a non-asymptotic convergence bound for stochastic heavy-ball methods with step decay scheduler on quadratic objectives, under the anisotropic gradient noise condition. As a direct implication, we show that heavy-ball momentum can provide $\tilde{\mathcal{O}}(\sqrt{\kappa})$ accelerated convergence of the bias term of SGD while still achieving near-optimal convergence rate with respect to the stochastic variance term. The combined effect implies an overall convergence rate within log factors from the statistical minimax rate. This means SGD with heavy-ball momentum is useful in the large-batch settings such as distributed machine learning or federated learning, where a smaller number of iterations can significantly reduce the number of communication rounds, leading to acceleration in practice.

Published as a conference paper at ICLR 2024
ACCELERATED CONVERGENCE OF STOCHASTIC HEAVY BALL METHOD UNDER ANISOTROPIC GRADIENT NOISE
Rui Pan";
Yuxing Liu"; Xiaoyu Wang',
Tong Zhang*
The Hong Kong University of Science and Technology
Fudan University
University of Illinois Urbana-Champiin
rpan@ccnnect. ust. hk, yuxinglius wiggplii220000 fudan. en, maxywanggust. .hk, tozhang@ill illinois edu
ABSTRACT
Heavy-ball momentum with decaying learning rates is widely used with SGD for optimizing deep learning models. contrast its empirical popularity, the understanding of its theoretical property is still quite limited, especially under the standard anisotropic gradient noise condition for quadratic regression problems. Although it is widely conjectured that heavy-ball momentum method can provide accelerated convergence and should work well in large batch settings, there is no rigorous theoretical analysis. In this paper, we fill this theoretical gap by establishing non-asymptotic convergence bound for stochastic heavy-ball methods with step decay scheduler quadratic objectives, under the anisotropic gradient noise condition. As a direct implication, we show that heavy-ball momentum
can provide (VK) accelerated convergence of the bias term of SGD while still achieving near-optimal convergence rate with respect to the stochastic variance term. The combined effect implies an overall convergence rate within log factors from the statistical minimax rate. This means SGD with heavy-ball momentum is useful in the large-batcl settings such as distributed machine learning or federated learning, where smaller number of iterations can significantly reduce the number of communication rounds, leading to acceleration in practice.
INTRODUCTION
Optimization techniques that can efficiently train large foundation models (Devlin al. 2019 Brown et al. 2020, Touvron et al.) [023a by Ouyang et al. [022] are rapidly gaining importance. Mathematically, most those optimization problems can be formulated as minimizing a finite sum
1 N
min ((w
Σ f((w),
where numerical methods normally applied find the minimum of the above form. Among all those methods, stochastic gradient descent (SGD) (Robbins and Monro 1951) and its variants can be regarded of the most widely used algorithms.
For instance, heavy-ball (HB) methods (Polyak] 1964), commonly referred as heavy-ball momentum, are one of those popular variants. Empirically, was extremely helpful for accelerating the training of convolutional neural networks (Szegedy et al. 20155 Simonyan and Zisserman, (005 He et al. 2015 Huang et al. 2017 Sandler et al. 2018). Theoretically, acceleration for gradient descent (GD) on quadratic objectives (Nemirovski 1995).
been shown to provide optimal
Nonetheless, when it comes SGD theory, things become much different. Despite its huge success in practice, most theoretical results of stochastic heavy ball (SHB) were negative, showing that the
Equul contribution.

---

Published as a conference paper at ICLR 2024
convergence rates of heavy-ball methods are better than vanilla SGD (Devolder et al.) 2013 Yuan ct al. 2016 Loizou and Richtárik, 2017; Kidambi et al., 2018 Jain et al., 201 Sa Li et al. 2022).
The existence of these gaps between GD and SGD, between practice and theory, is rather intriguing, which may make one wonder: Can stochastic heavy ball provide Φ accelerated convergence when the noise is small, such as under large-batch settings?
To answer this question, the first step is to find the missing pieces in those negative results. One key observation is that all those negative results assumed constant learning rates, while in practice, decaying learning rates are usually used instead. Those decaying learning rates, often referred as learning rate schedules, were demonstrated to be critical for improving the performance of a trained model in geal-world tasks (Loshchilov and Hutter 2017; [Hwwrd and Ruder 2018). Furthermore. if one only considers the vanilla SGD algorithm, the theoretical property of most schedules have already been well inspected (Shamir and Zhang, [003 [203]]aii et al. 2019 Ge et al., 2019 |Harvey et al.. 2019) Pan et al. [202]] Wu et al., (022a)). Briefly speaking, one can view learning rate schedules as a variance reduction technique, which helps alleviate the instability and deviation caused by stochastic gradient noise.
Since it has been pointed out by (Polyak 1987 that variance reduction is the key to improving stochastic heavy ball's convergence rate, it is then natural to ask: Are there proper learning rate schedules that can help us achieve accelerated convergence for SHB under large-batch settings?
Our paper gives positive answer to this question. As first step, we restrict ourselves to quadratic objectives. Although these problem instances are considered one of the simplest settings in optimiza- tion, they provide important insights for understanding model's behavior when the parameter is close to a local optimum. Furthermore, past literature on Neural Tangent Kernel (NTK) (Arora et al., 2019: Jacot al. 2018) suggests that the gradient dynamics of sufficiently wide neural networks resemble NTKs and can have their objectives approximated by quadratic objectives given specific loss functions.
Motivated by the empirical anisotropic behavior of SGD noises near minima of modern neural net- works (Sagun al. 2018 Chaudhari and Soatto 2018 Zhu al. 2019) and theoretical formalization of this noise property in least square regression (Jin et al. 201 8b a Pan et al. 2021). ). we conduct our analysis based on the assumption of anisotropic gradient noise, which is formally defined later as Assumption
3 in Section 3 Notice that the very same condition has already been adopted or
suggested by many past literatures (Dieuleveut al., 2017 Jastrzębski al. 2017 Zhang et al., 2018
Zhu et al. 2019 Pan et al. 2021
1.1
OUR CONTRIBUTIONS
1. We introduce novel theoretical techniques for analyzing stochastic heavy ball with multistage schedules, providing several key properties for the involved x 2 update matrix T, Specifi- cally, we show that | TT 22. TolTToll can be upper bounded by TF under certain conditions. This allows SHB with changing learning rates to exhibit similar theoretical properties as vanilla SGD: for each eigenvalue, SHB first exponentially decreases the loss with large learning rates, then retains the reduced loss with small learning rates.
2. As direct result of this technical innovation, we present a non-as -on-asympootic last iterate convergence rate for stochastic heavy ball with step decay learning rate schedule quadratic objectives, under the standard anisotropic gradient noise assumption. To the best of our knowledge, this is the first con-asymptotic result for SHB on quadratics that clearly expresses the relationship among iteration number T, condition number k and convergence rate with step-decay schedules.
3. Our results show that stochastic heavy ball can achieve near-optimal accelerated convergence under batch settings, while still retaining near-optimal convergence rate Ó(do² /T)) in variance (up to log factors away from the statistical minimax rate).
2
RELATED WORK
Large batch training: batch-batch training is realistic setting of its own practical interest, In several recent efforts of accelerating large model training, it has been observed that large batch sizes
2

---

Published as a conference paper at ICLR 2024
are beneficial for accelerating the training process (You al. 2017:
2018
2019 Izsak al. 2021
Pan et al.] 2022 Wettig et al. (022). On top of that, in distributed machine learning (Verbraeken
et al.] (2220) and federated learning (Kairouz et al.] 2021), one can normally support outrageous size of large batches by adding machines/devices to the cluster/network, but unable to afford large number of iterations due to the heavy cost of communication (Zheng et al. [2019] QQia et al. 2021).
This makes acceleration techniques even more tempting under those settings.
SGD + learning rate schedules: In contrast, the research SGD with learning rate schedules focused on more general settings without assuming constraints on the batch size. In (Ge et al. 2019). the convergence rate of step decay was proved to be nearly optimal on strongly convex linear square regression problems. (Pan et al. (2021) further pushed these limits optimal for some special problem instances and offered a tighter upper bound, along with a lower bound for step decay. Concurrently, (Wu et al., 2022a) extended the analysis of (Ge et al.) 2019) to dimension overparamterized settings, with tighter lower and upper bounds provided for step decay schedules. In (Loizou et al. 2021), the convergence rate of Polyak step size on strongly convex objectives was investigated. Nevertheless, all the bounds in above works require SGD to have least f2(1s log c) iterations to reduce the excess risk by any factor of c. There are also works with looser bounds but focus on more general objectives. Since we restrict ourselves to quadratics, we just list some of them here for reference: (Ghadimi and Lan 2013 3 Hazan and Kale 2014 Xu et al. [016 Yuan et al. 2019
version under
Vaswani et al. 2019. Kulunchakov and Mairal 2019 Davis et al. 2023 Wolf 2021
SGD + HB+ + constant learning rates: Opposed to the positive results of near optimality for SGD, most results of stochastic HB with constant learning rates were negative, showing that its convergence rate cannot be improved unless extra techniques like iterate averaging applied. In (Loizou and Richtárik 2017 2020). a linear convergence rate of SGD momentum on quadratic - f(w.)] was established, which
objectives for L2 convergence [T w.ll"] and loss requires at least S1(x: log c) iterations. A better bound for L.1 convergence [|E[wr - w.]ll² and B norm [E] / w.]ll was also proposed, but whether they are relevant to loss convergence is unclear. Here B is a positive definite matrix related to the problem instance and samples. In (Kidambi al. 2018) momentum was proved to be no better than vanilla SGD on worst-case linear regression problems. In (Jain et al. 2018a), both SGD and momentum are shown to require at least (xr)) single-sample stochastic first-order oracle calls to reduce excess risk by any factor of c, thus extra assumptions must be made to the noise. A modified momentum method using iterate averaging was then proposed on least square regression problems and achieves Ö(VRR) iteration complexity with an extra noise assumption. Here k S k is the statistical condition number. In (Gitman et al.) 2019), a last iterate rate of SGD momentum on quadratic objectives was presented, but the convergence rate asymptotic. Non-asymptotic linear distributional convergence was shown in (Can et al., 2019), where SHB with constant learning rates achieves accelerated linear rates of Wasserstein Distances between distributions. However, this does not imply linear convergence in excess risks, where the variance is still a non-convergent constant term. In (Mai and Johansson 2020), class of weakly convex objectives were studied and convergence rate of Okk/VT) was established for gradient L2 norm. In (Wang et al.) 2021), HB on GD is analyzed and shown to yield non-trivial speedup on quadratic objectives and two overparameterized models. However, the analysis was done in GD instead of SGD. In (Bollapragada et al.) 2022). SHB was shown to have a linear convergence rate 1/VK with standard constant stepsize and large enough batch size on finite-sum quadratic problems. Their analysis. however, was based on an extra assumption on the sample method. (Tang et al. 2023) proved SHB converges to a neighborhood of the global minimum faster than SGD on quadratic target functions using constant stepsize. In (Yuan et al., 2021). modified decentralized SGD momentum algorithm was proposed for large-batch deep training. Although achieves Ö(/TT) convergence rate on a T-smooth and jı-strongly convex objectives, it still requires at least (K)) number of iterations to converge, which is no better than SGD. Wang et al. (2023) also provided cases where SHB fails to surpass SGD in small and medium batch size settings, suggesting that momentum cannot help reduce variance. There are also other variants of momentum such as Nesterov momentum (Nesterov) 2003 Liu and Belkin 2019: Aybat et al. 2019), modified heavy ball, but since we only consider the common version of heavy ball momentum here, we omit them in context.
-T/vk) in terms
3

---

Published as a conference paper at ICLR 2024
SGD + HB learning rate schedules: As for SHB with learning rate schedules, only limited amount of research has been conducted so far. In (Liu et al. 2020), the convergence property of L-smooth objectives
SHB with multistage learning rate schedule the inverse relationship between the stage length and learning rate size was implicitly assumed, thus its convergence rate is actually ( /og/ T) for some constant a > 1. In (Jin et al. 2022). a convergence rate was derived for general smooth objectives. But the relationship between the convergence rate and T is still unclear, and the results were comparing SGD and SHB by their upper bounds. In (Li et al. 2022). worst-case lower bound of T/VTT) was found for SHB with certain choices of step sizes and momentum factors on Lipschitz and convex functions. A FTRL-based SGD momentum method was then proposed to improve SHB and achieve ( convergence rate for unconstrained convex objectives. Furthermore, in (Wang and Johansson 2021 (LIVVT) bound was derived on general smooth non-convex objectives. whose analysis supports more general class of non-monotonic and cyclic learning rate schedules. All these results only proved that SHB is no worse than SGD, or were comparing two methods by their upper bounds instead of lower bound against upper bound. Only until recently has SHB been shown to be superior over SGD in some settings. In (Zeng et al. 2003), a modified adaptive heavy-ball momentum method was applied to solve linear systems and achieved better performance than direct application of SGD. In (Sebbouh al. 20211). SHB was shown to have convergence rate arbitrarily close o(1 ((/vTT) on smooth convex objectives. However, the analysis stopped at this asymptotic bound and did not provide any practical implications of this result.
investigated. However,
In contrast to all the aforementioned works, we provide positive results in theory to back up SHB's superior empirical performance, showing that SHB yield accelerated convergence quadratic objectives by equipping with large batch sizes and step decay learning rate schedules.
3
MAIN THEORY
3.1 PROBLEM SETUP
In this paper, we analyze quadratic objectives with the following form,
min f())((w)))))) - Eg [f(w. [),, where f(w, {) 2 z H(E) b(E) w,
(3.1)
where & denotes the data sample. By setting gradient to O, the optimum f(w) obtained at
w. H Hb,,, where H = Eg H()), b = Egg [b(E)]
(3.2)
In addition, we denote the smallest/largest eigenvalue and condition number of the Hessian H to be
Amin H)), L Amax (H),
K - L/P.
(3.3)
where eigenvalues from largest to smallest are denoted as
= \i 2 22 -== 0.
We consider the standard stochastic approximation framework (Kushner and Clark 2012) and denote the gradient noise to be
([(w,) - V &w,,
(3.4)
Throughout the paper, the following assumptions are adopted. Assumption 1. (Independent gradient noise)
{nr} are pairwise independent.
(3.5)
Assumption 2. (Unbiased gradient noise)
E [n = 0.
(3.6)
Assumption 3. (Anisotropic gradient noise)
E n 3 HH.
(3.7)
4

---

Published as a conference paper at ICLR 2024
The anisotropic gradient noise assumption has been adopted by several past literatures (Dieuleveut et al. 2017; Pan et al. (022)) along with evidence supported in (Zhu et al. 2019 Sagun et al. 2018 Zhang et al. 2018 Jastzzbbkk et al., 2017 Wu et al. 2022b), which suggest that gradient noise covariance is normally close to the Hessian in neural networks training.
Let B, be the minibatch of samples iteration t. For simplicity, only consider the setting where all minibatches share the same batch size
B. = M, for VI 0,
1.
(3.8)
follows that the number of samples is N MT.
Remark 1. One may also employ the common assumptions strongly corvex least square regres- sions as (Bach and Moulines 2013 Jain et al. 2018a Ge et al. 2019):
min f(w), where f(w) def L (x))D [(y - (x, w)) and 1
(1) y w x + e, where B(x,)/D [ [xxxx] ]] 3 H
(3.9)
(2) E
X xx
R'H
which can also be translated into our settings under the compact set constraint w A, as suggested in (Jain et al. (018a]]
3.2 SUBOPTIMALITY OF SGD
We begin with the vanilla version of SGD,
V Fwf,Ww,))
W++1 W.
[B
(3.10)
CEB
whose theoretical property is well understood on quadratic objectives (Bach and Moulines 2013 Jain et al. 2018b] Ge et al.) 2019 Pan al. 202T). Here 7e means the learning rate at iteration
known that SGD requires least [2(x) iterations under the setting of batch size M 1 (Jain et al. 2018a), nevertheless, whether this lower bound still holds for large batch settings is not rigorously claimed yet. Here provide Theorem T to make things clearer.
Theorem 1. There exist quadratic objectives f(w) and initialization Wo. no matter how large the batch size is what learning rate scheduler is used, as long as ne S 2/ for Vt 0,1, l......; T running SGD for T iterations will result in
E [ff(wr) - f(w.)l f(wo) f(w.)
87 exp
K
The proof is available in Appendix B The existence of those counterexamples suggests that in the worst case, SGD requires least T 2 k/8 In(c/2) S(x ((<loo c) iterations to reduce the excess risk by a factor of c > 2, while in practice, can be quite large near the converged point (Sagun et al.) 2017
Arjevani and Field 2020: Yao al. 2020).
3.3
ACCELERATION WITH STOCHASTIC HEAVY BALL
To overcome this limitation, heavy-ball momentum (Polyak 1964) is normally adopted by engineers speed up SGD, equipped with various types of learning rate schedulers
Ve+1 = Bv, +
Σ Vwf(ww..
B.
EEB:
(3.11)
W2+1 W. Vr+1.
Despite its huge success in practice, the theoretical understanding of this method still limited, especially for quadratic objectives. Furthermore, although it was widely recognized that stochastic heavy ball should provide acceleration in large batch settings, positive theoretical results so far are still insufficient clearly account for that. We attempt to fill this gap.
5

---

Published as a conference paper at ICLR 2024
In this section, will show that SHB equipped with proper learning rate schedules can indeed speed up large batch training. The whole analysis is done in a general multistage learning rate scheduler framework, shown in Algorithm Specifically, in this framework, learning rates divided into stages, with each stages' learning rates and number of iterations being n/é and kg - K respectively, ie..
( (stan)
A K(E - 1); f A Ke
(3.12)
T)t = T/e for Vt 1 (
(start) 1, (cnd)
Algorithm Multistage Stochastic Heavy Ball with minibatch
Input: Number of stages learning rates 1 momentum 3, stage lengths K, minibatch size M, initialization Wo E Rd and Vo O.
1: 0
Iteration counter
2: for = .., do
3:
1/e <- 1/6
4:
for 1,2..... K do
5:
Sample minibatch B uniformly from the training data
6: 8: f M EçES (w,(w,
>> Mean gradient over a minibatch
7:
Ve+1 + 8V: + 1).
8:
W(+1 + We Vi+1
9:
+
10:
return W,
b. Last iterate
Given the above step decay scheduler, the following theorem states the convergence rate for SHB on quadratic objectives. To the best of our knowledge, this is the first non-asymptotic result that explicitly expresses the relationship between T and the convergence quadratic objectives.
of mutlistage SHB on
Theorem 2. Given quadratic objective f(w) and a step decay learning rate scheduler with (1 1/VR) with 2 4, and n. T/K with settings that
1. decay factor C
1 C<C
(3.13)
2. stepsize T/e
1 C
1/6 T,
(3.14)
3. stage length K
T
K =
(3.15)
loga (T VK)
4. total iteration number T
T
In (21478). In (26T4) logo(()) 20
(3.16)
then such scheduler exists, and the output of Algorithm 7 ssttisfies
27
E(((wT) - s(w-)] [E[[(w)) [[w=)] [[w=)]]eum) exp 15ln2 + 21n 2
Vx logo (T kR)
do In (2"T*) loge (TvR).
MT
Or equivalently, the result be simplified the following corollary.

---

Published as a conference paper at ICLR 2024
Given quadratic objective f(w) and step decay learning rate scheduler and momentum defined in Theorem
Corollary 3.
2
with T > (V and 4, the output of Algorithm satisfies
E[f(wT) f(w.)] S (f(w)) - f(w.)] exp (-
T VR
do²
+0
MT
where (O)) and (2.)) are used to hide the log factors.
Notice that the bias term [fw) f(w.)l ) (Xp(-Q((T/VR) is exponentially decreasing after T (VK) iterations, while the variance term can be bounded by ((/TT) This implies that under the large batch setting, the batch size is large enough to counteract the extra constant in the variance term, accelerated convergence will be possible as compared the iteration number of (x) required by SGD. It is worth noting that this é (k) acceleration is only log factors away from the optimal acceleration (Nemirovski] 1995) of Heavy Ball (Polyak) 1964) and Nesterov Accelerated Gradient (Nesterov 1983) in deterministic case.
The proof outline can be split into two major steps. The first step is bias- decomposition, which decomposes the expected excess risk [[r(wr)] ( -((w.) into two terms: bias and variance, where bias measures the deterministic convergence gradient noise. This step adapts the well-known vias-variance decomposition technique of SGD (Bach) and Moulines 2013 Jain et al.) 20186 Ge et al. 20T9 Pan et al., 202T) to SHB. Inside the adapted decomposition, critical "contraction" term (T 22-2...1 Tol is introduced in both bias and variance, where each matrix T, R2x2 depends on step size Ne and differs only by a diagonal matrix
and variance measures the effect of the
A, T, To diag(ór, 0).
The second major step is to bound the matrix product tightly. Notice that this term has a form of - (1 7:);) for SGD and is much easier to analyze. For the general form of IT Toll [T + AT + 2r Toll. the major difficulty arises from the non-commtativve matrix products of different T,'s s and A, ss. To overcome this obstacle, novel technique is proposed in our paper, which based on the special structure T,, The key observation is that product with form (T,, A Tu, A., ... T,, Ag R retains two important properties: 1) The first column is always nonnegative and second column is always nonpositive; 2) The absolute value of each entry is a monotonical increasing function of 8, ..,.... ST-1- Hence the sum of the exponential number of terms in the binomial-like expansion also retains those two properties, which leads to a bound [/T-iTTT-2- Toll < |T tight under certain conditions. This key technique. as rigorously stated in Lemma in Appendix, combined with subtle analysis of T, and learning rate schedule techniques in (Ge et al. 2019 Pan et al. 2021). gives birth to Theorem The full detail of the proof is provided Appendix
4 EXPERIMENTS
To verify our theoretical findings, two sets of experiments are conducted. The first one is ridge regression, which has quadratic loss objective and closer our theoretical settings. The second one is image classification
CIFAR- 10 (Krizhevsky et al. (2099) with ResNet18 (He et al. 2015). DenseNet121 (Huang et al., 2017) and MobilenetV2 (Sandler et al. 2018), which is more practical interest regarding our theory's potential applications.
4.1 RIDGE REGRESSION
In ridge regression, we consider the following setting
f(w) 1
Xw Y12 [//l
(4.1)
whose optimum has an analytic form
(X X nal) maa)a'''' x
7

---

Published as a conference paper at ICLR 2024
Therefore the optimum loss f(w.) can be directly computed. We a4 dataset CChagg and Lin 2011 Dua and Graff 2017) to realize this setting, which contains n = 4, 781 samples and 123 features.
In all of our experiments, we set the number of epochs to 100, so the total amount of data is N 478, 100. Besides, set different batch sizes M (2048, 512, 128), , and initialize Wo from a uniform distribution 1). The partial batch at the end of each epoch is not truncated, which means the total number of iterations T [N/M]]
Regarding hyperparameter choices for each scheduler & method, we do grid searches according to Table 3 in Appendix
A and report the best loss for each random seed. For all schedulers, 10- }.. As for the choice of momentum factor 8, we set B = 0.9 for
set E {10%, 10 10 10 stochastic heavy ball methods.
Table Training loss statistics of ridge regression in a dataset over runs.
Methods/Schedules
(f(w) - f(w.)) x 10-2
Batch size M = 512
M = 128
M = 32
M = 8
SGD + constant Y/:
2.10:10.46
1.17110.1
1.27.10.27
0.94.10.83
SGD + step decay
2.44:10.45
0.64:10.04
0.11110.11
0.04.10.04
SHB + constant 1/+
0,86+10.55
0.55+0.26
1.03:#0.35
0.97:10.58
SHB step decay
0.13:10.03
0.01+10.00
0.03:10.02
0.06:10.05
As shown in Table I one can observe that SHB are generally much better than SGD under large batch settings, and the step decay schedule always helps. The role of learning rate schedule and heavy-ball momentum is especially evident under the setting of M = 512, where SHB is able to greatly reduce the loss with a much smaller bias, but still has large loss due the existence of variance. This variance term is then further handled by step decay schedule and leads to fast convergence. As the batch size decreases, the variance term becomes dominant, which explains the closing gap between SGD and SHB.
4.2 IMAGE CLASSIFICATION ON CIFAR-10
In image classification, our key focus is still verifying the superiority of SHB over SGD, so no heavy tuning was done for 3. We follow the common practice of B 0.9 for our algorithm in Theorem 2 To simulate the practical settings of distributed learning and federated learning, we restrict the number of iterations to be few thousands (Kairouz et al., 2021), D, which roughly translated into # Epoch = 10 for batch size M = 128 and # Eppoch = 100 for batch size M = 2048. On top of that, for batch size M 2048, we replicate 16 nodes with micro batch size 128 on each node, hence the performance on distributed learning can be further simulated.
In this experiment, CIFAR-I0 (Kiizeessy al., 2009) dataset is adopted, which contains 50, 000 training samples and 10, 000 test samples. We use 5, 000 randomly chosen samples in the training set to form a validation set, then conduct grid searches by training on the remaining 45, 000 samples and selecting the hyperparameter with the best validation accuracy. The selected hyperparameter is then used for training the whole 50, 000 samples and testing on the test set. The final test results are thereby summarized in Table For grid searches, we choose learning rate 7o E (1, 0.0,,.01,, 0.001). with decay rate y E (1/2, 1/5, 1/10} and number of intervals n E {3, 4, 5, 6).
One can observe in Table 2 and Figure [ that under the large batch setting, SHB provides huge acceleration over SGD and achieves a significant performance improvement. This offers empirical evidence for our theory and suggests its practical value: Heavy Ball Momentum can provide true acceleration for SGD under large-batch settings.
'The dataset is accessible in htops: www.csie.... ntu wttpeettt......oii..iiee.....e.e wii.igiiiiiiii oolmtooll / html #a4a/
8

---

Published as a conference paper at ICLR 2024
Table 2: CIFAR-10: training losses and test accuracy of different methods over trials.
Setting
Method
Resnet18
DenseNet121
MobilenetV2
Crossent. Loss
Cmssent, Loss
Crossent.
Acc(%)
Aco(%)
Loss
128 (# Epoch 10)
SGD
0.46::0.01
81.19:00.93
0.22:1:0001
$8.58:1:0/23
0.45:20.000
82.90:20.37
SIBB
0.38.1..00 0.08
2..3
0,181
88.63:10.27
0.35:10.01
86.23....25
0.33
03.82:00.42
0.01+10.00
89.28:1:0233
0.32:4-002
84.37:10,77
(# Epoch = 100)
SIIB
0.01:10.00
89.78:10.23
0.00:10.000
92.46:10..15
0.07:10.01
89.57 0.18
momentum-0
momentum-
2.5
nomentum--
nomentum--0.9
300
1306
1530
2000
2000
30
1089
1306
2009
3500
momentum- momentum-0.9
momentum-0
momentum-0099
1300
1530
1063
390
momentum momontum-0.3
nomentum-0
nomentumm-0.9
TARDA 1530
2206
350
Figure 1: CIFAR- 10 training statistics of batch size M = 128 x 16 and # Epoch = 100 on Resnet18, 88, DenseNet121 and MobilenetV2 (from top bottom). Left: Training loss; Right: Test accuracy.
CONCLUSION
In this paper, we present non-asymptotic convergence rate for Stochastic Heavy Ball with step decay learning rate schedules on quadratic objectives. The proposed result demonstrates SHB's ' superiority over SGD under large-batch settings. To the best of knowledge, this the first time that the convergence rate of SHB is explicitly expressed in terms of iteration number T given decaying learning rates on quadratic objectives. Theoretically, our analysis provides techniques general enough to analyze any multi-stage schedulers with SHB on quadratics. Empirically, we demonstrate the practical benefits of heavy-ball momentum for accelerating large-batch training, which matches our theoretical prediction and explains heavy-ball momentum's 's effectiveness in practice to certain degree.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
Rui Pan acknowledges support from the Hong Kong PhD Fellowship Scheme (HKPFS).
REFERENCES
Yossi Arjevani and Michael Field. Analytic characterization of the hessian in shallow relu models: A tale of symmetry, 2020.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. Advances in neural information processing systems, 32, 2019.
Necdet Serhat Aybat, Alireza Fallah, Mert Gurbuzbalaban, and Asuman Ozdaglar. A universally optimal multistage accelerated stochastic gradient method. Advances neural information processing systems, 32, 2019.
Francis Bach and Eric Moulines. Non-stronglyycoovex smooth stochastic approximation with
convergence rate o (1/n). Advances in neural information processing systems, 26, 2013.
Raghu Bollapragada, Tyler Chen, and Rachel Ward. On the fast convergence of minibatch heavy ball momentum. arXiv preprint arXiv:2206. 07553, 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal,
Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Ben- jamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901, Curran Associates, Inc., 2020, URL httph: http::::s///pppeeeepeeeee newrips pdff
cc (c///aeerr/222//11e/111176660662 96 ootpe//ccccccctooo.ooaaoaaaaaeea htt.r//aaac12ffffaaaa
Bugra Can, Mert Gurbuzbalaban, and Lingjiong Zhu. Accelerated linear convergence of stochastic momentum methods in wasserstein distances. In International Conference on Machine Learning, pages 891-901. PMLR, 2019.
Chih-Chung Chang and Chih-Jen Lin. Libsvm: a library for support vector machines. ACM transactions on intelligent systems and technology (TIST), 2(3): 2011.
Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks. In 2018 Information Theory and Applications Workshop (ITA), pages 1-10. IEEE, 2018.
Damek Davis, Dmitriy Drusvyatskiy, and Vasilcios Charisopoulos. Stochastic algorithms with geometric step decay converge linearly sharp functions. Mathematical Programming, pages 1-46, 2023.
Jacob Devlin, Ming-Wai Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume I (Long and Short Papers), pages 4171-4188 2019.
Olivier Devolder, François Glincur, Yurüi Nesterov, et al. First-order methods with inexact oracle: the strongly convex case. CORE Discussion Papers, 2013016:47, 2013.
Aymeric Dieuleveut, Nicolas Flammarion, and Francis Bach. Harder, better, faster, stronger con- vergence rates for least-squares regression. The Journal of Machine Learning Research, 18(1): 3520-3570, 2017.
Dheeru Dua and Casey Graff. UCI machine learning repository, 2017. URL http:: ics. uci edduduum mml
10

---

Published as a conference paper at ICLR 2024
Rong Ge, Sham M Kakade, Rahul Kidambi, and Praneeth Netrapalli. The step decay schedule: A near optimal. geometrically decaying learning rate procedure for least squares. In Advances in Neural Information Processing Systems, pages 14977-14988, 2019. 2019.
Saeed Ghadimi and Guanghui Lan. Optimal stochastic approximation algorithms for strongly convex stochastic composite optimization, it: shrinking procedures and optimal algorithms. SIAM Journal on Optimization, 23(4):2061 2013.
Igor Gitman, Hunter Lang, Pengchuan Zhang. and Lin Xiao. Understanding the role of momentum in stochastic gradient methods. Advances in Neural Information Processing Systems, 32, 2019.
Nicholas JA Harvey, Christopher Liaw, Yaniv Plan, and Sikander Randhawa. Tight analyses for non-smooth stochastic gradient descent. In Conference on Learning Theory, pages 1579-1613.
PMLR, 2019.
Elad Hazan and Satyen Kale. Beyond the regret minimization barrier: optimal algorithms for stochastic strongly-convex optimization. The Journal of Machine Learning Research, 15(1): 2489-2512, 2014.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015.
Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classification. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pages 328-339, 2018.
Gao Huang. Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4700-4708, 2017.
Peter Izsak, Moshe Berchanshhhhheeeeeeeeeeeeeaeceee and Omer Levy. How to train bert with an academic budget, 2021. URL http:///wxvv..rrrr/aa///00000000
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
Prateek Jain, Sham M Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Accelerating stochastic gradient descent for least squares regression. In Conference On Learning Theory, pages 545-604. PMLR,
2a.
Prateek Jain, Sham M Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Parallelizing stochastic gradient descent for least squares regression: mini-batching, averaging, and model misspecification. Journal of machine leurning research, (8(223):1-43 2018b.
Prateck Jain, Dheeraj Nagaraj, and Praneeth Netrapalli. Making the last iterate of sgd information theoretically optimal. In Conference on Learning Theory, pages 1752-1755. PMLR, 2019.
Stanislaw Jastrzębski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint arXiv: 1.0622,, 2017.
Ruinan Jin, Yu Xing, and Xingkang He. On the convergence of msgd and adagrad for stochastic optimization, 2022.
220p.///....//.//21..... 11204
Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael G. D'Oliveira, Hubert Eichner, Salim El Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascón, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser, Zaid Harchaoui, Chaoyang He, Lie He, Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi. Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konecný, Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tancrède Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgür, Rasmus Pagh, Hang Qi, Daniel Ramage, Ramesh Raskar, Mariana Raykova, Dawn Song. Weikang Song, Sehastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramer, Praneeth Vepakomma,

---

Published as a conference paper at ICLR 2024
Jianyu Wang. Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and open problems in federated learning. Foundations and Trends@ in Machine Learning, 14(1-2): 1-2 2021. ISSN 1935-8237. 1561 /2600000000083
doi: 10.1561/22000000083. URL httt:://wx..o..o00/..0 www.pxoo.......
Rahul Kidambi, Praneeth Netrapalli, Prateek Jain, and Sham M. Kakade. On the insufficiency of existing momentum schemes for stochastic optimization, 2018. URL http:tts///xaii..orrrttt
abs/1803 05591
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
Andrei Kulunchakov and Julien Mairal. A generic acceleration framework for stochastic composite optimization. Advances in Neural Information Processing Systems, 32, 2019.
Harold Joseph Kushner and Dean S Clark. Stochastic approximation methods for constrained and unconstrained systems, volume 26. Springer Science & Business Media, 2012.
Xiaoyu Li, Mingrui Liu, and Francesco Orabona. On the last iterate convergence of momentum methods. In Sanjoy Dasgupta and Nika Haghtalab, editors, Proceedings of The 33rd International Conference on Algorithmic Learning Theory, volume 167 of Proceedings of Machine Learning mlr
Research, pages 699 717. PMLR, 29 Mar-01 Apr 2022. URL URLhhttps://procediii
Cress/v167/1122 html
Chaoyue Liu and Mikhail Belkin. Accelerating sgd with momentum for over-parameterized learning. In International Conference on Learning Representations, 2019.
Yanli Liu, Yuan Gao, and Wotao Yin. An improved analysis of stochastic gradient descent with momentum. Advances in Neural Information Processing Systems, 18261-118,,,,, 2020.
Nicolas Loizou and Peter Richtárik. Linearly convergent stochastic heavy ball method for minimizing generalization error. arXiv preprint arXiv:1710.10733 2017.
Nicolas Loizou and Peter Richtárik. Momentum and stochastic momentum for stochastic gradi- ent, newton, proximal point and subspace descent methods. Computational Optimization and Applications, 77(3):653-710, 2020.
Nicolas Loizou, Sharan Vaswani, Issam Hadj Laradji, and Simon Lacoste-Julien. Stochastic polyak step-size for sgd: An adaptive learning rate for fast convergence. In Arindam Banerjee and Kenji Fukumizu, editors, Proceedings The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pages 1306-133 1314, PMLR, 13-15 Apr 2021. URL https: https:::/proceedingg .mlr. pttps///////oooooo2222222
Ilya Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with warm restarts. In 5th International Conference Learning Representations, ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings, 2017.
Vien Mai and Mikael Johansson. Convergence of stochastic gradient method with momentum for non-smooth non-convex optimization. International Conference on Machine Learning, pages 6630-6639. PMLR, 2020.
Arkadi Nemirovski. information-basel complexity of convex programming. Lecture notes, 834, 1995.
Yurii Nesterov. A method for unconstrained convex minimization problem with the rate of conver- gence o (1/k") 2). Doklady an ussr, volume 269, pages 543-547, 1983.
Yurii Nesterov, Introductory lectures on convex optimization: A basic course, volume 87. Springer Science & Business Media, 2003.
Long Ouyang, Jeffrey Wu, Xu Jiang. Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35: 27730 2022.
12

---

Published as a conference paper at ICLR 2024
Rui Pan, Haishan Ye, and Tong Zhang. Eigeneurve: Optimal learning rate schedule for sgd on quadratic objectives with skewed hessian spectrums. In International Conference on Learning Representations, 2021.
Rui Pan, Shizhe Diao, Jianlin Chen, and Tong Zhang. Extremebert: toolkit for accelerating pretraining of customized bert. arXiv preprint arXiv:2211. 2022.
B.T. Polyak. Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 4(5):1-17, 1964. ISSN 0041-5553. doi: http::///..oogg///...000000111/1
URL wwt.tttsss.. www www.sccienceirectt
com/ http://nnn//rrrceeeeee/////11333
BT Polyak. Introduction to optimization, optimization software, inc. publ. division, 1987.
Xun Qian, Peter Richtárik, and Tong Zhang. Error compensated distributed sgd can be accelerated. Advances in Neural Information Processing Systems, 34, 2021.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pages 400-407, 1951.
Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond, 2017.
Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis the hessian of over-parametrized neural networks. CoRR, 044/1706600454, 2018.
Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen. Mo- bilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4510-4520, 2018.
Othmane Sebbouh, Robert M Gower, and Aaron Defazio. Almost sure convergence rates for stochastic gradient descent and stochastic heavy ball. In Conference on Learning Theory, pages 3935-3977 PMLR, 2021.
Ohad Shamir and Tong Zhang. Stochastic gradient descent for non-smooth optimization: Convergence results and optimal averaging schemes. In International conference on machine learning, pages 71-79. PMLR, 2013.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition, 2015.
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Du- mitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1-9, 2015.
Kejie Tang, Weidong Liu, and Yichen Zhang. Acceleration of stochastic gradient descent with momen- tum by averaging: finite-sample rates and asymptotic normality, arXiv preprint arXiv:2305.17665 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Maric-Anne Lachaux, Timothéc Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. / 13971, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint
2023b.
Sharan Vaswani, Aaron Mishkin, Issam Laradji, Mark Schmidt, Gauthier Gidel, and Simon Lacoste- Julien. Painless stochastic gradient: Interpolation, line-search, and convergence rates. Advances in neural information processing systems, 32, 2019.
Joost Verbraeken, Matthijs Wolting, Jonathan Katzy, Jeroen Kloppenburg, Tim Verbelen, and Jan S Rellermeyer. A survey on distributed machine learning. Acm computing surveys (csur), 53(2): 1-33, 2020.
13

---

Published as a conference paper at ICLR 2024
Jun-Kun Wang, Chi-Heng Lin, and Jacob D Abemethy. A modular analysis of provable acceleration via polyak's s momentum: Training a wide relu network and a deep linear network. In International Conference Machine Learning, pages 10816-10827. PMLR, 2021.
Runzhe Wang, Sadhika Malladi, Tianhao Wang. Kaifeng Lyu, and Zhiyuan Li. The marginal value of momentum for small learning rate sgd. arXiv preprint 2577 2023.
Xiaoyu Wang and Mikael Johansson. Bandwidthh-based step-sszes for non-convex stochastic opti- mization, 2021. htts:tps::/ //tpxiiv///xii.////aa/aa wtggabs//g/aaaeeo 02888
Alexander Wettig, Tianyu Gao, Zexuan Zhong, and Danqi Chen. Should you mask U5 https
//axxiv.org/ab//2222 08005
Florian Wolf. Stochastic gradient descent and its application for parametrized boundary value problems under uncertainties. 2021.
Jingfeng Wu, Difan Zou, Vladimir Braverman, Quanquan Gu, and Sham Kakade. Last iterate risk bounds of sgd with decaying stepsize for overparameterized linear regression. In International Conference on Machine Learning, pages 01280-24314 PMLR, 2022a.
Lei Wu, Mingze Wang, and Weijie Su. The alignment property of sgd noise and how helps select flat minima: A stability analysis. Advances in Neural Information Processing Systems, 35:4680 4693, 2022b.
Yi Xu, Qihang Lin, and Tianbao Yang. Accelerate stochastic subgradient method by leveraging local error bound. CoRR, abs/1667.00002) 2016.
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE international conference on big data (Big data), pages 581-590 IEEE, 2020.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arvi /708. 03888, 2017.
Yang You, Zhao Zhang, Cho-Jui Hsieh, James Demmel, and Kurt Keutzer. Imagenet training in minutes. In Proceedings of the 47th International Conference Parallel Processing, pages 1-10, 2018.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations, 2019.
Kun Yuan, Bicheng Ying, and Ali H Sayed. On the influence of momentum acceleration on online learning. The Journal of Machine Learning Research, 17(1):6602-666 2016.
Kun Yuan, Yiming Chen, Xinmeng Huang, Yingya Zhang, Pan Pan, Yinghui Xu, and Wotao Yin. Decentlam: Decentralized momentum sgd for large-batch deep training. In Proceedings of the IEFF/CVF International Conference on Computer Vision, pages 3029-3039, 2021.
Zhuoning Yuan, Yan Yan, Rong Jin, and Tianbao Yang. Stagewise training accelerates convergence of testing error over sgd. Advances in Neural Information Processing Systems, 32, 2019.
Yun Zeng, Deren Han, Yansheng Su, and Jiaxin Xie. On adaptive stochastic heavy ball momentum for solving linear systems. arXiv preprint arXiv:2305.05440 055-48 2023.
Yao Zhang, Andrew M Saxe, Madhu S Advani, and Alpha A Lee. Energy-entropy competition and the effectiveness of stochastic gradient descent in machine learning. Molecular Physics, 116 (21-22):3214 32233 2018
Shuai Zheng, Ziyue Huang, and James Kwok. Communicaiion-effcieenn distributed blockwise momentum sgd with error-feedback. Advances in Neural Information Processing Systems, 32, 2019.
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In International Conference on Machine Learning, pages 7654-7663. PMLR, 2019.
14

---

Published as a conference paper at ICLR 2024
A EXPERIMENTAL DETAILS
The detailed learning rate schedule forms and grid search hyperparameters for ridge regression in Section 4. are listed in Table
Table 3: Grid-search hyperparameter choices for ridge regression.
Scheduler
Form
Hyperparameter choices
Constant
Yu
11 70. ift [2,& 1).
E 3, 4,5] Y E 22 41 8J
Step decay
B PROOF FOR SECTION 3.2; SUBOPTIMALITY OF SGD
Theorem 1. There exist quadratic objectives f(w) and initialization Wo, no matter how large the
batch size is or what learning rate scheduler is used, as long as S 2/ L for Vt 0, .......? T 1, running SGD for T iterations will result in
E [[((T)) - f(w.)) f(wo) ((w.)
87 exp
Proof. Consider the case of x 2 wo = co I. d K 1, H() H
diag M.H. ,........ 14) and
b(E)= = 0, then according to SGD's ' update formula in Eqn. 3.10
Wi+1
Th B
V (w.,,
T/p H)w::
CEB.
we have the update of j-th (j 2 2) entry being
We+1, =(1 - wu. = (1 - m))))rr 2 (1 -1-1111
24
Wt. =
wtj 2 exp
Wt.j
L
where the first inequality comes from 7) S 2/L. L and the second inequality is entailed by 1 x 2 exp(-2x) for x = 2/K [o, 1/2], since for Vx E [0, 1/2]
Ög(x)
g(x) & Im(1 - => - In(exp(-22)) 2x)),
g = 0,
1 1 +220
dx
g(z) 2 g(0) = o In(1 - x) 2 - -2x))
7 - x 2 exp(-221) -22).
It follows,
d Lw + 11
1
f(wr) f(w.) =f(wr)
w. HwT
Σ w
2
2=2
d
d
T
WT
2
=
- m)
2 j=2
1=2
1=
11
d
87
D.j
exp
2 j=2
11 (d - 1) 2 exp
8T
pld 1). S(wo) - f(w.)
L + p(d -1) -xp (- 87
(-x K
(_87)
K
f(wo) f(w.)
87
m]d1) + exp
f(wo) -((w) exp
87
K
15

---

Published as a conference paper at ICLR 2024
C PROOF FOR SECTION 3.3, ACCELERATION WITH STOCHASTIC HEAVY BALL
We provide the broad stroke of our proofs here. The ultimate goal is bounding ( in terms of T, x and other fixed model parameters. To achieve this goal, we conduct the whole analysis in layer-by-layer fashion. each layer, we translate the terms in the previous layer into the terms of the current layer, specifically
Appendix C.I
[[ - f(w.)l +
| for some special 2 x 2 matrix Tr, via
bias-variance decomposition.
Appendix
IITetiTTe42.-. Trtkoll ((T++RJ)* via analyzing the property of | and 1411.:T442... T(tkj]
C.2
Appendix
C.3
= {y. [3, 25} via analyzing the property of p(Tx,)) p(Tx,))
Appendix C.4
8,, [,,) = fite 8, A3} (T, K........ via specializing the scheduler with
multistage scheduler step decay.
C.I BIAS VARIANCE DECOMPOSITION: [[f(w) - f(w.)) = |
The section presents the lemmas that decompose the loss into bias and variance, expressing them in terms of the norm of product of series of special 2 x 2 matrices IIII, Te, Il's, specifically [ Tr-2,5 T for some T 0, ......T 1. For simplicity, we only consider the case of batch size M 1. The case of larger batches equivalent replacing noise term o² with o² /M..
Lemma 1. (Quadratic excess risk is H-norm) Given quadratic objective f(w), we have
f(w) f(w.) 2 ||w -w,
w.ll H 2
A (w -.. [[(w-w) -
(C.1)
Proof. It holds that
1 W. H(g)w. D(E) w.
S(w) f(w.)
EE
H(E) T6E) E.
[[)) R[b(E)]
W. (HH)))) E[b(E)] w.
Hw bTw)
2 w. Hw. T w_
2
Hw (Hw.)
w. Hw. (Hw.) Tw.
w Hw w Hw. W. Hw + w Hw.]
(w w.) H(w W.).
Here the fifth equality comes from H H being symmetric matrix, and w. Hw (w. Hw)" =w Hw, being a scalar.
Lemma 2. (Bias-variance decomposition) Given a quadratic function f(w), if batch size M 1, after running iterations of Algorithm 7 we have the last iterate Wr satisfying
E
WT w.lii
Wo
w. w.
M M.
MT
M Mr- My
Wo - w.
O
O
w-1 w.
(C.2)
T-1
+
2 n
M M. ......
M
H
O
0
0
Mr- M
n
0=0
16

---

Published as a conference paper at ICLR 2024
where W-1 Wo
Vo and
M -
(1
+
8)I
H
(C.3)
0
Proof. Denote
W. We
Rd2d
n,
n 0
R2dd
H
H
0 O
R 2dx2d
(C.4)
W w.
O
as the extended parameter difference, extended gradient noise and extended Hessian matrix, then according to momentum's 's update formula in Eqn. (3.11). or line 78 8 in Algorithm with batch size M 1, we have
V++1 =BV; + 71 V (w.:g)
W1+1 =W,- - Ve+1
WH11 W. -W; - W. V4+1
=W; W. (5v, + 7/t w/ /w/(ww (ww,,()
W. B (wr-1 w.) T/ V fw.,,
We w. B (Wi-1 1 wx) 71 (Vwj(w,) nn))
W; w. B (w,-1 w,) 1/1 (H (w, w.) n2))
We W. - B [[ (w,-1
w.) - (w, w.)l - 1/: (H (w:- - ww.) n.)
[11 + 3) The (w, w.) 31 (W,-1 w.) n/mn
W++1 W.
+ B)I H 3I
W. W.
n,
Wx41
0
Wr-1 w.
+ 1/
M.W. 7/.
It follows
W++1 -M,W; + n/
-M,M, W. Th- M + +
M.M + Th- M.MM, + 77
-M,M ... Mowo M,M (..M.-,...M,11..)
where We
Wo w. W-1 w.
and w-1 Wo + Vo which is associated with iteration = 0.
We decompose the above process into two parts
w ++1 -M,w
(b)
with (b) Wu + 1 with w (e) O,
(C.5)
w f+1 -M,w
(0)
since
w W41 = -M,M--1 Mow! = Mo
w (0)
(1.MM, M.+1
(C.6)
[11
(8)
(v)
W41
2+1
17

---

Published as a conference paper at ICLR 2024
Furthermore, we have
E
w 1+1
T
Hw c+1
[+1
(b)
W 1+1
C6
w ++1
(b)
(1, M..MM
(C.7)
(w+11
H
(n,MM, B
r=0
and
E
(v] [41
Hw
i+1
(0)
M,,M I.M.M.M....MM..... M-H1 )
E
H Σ
M,, .......-1hr.))
r'=0
E T=0,7'=0
M M M ...
(C.8)
Pu. M ........,HMMM HM ...
E B M 122 M, HM. M
-=0,r'=0,
n M. M 1 ... M. HMM ... ...MM,,,..M,,1,,, +0,
where the last equality is because
n- and i,."'
n-
are pairwise independent and have mean O given Assumption I and
Then obtain
E WT w.
WT
w.
H (WT W.) =
WT W.
H O
WT W.
WT-1 w.
WT-1 W.
W HWT
C.6
E
r
+
T
(v)
w T
(6)
(v)
E
w T
W T
+
ww(""
H
W T
+ 2
(T
H
w T
T
E
r
H
T
+
E
(w
H
w
E
W
M
M.
HMM MT ... .Mawo
M 42 M. HMr M..M....MMM11,,
Here the fourth equality is because (w extended terms w, ñ, H with their definitions in Eqn. (C.4) we get the desired form in Eqn. CC.
(8)
H (1) is scalar and H is symmetric. Replacing the
(+1
18

---

Published as a conference paper at ICLR 2024
Lemma 3. (Decomposing M, into block diagonal matrix) Given matrix M. R2xx2d defined in Eqn. (C.3) we have
Tr,1
M, VII
II VV
(C.9)
T r
where
1+ 3 - -
-8
T
E R2xx
(C.10)
1
0
and orthogonal matrices
0
E 2dd
V
U 0
E R
2d22d
(C.11)
C1
C2
ed
given the eigendeccompoition of being
H = UAU E Bddd
(U'U=))
(C.12)
and standard unit vectors/stannarrd basis being
e; = 00 0 th
0...0
00
E Rdx1
(C.13)
Proof.
II V M, VII
(1 + 3)1 - H
-81
VII
I
0
UT
(1 +
3)1
7
U
0 0
0
UT
I
0
0
(1+ B)I
7
-31
I
0
e
T
C11
(1+ + 3)I
n
-33
C1
C2
0
Cd
I
0
0
91
0
C2
0
Fd
cd
0
(1 + 8 n1d1)01
-0e1
(1 +
B
-nAaeed
-Bea
01
ed
0
e T
e
SS.1
S1.2
Si.d S2.
S2,1
S22
[1 3 m,dyee
jeTe;
-seTe,
where Sijj
e, ej
0
Sa,1
S
Saa
Tu.1
T2
C110
T T2d
19

---

Published as a conference paper at ICLR 2024
Since
0
e
CI
e
0
C.B
12dx2d
0
62
ed
0
(C.14)
0
e
MI
U
U
12xx2d
U
0
both orthogonal matrices, thereby have
T1,1
M, VII
II vv
T T
Lemma 4. (Bound Variance with | Assuming batch size 1, we have
T
T
n
M
M
H
0 O
My
Mr
M
n2
722
O
==0
(C.15)
2 TT-TTTTT-2
j=1
7=0
where M, defined in Eqn.
C.3 and Tei E 1g22 is defined in (C.10
Proof. Similarly, we define n and H as Eqn. matrix since if the eigenvalue decomposition of Hessian H
Notice that H a positive semi-definite
UAUT, have
H 0
O
0
U
0
H
0
U
0
0
0
U
Therefore, 1/2 well-defined. Denote
A, H M M,41,
(C.16)
then
A = M M
M
[ (2) = M -M,,M M ,
(C.17)
where the second equality is because 1/ is symmetric. follows
T
M.
MT
H 0 Mr- Mr+1
n-
0
T-
" M M M M MM M T.n-.....MM.... M.+1..)
==0
T-1 Σ
C.16
[ A A n.]
E tr ( A A.)
Trace of scalar is itself
Σ
E [ir (A,,,,,,AA)
Cyclic property of trace
Σ
nrt (E (E[A,ú,
> Linearity of expectation
20

---

Published as a conference paper at ICLR 2024
D1
E
A
T-1
7
A,E
n O
0 0
A
7=0
T-1
A
H2 0
A
T
0²
T-1
Σ
1
A
H
0 0
A
0
T=0
==0
H
0 0
1/
A A
Cyclic property of trace
0
T-
C
C
1/1
M M ... M M7- M,+1
0
0
0
T-1
tr
MM M M
7=0
where the inequality because
n
0
0 H
0
E
given E n 3
o 2HH in Assumption 3
along with basic properties of Loewner order in
Lemma 13 and 15
Let
Tt,1
Te,2
T, A
(C.18)
T
we have the variance term being
E
F/
n
0
M
M 2
M
H 0
O O
Mr
M
n, 0
T-1
<a²
1/r
(M.,M, . M H M M,
7
T-1
2
1//tr
(VIIT --1 T T T II V H T v H
T=0
tr
( (T
TT 7-2... ... T I v Hvπ) Tr-iTr-2.. T (I VVVV))
21

---

Published as a conference paper at ICLR 2024
Here the last equality comes from the eyclic property of trace. Given the definition of II, V and H in Eqn. (C.11) and Eqn. (C.4) we have
e
el
IVII
U
H
0
-0
0
62
ed
o
U
0
0
ed
O
e
A 0
C1
0
O 0
0
ed
c T
T1
0
22
0
A
0
Ad
(C.19)
Here & is the Kronecker product. Then the variance term is simplified to
E
n,
M
11+1 M 72 ... M
H
O
Mr-MMr 2.. Mr+1
n,
0
O
O
0
T-1
1
0
S 2
Σ
1r
T
T
T-
Tr- T TT ........ Tr+1
A
0
7+1
+22
0
0
C18
T
1-
tr
T T ... T
1.
A
0
TT- TT- .T..IT
0)
==0
/=1
All are block diagonal matrices
7F
T F
I
0
TT-1ŞTT-2 ... T
0
Ø=1
Σ
tr
(b]
TI
T +22 ... T
[
Tr-1 Tr TT- ... T,+15
[]]
01
0
Cyelic property of trace
T-
(
T
T (
T T
T
0
TT- Tr-2.
-=0
1=1
Notice that the term inside the trace is a scalar
T-1 [
PT-ITTT-2 Ti+1s 11]
2
2
X3
T=0
1=1
Transpose of scalar is itself
T
<
Σ
T
x3 IT- TT-25 T
22

---

Published as a conference paper at ICLR 2024
T-1
=o²
Σ j=1
2 Σ o
TI--1TT-22.
Trtisl. Trirslll
T=0
Here the last inequality is entailed by the fact that for vC E R2x22
( < (,) ())
(c !)
| 01 - <||||,
with (x)1, (x)2 standing for the first and second element of vector
Lemma 5, (Bound Bias with |TT
Wo w. w. M
M
H O
Mr-IMr-2
My
Wo W.
O
O
W-1 W.
(C.20)
IT IT- ... To
II V
Wo
W-1
W.
2j- 2:21:
where M, is defined in Eqn. CC3).
T E g2x2 is defined in C.10 and II, V are orthogonal
matrices defined in (C.1l. Here notation Zjzrjz means
21
2j
22
2j+1
For Vz
Rd i.sijsj sd
Zj:j' &
(C.21)
Proof. The proof is similar to a simplified version of the variance case, so we will reuse some of its notations to shorten the proof.
w. W w. 1 M, MT
M,M,
M
H 0 MT
Mo
w. W-1 w,
O
0
M HM, MT (.r-2...iw)) Miw0
VIIT
T II V THVIIITIITTT
T V wo]
=E
(wo VII
II
V
HVII
TI-ITr-
To (II V wo)
(wo VII To T
T
A
1
0
T
Tr-2
To III V wo)
0
0
Wo
S1
V wo III
1:2
122
wo
S
wo
C.I
V Wu
Sa
(II' T V Wo
2d 1:24
2d-1:22
D 0
where S. TT,
T
1.j
TT-15 15
0
TT 1, TT- ... To.
=E
Wo
21-1:22
S, (II
V
Wo) wo)2
A
0
where S,
T
T
15 T
Tr
T Tr-25- ... To
0
0
=
Wo
(M,s,s,)
[II
wo) 25-1:25
1:25
23

---

Published as a conference paper at ICLR 2024
where 8, = Tū, T
TT-1
1
=E Σ N
(IT'
Wo) 2j 1:2j 81) (s] III' V wo) 20)22-1123)
where s,
Toi
TT 0 [II
T
A
5
III
Wo 23 1:21
where S; -To.j -To,TII ... T I--
1
=E
Σ X
0
T TT- TT-2------. 2.... Tos (II'V wo)
<B
Σ
[T--, T
Tu.ll
II V Wo) 2j
A STr-I-T-22--1 ... To, 2 III V
wo) 20)23-1:23
> Linearity of expectation
C4
1T--ITTT-22 To
V
Wo W,
W-1 w. 2:-1:2)
Here the inequality is entailed by the fact that for vC E Rx2 z E R²
I
Cz (Cz), S /(zz) +Cz), = |Czz S
0
C.2 BOUNDING [TT4k. Te415 WITH p p(T(+k.))
This section upper bounds the matrix product (TTtk.j..TTT-.. with the spectral radius (T++1,5) Similar results for bounding ||TH), have been shown in Wang et al. 2021) (Theorem 5), but our result more general, so we still put our proofs here. In the following proof, we use ||-1lp to denote the Frobenius norm of matrices.
Lemma 6 (Bounding |Tbylle with ((T Given momentum matrices Tej that are defined in Eqn. (C.10) and 144, for all positive integer k 2 it holds that
8
| min
Sk:
(TT.))
(C.22)
VI(I+ + 3 m,d1)² 43
Proof.
According to the definition of Tí, in Eqn. C.100..
T 1 + 3 2 3
We can directly analyze the product by Jordan decomposition that there exists P E C2x2 such that
Te PJP-1
where J can have the following two cases
71 0
J =
72
or
0
72
24

---

Published as a conference paper at ICLR 2024
depending whether Tu,j is diagnolizable. Here y1, "/2 the eigenvalues of Te, and assume without generality that "1 2 72 if Y1, 12 E R. And when V1. 72 & R, it holds that 71 and 72 are conjugate thus In h/2l- Therefore |m| is the spectral radius of Tr.j ie. p(Te.) In Then discuss case by case.
i) If 71 o 72F In this
can verify that
P
72
71
P-1
72
0
71 72
where 71, 72 eigenvalues of Te.j- And the characteristic polynomial of Tr.j
det (Te. -1) 1+ 3
-3 -1+
-7
entails 7/ + 72 = 1 + -m/-2j, 71 3. Thus in this case, it holds that
T PJpp-1
7/1
72
-72
72
71
72
-1
A++1 72
-8
12
71
21
8
7 72
71 72
It holds that
[ k+1
k+1
-3 (1)
IT T
72
1 72
7 72
3 (1)
7/2
(T() 2p(TT.5²
-28p(Tr.))
m 72
-23p(Te.)) 66-1
(Tg)T (Tp()
-28p(Tes)))
2p T k
1 -3
11 72
-43p(Tr.a)"
m
72
[
-23
1n1 72l
VI(1+ 3 722))² 43
where the first inequality holds as I4a = -| P(T,,) and the third inequality holds as p(T.))
712 = 3 2 1/4 according to Lemma and thus (T,)) E (1/2, 1). For more details about the We also analyze the norm from another
properties of ((Tr.)) one refer to Appendix C.3 point: it holds that for all k.
f 2
A-
k-1
8-1-1
Σ 1=0
k-1
p(TT,)* h-1 kp(TT,,)**
?1
7/2
i=0
j=0
Substituting we have
31 kt1 kt1
B -
31 78
7
-32
71 72
711 72
(k + 1)P((T.)))
-Bkpp(T x.
kp Tr.j)
-3(k 1)pl T k-2
2kp (Tp Te.)) 2kp (Tt,))
-2k
S 8kpp ( T TT.)
-4kp(Teg)**
Thus we prove the lemma when 71 * 7/2-
ii) If 72 = 7: In this case, Tr, can not be diagonalized, which means that
7 1
T
0
7
25

---

Published as a conference paper at ICLR 2024
can verify that
P
J
p-1
0
->
1
0
[
where 71 72 = R. In this case, it holds that
T PJ
?
k
koyk-1]
0
0
1
[k + 1)7t
-kght1
ky
- (Ks 1)/
Then it holds that
(I 1)1 1ork-1
-ky h+1 -(k 1)
2kg
-koy
|Ttsll -
S 8k7" = 8kp(TT.))
2k
- 1og
Therefore, combining the two cases, we obtain the conclusion.
Lemma 7 (Bounding (T (T,, + AT)...
Ax)l with matrix power). Given ma-
trices Túj defined in Eqn. (C.10) and A, A defined as
0
6 0
A.
A
0 0
0
0
where 6, 2 and maxigisa 6, if + m)2g)²
48 2 0, it holds that
[[T, + (T) A2)... (T. + Ax) S (T + ))
(C.23)
Proof. As assume +1+33 3 the same method as the proof of Lemma momentum matrix Tt, can be written as
43 0, the eigenvalues of Tr 71, 72 R. Following
6 to apply Jordan decomposition to T the power of
B - B( 22
T
71
2 k k-13
i1 + 72
E.
21 72
(k + 1)7* kn k-
- k+1
T
if 71 = 77
-(k: - 1)2
We can observe that in this case, the first column of TÝ, is nonnegative and the second column is nonpositive as 71, 7 E R. For simplicity, in the following proof we use II to denote product from
& orderly from left to right, namely, IT TH4LK Tetk.Tet Te+1j- We first consider the combination product form of II. (Tes + A,) that
k
II (T., + A,)
Σ
111... 11 Th A21
21,
Tha f.j
A Au, Thun
which is similar to the binomial expansion but without the commutativity of T and A. Now consider one arbitrary combination term S that
S = A A11... T 21
T A.1... A.. T t.
We first prove by induction that S has the following properties:
1. the first column of S is nonnegative and the second column of S is nonpositive;
2. the absolute value of each entry of S is monotonically increasing with respect to 61, Sz.
We call product of the first p multiple component of S in the following proof. The first multiple component of S can be S Thj or S Aur which satisfies the two desired properties naturally. Then we assume that the product of the first p multiple component S,, satisfies the two properties. We discuss Sp+1 in cases that
f, or Aij one multiple component of S in the following proof. And we denote Sp the
26

---

Published as a conference paper at ICLR 2024
1. if the p-th multiple component is Tí.j where represents arbitrary integer, then the (p 1)-th multiple component should be Ay, where i' also represents an arbitrary integer, or the p-th and (p + 1)-th component can be merged. Then it holds that
[ p
p.
6
0
0
S p.11 0
Sp.2
Sp.2
0
0
8 Sp.21 0
Thus if the two properties hold for p, it also holds for p + in this case.
2. if the p-th multiple component is A where represents an arbitrary integer, and the
(p 1)-th multiple component is
++1 then it holds that
S p.
pp,12
6.+1
0
6 p.11 0 S p.21
0
Sp+1 Sp4++1
S p,21
p.22
0
Thus if the two properties hold for p, it also holds for p + this case.
3. if the p-th multiple component is A where represents an arbitrary integer, and the then it holds that (p + 1)-th multiple component T A.j
Sp 1,11
1,-1, p
0 0
S.Sp 1,11 0
S
1,21
0
5.S 1-21
which implies that Sp, Sp.2 = 0, thus we can substitute that
Sp
41
412
S Sp
Sp,1/12 Sp.21/1/
Sp.21
(21
122
As f11, Sp.11, Sp.21 are nonnegative, f12 is nonpositive, the two properties also hold for Sp+1-
Therefore, the two properties hold for Sp+1 and thus for any combination term S by induction. And one can verify that the properties also hold for their summation II (T + A.). Because of the definition of frobenius norm, when the two properties hold, | (T + A is monotonically increasing with respect to ò,, as well. Therefore, it holds that
b II (T + +
(TT++
[|((TT) + A)
which concludes the proof.
Then combining Lemma and Lemma 7 we can obtain a conclusion.
Lemma 8. Given 3 E [1/4, 1), Tú, defined as Eqn. (C.10), if only has real eigenvalues, which is equivalent to that the discriminant of TT,j satisfies that (1 + 3 - mT: 43 2 0, it holds that
8
S <min| 8k,
P ((Te+ke)
(C.24)
(11 + )= 43
Proof. The difference of two momentum matrices Tú, and Try that ny n is
Te Tt,j = (Tw 7))
which has the same structure with A; in Lemma Lemmapand and Lemmal One can verify that
Thus Lemma is natural combination of
T T1.....TTT.... < [ T44.j... T.ta T.2....TTrklllll
|| | |
8
min
8k, VI(1 + )² 43
o(T
which concludes the proof.
27

---

Published as a conference paper at ICLR 2024
C.3 KEY PROPERTIES OF P(T)
This section offers some useful property of spectral radius p(Tx.) ) in terms of different vt, 3 and Xz-
Lemma (The exact form of p and its relationship with {8, A)). Given momentum matrix
1
-3
0
if only has real eigenvalues, which is equivalent to that the discriminant (1 + -1+--1A)= 4,3 the spectral radius of is
1 P(T) = 1+3-
- n v(1 n2)2
Else the spectral radius Tf is
(T) VB. Thus under the assumption that 1n < 1, we have (TT) is monotonically decreasing with respect to nd.
Proof. To find the spectral radius, we first derive the eigenvalues of T,, which is equivalent to solving the equation
1 + 8 - 1/d - -3 -7 0.
det T -I| = det
After rearrangement have
2² - 31 - + 3 0,
which is quadratic equation. The discriminant is that
A = [1+3 - - - - 43 - ((1-V)) ((+vV)) -mx)
Under the assumption that 17 S 1, the positivity of A depends on the positivity of (1 - /3)² - If A 2 O, then the eigenvalues 71, 72 E R. If < 0, then the eigenvalues 71. 72 & R. We then cases.
discuss these
Case 1: A 2 0 In this case, we have
+ 7 (11 3 n))2 48
71
2
1 1+3--1-- -
72- 2
Then we have 2 thus the spectral radius
p T = = [1+3- nd + v(1+3 1/)2 48
Then we justify the monotonicity in this case. We have
0p(T T (n))
B
-1
0,
V11 7A)2 48
thus P(T) is monotonically decreasing with respect to 17A.
Case 2: A < 0 In this case, we have
71
1 - 3 - n 48 (1 3
1 + - - /48 - (1 +1+ -
2
where is the imaginary unit. One can observe that 71 is the complex conjugate of 72. Thus we have the spectral radius is
(Tt,)) ml 121 = 7171 7772
3.
In this case, ((Tr,)) does not depend 17λ. Thus finish the proof.
28

---

Published as a conference paper at ICLR 2024
Lemma 10 (p bounded by 17, A, 8 in real-eigenvalue case). For momentum matrix
1
17
-5
0
iT only has real eigenvalues, which is equivalent that the discriminant (1+ (1+8-11)22-1 1/)) it holds that
7d
nn
p (T) 2
(C.25)
2 4 - vB)
Proof. According to Lemma 9 the spectral radius of is
(T)) 1 B nn + v(1 (1 + B mA)2 48
It holds that
(T))
1 + B - n v(1 n/122 43
2
1 1+3- ax+ + v(1+
(1+ v/)
m7t 1(1- v3) 2 mn
2
1+3 au+ + (1+V3) / (1- /3) mn
<
7d
1+B A + (1- B) /)/
V
(1- 1=-
7d1
1 + AA + (1-
1
2 2(1 - /3)²
(1 3)md
2(1 7/
=
2
17À
2(1 vB)²
s 2 y/1
v3)
1
2 4 (1 v3)
C.4 PROOF OF THEOREM 2
In this section, we specify the schedule to step decay and prove Theorem We denote the step size of the f-th stage 7/c and its corresponding momentum matrix TV, specify the stagewise case.
We first present a lemma to simplify our proof in the stagewise case. Lemma 11. For all stage > given matrices T7, defined in Eqn. (C.IO) if > (1 vB)²
and the length of the stage
K 2 (8ln(8T),
holds that
(TT2,) S..
(C.26)
Proof. In this case, the cigenvalues of TV E. are not real and thus the spectral radius p p(TV.,) = /B as Lemma 9 suggests. Thus it holds that
(T | < (TT.)
8KP K 8K (vB) K
F
(u()))
<ST
1
Vx
29

---

Published as a conference paper at ICLR 2024
D.
ST In(87)
which concludes the proof.
Then we are ready to prove the convergence of step decay schedule.
Theorem 2. Given quadratic objective f(w) and step decay learning rate scheduler with
(1 1/vm) with n. 2 4, and n T/K with settings that
1. decay factor C
1 I<CSTVR
(3.13)
2. stepsize r/e
1/6 = L C
(3.14)
3. stage length K
K T
(3.15)
logc (T /k)
4. total iteration number T
T
2 20
(3.16)
In (214718 In (2674) logo (T2)
then such scheduler exists, and the output of Algoritth 7 satiisies
27
E's(wT) - [(w-)] SE (s(wo) [[ww)]
-xxp exp 15ln2 + 21n 21nx
vk logo (T RR)
4096C3 In? (29T") logo (TVR).
MT
Proof. From 3.16 the total iteration number T satisfies that
T 2C 2C In (2447(m) In (29T4) logo (TVR).
(C.27)
and we define an auxiliary constant for our proof that
h = hrT,, x) = 4nn (2fr*)- logo (TVR) 2 1.
(C.28)
From (.16), we know that 2 20. Then with following requirements satisfied.
CC.7 and CC28)
we can verify that
1. From Lemma 11
K 2 Valn (8T).
(C.29)
Verify: As T 2 /R. holds that
K
3.15
3.16
20 2C (n(2))))))))))))))) 2 VIIn(8).
log
T
2. From variance case 1.1, the final stage needs the variance be small enough that n7, L < hT)]:
T
K
K<
(C.30)
logo
TvR
T
h
Verify: As T 2 VK and h 2 1, holds that
T
T
loga (T FR)
logo
Tv/
30

---

Published as a conference paper at ICLR 2024
3. From variance case 1.1:
K
(C.31)
logc
1x 3
Verify: As 2 20 it holds that
3.15
T
T
logg (T RR)
logo (3
+
4. From variance case 1.2:
4T K 2
In (20T4).
(C.32)
h
Verify: holds that
C.28
4nn (2
logg (TVR)
3.15
47
In (2
K
5. From variance case 2.1:
2 2C (21210).
(C.33)
Verify: It holds that
K
31S
3.16
2C 20 (21276).
logo(TVl)
6. From variance case 2.2:
VE
K
In (21*T(16).
(C.34)
2
Verify: A T 2 V K, it holds that
K
3.15
T
3.16
2C 2CVeln(214T)) (1((22T /i In (2)*T²x)
logo
Now we are ready to start our main analysis. From the former lemmas, it holds that
2E (f((T) f(w.)]
WT w.)
H (WT - W..
Wo w.
M
M
O
Mr MT Mo
Wo - w. w-1 W.
W-1 W.
O
O
+
T-1
n/
n1
T
M. M T+2
M
O
nn My Mr
M,+1
O
0
r=0
015
TT T 1. Tr- 2 To 2
IITVT
Wo w. w-1 w.
/=1
2:-1:22
Mr IT TT- "..T. Tr
j=1
If we denote bias and variance as
2
B
A
(TT-1
TT. To. 12 PB VVVV
Wh - W.
w-1 W.
j=1
2j 1:29
V
A
2
Σ
TT- Ty 2, T T++1,112
31

---

Published as a conference paper at ICLR 2024
have the following results. (1) Bounding bias term:
It holds that
B V
wo w. w- Wx
B
j=1
2j 1:25
Xy |T- Tr Tooll 1 (IIVV wo) Wo)2 25-1:
j=1
Σ
(TT..) ( (T (Tn--1)) Amg-s )
... (Ti,) A
E
V wi
2j 1:20
j=1
2
C.26
> (TT.)))))
8
E
V w.) w 25-1:21
V(1+ 3 - 11 12,) - 4,3
256
Σ
> (T 11 E Vii'vi vw)) 2j
2
1/11 3=1
d
2.
256
A (8kip (Ti.) E ww)
29-1:25 2
1/144
/=1
Lem 9256 d
X (skip (Ti,)) E | V ww)
25-1:25 w0)2311225
2
1/ /=1
256
(8kip (T' (Tl'"' k:)
2
1
(II VTo) 25-1:2j
7/11
d
256
8kip (Tr.d) kk)
E
Σ
Ay
wo) T
V wo 25-1:21
1/11
2:-1:22
256
[8hi (Ti.)) ) )
E
Wo)
Wo)2y- 1:21
Aj
[II V wo
1/11
A3
21-1:
256
(8kip (T 1h))
7/11
II
Wo)
Si
wo 1.2
II
Wo) 34
S2
wo 3:4
-E
Wo
2:-1
S
Wo
2d-1-2a
0
with
Sy
0
Aj
256
(8kp (T )1)
E
[II
V V T
A
+
0
1
III V ww)
256
(8kip (T /1)
0
E
VII
A
wo
7/
where the first inequality is because that Eqn. C.26 can be applied for the stages E where TV has only complex eigenvalues and C.2.4 can be applied to the matrix product for all stages E that Te, has only real eigenvalues. And the second inequality holds
(1 + B ,2,,)² 43 v( vBj² nhh9 (1 + VB22 1//bj
32

---

Published as a conference paper at ICLR 2024
, -V .1/2g
V (1 /BR
3
4 (1 /392 = 2 vB) = V 1/1.
Notice that
A & 1 1 II
V 111 :
VII
since
H
II
VII
H
0
T
U
0
0
U
e1
ed
T
e
0
el
A O
e:
0
02
0
A
0
C1
0
T
T
e
A1
A1
A2
X2
1
0 1
6
0
Ad
Ad
and II, V are orthogonal matrices (Eqn.
CC.4
We then further have
B < 7/
256
[81p (T
11)
E
W
6
H
O
O
H
Wu
256
(8ki (T
E
Н
Wo
W-1
w,
O
W-1 - w.
256
(Skip (Tí,a)
k1
2
(wo - w.) H (Wo- w.) + (w-1 - w.) H (W - w.)
1/
256
2
(skip (TH,a)
E [[f(w s(wo) 2f(ww)]
7/14
256
8h1 (VB) E[f(w f(wo) - 2f(w)
256
8k1 (1 (1- (1- --Vii 11 -E(((w + f(wo) 2f(ww]]
7/
< <exp 114 + 21 21n (7/11)
2k v 1/1)))
(WW + f(wo) - -2((w.)]
14 + k - 21n 1
2T
=E[[(w-1) 1)+ + s(wo) - 2/(ww)) exp
Vk logg (TVR)
2T
<E[ f(w)) - 2f(w.) exp
141n + 21n 2ln
Vx logog (T (K)
33

---

Published as a conference paper at ICLR 2024
2T vx logc (TVK)
-B - [(w.)] exp 15 2n 2
where the second last inequality is because of the fact that 1 S exp(-z) for 2 0 and the last cquality is because of the setting Algorithm| that Vo = and thus W-1 = Wo.
(2) Bounding variance term:
denote
Vi,
in the following analysis. We first assume that the batch size M 1 in the main analysis and we will transfer the result to the general M
case.
We make use of Lemma 9 and 10 and divide the analysis of Vi, into cases with respect to m and cquivalently the corresponding cigenvalues. The division of the cases is due to two major boarders: 61 (1 2 and b2 = (K).
For momentum matrix Tri with T)t > bi. T has complex eigenvalues and large spectral radius that supports sufficient geometric decay of Vi.j, the variance generated at iteration on Ay. Case 2.1 and Case 2.2 discuss this case.
For momentum matrix Tej with b2 < n S bi. Tr,j has real eigenvalues but its spectral radius still large enough support sufficient geometric decay of Vv.j- Case 1.2 discusses this case.
For momentum matrix Te with m < b2, the corresponding step size neyy is small enough to ensure the generated variance Vi.j is small. Case 1.1 discusses this case.
Among this cases, we use requirement CC.31 to ensure that Case 1.1 exists and thus the variance can be small enough. Let's discuss Vr,j case by case then.
Case 1.1: Real eigenvalues with small step size. We first consider the case that Tt only has real eigenvalues. This case is equivalent to that
(1 + 3 - T/A)))² - 43 2 0.
is also equivalent that
n S 51 2
(C.35)
And in this case we further assume that n/ small enough that
Ch
mo - TVK'
(C.36)
where h is defined in (C.28) We also discuss From requirement (C.30). this case exists. Then from the fact that Frobenius norm of matrix is always larger than norm and using Lemma we can obtain that
|T TT 2 T < TT TT Tr+1sll Tr.
C224
min
8(T (8T'-11 1-7), -),,
p (TT--1))T-1-
(1 + B 1/1 1 43
(T --1-7
V + AT IT- 12, 43
T T-1-1
(1+ Th , 2 43
(C.37)
34

---

Published as a conference paper at ICLR 2024
Then analyze the right hand side term by term. First it holds that
(1 V1+ B w/>a)² 43 V (1 vB2 2/123 ((1 + vB2 1//
2 (1 vB)² 1// Aj = / vB2² - C- Thi
(C.38)
C.31
3 (1-
v32²
vB) = 7/H.
From Lemma 9
(T,, < 1 holds for all t,j. Combining all above we can obtain that for j
satisfying 7): Ag h// (T V / F it holds that
V1j 23 [T 1.5
TT 2, T T+
Ch
8
TVK
(I+ B - 1
1/,
43
2²
64
256C
25
T
TRIT/H T
2 v
7
where the last equality holds as v7f = 1/L.
Case 1.2: Real eigenvalues with large step size. Then continue to consider the case that Te. only has real cigenvalues. Similar to case 1.1, this case is equivalent to that
(1 + 3 TxA,)² - 43 2 O.
It is also equivalent to that
n, (1 11-1-/9 23)² = 1/2.
(C.39)
But in this case we further assume that 1/:11 is large enough that Ch
TVK
as the case that n S Ch/(TVF) has been discussed in case 1.1. From requirement (3.66) we know that such stage exists. Then according Lemma 10 with n.dy h/(Tww). it holds that
(TT,,) 2
1 3 -n + V (1 + B )2 4,3
4 (1 m
=1 nhy
vB)
4 1/
h
4T
Thus if we denote f, to be the first iteration that T/. X3 < h/ (T VK), it holds that
Vij = 2/ T 1.j Tr 2.j T T
< 1 1²²
TT- T-
T
TT. T 22- 11
2.23
n
TT-11 Tx -2g.....
Te,
TT. 1
58
8
2
h
2(t.-t-1)
n2
(8(t. - -t- 1))²
1
1(1
T/ )2 4B
4T
C 38
222
(1- (1- 4T
h
2(1, -1-1)
2127
2K
S n/ {3
2
4T
7/11
212722
h
4T. 1n(""T*
D.4 256
1 (/T
2
4T
T2
35

---

Published as a conference paper at ICLR 2024
where the third last inequality holds mı 2 Ch/(TVR) which suggests that there is at least one stage between t and t,
Case 2.1: Complex eigenvalues with small step size. Then consider the case that T2 has complex eigenvalues but TV 111, has real eigenvalues, which is equivalent to that
C(1 /3)² Z mdj > (1 VB)² = T/1..
(C.40)
Thus we denote i, to be the first iteration in the + stage, consider the spectral radius of the l + stage, holds that
P(T..)) ((a))=((T(.)))
1 + 3 Net12 + / (1+ - 43
1 = 16t12,
4(1 v3) 4V1//H
(C.41)
n/1 = TK
4C 1C/1/11 4C
Thus have
Ves = 12 T.-2.... T
C.35
(C11+222))
[ - ...... T+t1ol
C² (T/1)) 1T T2..... 1et1
2 -1 T 2.... T,, K.j | T T -2-2j.. 2.... T..ll -T, Te.- 25. 2...TTTTT Tet1zl?
256
T TT 2 T ?1? TT. Tr, 2..... T..1lll
T/2u
2 -I 1|2
T (T. Te, 22. 2....T.411. T++1,11
K+1
C:41
V
256C ni
8T
[T, T T
4C
K
256C
8T
1
1
max oSiSk
(T' K-il xxi1)
4C
v
V
K
v 1/11
2K
256C2
8T
(8K)² S 2
C2T
1
4C
4C
4C
40
v Vhh
Th
256
20 ThP
1
4C'
722
where in the fifth inequality we introduce max, because the leading stage may be incomplete.
Case 2.2: Complex eigenvalues with large step size. Finally we consider the case that Tú, has complex eigenvalues which equivalent to Equation (C.40) and also T,, is far away from the boundary, namely, the step size of the next stage can still satisfy Equation (C.40). In this stage, it holds that
M1 > Crtp..
(C.42)
Thus if we denote f., the first iteration of the f,, + 1 stage where f, is the last stage that TV has complex eigenvalues, namely, 1/6 XÚ E (The, C1/11). We denote & the stage where iteration / is in. Thus holds that l, < l. Then we consider that
V
x
T T 1, TT
2. - 111
3.
14
(T-- Tl (T x1| (T 111e_ ... (T xI2
36

---

Published as a conference paper at ICLR 2024
(T (T
max 0sisk
C.37
8
(TT (1' (TT
(T
x11 ... (Ta+io)) x
- 1/2 2,) 48
(TT K -1=
max
0SISK
C.38
256
(T% K|' (T% M11 ...| (Th, (T2)t.))
K
(Tam.) K-i
max
1/11
osisk
C.26
256
(T
$11²
(T2 10=
max OSisk
1/
256
8K (V8) K)
2²0O2 K
(8K)²
[1
v 7/
1/1P
7/14 u
C.34
042
220T1
(1-Vih
Valu(2"'
D4
256
1//10
T
where in the first inequality we introduce max, because the leading stage & may be incomplete. Therefore, combining the four cases, have the result
V= ΣΣ
d T-1
d T-1
25
22
4096C²
In² In? (2 logg (TVR),
T
T
1=
which concludes the proof in the case that batch size M 1.
For general batch size M 2 1, the gradient noise
n/ -Vv[((w) (wi))
Vw Vwi((w,,))
Be
EEB.
M SEB.
M-1 1
nea++i
M
1=0
satisfies
E
n
(n/)
T
M
1 M-1
I2
neM+i
M
M
===
M-1
1
E
D + Σ
M2
1
E [D (eM+)DMM4 + DeM+endee
M
Egtts
[neM4 +
[neare]]
mostel
D> Assumption
M2
M-1
1
E
(DEM+aNTMAl
Assumption
2
M
t=
-M
Assumption 3
M2
37

---

Published as a conference paper at ICLR 2024
H
M
So for general M 2 is equivalent to replacing the noise term o² with o² /M. Thus it holds that
do² In (2 loge (TVR).
V
MT
Combining the bias and variance term, we can verify that Theorem 2 holds.
D PRELIMINARY: USEFUL. LEMMAS
We provide some basic mathematical tools relevant our proof. serves manual section and can be skipped if one is already familiar with those tools.
D. RANDOM VARIABLES
Lemma 12. each pair of entries of matrix and matrix Y are independent from each other, then
=
(D.1)
Proof. For Vi,j.
XY
=E
(XY)s
Σ Yk.j
Σ E E(
Σ E A
E X
E
[Y] Je.
(E
E (Y))
Thus IE = E
D.2 LOEWNER ORDER
In Loewner order, X Y if and only if V2, x. T -)) 0. Lemma 13. Given two x block matrices X and Y where
Xn 3 Yn,
x12= = Y12, X = Y21, X = Y
then X Y.
Proof.
For vz
Z1 //
2 T
X)z = Zi
Z
Yn X11
O
0 Z1 2 (Yn X 20
22
where the inequality comes from X Y Therefore X < Y according to the definition of Loewner order.
Lemma 14. 1fX YE Raxn, then for VC Raxm
XC YC
Proof. For Vz Rm
z YC- XC) TY X)) 20 = c CTCC
Lemma 15. XfX < Y, then tr(X)tx(xx)) S ir(Y).
Proof. Denote
0 0...0 00 E Rdd
e;
i-th
Y, - X, X)e 20 0
for Vi
tr(Y) tr(X) = Σ YY xx) 20
38

---

Published as a conference paper at ICLR 2024
D.3 QUADRATIC EQUATIONS
Lemma 16. (Roots of quadratic equations) If 1, 32 E C are roots of equation x² + Bx + C = 0,
where B, C I., then
(1) 31 + 2/2 = -B,
(2 21 C
[2C,
if a 22 are real
(3) 21J2 + -1022 =
(D.2)
1B²-20,
if a1 are imaginary
[Ba-2C,
if W2 are real
(4)
[1
x2
12C,
if #1 22 are imaginary
Proof. (1) and (2) are a special case of the famous Vieta's formulas, which can be directly obtained from
+ Bx + (x - x1)(x x2) = x² (x) + x2)x + (2)22))
(3) can be derived from the form 2ı and z2. The roots of x² + Ba = is
12 VB
x =
21 and 22 are both real, then
31 + 1142
I + -' = 2x 2C.
Otherwise Tı and 22 are both imaginary. then #1 = T2, 22 = Ij. which follows
T2 a + zي (x) + x22² 201 B² -2-
(4) can be obtained from (1) and (3) by
|x 2 + |x2
-i + D2 = (01 + 3:2) 11 + T22
(x1 + +1122) B² (x1 + x1022
B² -2,,
if æ1 22 are real
2C, if 11, 32 imaginary
D.4 BOUNDING SPECIAL FUNCTIONS
Lemma 17.
V1-x 2
holds for Vx E [0.1]. 1].
(D.3)
Proof. Only this lemma, denote
f(x) A V1-x -
then for Vo E [0, 1).
1
1 1
1
S(0) = 0, dr (2) =
1
0
2 VI
2
2
vI
+ f(x) < f(0) =0
x 21- I 2
For x = 1, we also have v1 - 0 052 x/2.
Lemma 18. For r E [1, +00), it holds that
(1 -
f(x) =
(D.4)
39

---

Published as a conference paper at ICLR 2024
Proof. The lemma is equivalent to that
1 - - -1 x In
Denote g(t) = In(1 - t) +t, e (0, 1), then g(t) is monotonically decreasing
g((() 1-t +1<0
Therefore, g(t) S g(0) = 0 when t (0,1), 1). Thus substituting = 1/2,, it holds that
In 1
1
5.
x
x
After rearrangement, we can obtain the result.
40


Reviews:
[{"review_id": "l04CvOu2G0", "reviewer": "Reviewer_Uq7v", "summary": "This paper analyzes the performance of SGD and SHB (Stochastic Heavy Ball) in stochastic optimization with a quadratic objective. A lower bound of the convergence rate of SGD is given, which argued to be strictly worse than the upper bound of convergence rate of SHB also proved in the paper. More precisely, SGD is shown to take at least $O(\\kappa)$ iterations to converge, while SHB requires only $\\tilde{O}(\\sqrt{\\kappa})$ iterations with proper learning rate schedule.", "questions": "My foremost concern is the comparison with previous literature, which I discussed in the `Weaknesses` part.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The contrast between SGD and SHB does seem plausible. I checked the proof briefly and found no obvious mistakes.", "weaknesses": "1. The upper bound of SHB seems very similar (and worse by a logarithmic factor) than the quadratic setting in [Can et al. (2019)](https://arxiv.org/pdf/1901.07445.pdf). I would suggest adding a detailed comparison. Also the convergence rate is sublinear, which is not very satisfactory.\n2. The batch size is required to be of order $\\Omega(1/\\epsilon)$, which does not seem realistic, as $\\epsilon$ is usually exponentially small.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ds1Cg85eax", "reviewer": "Reviewer_dqiw", "summary": "This study presents an analysis of the convergence properties of the stochastic HB method when applied to quadratic objective functions. Specifically, when we make an assumption about the presence of anisotropic noise, it demonstrates that the stochastic HB method exhibits a significantly faster rate of convergence. Earlier research has already established that this anisotropic noise model can manifest in neural networks when their parameters are in proximity to the optimal values.", "questions": "In the paper you mentioned that this rate achieved for the objective in the large enough batch regime. However the acceleration is oblivion to the mini-batch size. So I was wondering why you emphasise on this in the paper.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper for the first time shows that in the anistorpic setting for the quadratic objective function, the stochastic HB method can achieve accelerated rate. This is a novel contribution.", "weaknesses": "The paper should delve further into the circumstances under which the anisotropic setting is applicable to machine learning models. The content of Theorem 1 may seem redundant as its proof is not contingent on stochasticity; it essentially establishes a lower bound for gradient descent (GD), a well-known result. Therefore, its inclusion in the paper is somewhat unclear. It would be more appropriate to allocate space to discuss the technical innovations of the main theorem's proof within the main body of the paper. Additionally, there are several typographical errors in the proofs, including indexing, which should be rectified.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8C7QAIjkK5", "reviewer": "Reviewer_Vpgj", "summary": "The paper investigates effect of Polyak-type momentum in SGD with large batch size. The authors show that momentum, combined with a stage-wise geometric decay stepsize schedule, improves the linear convergence part of SGD. Although the analysis is done on convex quadratic functions, it takes a step towards better understanding of SGD with momentum.", "questions": "1. Assumption 3 seems to suggest \"noise is small in the coordinates that result in ill-conditioning\". What do you think is the difficulty in extending it to commonly used assumptions (e.g., bounded variance)?\n2. In your first experiment, It seems that SGD performs worse when $M$ gets large. To me it feels that it is because your range of grid search for $\\eta_t$ is invariant for all batchsizes. I'm curious whether SGD performs better when you multiply the stepsize by $\\sqrt{M}$.\n3. Nonconvex models are actually used in the second experiment. Although the authors claim that the strongly convex quadratic optimization model approximates the landscape near local optimum, Figure 1 actually suggests that benefits of momentum are significant at the beginning of the training procedure. Could you elaborate more on this?\n4. Some recent researches show that even without decay of learning rate [1], momentum SGD also outperforms SGD using large batch size. What do you think might contribute to this phenomenon?\n5. I would suggest the authors remove (or delay to the end of the appendix) proofs for auxiliary results, most of which are well-known results. Some notations should be properly defined (such as $\\mathbf{X} \\succeq \\mathbf{Y}$ and Hessian norm $\\\\|\\cdot\\\\|_\\mathbf{H}$)\n\n**Minor typos and stylistic issues**\n\n1. Algorithm 1. Line 6\n\n   $m$ => $M$\n\n2. Experiment\n\n   initialize $\\mathbf{w}^0$ from $(-1,1) \\Rightarrow (-1,1)^d$.\n\n3. Proof of Lemma 3 \n\n   $\\succ$ should be $\\succeq$\n\n4. Page 29\n\n   genarality => generality\n\n**References**\n\n[1] Wang, R., Malladi, S., Wang, T., Lyu, K., & Li, Z. (2023). The marginal value of momentum for small learning rate sgd. *arXiv preprint arXiv:2307.15196*.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well-written and easy to follow. The motivation and main results are clearly presented, with enough intuition given after each result. The technical results seem solid and are supported by empirical evidence.", "weaknesses": "There is still a gap between real-life SGD applications and the analysis. Most SGD applications are for nonconvex (even nonsmooth) problems, where Hessian may even not exist and the notion of condition number is uncommon.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pVrkcvqkve", "reviewer": "Reviewer_JMNK", "summary": "Motivated by empirical successes of stochastic heavy ball (SHB) method and the prevalence of the anisotropic gradient noise in training neural networks, the paper presented last iterate convergence rate for the SHB method on quadratic objectives under the anisotropic gradient noise assumption. The results implies SHB with polynomially decaying learning rate provides $\\tilde{O}(\\sqrt{\\kappa})$ acceleration rate with respect to stochastic gradient descent (SGD) when batch size is large, where $\\kappa$ denotes the conditional number.", "questions": "No", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well-written, and proofs are neat and easy to follow. \n\nThe paper formally established a lower bound for SGD tailored for the scope of investigation, which served as a concrete comparison benchmark for SHB method.\n\nThe paper presents novel theoretical results of SHB for quadratic loss. The result implies $\\tilde{O}(\\sqrt{\\kappa})$ acceleration in comparison to SGD when batch size is big. The results is also nearly optimal in comparison to heavy ball method in deterministic setting. \n\nThe experiment results matches with the theoretical bound: acceleration guarantee with large batch size. The advantage of SHB with large batch size is also well-motivated in practice. \n\nThe analysis follows the classical bias variance decomposition paradigm. The novelty is to quantifying bias and variance with some linear operator, and then analyze the property of those linear operators.\n\nAlthough the scope of theoretical investigation is limited at quadratic loss with anisotropic gradient noise, the author justified the broad implication of such assumptions to real applications.", "weaknesses": "Potential minor typos:\n\npg5: eqn 3.11: missing - sign before $\\beta$. (Same typo appears at pg 21 under the proof of Lemma 13 at two places)\n\npg6: Algorithm 1 line 6, batch size $M$ instead of $m$ being consistent \n\npg27: the equality applies (C.18) and (C.21): the last row of the matrix $2d$.\n\n\nEverything is rigorous, but the notation might be a bit overly complicated for what was actually being used in order to establish Theorem 2. For example: number of stages $n_{\\ell}$, the stage lengths $\\set{ k_1, \\cdots,  k_{n_{\\ell}} }$. Theorem 2 actually set all stages the same length, and each stage has the same learning rate.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)