Provable Acceleration of Nesterov’s Accelerated Gradient Method over Heavy Ball Method in Training Over-Parameterized Neural Networks

1 Introduction

Nowadays, the first-order gradient method is widely adopted as the standard choice for training neural networks [Schmidhuber, 2015]. The gradient descent (GD) method is the most famous first-order method, which only accesses the current gradient to update the model. To accelerate its convergence, the momentum method is proposed by utilizing the historical gradients. The idea of incorporating momentum term into GD can be traced back to 1964, when Polyak introduced the heavy ball (HB) method for solving the quadratic convex problem [Polyak, 1964]. Compared to GD, HB achieves a local accelerated convergence rate in a neighborhood around the global optimal. Subsequently, Nesterov proposed a different type of momentum method known as Nesterov's accelerated gradient (NAG)method [Nesterov, 1983]. In the convex setting, NAG achieves the optimal global convergence rate [Nesterov, 2013], which improves the local convergence of HB.

∗Corresponding Authors.

Nowadays, neural networks have gained remarkable performance in a various tasks owning to their overparameterization, such as GPT-3 with 175 billion parameters [Brown et al., 2020]. As a result, the training process of neural networks suffers from significant computation costs. To save the expensive computational resources, momentum methods are widely adopted in the deep learning community due to their acceleration over GD. Besides, NAG converges faster than HB in practical neural network training [Schmidt et al., 2021; Nado et al., 2021]. Unfortunately, this phenomenon lacks sufficient understandings mainly due to the non-convexity faced by the optimization with neural networks. Recently, some advances have been achieved in providing explanations for the convergence of GD in training over-parameterized neural networks. Later, some works extended these findings to momentum methods [Wang et al., 2021; Bu et al., 2021; Liu et al., 2022b]. However, their analysis only establishes convergence guarantees, but stops short of distinguishing the convergence difference between HB and NAG. Specifically, existing works can be roughly categorized into two parts: discrete-time and continuous-time. For the discrete-time, [Wang et al., 2021] proved that HB can reach a global minimum when training two-layer neural networks. Likewise, [Liu et al., 2022b] established the global convergence guarantee for NAG. These two researches only provide convergence results with specific hyper-parameters, but it is still unclear whether NAG theoretically converges faster than HB in training neural networks. From the continuous-time, [Bu et al., 2021] exploited the connection between the limiting ordinary differential equations (ODEs) and the momentummethods. However, their result relies on a low-resolution ODE, which fails to differentiate the convergence behavior between HB and NAG. As a result, there still exists a gap between practice and theory.

Motivated by the fact that the high-resolution ODEs provide more accurate characterizations than low-resolution ODEs, we study the training dynamics of momentum methods from a high-resolution perspective and establish an explicit convergence difference between HB and NAG in training neural networks. For the architecture of the neural network, we focus on an over-parameterized two-layer neural network with ReLU activation, which is extensively studied in prior works [Du et al., 2019c; Wang et al., 2021; Liu et al., 2022b]. Our contributions are summarized as follows:

• Firstly, we derive the residual dynamics of HB and NAG by applying high-resolution ODEs and neural tangent kernel (NTK) theory, which reveals that the original non-convex weight dynamics corresponds to a strongly convex residual dynamics under over-parameterized assumption.
• Secondly, with the Lyapunov-based approach, we investigate the residual dynamics via a sharper Young's inequality instead of Cauchy-Schwarz inequality. Compared to existing works, we improve the convergence rates of HB and NAG with tighter upper bounds. Besides, our result provides the first theoretical guarantee for the acceleration of NAG over HB, which is mainly caused by the extra gradient correction term of NAG.
• Finally, we validate the findings through a convergence comparison between HB and NAG under different settings of hyper-parameters on three benchmark datasets. The maximum distance between the training iterate and its initial value is also tested to support our theory.

2 RelatedWorks

2.1 MomentumMethods

