Exploring the Inefficiency of Heavy Ball as Momentum Parameter Approaches 1

1 Introduction

Training a given machine learning model y = g(x,w), parameterized by w, is typically formulated as solving the following empirical risk minimization (ERM) problem

where L denotes the loss function, (xi, yi) represents the ith data-label pair, and n is the size of the training dataset S. The main workhorse for solving the ERM problem is stochastic gradient descent (SGD) [Robbins and Monro, 1951], especially when the dimension of w is very high, and the number of training data n is large. Starting from an initial point w1 ∈ Rd , SGD iterates as follows

∗Corresponding authors

MNIST Classification

We conduct experiments employing several machine learning models, including the ℓ2-regularized multi-class logistic regression and a multi-layer perceptron (MLP) for MNIST classification [LeCun and Cortes, 2010]. The dataset contains 60,000 and 10,000 grayscale images in the training and test sets. Note that the ℓ2-regularized logistic regression model is strongly convex, while the other is nonconvex. The training batch size is set to 256 for all the following tasks.

ℓ2-regularized multi-class logistic regression. We apply both SGD and SHB with a learning rate of 0.01 and use different momentum parameters β for SHB to train the multiclass logistic regression classifier with ℓ2-regularization (set as 0.01) on the weights (note that this regularization guarantees the loss function to be strongly convex). We first run the full batch gradient descent for 10,000 iterations with a learning rate of 0.01 to find the global minimum, denoted as w∗ , of the model. Figure 1 illustrates the behavior of the minibatch approximation of E[RS(wk ) - RS(w∗ )] (the expectation is obtained via averaging over five independent runs) with wk obtained from SGD or SHB with different β. It is evident that SHB converges faster than SGD at the early training stage, particularly as β increases. However, a large value of β, such as 0.9, increase oscillation and deteriorates generalization. Figure 1 also shows the test loss (b) and accuracy (c) of the model trained by different optimization algorithms over iterations, affirming that the heavy ball momentum can improve generalization.

Multi-layer perceptron (MLP). We next compare the performance of SGD and SHB with different momentum hyperparameters β in training a MLP for MNIST classification. The MLP architecture used in this study follows the same structure as described in [McMahan et al., 2017], consisting of two hidden layers with 512 units, and each is activated using ReLU. We run both SGD and SHB with different β for 60 epochs and employ a learning rate of 0.01. Figure 2 shows that the heavy ball momentum can accelerate SGD and improve its generalization for solving nonconvex optimization problems, and a larger β tends to yield better performance during the early training stage. However, SHB with a large β, such as 0.99, performs significantly worse than SGD or SHB with a smaller β in both optimization and generalization. Specifically, the training loss and test accuracy of SHB with β = 0.99 plateau at much worse results than that of SGD and SHB with a smaller β.

CIFAR10 Classification

We further consider using the ResNet18 [He et al., 2016a] model to classify the CIFAR10 [Krizhevsky et al., 2009] dataset. This dataset consists of 50,000 and 10,000 colored images in the training and test sets, respectively. We run both SGD and SHB with different β for 90 epochs, and the batch size is set to 256. For both algorithms, we use the same learning rate of 0.01. Figure 3 plots the training loss, test loss and test accuracy versus iterations. These results confirm the previous findings. Here, we notice that when β = 0.9, SHB significantly outperforms SHB with a smaller β and SGD, which resonates with the practical choice of momentum for training ResNets [He et al., 2016a]. But if we use an even bigger β, such as 0.99, SHB performs worse than SGD and much worse than SHB with a smaller β.

Numerical Observations Summary: We summarize four empirical findings from the above experiments:

• When β is very close to 1, SHB exhibits slower performance compared to SGD after several iterations, and the training curve of SHB becomes highly oscillatory.
• SHB still converges faster than SGD in the early training stage, even when β is very close to 1.
• SHB with a larger β converges faster than SHB with a smaller β, provided that β is not very close to 1.
• The test accuracy of the model trained by SHB with a β that is very close to 1 is also worse than that trained with a smaller β. That is, SHB with a β close to 1 generalizes worse than SHB with a smaller β.

Although SHB has been studied in various settings, we noticed that the theoretical interpretations for the aforementioned numerical observations have not been established yet. The focus of this paper is to provide theoretical interpretations for these empirical phenomena.

1.2 Contributions

