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

TL;DR

This work analyzes training of over-parameterized two-layer ReLU networks through high-resolution ODEs and Neural Tangent Kernel theory, establishing global convergence for both Heavy Ball and Nesterov's Accelerated Gradient methods. It derives residual dynamics that map non-convex training to strongly convex behavior in the prediction error, enabling explicit linear convergence rates. Crucially, it provides the first theoretical guarantee that NAG accelerates over HB in this non-convex, NTK-regime, with a rate exponent ρ*_{NAG}(α) exceeding the HB rate (2−√2) under appropriate parameter choices. The results are validated on MNIST, FMNIST, and CIFAR-10, showing consistent NAG superiority and revealing how network width controls weight excursion, offering guidance for momentum-based optimization in deep learning.

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.

Figures (11)

• Figure 1: The value of $\rho_{NAG}^*(\alpha)$ with respect to $\alpha$.
• 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 $\eta$ and momentum parameter $\beta$.
• Figure 3: The maximum distance $\max_{r \in [m]}\|\bm{w}_r(t) - \bm{w}_0(t)\|$ of HB in training two-layer ReLU neural networks on the MNIST dataset during training process.
• Figure 4: The maximum distance $\max_{r \in [m]}\|\bm{w}_r(t) - \bm{w}_0(t)\|$ of NAG in training two-layer ReLU neural networks on the MNIST dataset during training process.
• Figure 5: The feasible region (grey area) of case 1.

Theorems & Definitions (11)

• Lemma 1: Theorem 3.1 in Du2019
• Lemma 2: Lemma 3.1 in Du2019
• Lemma 3: Lemma 3.2 in Du2019
• Theorem 1
• Lemma 4
• Lemma 5
• Theorem 2
• Lemma 6
• proof
• proof