Momentum methods apply the historical information of gradients to accelerate GD. For the quadratic convex problem, Polyak proved HB achieves a local iteration complexity O(√κ log(1/ϵ)) [Polyak, 1964], where κ > 1 is the condition number of the quadratic problem. Compared to the iteration complexity O(κ log(1/ϵ)) of GD, HB achieves an acceleration. For L-smooth and µ-strongly convex problems,[Ghadimi et al., 2015] proved that the average iterates of HB converge to a global optimal with O( Lµ log(1/ϵ)) iteration complexity, similar to the convergence of GD. Nevertheless, [Lessard et al., 2016] discovered a counterexample that HB fails to converge in tackling some smooth and strongly convex problems. In contrast, NAG enjoys a global iteration complexity O(√κ log(1/ϵ)) in the same setting [Nesterov, 2013].

Meanwhile, researchers investigated the connection between discrete momentum methods and ODEs. With the step size towards infinitesimal, [Wilson et al., 2021] derived HB and NAG share the same ODE formulation. [Su et al., 2014] obtained the limiting ODE for NAG with a time-varying momentum parameter. Subsequently, [Shi et al., 2021] applied dimensional analysis to establish high-resolution ODEs for momentum methods, where HB and NAG exhibit different limiting ODEs. However, in the continuous viewpoint, [Shi et al., 2021] only proved HB and NAG have the same convergence rate O(e-√µt/4 ) for L-smooth and µ-strongly convex problems. By discretizing the Lyapunov function and the high-resolution ODEs, Shi further proved that NAG converges faster than HB in the strongly convex setting. Additionally, [Sun et al., 2020] improved above continuoustime result of NAG with a faster O(e-3√µt/7 ) convergence rate. Recently, [Zhang et al., 2021] further derived a shaper O(e-√µt/2 ) result.

2.2 Neural Tangent Kernel

[Jacot et al., 2018] established the connection between the optimization of infinite wide neural networks and NTK. Later, this result was further extended in providing convergence guarantees of GD for different architectures of neural networks, including convolutional neural network [Arora et al., 2019a], recurrent neural network [Alemohammad et al., 2020], residual network [Huang et al., 2020] and graph neural network [Du et al., 2019a].

With the NTK theory, there is an increasing interest in explaining the acceleration phenomenon of momentum methods in training neural networks. From the discrete-time view, [Wang et al., 2021] established the first convergence result of HB in training two-layer ReLU neural networks and deep linear neural networks. Besides, their result theoretically shows the acceleration guarantee of HB over GD. In addition, [Liu et al., 2022b] presented the convergence result of NAG in training two-layer neural networks. From the continuous-time view, [Bu et al., 2021] exploited the limiting ODEs in analyzing the convergence of momentum methods for training neural networks. However, their result relies on low-resolution ODEs, which fails to capture the convergence difference between HB and NAG. At the same setting, our work differentiates the convergence behaviors between HB and NAG through sharper inequality and analysis. Specifically, our result shows that HB converges at a O(e-(2-√2)√µt ) convergence rate and NAG enjoys a faster O(e-ρNAG√µt ) rate, where 2 - √2 < ρNAG ≤ (11 - √73)/4. The comparison of existing results for HB and NAG from the high-resolution view can be founded in Table 1.

3 Preliminary

3.1 Problem Setting

In this paper, we focus on the supervised learning problem for training a two-layer ReLU neural network with m hidden nodes:

where x∈Rd denotes the input feature, σ(z)=max{0, z} denotes the ReLU activation function, W = (w1, · · · ,wm) ∈ Rd×m is the parameter of the hidden layer and a ∈ Rm is the parameter of the output layer. For simplicity, we assume the input feature has unit norm, i.e., ∥x∥2 = 1, which follows [Bu et al., 2021; Du et al., 2019c]. For initialization, wr(0) ∼N(0, Id) and ar ∼Rademaher(0.5) for any r ∈ [m]. According to [Wang et al., 2021; Bu et al., 2021; Du et al., 2019c], we use the square loss as the loss function. Then, it can find the optimal parameter by minimizing the empirical risk ℓ to :