In this paper, we aim to develop a theoretical understanding of the empirical observations above. Our contribution lies in theoretically analyzing the detriment effect of the heavy ball momentum on SGD and proposing a more efficient heavy ball momentum acceleration algorithm based on this finding. Specifically, we summarize our main contributions into the following fourfold.

• When β is close to 1, we have proven that the upper bound complexity of SHB is dominated by a smaller constant geometric factor compared to SGD in the early iterations. This result elucidates why SHB outperforms SGD in the initial training stage. For more detailed information, please refer to Section 3.1.
• We demonstrate why SHB can be slower than vanilla SGD when β is very close to 1 after a certain number of iterations. Specifically, we establish that SHB has a larger lower bound in complexity for two consecutive iterates compared to SGD. This lower bound helps explain why SHB experiences slower convergence, reduced accuracy, and decreased stability when β approaches 1. Please refer to Section 3.2 for a more detailed explanation.
• When β is not near 1, we provide a theoretical explanation for why SHB with a larger value of β converges faster than SHB with a smaller value of β. More detailed information can be found in Section 3.3.
• Building upon previous understandings and insights of momentum, we introduce a descending warmup method for SHB. This technique accelerates convergence, enhances accuracy, and improves the stability of SHB. For more detailed information, please refer to Section 4.

RS(w k ) - RS(w * )

1.3 RelatedWork

We will briefly review some of the most relevant literature regarding the development of the heavy ball momentum algorithm and its theoretical foundations. The heavy ball momentum algorithm was initially proposed by Polyak in 1964 [Polyak, 1964], and since then, its convergence properties have been extensively studied in both convex and nonconvex settings [Ochs et al., 2014; Ghadimi et al., 2015; Yuan et al., 2016; Sun et al., 2019c]. In a notable paper by Sun et al. [Sun et al., 2019a], it was proven that heavy ball momentum can effectively escape saddle points with larger learning rates compared to traditional gradient descent. SGD with heavy ball momentum has become a widely used and powerful algorithm for machine learning tasks [Krizhevsky et al., 2009; Sutskever et al., 2013; Sun et al., 2021]. Recent research efforts have focused on developing efficient heavy ball-style algorithms for training machine learning models, particularly in the field of deep learning [Yan et al., 2018; Ma and Yarats, 2019; Sun et al., 2019b; Gitman et al., 2019; Sun et al., 2020]. One notable extension of the heavy ball momentum algorithm is Nesterov's acceleration, which incorporates lookahead momentum scaled by an iteration-dependent weight [Nesterov, 1983]. This variant has proven to be efficient in accelerating gradient descent [Nesterov, 2005; Nesterov, 2013]. However, the stochastic version of Nesterov's acceleration [Wiegerinck et al., 1994] faces challenges such as error accumulation and potential non-convergence [Kulunchakov and Mairal, 2019; Vaswani et al., 2019; Aybat et al., 2020; Wang et al., 2022]. [Kidambi et al., 2018; Liu and Belkin, 2020; Ganesh et al., 2023] studied the inefficiency of SHB in specific cases, but failed to explain the inefficiency phenomenon when the momentum parameter approaches 1. To assess the generalization error of empirical risk minimization, [Bousquet and Elisseeff, 2002] introduced the concept of uniform stability. Building upon this notion, researchers such as [Ong, 2017; Chen et al., 2018] have further investigated the bounds of generalization error for the specific cases of SHB. Additionally, [Ramezani-Kebrya et al., 2024] have established a generalization bound for SHB in more general scenarios. These studies contribute to our understanding of the performance of SHB and shed light on its potential for generalization in machine learning tasks.

2 Preliminaries

2.1 Notation

We denote scalars and vectors by lowercase and bold lowercase letters, respectively, and denote matrices by uppercase boldface or uppercase curly letters. For a vector x = (x1, · · · , xd) ∈ Rd , we denote its ℓ2 norm by ∥x∥. The transpose of a matrix A is denoted as A⊤, its spectral norm and spectral radius are denoted as ∥A∥ and ρ(A), respectively. λmax(A) and λmin(A) represent the largest and smallest eigenvalues of matrix A, respectively, and the trace of a matrix is denoted by Tr(·).

Given two sequences {am} and {bm}, we write am = O(bm) if there exists a positive constant 0 < C < +∞ such that am ≤ Cbm, and we write am = Θ(bm) if there exist two positive constants C1 and C2 such that am ≤ C1bm and bm ≤ C2am. We use am = ˜O(bm) and am = ˜Θ(bm) to hide the logarithmic factor on top of O(·) and Θ(·), respectively.

