Paper: Accelerated Convergence of Stochastic Heavy Ball Method under Anisotropic Gradient Noise (ICLR-2024)

Reviewer comment (Writing):
pg5: eqn 3.11: missing - sign before $\beta$. (Same typo appears at pg 21 under the proof of Lemma 13 at two places)

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

Reviewer comment (Theory):
pg27: the equality applies (C.18) and (C.21): the last row of the matrix $2d$.

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

Reviewer comment (Experiments):
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 batch sizes. I'm curious whether SGD performs better when you multiply the stepsize by $\sqrt{M}$.

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
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.

|| ResNet-18 || MobilenetV2 || DenseNet-121 ||
| -- | -- | -- | -- | -- | -- | -- |
| Method |  Training loss | Test acc (%) | Training loss | Test acc (%) | Training loss | Test acc (%) |
| SGD | $4.13 \times 10^{-2}$  | $90.14$ | $3.40 \times 10^{-2}$ | $90.29$ | $3.93 \times 10^{-3}$ | $92.91$ |
| SHB  | $2.08 \times 10^{-3}$ | $91.10$ | $6.32 \times 10^{-3}$ | $91.07$ | $2.51 \times 10^{-4}$  | $93.58$ |

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.

Reviewer comment (Writing):
Minor typo in Algorithm 1, Line 6: $m$ ⇒ $M$.

Author response:
Thank you very much for pointing out these issues! We have fixed them accordingly in the revision.

Reviewer comment (Experiments):
Minor issue in Experiment description: initialize $\mathbf{w}^0$ from $(-1,1) \Rightarrow (-1,1)^d$.

Author response:
Thank you very much for pointing out these issues! We have fixed them accordingly in the revision.

Reviewer comment (Theory):
Minor issue in Proof of Lemma 3: $\succ$ should be $\succeq$.

Author response:
Thank you very much for pointing out these issues! We have fixed them accordingly in the revision.

Reviewer comment (Writing):
Minor typo on Page 29: genarality ⇒ generality.

Author response:
Thank you very much for pointing out these issues! We have fixed them accordingly in the revision.

Reviewer comment (Theory):
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.

Author response:
**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)](https://arxiv.org/pdf/1901.07445.pdf) 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)](https://arxiv.org/pdf/1901.07445.pdf) 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.


[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.

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

Author response:
**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.

Reviewer comment (Novelty):
My foremost concern is the comparison with previous literature, which I discussed in the `Weaknesses` part.

Author response:
**Summary: My previous evaluation is based on ... Considering that Theorem 2 also seems incremental compared with Can et al. (2019), I will keep my score unless there is strong evidence to refute this point.**

Theorem 2 is not incremental at all. It is the first work that achieves near-optimality statistical lower bound for SHB with step decay schedules on quadratic objective functions. To achieve this, we have invented new proof techniques that can be beneficial for future researches in this direction. It also clearly illustrated SHB's superiority over SGD under large batch settings.

The reviewer's claim of incremental is based on unsubstantiated claim that [Can et al. (2019)](https://arxiv.org/pdf/1901.07445.pdf) covers our result. In fact, that paper has no convergence result for stochastic optimization, not to mention matching statistical lower bound as we do. Moreover, we do not see any easy adapation of their results to derive our result proved in this work.