where (xi, yi) ∈ Rd × R denotes the i-th training instance. Following [Bu et al., 2021; Du et al., 2019c], we only update W and keep a fixed during training. Note that the problem in Eq.(2) is both non-convex and non-smooth. Recently, some works provided global convergence guarantees for gradient-based methods in optimizing Eq.(2) [Bu et al., 2021; Du et al., 2019c]. GD is the most widely used gradient-based method due to its simplicity and efficiency, which updates as

where η > 0 is the learning rate.

For accelerating GD, momentum methods blend history information of gradients into current updates. HB starts from the initial parameter wr(-1)=wr(0) and updates as follows

where β ∈ [0, 1) denotes the momentum parameter. According to [Nesterov, 2013], NAG has several variants with different schemes of hyper-parameters. In this paper, we mainly focus on NAG with a constant momentum parameter β. Given the initial parameters wr(-1) = wr(0), NAG involves the update procedure as

Compared to HB, NAG has an extra term βη( ∂ℓ(W(t),a) ∂wr(t) - ∂ℓ(W(t-1),a) ∂wr(t-1) ), which is associated with the difference between two consecutive gradients and is referred to as gradient correction [Shi et al., 2021]. This term is crucial for the provable acceleration of NAG over HB in the convex setting, but lacks enough understandings in the non-convex setting.

3.2 Continuous View of MomentumMethods

Continuous view of GD. Before introducing the limiting ODEs of momentum methods, we present existing result of GD in training two-layer ReLU neural networks.

When learning rate tends to infinitesimal, GD corresponds to the gradient flow

Based on Eq.(2), the time derivative of the hidden layer has

where f(t) = (f1, · · · , fn)⊤ ∈ Rn with fi = f(W(t), a; xi) and y = (y1, · · · , yn)⊤ ∈ Rn . Denotes ∆= f -y as the residual vector, it has

where H(t) is a Gram matrix induced by the two-layer ReLU neural network with following formulation

When m goes to infinity, H at the initialization has

which holds for any i, j ∈ [n]. Note that H∞ is positive definite with following lemma.

Lemma 1 (Theorem 3.1 in [Du et al., 2019c]). Suppose xi is not parallel with xj for any i̸=j, then λ0 := λmin(H∞)>0.

When m is large enough, H(0) is also positive definite.

Lemma 2 (Lemma 3.1 in [Du et al., 2019c]). If m = Ω ( n2 λ20 log ( n2 p )), with probability at least 1 - p, it has ∥H(0) -H∞∥2 ≤ λ0 4 and λmin(H(0)) ≥ 3λ0 4 .

Next, we introduce a lemma that shows for any t, if wr(t) is close to wr(0), then H(t) is identical to H(0), which ensures the positive definite of H(t).

Lemma 3 (Lemma 3.2 in [Du et al., 2019c]). Assume wr(0) ∼ N(0, I) for r ∈ [m] and ∥wr(0) -wr∥2 ≤ cpλ0 n2 =: R for some small positive constant c, with probability at least 1 - p, it has ∥H(t) - H(0)∥2 < λ0 4 , λmin(H(t)) > λ0 2 and λmax(H(t)) < λm := λmax(H∞) + λ0 4 .

Therefore, under the over-parameterized assumption, H(t) stays positive with the condition number κ := λm/(λ0/2) > 1 and wr(t) remains close to its initialization during the training process.

Low-resolution view of momentum methods. From a continuous view, [Bu et al., 2021] derived the convergence guarantees of HB and NAG in optimizing Eq.(2). They consider a non-linear dissipative dynamical system with b > 0

which corresponds to the limiting ODE for both HB and NAG when the learning rate tends to infinitesimal [Wilson et al.,

2021]. Consequently, Eq.(11) fails to capture the convergence difference between HB and NAG.

High-resolution view of momentummethods. In this subsection, we first present the high-resolution formulations of HB and NAG, then derive their corresponding residual dynamics when training the two-layer ReLU neural network.

On the one hand, according to [Shi et al., 2021], HB has the following high-resolution ODE

where α ≥ 0 and s > 0. For each neuron r ∈ [m], HB has the corresponding ODE

On the other hand, NAG has a different ODE representation