For a function f(w) : Rd → R, we denote its gradient and its Hessian as ∇f(w) and ∇2 f(w), respectively. The minimum value of f(w) is denoted as minf, and the minimizer of f(w) is denoted as w∗ . We use E[·] to denote the expectation with respect to the underlying probability space.

2.2 Assumptions

Before presenting our theoretical results, we first collect several necessary assumptions for the subsequent analysis.

Assumption 1. The stochastic gradient is an unbiased estimate of the true gradient and has bounded variance, i.e.,

and E∥ek ∥2 ≤ σ2 for some σ > 0.

Assumption 1 is a widely accepted fundamental principle in the stochastic optimization community. It is worth noting that in the finite-sum minimization, the gradient is typically approximated using the mini-batch gradient, denoted as gk = (∑mi=1 ∇L(g(xki , wk ); yki )) /m, where m ≪ n represents the batch size (commonly referred to as mini-batch SGD). This stochastic gradient is unbiased. However, relying solely on this assumption is insufficient to provide a theoretical explanation for all the experimental phenomena reported in the introduction. Therefore, to further investigate and elucidate these phenomena, additional assumptions need to be employed.

Assumption 2. The noise ek satisfies some distribution E, whose covariance matrix

is positive semi-definite and Tr(Σ) > 0.

In [Jastrze¸bski et al., 2018], the authors assume that the covariance matrix of the stochastic noise is constant, and they provide an explanation for this assumption within the context of mini-batch SGD. It is important to note that Assumption 2 indicates Assumption 1, but not vice versa. A specific case of E in Assumption 2 is the normal Gaussian distribution, such as gradient Langevin dynamics [Welling and Teh, 2011; Stephan et al., 2017; Gitman et al., 2019], where Σ reduces to the identity matrix. In our paper, there is no need to assume a normal Gaussian distribution. Instead, we only require the covariance matrix of the distribution to be positive and semidefinite, which is a much weaker condition and holds for almost any distribution. The positivity of the trace of the covariance matrix ensures that E∥ek ∥2 = Tr(Σ) > 0.

3 Heavy Ball Slows Down the Convergence of SGD When β Is Close to 1

In this section, we analyze the convergence of SHB when RS is quadratic and strongly convex, i.e., ∇2RS(w) ≡ A is positive definite.

Main Techniques. First, we briefly introduce the main techniques that we used to establish the theoretical results of this study. Specifically, we reformulate SHB as follows

where

and I ∈ Rd×d is the identity matrix. Let

We then have yk+1 = T yk - γek , and the iteration can also be formulated as the sum of a power series, given below

Since the noises are independent and have zero mean, taking expectations on both sides of the equation yields

Therefore, the problems under study turn to estimating the lower and upper bounds of E∥T k-iei ∥2 . The bounded variance of the noise directly gives us

The remaining problem lies in providing upper bounds of ∥T k ∥ for a fixed k ∈ Z+.

Lemma 1. Assuming that the matrix A ∈ Rd×d is symmetric positive definite, and 0 < ν ≤ λmin(A). Let ϵ > 0 be small enough and γ = Θ(ϵ), and 0 ≤ β = (1- √γν)2 +ϱ < 1 for 0 < ϱ ≪ ϵ. We then have

where C1 > 0 is a constant independent of k and γ.

Here, the constant ϱ is introduced to ensure that T has distinct eigenvalues, making it similar to a diagonal matrix and simplifying the analysis. Determining the lower bound requires a non-trivial analysis, which relies on Assumption 2. Assuming that Assumption 2 holds, we can derive the following result regarding the matrix T .

Lemma 2. Let E be the stationary distribution of stochastic noise and ξ ∼ E ⊕ 0d. The momentum parameter satisfies β = 1 - Θ(γ τ ) with 1 ≤ τ ≤ 2. Under Assumption 2, we can establish that

for some C2 > 0 which is only dependent on A.

With this in place, we can now present our main results.

3.1 Why Faster in the Early Training

Theorem 1. Consider a quadratic function RS with its Hessian defined as ∇2RS(w) ≡ A. The sequence {wk }k≥0 is generated by (1). Assume that 0 < ν ≤ λmin(A) ≤ λmax(A) ≤ L and Assumption 1 holds. For any sufficiently small ϵ > 0, let γ = Θ(ϵ) be chosen such that 0 ≤ β = (1 - √γν)2 + ϱ < 1, where ϱ satisfies 0 < ϱ ≪ ϵ. Then the output of SHB satisfies

