The Marginal Value of Momentum for Small Learning Rate SGD

Thank you for your comments.

---

Q1: I did not recall a theoretical result that claims acceleration of momentum for GD with a learning rate.

A1: The result quoted by the introduction comes from [1] (Theorem 9). Specifically it shows that the heavy-ball momentum with optimal hyperparameter constants enjoys a faster convergence rate than GD in the neighborhood of a local minimizer, and when the hessian is ill-conditioned the optimal learning rate with momentum can exceed $2/L$.

• [1] 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.

---

Q2: Is it surprising that momentum does not help when the LR is sufficiently small? The benefit of momentum over standard GD has only been claimed in the large LR regime

A2: There are works that claim the role of momentum (for instance [2]) for GD in terms of the implicit regularization in the setting where the LR is sufficiently small, and we are writing to further show that when there are enough gradient noises, then adding momentum does not help implicit regularization.

• [2] Avrajit Ghosh, He Lyu, Xitong Zhang, and Rongrong Wang. Implicit regularization in heavyball momentum accelerated stochastic gradient descent. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id= ZzdBhtEH9yB

---

Q3: The condition on $1-\beta$ in Definition 3.1 says that the past gradient will be quickly forgotten, so it is not surprising that SGDM is close to SGD in that case. What happens when $\alpha > 1$?

A3: When $\alpha>1$, there will be no limiting dynamics as $\eta\to 0$. For any non-zero initial momentum, the parameter will follow the initial momentum up to the scale $O(\eta^{1-\alpha})$, so when $\eta$ is small enough the trajectory will tend to explode as the parameters go unbounded. In general when $\alpha>1$, it takes longer time for the momentum to adapt to local curvature than for the parameter to go out of local areas, therefore conceptually such a momentum choice makes optimization uncontrollable.

Even when the initial momentum is zero and the initial point is near a strongly-convex local minimizer, the asymptotic convergence rate will still be $O(1/\eta^{(1+\alpha)/2})$ for $\alpha>1$, strictly slower than the rate $O(1/\eta)$ for the case when $\alpha<1$, which is not preferred for optimization.

Empirically, setting $\alpha>1$ at LR annealing also produces very bad performance. We believe the above are the reasons.

---

Q4: How does the scaling of the noise work? Lemma 2.4 and Theorem 3.5 appear to contradict. Can Theorem 3.5 be applied to full-batch GD with and without momentum?

A4: Yes Theorem 3.5 can be applied to full-batch GD where $\sigma = 0$. We write $\sigma\leq \eta^{-1/2}$ to make the theorem more general, but the most interesting case is $\sigma= \eta^{-1/2}$. As Lemma 2.4 shows, when $\sigma= \eta^{-1/2}$, the gradient noise has non-negligible impact on the loss curve, while when $\sigma\ll \eta^{-1/2}$, the impact of the gradient noise is negligible and SGD will have a similar trajectory as GD.

We realized that the paragraph after Lemma 2.4 may cause confusion, as our results actually hold for a range of $\sigma$ as stated in the theorems. We will fix the confusion in the revision, and thank you for the suggestion.

---

Q5: Does the learning rate for SGD need a rescaling in the CIFAR-10 experiments?

A5: Yes in all the experiments the learning rate needs to be rescaled so that $\eta$ matches for SGD and SGDM (or $\eta = \gamma/(1-\beta)$). We will add that into the experiment details in the revision.

Thank you for your comments.

---

Q1: The assumptions that the paper considers to show that SGDM approximate SGD is too restrictive. It would be great if they can justify these too restrictive assumptions.

A1: The assumptions we make are standard from a series prior works [1-4] on the trajectory analysis of stochastic updates, and theoretical findings from those works have been predictive of complex real-world settings (e.g., training vision models and fine-tuning language models, see [2,5] for instance).

Specifically, for training smooth loss on a neural network with smooth activations (Swish, sigmoid, GeLU etc) and weight decay, our assumptions will hold. In practical settings, all the data points and networks weights are actually mathematically bounded for a fixed number of steps, and the corresponding assumptions mostly serve for mathematical rigor.

References:

• [1] Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. Journal of Machine Learning Research, 20(40): 1–47, 2019.
• [2] Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling SGD with stochastic differential equations (SDEs). Advances in Neural Information Processing Systems, 34:12712– 12725, 2021a.
• [3] Li, Zhiyuan, Tianhao Wang, and Sanjeev Arora. "What Happens after SGD Reaches Zero Loss?--A Mathematical Framework." In International Conference on Learning Representations. 2021.
• [4] Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the SDEs and scaling rules for adaptive gradient algorithms. Advances in Neural Information Processing Systems, 2022.
• [5] Elkabetz, Omer, and Nadav Cohen. "Continuous vs. discrete optimization of deep neural networks." Advances in Neural Information Processing Systems 34 (2021): 4947-4960. Blog http://www.offconvex.org/2022/01/06/gf-gd/.

Thank you for your comments.

---

Q1: Although the implications of Theorem 3.5 is clear, the statement of the theorem is a bit confusing. Specifically, the averaged learning schedule is introduced, and it does not appear again until the appendix.

A1: The averaged learning rate schedule is defined in our theorem as the formula in the parentheses directly followed. We will find a better way to introduce this notion in the revision to eliminate the confusion.

