---
title: "Provable Acceleration of Nesterov’s Accelerated Gradient Method over Heavy Ball Method in Training Over-Parameterized Neural Networks"
authors:
  - "Xin Liu"
  - "Wei Tao"
  - "Wei Li"
  - "Dazhi Zhan"
  - "Jun Wang"
  - "Zhisong Pan"
doi: "10.24963/ijcai.2024/508"
url: "https://ijcai.org/proceedings/2024/508"
pdf: "https://ijcai.org/proceedings/2024/0508.pdf"
venue: "Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence"
year: 2024
pages: "4596-4604"
track: "Main Track"
keywords:
  - "Machine Learning: ML: Optimization"
  - "Machine Learning: ML: Theory of deep learning"
---

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

Xin Liu, Wei Tao, Wei Li, Dazhi Zhan, Jun Wang, Zhisong Pan

*Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence*, pages 4596-4604. DOI: [10.24963/ijcai.2024/508](https://doi.org/10.24963/ijcai.2024/508)

## Abstract

Due to its simplicity and efficiency, the first-order gradient method has been extensively employed in training neural networks. Although the optimization problem of the neural network is non-convex, recent research has proved that the first-order method is capable of attaining a global minimum during training over-parameterized neural networks, where the number of parameters is significantly larger than that of training instances. Momentum methods, including the heavy ball (HB) method and Nesterov's accelerated gradient (NAG) method, are the workhorse of first-order gradient methods owning to their accelerated convergence. In practice, NAG often exhibits superior performance than HB. However, current theoretical works fail to distinguish their convergence difference in training neural networks. To fill this gap, we consider the training problem of the two-layer ReLU neural network under over-parameterization and random initialization. Leveraging high-resolution dynamical systems and neural tangent kernel (NTK) theory, our result not only establishes tighter upper bounds of the convergence rate for both HB and NAG, but also provides the first theoretical guarantee for the acceleration of NAG over HB in training neural networks. Finally, we validate our theoretical results on three benchmark datasets.

## Full text

## 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 κ &gt; 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.

Table 1: Comparison of existing results for HB and NAG from the high-resolution view, where SC, NC and Obj represent stronglyconvex, non-convex and Objective function, respectively. In addition, 2-√2&lt;ρNAG≤ 11-√73 4 .

|                      | Obj   | HB              | NAG           |
|----------------------|-------|-----------------|---------------|
| [Shi et al., 2021]   | SC    | O(e- √µt 4 )    | O(e- √µt 4 )  |
| [Sun et al., 2020]   | SC    | -               | O(e- 3√µt 7 ) |
| [Zhang et al., 2021] | SC    | -               | O(e- √µt 2 )  |
| This work            | NC    | O(e-(2-√2)√µt ) | O(e-ρNAG√µt ) |

## 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 &lt; ρ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:

<!-- formula-not-decoded -->

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 :

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

where η &gt; 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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

When m goes to infinity, H at the initialization has

<!-- formula-not-decoded -->

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∞)&gt;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 &lt; λ0 4 , λmin(H(t)) &gt; λ0 2 and λmax(H(t)) &lt; λm := λmax(H∞) + λ0 4 .

Therefore, under the over-parameterized assumption, H(t) stays positive with the condition number κ := λm/(λ0/2) &gt; 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 &gt; 0

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

On the other hand, NAG has a different ODE representation

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

<!-- formula-not-decoded -->

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

Then, it has

<!-- formula-not-decoded -->

<!-- formula-not-decoded -->

and

<!-- formula-not-decoded -->

<!-- formula-not-decoded -->

<!-- formula-not-decoded -->

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)

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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′ &lt; 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 &lt; 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′ &lt; 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′ &lt; 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.

Figure 1: The value of ρ∗NAG(α) with respect to α.