Theorem 1 provides an upper bound when the momentum hyperparameter β is close to 1, and it solely relies on Assumption 1. In the early training stage, i.e., when the error ˆ ϵ is large, we can set ˆ ϵ = √ϵ/ν 32 , then achieving E∥wK - w∗ ∥2 = O(ˆ ϵ) for a given K only requires (1 - √γν)2K /γν = O(ˆ ϵ). Therefore, the worst-case bound for K of SHB is given by

In comparison to SGD, which requires ˜ O(1/(ϵν)) iterations to achieve the desired solution with error ˆ ϵ in the strongly convex case [Rakhlin et al., 2012], SHB significantly accelerates the convergence of SGD in the early training stage. Consequently, when the desired error is on the order of O(ˆ ϵ), SHB exhibits faster convergence than SGD. Thus, even when the momentum hyperparameter β is close to 1, SHB remains faster than SGD in the early training phase for large desired errors.

Algorithm 2 Stochastic Heavy Ball with Descending Warmup (SHB-DW)

Require: parameters γ > 0, 0 ≤ β < β < 1, 0 < α < 1 Initialization: w0 = w1 , β = β

for k = 1, 2, . . .

step 1: update βk according to (7)

step 2: get gk being an unbiased sample of ∇RS(wk )

step 3: update wk as (1) with β ← βk

end for

3.2 Why Slower and Unstable after Enough Iterations

Theorem 2. Let conditions of Theorem 1 and Assumption 2 hold, and β = 1 - Θ(γ τ ) with 1 ≤ τ ≤ 2. Then the output of SHB satisfies

for any integer k ≥ 1.

If τ > 1, Theorem 2 indicates that no matter how many iterations are done, SHB will never achieve an error of O(ϵ). On the other hand, the SGD algorithm can reach an error of O(ϵ) after O(1/(ϵν)) iterations with a suitable learning rate γ = Θ(ϵ). This means that given enough iterations, SHB will be slower than SGD.

Theorem 2 also provides an explanation for the instability of the SHB curves. If E∥wK - w∗ ∥2 ≪ ϵ2-τ for some K ∈ Z+, then

3.3 Why SHB Converges Faster as β Increases Provided That β Is Not Close to 1

Given a fixed positive number β0 such that 1 - β0 ≫ ϵ, we study the convergence of the SHB algorithm. We first present a lemma as follows.

Lemma 3. Let A ∈ Rd×d be a symmetric positive definite matrix. Assume that ϵ > 0 is small enough and γ = Θ(ϵ). Let 0 ≤ β ≤ β0 < 1 and 1 - β0 ≫ ϵ, we can establish that

where constants C3, C4 > 0 are independent of k and γ.

With Lemma 3, we can prove the convergence of SHB.

Theorem 3. Let RS be a quadratic function (∇2RS(w) ≡ A), and {wk }k≥0 be the sequence generated by (1). Assuming that 0 < ν ≤ λmin(A) ≤ λmax(A) ≤ L and Assumption 1 holds. Given any ϵ > 0 small enough, let γ = Θ(ϵ) such that 0 ≤ β ≤ β0 < 1 and 1 - β0 ≫ ϵ. Then to achieve E∥wK -w∗ ∥2 ≤ ϵ, the worst-case iteration complexity is

Theorem 3 provides an explanation for why SHB can enjoy faster speed when β increases but not close to 1. As β ∈ [0, β0], the worst-case complexity decreases if β increases.

RS(w k ) - RS(w * )

4 DescendingWarmup

Our results above demonstrate that SHB with a large β accelerates convergence during the early training stage but causes instability and reduces accuracy after sufficient iterations. Based on this insight, we propose an initialization method that employs a large momentum hyperparameter β0 := β ∈ (0, 1), which is then reduced to a fixed number β ∈ [0, 1) according to the following rule

where 0 < α < 1 is the descending factor. It is straightforward to see that {βk }k≥0 will be unchanged and remains as β after the following number of iterations

Therefore, (7) is actually a warmup process, and called Descending Warmup (DW) in this study. We formulate SHB with DW (SHB-DW) in Algorithm 2. In the early iterations of SHB-DW, relatively large momentum hyperparameters are utilized to achieve fast convergence, while smaller momentum hyperparameters are employed in the subsequent iterations to produce stable and accurate solutions.