The third term in Eq.(14) is the main difference between NAG and HB, which is induced by the gradient correction term in Eq.(5) [Shi et al., 2021]. Accordingly, NAG has the following ODE for any wr

When s → 0, Eq.(13) and Eq.(15) degenerate to the same low-resolution Eq.(11). However, ℓ is non-convex while [Shi et al., 2021] derived the convergence rate only in the convex setting. For the µ-strongly convex and smooth problem ℓ, Shi et al. sets α = µ for both Eq.(13) and Eq.(15). However, their result only shows that HB and NAG converge at the same rate O(e-√µt/4 ), which fails to distinguish them in the continuous view.

Instead of the convex setting, we mainly focus on a nonconvex ℓ. Following [Bu et al., 2021], it has

where 0d×d represents the zero matrix with dimension d × d, ∂I{x≥0} ∂x = δ(x) and δ denotes the delta function that

Then, it has

and

When multiplying Eq.(13) with ∂f ∂wr and summing over all r ∈ [m], it has the following ODE based on Eq.(9) and Eq.(20)

Then, the dynamics of the residual vector ∆ = f - y has

Instead of directly analyzing the non-convex ℓ over W, we show that optimizing ℓ corresponds to the pseudo-loss ˆ ℓ(t) := 1 2 ∆⊤(t)H(t)∆(t), which is λ0 2 -strongly convex over ∆(t) according to Lemma 3. With µ = λ0/2, Eq.(13) can be transformed to

where α = λ0/2 according to [Shi et al., 2021]. Similarly, it can derive following residual dynamics of NAG as

where the third term in Eq.(27) is derived by

Note that the third term in Eq.(27) is derived from the gradient correction term.

4 Main Results

In this section, we provide the convergence analysis and comparison of HB and NAG.

4.1 Convergence Analysis of HB

Theorem 1. Suppose m= Ω( n6 p3 λ40 ), wr(0) ∼ N(0, Id) and ar ∼ unif{-1, 1} for any r ∈ [m], with probability at least 1-p, it has

Remark 1. The above theorem implies HB can attain a global minimum of Eq.(2) at a linear rate O(e-(2-√2)√ λ0 2 t ) in training two-layer ReLU neural networks. It is worth noting that [Shi et al., 2021] studied the same high-resolution ODE of HB and proved HB can attain a global convergence rate O(e- 14 √µt ) for µ-strongly convex problems. In contrast, our focus is on a non-convex optimization problem. We show that optimizing the non-convex loss ℓ corresponds to an error dynamics, as if we were actually optimizing a strongly convex loss ˆ ℓ over the prediction error. Compared to [Shi et al., 2021], Theorem 1 further provides a tighter upper bound for the convergence rate of HB in the high-resolution continuoustime view.

Remark 2. We are aware of some recent works that improve the requirement of the over-parameterization [Song and Yang, 2019; Nguyen, 2021]. In this work, we mainly focus on providing the theoretical guarantee for the acceleration of NAG over HB in training over-parameterized neural networks. Therefore, we incorporate the first and also the most well-known result about the requirement on m that ensures H(t) positive [Du et al., 2019c]. We believe incorporating these improved results on m into our analysis can also improve the bound on m in our work.

The proof of Theorem 1 mainly consists of three parts:

1. Firstly, based on the positive definite assumption of H, we demonstrate the training loss decays to zero at a linear rate and the distance between w and its initial value has a upper bound R′ .
2. Secondly, we show that if R′ is smaller than the distance R defined in Lemma 3, our conclusions from the first procedures hold true and the positive definite of H is guaranteed.
3. Finally, the width m is determined by satisfying R′ < R, which completes the proof of Theorem 1.

To start with, we present the first procedure of the proof, which involves the smallest eigenvalue λ0 2 ofH as introduced in Lemma 3.

Lemma 4. Assume λmin(H(i)) ≥ λ0 2 for 0≤i≤t. With 0 < s ≤ 1/λm and ˙wr(0) = 0 for any r ∈ [m], HB has ℓ(t) ≤ 6ˆ ℓ(0) λ0 e-(2-√2)√ λ0 2 t and ∥wr(t)-wr(0)∥2 ≤ R′ := 10√6ˆ ℓ(0)n/(λ30m).