![Image](https://ijcai.org/proceedings/2024/fulltext/508/image_000000_06df7398b77bf3148063bdee93c923675117d12e4bf2467f1e55ebd06e3c5816.png)

## 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

<!-- formula-not-decoded -->

where

<!-- formula-not-decoded -->

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 &lt; s ≤ 1/λm, it holds that 0 &lt; α ≤ 1/(2√κ) ≤ 0.5. From Fig. 1, it can observe that ρ∗NAG(α) is a monotonic increasing function with respect to α when 0 &lt; α ≤ 0.5. Thus, it has

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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.

Figure 2: Convergence comparison between NAG and HB in training two-layer ReLU neural networks on the MNIST dataset with different width m, learning rate η and momentum parameter β.

![Image](https://ijcai.org/proceedings/2024/fulltext/508/image_000001_5798f57a18908d4e1d18be2da294f5a7887b2955929ca56ebeae1b00584afdb1.png)

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 &lt; 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.

Figure 3: The maximum distance maxr∈[m] ∥wr(t) - w0(t)∥ of HB in training two-layer ReLU neural networks on the MNIST dataset during training process.

![Image](https://ijcai.org/proceedings/2024/fulltext/508/image_000002_bbcad78895c31c4ef6055741d18dfe1cbd350507d16f0cc0a10de19b6124a397.png)

Figure 4: The maximum distance maxr∈[m] ∥wr(t) - w0(t)∥ of NAG in training two-layer ReLU neural networks on the MNIST dataset during training process.

![Image](https://ijcai.org/proceedings/2024/fulltext/508/image_000003_60d86c9b605fc74b22db995acbc5528e04d58f3390a880e8aea49bb27553f5bd.png)

## 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).

## References

1. [Alemohammad et al., 2020] Sina Alemohammad, Zichao Wang, Randall Balestriero, and Richard Baraniuk. The re- current neural tangent kernel. In International Conference on Learning Representations, 2020.
2. [Arora et al., 2019a] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems, pages 8141-8150, 2019.
3. [Arora et al., 2019b] Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pages 322-332, 2019.
4. [Bradbury et al., 2018] James Bradbury, Roy Frostig, Pe- ter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: Composable transformations of Python+NumPy programs, 2018.
5. [Brown et al., 2020] Tom Brown, Benjamin Mann, Nick Ry- der, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Sys- tems, pages 1877-1901, 2020.
6. [Bu et al., 2021] Zhiqi Bu, Shiyun Xu, and Kan Chen. A dynamical view on optimization algorithms of overparam- eterized neural networks. In International Conference on Artificial Intelligence and Statistics, pages 3187-3195, 2021.
7. [Du et al., 2019a] Simon S. Du, Kangcheng Hou, Ruslan Salakhutdinov, Barnabás Póczos, Ruosong Wang, and Keyulu Xu. Graph neural tangent kernel: Fusing graph neural networks with graph kernels. In Advances in Neural Information Processing Systems, pages 5724-5734, 2019.
8. [Du et al., 2019b] Simon S Du, Jason D Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. International Conference on Machine Learning, pages 1675-1685, 2019.
9. [Du et al., 2019c] Simon S. Du, Xiyu Zhai, Barnabás Póczos, and Aarti Singh. Gradient descent provably op- timizes over-parameterized neural networks. In Interna- tional Conference on Learning Representations, 2019.
10. [Ghadimi et al., 2015] Euhanna Ghadimi, Hamid Reza Feyzmahdavian, and Mikael Johansson. Global conver- gence of the heavy-ball method for convex optimization. In European Control Conference, pages 310-315, 2015.
11. [Huang et al., 2020] Kaixuan Huang, Yuqing Wang, Molei Tao, and Tuo Zhao. Why do deep residual networks gen- eralize better than deep feedforward networks -A neural tangent kernel perspective. In Advances in Neural Infor- mation Processing Systems, pages 2698-2709, 2020.
12. [Jacot et al., 2018] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing System, pages 8571-8580, 2018.
13. [Krizhevsky et al., 2009] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny im- ages. 2009.
14. [Laborde and Oberman, 2020] Maxime Laborde and Adam M. Oberman. A lyapunov analysis for accelerated gradient methods: from deterministic to stochastic case. In International Conference on Artificial Intelligence and Statistics, pages 602-612, 2020.
15. [LeCun et al., 1998] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning ap- plied to document recognition. Proceedings of the IEEE, 86(11):2278-2324, 1998.
16. [Lessard et al., 2016] Laurent Lessard, Benjamin Recht, and Andrew Packard. Analysis and design of optimization al- gorithms via integral quadratic constraints. SIAM Journal on Optimization, 26(1):57-95, 2016.
17. [Li et al., 2019] Zhiyuan Li, Ruosong Wang, Dingli Yu, Si- mon S Du, Wei Hu, Ruslan Salakhutdinov, and Sanjeev Arora. Enhanced convolutional neural tangent kernels. arXiv preprint arXiv:1911.00809, 2019.
18. [Liu et al., 2022a] Jinlan Liu, Jun Kong, Dongpo Xu, Miao Qi, and Yinghua Lu. Convergence analysis of adabound with relaxed bound functions for non-convex optimization. Neural Networks, 145:300-307, 2022.
19. [Liu et al., 2022b] Xin Liu, Zhisong Pan, and Wei Tao. Provable convergence of nesterov's accelerated gradi- ent method for over-parameterized neural networks. Knowledge-Based Systems, 251:109277, 2022.
20. [Nado et al., 2021] Zachary Nado, Justin M Gilmer, Christo- pher J Shallue, Rohan Anil, and George E Dahl. A large batch optimizer reality check: Traditional, generic optimizers suffice across batch sizes. arXiv preprint arXiv:2102.06356, 2021.
21. [Nakakita, 2022] Shogo Nakakita. Parametric estimation of stochastic differential equations via online gradient de- scent. arXiv preprint arXiv:2210.08800, 2022.
22. [Nesterov, 1983] Yurii Nesterov. A method of solving a con- vex programming problem with convergence rate o (1/k2). In Sov. Math. Dokl, 1983.
23. [Nesterov, 2013] Yurii Nesterov. Introductory lectures on convex optimization: A basic course. Springer Science, 2013.
24. [Nguyen, 2021] Quynh Nguyen. On the proof of global con- vergence of gradient descent for deep relu networks with linear widths. In International Conference on Machine Learning, pages 8056-8062, 2021.
25. [Polyak, 1964] Boris T Polyak. Some methods of speeding up the convergence of iteration methods. USSR Computa- tional Mathematics and Mathematical Physics, 4(5):1-17, 1964.
26. [Schmidhuber, 2015] Jürgen Schmidhuber. Deep learning in neural networks: An overview. Neural Networks, 61:85- 117, 2015.
27. [Schmidt et al., 2021] Robin M. Schmidt, Frank Schneider, and Philipp Hennig. Descending through a crowded val- ley -benchmarking deep learning optimizers. In Inter- national Conference on Machine Learning, pages 9367- 9376, 2021.
28. [Shi et al., 2021] Bin Shi, Simon S Du, Michael I Jordan, and Weijie J Su. Understanding the acceleration phe- nomenon via high-resolution differential equations. Math- ematical Programming, pages 1-70, 2021.
29. [Song and Yang, 2019] Zhao Song and Xin Yang. Quadratic suffices for over-parametrization via matrix chernoff bound. arXiv preprint arXiv:1906.03593, 2019.
30. [Su et al., 2014] Weijie Su, Stephen Boyd, and Emmanuel Candes. A differential equation for modeling nesterov's accelerated gradient method: Theory and insights. In Ad- vances in Neural Information Processing Systems, pages 2510-2518, 2014.
31. [Sun et al., 2020] Boya Sun, Jemin George, and Solmaz S. Kia. High-resolution modeling of the fastest first-order op- timization method for strongly convex functions. In Con- ference on Decision and Control, pages 4237-4242, 2020.
32. [Wang et al., 2021] Jun-Kun Wang, Chi-Heng Lin, and Ja- cob D. Abernethy. A modular analysis of provable ac- celeration via polyak's momentum: Training a wide relu network and a deep linear network. In International Con- ference on Machine Learning, pages 10816-10827, 2021.
33. [Wilson et al., 2021] Ashia C Wilson, Ben Recht, and Michael I Jordan. A lyapunov analysis of accelerated methods in optimization. Journal of Machine Learning Research, 22:1-34, 2021.
34. [Xiao et al., 2017] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: A novel image dataset for bench- marking machine learning algorithms. arXiv preprint arXiv:1708.07747, 2017.
35. [Xu et al., 2021] Dongpo Xu, Shengdong Zhang, Huisheng Zhang, and Danilo P. Mandic. Convergence of the rm- sprop deep learning method with penalty for nonconvex optimization. Neural Networks, 139:17-23, 2021.
36. [Zhang et al., 2021] Peiyuan Zhang, Antonio Orvieto, Hadi Daneshmand, Thomas Hofmann, and Roy S. Smith. Re- visiting the role of euler numerical integration on acceler- ation and stability in convex optimization. In International Conference on Artificial Intelligence and Statistics, pages 3979-3987, 2021.