Let the sequence {wk }k≥0 be generated by the SHB-DW algorithm. Since the iterates of SHB-DW after K(β, β, α) iterations can be regarded as SHB with an initialization set as wK(β,β,α) , the stationary analysis of SHB-DW is then trivial and will not be repeated. To get the fine-grained convergence result for SHB-DW, we just need to bound ∥w∗ - wK(β,β,α)+1 ∥ under the iteration-dependent momentum hyperparameters.

5 Numerical Experiments

5.1 DescendingWarmup Improves SHB in Optimization and Generalization

In this section, we replicate the experiments conducted in Section 1.1 using the proposed SHB-DW algorithm. Our theoretical findings indicate that a larger momentum parameter can be employed to speed up the training process in the early stages, and subsequently, a β value not close to 1 is required to ensure optimal performance. By default, we set the initial momentum hyperparameter β0 to 0.999 for warmup and decrease it by a factor of α = 0.999 after each iteration. For the ℓ2-regularized multi-class logistic regression task, the lower bound of the momentum hyperparameter β for SHB-DW is set to 0.1. While for the other two tasks, Figures 2 and 3 show that SHB achieves favorable results with β = 0.9. Therefore, we set β of SHB-DW to 0.9 for the two tasks. Additionally, we compare the widely used Adam algorithm [Kingma and Ba, 2015]. Since Adam requires a relatively small learning rate to fully utilize its performance [Kingma and Ba, 2015], we set the learning rate of this algorithm to 0.0001, while keeping all other parameters unchanged. In Figures 4, 5, and 6, the performance of SGD, Adam, SHB (β = 0.9, 0.99), and SHB-DW is compared in the three benchmark experiments described in Section 1.1. These results confirm the superiority of SHB-DW: the decreasing warmup approach accelerates training in the early stages, and compensates for the performance degradation caused by the use of a large β.

5.2 SHB-DW for Deep Learning

To further validate the effectiveness of the descending warmup technique in training deep neural networks, we conduct more detailed experiments on CIFAR10 using ResNet18 and ResNet34 [He et al., 2016b]. We compare four optimization algorithms, SGD, Adam, SHB, and SHB-DW. The momentum parameter β for SHB is set to 0.9, a commonly utilized value in practice [He et al., 2016a], and the settings for SHB-DW remained the same as in Section 5.1. The algorithms are run for 200 epochs with a batch size of 256. The initial learning rate for Adam is set to 0.001, while the others are set to 0.1. All algorithms use a decreasing learning rate strategy, i.e., decreasing by a factor of 10 at the 60th, 120th and 180th epochs, respectively. Table 1 presents the test accuracy of the two models trained by the four different optimizers, and the training process is illustrated in Figure 7 (ResNet34). Compared with other optimization algorithms, SHB-DW exhibits greater stability and efficiency, outperforming SGD, Adam, and SHB in test accuracy.

6 Concluding Remarks

In this paper, we have provided an explanation for the deterioration of convergence speed in the heavy ball momentum when the hyperparameter β approaches 1. Our analysis is based on novel reformulations of the heavy ball algorithm, allowing us to utilize non-Lyapunov analysis to establish the lower bound of SHB. To the best of our knowledge, our results are the first to explain how a hyperparameter close to 1 in heavy ball momentum can negatively impact convergence and lead to instability in SHB.

There are several potential directions for future research: 1) Can we extend our analysis to Nesterov's momentum, another widely utilized technique in deep neural network training? 2) Can we generalize our analysis to Nesterov's acceleration with restart, which has been empirically observed to significantly enhance the performance of SGD [Wang et al., 2022]? 3) Can we establish similar theoretical results for heavy ball momentum when combined with adaptive learning rate algorithms, such as Adam?

Ethical Statement

This study focuses on the fundamental optimization algorithm and its analysis. No ethical issues are discussed or related to the research.

Acknowledgments

This work was supported in part by the National Science Foundation of China (No. 62025208, No. 62376278), Hunan Provincial Natural Science Foundation of China (No. 2022JJ10065), Young Elite Scientists Sponsorship Program by CAST (No. 2022QNRC001), Continuous Support of PDL (No. WDZC20235250101), and Xiangjiang Laboratory Fund (No. 22XJ01012).