The proof is provided in the supplementary materials. Then, the conclusions in Lemma 4 hold when the distance R′ defined in Lemma 4 is less than the distance R given in Lemma 3.

Lemma 5. Assume R′ < R, it has (1) λmin(H(t)) ≥ λ0/2, (2) ∥wr(t) -wr(0)∥2 ≤ R′ for all r ∈ [m] and (3) ℓ(t) ≤ 6ˆ ℓ(0) λ0 e-(2-√2)√ λ0 2 t .

The proof can be founded in the supplementary materials. To satisfy the assumption R′ < R in Lemma 5, it requires m = Ω( n6 p3 λ40 ) based on ˆ ℓ(0) = Ω(nλ0/p) as proved in [Bu et al., 2021], which completes the proof of Theorem 1.

4.2 Convergence Analysis of NAG

In this subsection, based on the similar analysis framework as the proof of HB, we present the convergence guarantee of NAG and show its acceleration.

Theorem 2. Suppose m = Ω( n6 p3 λ40 ), wr(0) ∼ N(0, 1) and ar ∼ unif{-1, 1} for any r ∈ [m], with probability at least 1-p, it has

where

Remark 3. The above theorem indicates that NAG can reach a global minimum at a linear rate O(e-ρ∗NAG(α)√ λ0 2 t ). When 0 < s ≤ 1/λm, it holds that 0 < α ≤ 1/(2√κ) ≤ 0.5. From Fig. 1, it can observe that ρ∗NAG(α) is a monotonic increasing function with respect to α when 0 < α ≤ 0.5. Thus, it has

Compared to the convergence rate O(e-(2-√2)√ λ0 2 t ) of HB, our theoretical result demonstrates that NAG converges faster than HB.

In addition, according to Eq.(44) in the supplementary materials, the derivative of the corresponding Lyapunov function V of NAG has

where the third term on the right hand side of Eq.(31) is induced from √sH(t) ˙∆(t), which corresponds to the gradient correction term as mentioned in Section 3.2. If neglecting this term, based on the upper bound of V (t) (Eq.(43) in the supplementary materials), the corresponding convergence rate has

which is the same as the convergence rate of HB. Therefore, we provably show that the extra gradient correction term of NAG in Eq.(27) leads to the acceleration of NAG over HB in training neural networks.

Remark 4. It is noted that [Shi et al., 2021] conducted an analysis on the same high-resolution ODE of NAG for µstrongly convex problems. Their result demonstrated that the convergence rate of NAG is upper bounded by O(e- 14 √µt ). Later, [Sun et al., 2020] and [Zhang et al., 2021] improved the convergence rate to O(e- 3 7 √µt ) and O(e- 12 √µt ), respectively. Comparing to [Shi et al., 2021; Sun et al., 2020; Zhang et al., 2021], our study focuses on a non-convex ℓ over the weights. Furthermore, we demonstrate that ℓ corresponds to a λ0 2 -strongly convex problem by utilizing the NTK theory. Additionally, we provide a tight convergence rate of NAG with a more refined inequality and analysis.

Remark 5. Noted that directly studying the convergence of discrete algorithms for non-convex ℓ is far too complex and general to offer insightful explanation about their efficiency in deep learning. Instead of the low-resolution approach in previous work [Bu et al., 2021], we study from a high-resolution continuous viewpoint. Our result provides the first acceleration guarantee of NAG over HB with a simple and intuitive proof. Furthermore, Proposition 1 in [Shi et al., 2021] shows that the discrete iterates and the high-resolution ODE trajectories remain close for both HB and NAG when the learning rate is small enough. It would be interesting to consider the discrete setting in the future work.

Lemma 6. Assume λmin(H(i)) ≥ λ0 2 for 0 ≤ i ≤ t. With 0 < s ≤ 1/λm and ˙wr(0) = 0 for any r ∈ [m], NAG has ℓ(t) ≤ 26ˆ ℓ(0) 3λ0 e-ρ∗NAG(α)√λ0/2t and ∥wr(t) -wr(0)∥ ≤ 67√ˆ ℓ(0)n/(λ30m).

