Accelerated Convergence of Stochastic Heavy Ball Method under Anisotropic Gradient Noise

We would like to extend our sincere appreciation to the reviewer for all the constructive and positive feedbacks on our contribution, which definitely helps us improve the paper concretely and motivates us to do our best. Thanks!

Overall:

We have made corresponding changes in our revised paper to address the mentioned potential issues.

Weaknesses

Thanks very much for pointing out those typos and organization issues! We have uploaded a revised version of our paper to reflect the changes.

We would like to offer our sincere thanks for the reviewer's constructive comments and recognition on our contributions.

Overall:

Regarding the raised questions, We have provided additional experimental results and relevant discussions to address the concerns. The details are as follows.

Q1: What do you think is the difficulty in extending Assumption 3 to commonly used assumptions (e.g., bounded variance)?

That's an interesting question. I think one major difficulty arises from bias-variance decomposition.

First, for SGD, a more general assumption $E[n_t n_t^\top] \preceq \sigma^2 D$ can be well handled, where the rest standard proofs for quadratics lead to $\sigma^2 D_{jj} / (\lambda_j T) \le \sigma^2 D_{jj} / (\mu T)$ for each dimension, summing them up across dimensions and we obtain the common upper bound $O(\sigma^2 / (\mu T))$ given $\sum_j D_{jj} = 1$.

For SHB in our proof, one thing is that bias-variance decomposition requires the noise and the hessian to be simultaneously diagonalizable, as shown in (C.19). If this part is resolved, we shall obtain similar results as SGD and recover the $~O(\sigma^2 / (\mu T))$ bound, given the rest part of the proof in our paper providing similar guarantees.

Also, it is worth noticing that generally our setting is technically harder than normal settings of bounded variance. For bounded variance, standard techniques such as [1] can be considered to apply, which we believe to lead to a results similar to existing literatures. We choose Assumption 3 because the corresponding result is still unknown and the setting is technically more difficult.

Q2: In your first experiment, ... I'm curious whether SGD performs better when you multiply the stepsize by $\sqrt{M}$.

Thanks for the question! Actually we tried larger learning rates before, such as initial learning rate $\eta_0 = 10.0$, but the loss soon exploded. This is no surprise given it matches the theoretical insight that learning rate $\eta_t$ must be no greater than $2/L$, no matter how large the batch size is.

We also conducted additional experiments of multiplying the stepsize by $\sqrt{M}$. The loss also exploded for $M \ge 32$. According to our experience, increasing the stepsize generally helps SGD to converge faster, but there is an upper limit for this increased stepsize.

Q3: 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?

That's a very good question. We provide further experimental results on cifar10 to verify SGD and SHB's behavior near optimum, where we first run SHB to for 50 epochs (batch size 2048, constant learning rate $\eta = 0.1$) to approach the optimum. We then conduct the exactly same hyperparameters search process as our original experimental settings, with batch size 2048 and number of epochs 100.

As shown in following table, both SGD and SHB has a much lower loss as it is near the optimum, but SHB's loss is an order of magnitude lower than SGD, along with a higher test accuracy.

In addition, our intention of the original set of experiments is to demonstrate the effectiveness of momentum in common non-convex settings, whose assumption may not be perfectly aligned with our theory as the reviewer implied.

[1]: Aybat, Necdet Serhat, et al. "A universally optimal multistage accelerated stochastic gradient method." Advances in neural information processing systems 32 (2019).

Thank you very much for all the constructive comments and recognition on our contributions. We really appreciate it.

Overall

Regarding the mentioned concerns, we have provided relevant discussions and a revised version of our paper to address the issues.

Weakness

Thanks for all the constructive suggestions! We have made corresponding revisions based on the mentioned valuable advices, including

1. Shorten the paragraph for Theorem 1.
2. Add two short paragraphs after Corollary 3 to discuss about the intuition of this paper's technical innovation in proofs.
3. Rectify the typos pointed out by the reviewer.

Any further comments would be greatly appreciated.

Q: 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.