---

Q2: It would be interesting to see other concrete types of hyperparameter schedule apart from the constant ones, both theoretically and empirically. It would also be interesting to see experiments demonstrating the main theoretical contribution of this paper, i.e., the trajectories rather than the loss / accuracy stay close.

A2: Our theory is proved for a general hyperparameter scheduling (as long as its scale is controlled), so the conclusions hold for non-constant schedules. Also, for the experiments, we used the 3-stage annealing learning rate schedule and added a linear warm-up, which is technically not constant. In the future, we can surely provide more experiments with more varying hyperparameter schedules.

Our theory prescribes that the trajectories are close according to any test function (Definition 3.4), and we used training and test curves which are the most empirically concerned test functions. In the future we can add more test functions to show the closeness of the trajectory distributions.

---

Q3: Can we understand what’s going on when $t=\Omega(1/\eta)$ and $t=O(1/\eta^2)$? That is, between the two regimes?

A3: The two regimes are separate in our results because they entail different assumptions on the landscape. The $\Omega(1/\eta)$-step result holds for general smooth landscapes, while for the trajectory to be tractable for $O(1/\eta^2)$ steps, we need stronger assumptions on the landscape (i.e. the existence of a minimizer manifold). The two regimes also have different levels of tolerance of the noise scale $\sigma$.

When the $O(1/\eta^2)$-trajectory is tractable, the first $O(1/\eta)$ steps in the trajectory actually follows an SGD trajectory that decrease the loss from initial to almost zero. At the end of the first $O(1/\eta)$ steps, the parameter will be $o(1)$ close to the minimizer manifold, and it stays in the neighborhood of the manifold for $O(1/\eta^2)$ steps and do a $O(\eta)$ progression along the manifold every $O(1/\eta)$ steps.

---

Q4: Why is a class of hyperparameter schedules used in the main theorems?

A4: For a specific hyperparameter scheduling, we can indeed use our theorem to bound the distance between SGDM and SGD trajectories.

We write a class of schedules in our theorem to make the theorem more general and to emphasize the dependence of our result on the scale of $\eta$, e.g. when $\eta\to 0$ the trajectories indeed have limiting distributions at a rate $O(\eta^{(1-\alpha)/2})$ for some limiting hyperparameter schedule.

Thank you for your comments.

---

Q1: The analysis is based on the fact that the momentum effect is dominated by other factors when the learning rate is negligible. This is somewhat intuitive and needs to dive deeper into what the special role of momentum is.

A1: We are not sure about what is excluded from “other factors”, and we appreciate it if more clarifications are provided.

The main message we wish to convey through this paper is that momentum has no special role when the learning rates are small, and tuning the momentum parameter has very limited benefits as model performance is agnostic to the choice of momentum parameters over a range of hyperparameter scales. This conclusion is supported by both our theoretical and empirical findings.

For some clarifications of our results, first, we believe that it is incorrect to claim that the momentum effect is dominated by changing the curvature-induced instability, rather than changing the noise-induced instability (defined in Lemma 2.4). Adding momentum actually modifies the noise part of parameter updates more than the curvature part. Second, simple intuitions like Section 3.1 cannot offer any suggestions on the number of steps when adding momentum accumulates non-negligible deviations. For instance, as our setting involves a $O(\sqrt{\eta})$-normed update per step, two adjacent SGD updates may differ by $O(\sqrt{\eta})$, therefore after $O(1/\sqrt{\eta})$ steps an SGDM trajectory will deviate from an SGD trajectory by the above intuition, which only accounts for $O(\sqrt{\eta})$ training loss decrease. Therefore even when the learning rate is small enough, it is not trivial to conclude that momentum has a negligible effect over the whole trajectory of $O(1/\eta)$ or $O(1/\eta^2)$ steps, when the training loss goes from initial to almost optimal and plateaus.

---

Q2: There needs to be clear examples of what scenarios the theory fits in. For example, with what class of loss and neural network does the theory fit in?

A2: For a short clarification, as our theory is a proof of the general non-convex landscape, it aims to model model training with any smooth loss, any neural networks architecture with any smooth activations (with polynomial growth) and small learning rates. For instance training with square loss for regression or logistic loss for classification falls into the assumptions. The boundedness of noise covariance and hessian are satisfied when the weights of the network have bounded moments during the course of the training, which is satisfied when weight decay is involved for general neural networks.

The assumptions of the paper are standard in a series of trajectory analysis works [1-4].

References

• [1] Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. Journal of Machine Learning Research, 20(40): 1–47, 2019.
• [2] Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling SGD with stochastic differential equations (SDEs). Advances in Neural Information Processing Systems, 34:12712– 12725, 2021a.
• [3] Li, Zhiyuan, Tianhao Wang, and Sanjeev Arora. "What Happens after SGD Reaches Zero Loss?--A Mathematical Framework." In International Conference on Learning Representations. 2021.
• [4] Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the SDEs and scaling rules for adaptive gradient algorithms. Advances in Neural Information Processing Systems, 2022.

Furthermore, our experiment results further corroborate that our conclusion holds for general small-batch training from image tasks to natural language tasks, even when the learning rate is large. In practice, this implies that model performance is agnostic to the choice of momentum parameters when the batch sizes are scaled down.