The proof is presented in the supplementary materials. Similarly, it can establish the conclusion of the Theorem 2 following Lemma 5. Moreover, the required width should meet m = Ω( max{√sn/√λ0, n6 p3 λ40 }) = Ω( n6 p3 λ40 ) using λ0 ≤ 1/2 as proved in [Bu et al., 2021].

5 Numerical Experiments

5.1 Setup

We use three commonly used datasets: MNIST [LeCun et al., 1998], FMNIST [Xiao et al., 2017] and CIFAR10 [Krizhevsky et al., 2009], where the pre-processing procedure follows the existing work [Arora et al., 2019b]. Specifically, we use the first two classes of instances, assigning the label of the first class to +1 and the second to -1. In addition, each instance is normalized with the unit norm.

Regarding the architecture and initialization of the neural network, we adopt the setting in Section 3.1. Besides, we set the width m ∈ {1000, 2000, 4000, 8000}, learning rate η ∈ {1/(10λm), 1/(20λm)} and momentum parameter β ∈ {0.97, 0.99} , where λm is calculated with the eigenvalues of the NTK matrix in Eq.(10). For each setting of hyper-parameters, we conduct 5 independent runs with different random seeds for initialization. The average training loss over iterations is presented by the dashed line. We use JAX [Bradbury et al., 2018] library for code implementation and all experiments are conducted on 4 NVIDIA V100 GPU.

5.2 Results Analysis

Due to the space limitation, only the results on the MNIST dataset are presented here. The remaining results can be founded in the supplementary materials. From Fig. 2, it can observe that the training losses of NAG and HB decay at a almost linear rate, which demonstrates the linear convergence of HB and NAG in training the two-layer ReLU neural networks. Additionally, the loss curve of NAG lies below that of HB for different hyper-parameter settings, which indicates NAG achieves a smaller training loss compared to HB at the same iteration. Thus, NAG achieve an acceleration over HB. These empirical findings validate the conclusion in Section 4 that NAG theoretically obtains a faster global convergence rate compared to HB.

To further validate our findings, we present the maximum distance maxr∈[m] ∥wr(t)-w0(t)∥ of HB and NAG during the training process with different widths. According to Lemma 4 and Lemma 6, the theoretical maximum distances of HB and NAG both satisfy maxr∈[m] ∥wr(t)-w0(t)∥ ∼ O(1/√m), which indicate that the maximum distance decreases when the width of the neural network increases. The corresponding empirical results are shown in Fig. 3 and Fig. 4, which is in accordance with our theoretical findings.

6 Conclusion and FutureWork

In this paper, we study the convergence of momentum methods in training an over-parameterized two-layer ReLU neural network. By exploiting high-resolution ODE formulations and the NTK theory, the global convergences of HB and NAG are established. In contrast with previous works, we exploit high-resolution ODEs instead of the low-resolution. In addition, our work yields tighter upper bounds via a more refined inequality and analysis, which provides the first theoretical guarantee for the acceleration of NAG over HB in tackling non-convex optimization problem on neural networks.

There are several directions for future research. Firstly, it is natural to further extend our result to deep neural networks, where the key technical challenge lies in deriving and analyzing the spectral properties of the corresponding H. We note that there is some works about the eigenvalues of H, including deep fully-connected neural networks [Du et al., 2019b], graph neural networks [Du et al., 2019a] and convolutional neural networks [Li et al., 2019]. Secondly, stochastic optimization methods are widely used in training neural networks due to relatively low computational costs. Therefore, it would be interesting to extend existing stochastic convex works [Nakakita, 2022; Laborde and Oberman, 2020] to analyze the stochastic optimization problem in neural networks. Finally, the understanding of the adaptive first-order method is still lacking, such as AdaBound [Liu et al., 2022a] and RMSProp [Xu et al., 2021]. These adaptive methods also incorporate the idea of momentum and our work may inspire the development of new tools for understanding their convergence behaviors.

Acknowledgements

This work was supported by National Natural Science Foundation of China (No.62076251 and No.62106281).