That's a very good question. It is because without a sufficiently large batch size, the acceleration effect in the bias term will become much less obvious, and may even be completed covered by the large variance term. This phenomenon is similar to the one in [1] mentioned by reviewer Vpgj, where under small or medium batch size settings, momentum's empircal benefit is not evident.

[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

I would like to express my appreciation to the authors for their very patient response to my review, and my apologize for a few comments that are unclear or incorrect. The authors did make an impressive effort to address my concerns, but some questions persist.

1. I appreciate the efforts of the authors to provide a very careful comparison with Can et al. (2019). This helps to clarify the situations, but aside from the obvious differences the authors mentioned, there are some deeper connections that need to be addressed. For example, the point that there is an extra non-convergent term in Can et al. (2019), which uses a constant learning rate, seems easily remedied by using a decaying learning rate (hinted in Remark 23 therein).

2. By sublinear convergence I meant the first term $\mathbb{E}(f(w_0)-f(w_*))\cdot\exp(2\log(T)-\Omega(T/\log T))$ rather than the variance term (which, of course, cannot be linear). Usually something like $\exp(-\Omega(T))$ is more welcomed here, so that the algorithm converges linearly before reaching the variance-dominated regime. Anyway, this point is only a minor complaint which does not affect my assessment of this paper. But I greatly appreciate the authors' patience in answering this point which I did not word properly in detail.

3. I apologize for the careless mistake claiming that $M=\Omega(1/\epsilon)$. I appreciate again the authors' patience for forgiving my carelessness.

My previous evaluation is based on the thought that Theorem 1 is novel, which, as the other reviewers pointed out, does not depend on stochasticity and follows from the corresponding (well-known) result for GD. Considering that the contribution of Theorem 2 is fair but not significant enough compared with Can et al. (2019), I will keep my score unless there is strong evidence to refute this point.

Thanks to the authors for their detailed response. Concerning the interpretation of Can et al. (2019), I would like to remind that distributional convergence is stronger than convergence in expectation, as $\mathbb{E}f(w) - f(w_*)\le \frac{L}{2}\mathbb{E}\|w-w^*\|^2\le \frac{L}{2}W_2(\nu(w), \delta_{w_*})^2$.

We would like to appreciate for all the time and effort the reviewer has spent on checking our proofs, along with recognition on our writing. Thanks!

Overall:

According to the reviewer's suggestion, we have included the mentioned paper in related works, together with a corresponding comparison and relevant discussion.

On the other hand, we have noticed that there maybe some misunderstandings or incorrect points in the review, so we would like to provide further explanations to clarify our contributions.

Weakness 1.1: The upper bound of SHB seems very similar (and worse by a logarithmic factor) than the quadratic setting in Can et al. (2019). I would suggest adding a detailed comparison.

Thanks for the suggestion! We've added corresponding discussions and comparisons with the aforementioned paper in the "Related Work" section.

To make things clearer, we also summarize the comparisons here. Briefly speaking, the proofs in Can et al. (2019) has different settings as ours. Among all the results related to quadratic objectives:

• Theorem 3, 5, 9 focus on the deterministic version of Stochastic Heavy Ball (SHB) and Nesterov's Accelerated Gradient (NAG), which has no stochastic gradient noise and thus can achieve accelerated linear convergence rate with no surprise.
• Theorem 4, 7, 11 measure the distributional convergence, i.e. convergence speed of parameter distributions towards a stable distribution. Notice that this does NOT imply parameter convergence to optimum or expected excess risk $E[f(w)] - f(w_*)$ convergence to 0.
• Theorem 8, 12 provides an upper bound in terms of the expected loss/expected excess risk, which has an extra non-convergent term thus is much worse than our bounds.

Weakness 1.2: Also the convergence rate is sublinear, which is not very satisfactory.

We would love to obtain linear convergence if possible. Regrettably, it is theoretically non-achievable under stochastic gradient settings, as pointed by lower bounds mentioned in [1] and [2]. This lower bound is called statistical minimax rate, which not only applies to SHB and NAG, but also all possible optimization methods based on stochastic gradient oracles. Briefly speaking, all of them have a $\Omega(1/T)$ lower bound for convergence rate in terms of expected excess risk $E[f(w)] - f(w_*)$.

It is worth noticing that the provided convergence rate in our paper is already near-optimal when compared with the aforementioned lower bounds (up to log factors).

The mentioned paper Can et al. (2019) may give readers a false impression that linear convergence is achievable for SHB or NAG, where they are actually only achievable in distributional convergence or deterministic settings, but not for stochastic optimization, as discussed in our responses for Weakness 1.1.

Weakness 2: The batch size is required to be of order $\Omega(1/\epsilon)$, which does not seem realistic, as $\epsilon$ is usually exponentially small.

We would like to kindly remind the reviewer that there is no such requirement in our paper. Notice that in Corollary 3, we only need the number of training samples $MT$ to be $\Omega(1/\epsilon)$, where $M$ is the batch size and $T$ is the number of iterations. This is achievable since we adopt decaying learning rates instead of constant learning rates, where the generated variance decreases when learning rates gradually diminish during training.

[1]: Ge, Rong, et al. "The step decay schedule: A near optimal, geometrically decaying learning rate procedure for least squares." Advances in neural information processing systems 32 (2019).

[2]: Wu, Jingfeng, et al. "Last iterate risk bounds of sgd with decaying stepsize for overparameterized linear regression." International Conference on Machine Learning. PMLR, 2022.

Thanks for your prompt reply. Regrettably, this claim is simply wrong.

Convergence in distribution generally does not imply convergence in optimization and the two objectives measuring the convergence are not the same. If the reviewer insist on claiming one is stronger, we would encourage the reviewer to write down a step-by-step proof.

There are two mistakes the reviewer is making here.

First, in advanced probability theory, convergence in distribution is indeed stronger than convergence in expectation. However, in the context of optimization, this is totally different. The tool is simply applied in a wrong way here. In Can et al. (2019), it proved that the initial distribution converges to a stable distribution in terms of Wasserstain Distance. Let's omit the details and say that if the tool could be applied here, then we get:

• The expectation of the distribution, i.e. $\mathbb{E}[w_t]$, converges to the expectation of the stable distribution $\hat{w}$.

Whether $\hat{w}$ is $w_*$ remains unknown. Also, even $\hat{w} = w_*$ and $\mathbb{E}[w_t] \to w_*$, it doesn't imply $\mathbb{E}[f(w_t)] - f(w_*) \to 0$, just consider the simple case of $w_t \in U(-1,1)$ (uniform distribution) and $f(w) = w^2$.

The second mistake is the right hand side of the second inequlity $W_2(\nu(w), \delta_{w_*})^2$. Notice that in Can et al. (2019), it only says that the distribution converges to a stable distribution, whether this stable distribution is Dirac delta distribution at optimum $w_*$ still remains unknown.

The prompt reply is greatly appreciated. As the reviewer repetitively referred to Remark 23 in Can et al. (2019), we would like to point out some key aspects that the review missed in the remark.

• First, this is for Nesterov's Accelerated Gradient (AG), not Stochastic Heavy Ball (SHB).

• Also, we highly recommend the reviewer to double-check the proof for the remark, is $\alpha = \log^2(k) / (\mu k^2)$ a decaying learning rate?

As we provide techniques for SHB + decaying learning rates, differences in those key aspects matter a lot.

We appreciate your prompt reply.

Again, as we pointed out in Response 2 (R2) to the reviewer, there is absolutely no evidence that Can et al. (2019) implies $E[f(w_T)] - f(w_*)$ matching stochastic lower bound. The reviewer's claim is completely unsubstantiated.

Thank you for the prompt reply. We encourage the reviewer to write down the full proof rigorously and step-by-step.

Among various issues in the above proof idea, one is that the reviewer failed to take into account the effect of small learning rates, where the convergence is a logarithmic number of iterations only when the learning rate is large, such as constant learning rate $\alpha_{HB}$ defined in (12) of Can et al. (2019).

I wish I could have the privilege of omitting most of the details, since most of them are contained in Can et al. (2019). An argument that is completely similar to my previous comments will show $\mathbb{E} f(w_t) - f(w_*) \lesssim W_2^2(\nu(w_t), \nu_\infty) + \mathbb{E} \|w_\infty - w_* \|^2$, where $w_\infty\sim\nu_\infty$ is the stable distribution. The MSE $\mathbb{E}\|w_\infty-w_*\|^2$ is computed explicitly in (87), which can be adapted to (stagewise) decaying learning rate with only notational change. The authors may also notice that the proof in Appendix C.2 of Can et al. (2019) is actually very similar to this paper.

We thank for the reviewer's active participation during the discussion session.

Modifications

To reach a consensus, we have made corresponding modifications in a new revised version of our paper (version 3/V3) asked by the reviewer. Specifically,

• In the "Related Work" section, we provide the explicit form of accelerated distributional convergence for Can et al. (2019).
• In our contribution (Section 1.1), we emphasize more on the variance term $\tilde{O}(d \sigma^2/T)$.
• In paragraphs after Corollary 3, which explains the intuition behind one of our key technical contributions, we add a specific reference to the Lemma that provides the contribution.

Consensus

Regarding the reviewer's proposed two points, we think the reviewer now understands most parts correctly. However, there are still some missing points.

To avoid confusion, we would like to provide more clarification about the background and context:

• Lemma 19 in V1 and Lemma 8 in V2/V3: For Lemma 19, we assume the reviewer is referring to the version before rebuttal. This lemma corresponds to Lemma 8 in our revised versions during the rebuttal.
• Accelerating bias is easy: Showing the bias term to be converged linearly in $O(\sqrt{\kappa})$ for constant learning rates is straightforward, as it is essentially the same for deterministic HB, which is covered in almost every textbook about optimization. The main contribution of Can et al. (2019) is that it proves accelerated linear convergence rates for distributional convergence.
• Balancing bias and variance is hard: It is highly non-trivial to obtain accelerated linear convergence in bias while still achieving near-optimal rates in variance for SHB + decaying learning rates on quadratics, as to the best of our knowledge, no such work existed before.
• Other technical contributions: The techniques of using schedulers are also adopted in our proofs to balance bias and variance. Combining it with SHB is non-trivial as well.

About Top Conferences

We would also like to provide the reviewer with a list of past papers published in top conferences for reference, serving as a background for this subfield and top conferences' attitude towards it:

• COLT 2018: Modified SHB + averaging + least square regression (implies anisotropic noise on quadratics), achieving optimal rate and acceleration in certain types of noises [1]
• NeurIPS 2019: SGD + step decay + least square regression (implies anisotropic noise on quadratics), achieving near-optimal rate [2]
• ICLR 2022: SGD + step decay/improved scheduler + strongly convex least square regression & anistropic noise on quadratic, achieving near-optimal rate generally and optimal-rate for certain types of instances [3]
• ICML 2022 (Oral): SGD + step decay/polynomial decay + overparameterized linear regression (implies anisotropic noise on quadratics), a detailed analysis of lower bounds and upper bounds [4]

As SHB covers SGD, our result is technically much harder. Also to the best of our knowledge, this is the first time that an accelerated convergence rate is obtained while still achieving near-optimal variance for SHB + decaying learning rates on quadratics. It also implies SHB benefits in practice, which is mostly negative results in past literature.

We really appreciate the reviewer's prompt replies during the discussion and hope this can clarify the significance of our contributions.

[1]: Jain, Prateek, et al. "Accelerating stochastic gradient descent for least squares regression." Conference On Learning Theory. PMLR, 2018.

[2]: Ge, Rong, et al. "The step decay schedule: A near optimal, geometrically decaying learning rate procedure for least squares." Advances in neural information processing systems 32 (2019).

[3]: Pan, Rui, Haishan Ye, and Tong Zhang. "Eigencurve: Optimal Learning Rate Schedule for SGD on Quadratic Objectives with Skewed Hessian Spectrums." International Conference on Learning Representationsk(2022).

[4]: Wu, Jingfeng, et al. "Last iterate risk bounds of sgd with decaying stepsize for overparameterized linear regression." International Conference on Machine Learning. PMLR, 2022.