Provable Accelerated Convergence of Nesterov's Momentum for Deep ReLU
  Neural Networks
Current state-of-the-art analyses on the convergence of gradient descent for
training neural networks focus on characterizing properties of the loss
landscape, such as the Polyak-Lojaciewicz (PL) condition and the restricted
strong convexity. While gradient descent converges linearly under such
conditions, it remains an open question whether Nesterov's momentum enjoys
accelerated convergence under similar settings and assumptions. In this work,
we consider a new class of objective functions, where only a subset of the
parameters satisfies strong convexity, and show Nesterov's momentum achieves
acceleration in theory for this objective class. We provide two realizations of
the problem class, one of which is deep ReLU networks, which --to the best of
our knowledge--constitutes this work the first that proves accelerated
convergence rate for non-trivial neural network architectures.
Fangshuo Liao, Anastasios Kyrillidis
Accelerated Convergence of Nesterov's Momentum for Deep Neural Networks under Partial Strong Convexity

###### Abstract

Current state-of-the-art analyses on the convergence of gradient descent for training neural networks focus on characterizing properties of the loss landscape, such as the Polyak-Lojaciewicz (PL) condition and the restricted strong convexity. While gradient descent converges linearly under such conditions, it remains an open question whether Nesterov's momentum enjoys accelerated convergence under similar settings and assumptions. In this work, we consider a new class of objective functions, where only a subset of the parameters satisfies strong convexity, and show Nesterov's momentum achieves acceleration in theory for this objective class. We provide two realizations of the problem class, one of which is deep ReLU networks, which -to the best of our knowledge-constitutes this work the first that proves accelerated convergence rate for non-trivial neural network architectures.

## 1 Introduction

**Theory of gradient descent in neural networks.** Training neural networks with gradient-based methods has shown surprising empirical success (Lecun et al., 1998; LeCun et al., 2015; Zhang et al., 2017; Goodfellow et al., 2016); yet, it is not lucid why such a simple algorithm works (Zhang et al., 2018; Li et al., 2020). Overall, it has been a mystery why such routines designed originally for convex problems can consistently find a good minimum for complicated non-convex objectives, as that of neural network training (Yun et al., 2019; Auer et al., 1995; Safran and Shamir, 2018).

Focusing on the case of neural networks, a major advance in understanding is the analysis through the so-called Neural Tangent Kernel (NTK) (Jacot et al., 2020). The use of NTK shows that, when the width of the neural network approaches infinity, the training process can be treated as a kernel machine. Inspired by the NTK analysis, a large body of work has focused on showing the convergence of gradient descent for various neural network architectures under finite over-parameterization requirements (Du et al., 2019, 2019; Allen-Zhu et al., 2019; Zou and Gu, 2019; Zhang et al., 2019; Awasthi et al., 2021; Ling et al., 2023; Allen-Zhu et al., 2019; Song and Yang, 2020; Su and Yang, 2019). In the meantime, the analysis was also extended to other training algorithms and settings that differ from
gradient descent, such as stochastic gradient descent (Oymak and Soltanolkotabi, 2019; Ji and Telgarsky, 2020; Xu and Zhu, 2021; Zou et al., 2018), drop-out (Liao and Kyrillidis, 2022; Mianjy and Arora, 2020), federated training (Huang et al., 2021), and adversarial training (Li et al., 2022).

A recent line of work follows a different route and characterizes the loss landscape of neural networks using the local Polyak-Lojaciewicz (PL) condition (Song et al., 2021; Liu et al., 2020; Nguyen, 2021; Ling et al., 2023). Building upon the well-established theory of how gradient descent converges under the PL condition (Karimi et al., 2020), this line of work decouples the neural network structure from the dynamics of loss function, along the optimization trajectory. This way, these works were able to show a more fine-grained analysis of the relationship between regulatory conditions, such as the PL condition, and the neural network structure. Such analysis not only resulted in further relaxed over-parameterization requirements (Song et al., 2021; Nguyen, 2021; Liu et al., 2020), but was also shown to be easily extended to deep architectures (Ling et al., 2023), suggesting that it is more suitable in practice.

**Theory of accelerated gradient descent in neural networks.** In contrast to the fast-growing research devoted to vanilla gradient descent, there is limited work on the behavior of momentum methods in training neural networks. In particular, the acceleration of both the Heavy Ball method and Nesterov's momentum is shown only for shallow ReLU neural networks (Wang et al., 2021; Liu et al., 2022; Zhang et al., 2022) and deep linear neural networks (Wang et al., 2021; Liu et al., 2022; Wang et al., 2022). For these architectures, (Wang et al., 2022) shows that training reduces to a case very similar to optimizing a quadratic objective. Thus, it is not clear whether these analyses can be extended to neural networks with more layers, or more complicated structure.

In addition, showing an accelerated convergence rate under only the PL condition has been a long-standing difficulty. For the Heavy Ball method, Danilova et al. (Danilova et al., 2018) established a linear convergence rate under the PL condition, but no acceleration is shown without assuming strong convexity. Wang et al. (Wang et al., 2022) proved an accelerated convergence rate; yet, the authors assume the \(\lambda^{\star}\)-average out condition, which cannot always be easily justified for neural networks. To the best of our knowledge, any proof of convergence for Nesterov's momentum under the PL condition in non-convex settings is currently missing. In the continuous limit, acceleration is proved in a limited scenario (Apidopoulos et al., 2022), and the analysis does not easily extend to the discrete case (Wang et al., 2022; Shi et al., 2022). Finally, Yue et al. (Yue et al., 2022) shows that gradient descent already achieves an optimal convergence rate for functions satisfying smoothness and the PL condition. This suggests that, to prove the acceleration of momentum methods on a wider class of neural networks, we need to leverage properties beyond the PL condition.

Based on prior work (Liu et al., 2020) that shows over-parameterized systems are essentially non-convex in any neighborhood of the global minimum, we aim at developing a relaxation to the (strong) convexity that is suitable for such systems and enables the momentum methods to achieve acceleration. In particular, we consider the minimization of a new class of objectives:

\[\min_{\mathbf{x}\in\mathbb{R}^{d_{1}},\mathbf{u}\in\mathbb{R}^{d_{2}}}f( \mathbf{x},\mathbf{u}), \tag{1}\]

where \(f\) satisfies the strong convexity with respect to \(\mathbf{x}\), among other assumptions (more in text below). Intuitively, our construction assumes that the parameter space can be partitioned into two sets, and only one of the two sets enjoys rich properties, such as strong convexity. In this paper, we focus on Nesterov's momentum. Our contribution can be summarized as:
* Focusing on the general problem class in (1), we prove that Nesterov's momentum enjoys an accelerated linear convergence where the convergence rate is \(1-\Theta\left(\nicefrac{{1}}{{\sqrt{\kappa}}}\right)\), compared to \(1-\Theta\left(\nicefrac{{1}}{{\kappa}}\right)\) in the case of gradient descent. Here, \(\kappa\) is a variant of the condition number of \(f\). Our result holds even when \(f\) is non-convex and non-smooth.
* We provide two realizations of our problem class. In Section 5.1, we first consider fitting an additive model under the MSE loss. We prove the acceleration of Nesterov's momentum as long as the non-convex component of the additive model is small enough to control the upper bounds of Lipschitz-type of assumptions.
* Next, we turn to the training of a deep ReLU neural network in Section 5.2. We show that when the width of the neural network trained with \(n\) samples is \(\Omega\left(n^{4}\right)\), under proper initialization, Nesterov's momentum converges to zero training loss with rate \(1-\Theta\left(\nicefrac{{1}}{{\sqrt{\kappa}}}\right)\). To the best of our knowledge, this is the first result that establishes accelerated convergence of deep ReLU networks.

## 2 Related Works

**Convergence in neural network training.** Recent studies on the convergence guarantees of gradient-based methods in neural network training mainly focus on two techniques: the Neural Tangent Kernel (Jacot et al., 2020) and the mean-field analysis (Mei et al., 2019). The NTK-based analysis builds upon the idea that, when the width approaches infinity, training neural networks behaves like training a kernel machine. Various techniques have been developed to control the error that occurs when the width becomes finite. In particular, (Du et al., 2019) tracks the change of activation patterns in ReLU-based neural networks, and often requires a large over-parameterization. Later works improve the over-parameterization requirement by leveraging matrix concentration inequalities (Song and Yang, 2020) for a more fine-grained analysis on the change of Jacobians (Oymak and Soltanolkotabi, 2019), and building their analysis upon the separability assumption of the data in the reproducing Hilbert space of the neural network (Ji and Telgarsky, 2020).

A noticeable line of work focuses on establishing that the PL condition is satisfied by neural networks, where the coefficient of the PL condition is based on the eigenvalue of the NTK matrix. (Nguyen, 2021) shows the PL condition is satisfied by deep ReLU neural networks, by considering the dominance of the gradient with respect to the weight in the last layer. (Liu et al., 2020) proves the PL condition by upper bounding the Hessian for deep neural networks with smooth activation functions. (Song et al., 2021) further reduces the over-parameterization, while maintaining the PL condition, via the expansion of the activation function with the Hermite polynomials. Lastly, (Banerjee et al., 2023) establishes the restricted strong convexity of neural networks within a sequence of ball-shaped regions, centered around the weights per iteration; yet, the coefficient of the strong convexity is not explicitly characterized in theory.

**Convergence of Nesterov's Momentum.** The original proof of Nesterov's momentum (Nesterov, 2018) builds upon the idea of estimating sequences for both convex smooth objectives and strongly convex smooth objectives. Later work in (Bansal and Gupta, 2019) provides an alternate proof within the same setting by constructing a Lyapunov function. In the non-convex setting, a large body of works focuses on variants of Nesterov's momentum that lead to a guaranteed convergence, by employing techniques such as negative curvature exploitation (Carmon et al., 2017), cubic regularization (Carmon and Duchi, 2020), and restarting schemes (Li and Lin, 2022). For neural networks, (Liu et al., 2022, 2022) are the only works that study the convergence of Nesterov's momentum. However, with the over-parameterization requirement considered, the objective is similar to a quadratic function.
Deviating from Nesterov's momentum, (Wang et al., 2021) studies the convergence of the Heavy-ball method under similar over-parameterization requirements. A recent work (Wu et al., 2023) proves the convergence of the Heavy-ball method under the mean-field limit; such a limit is not the focus of our study in this paper. Lastly, (Jelassi and Li, 2022) shows that momentum-based methods improve the generalization ability of neural networks. However, there is no explicit convergence guarantee for the training loss.

## 3 Preliminary

**Notations.** We use bold lower-case letters (e.g. \(\mathbf{a}\)) to denote vectors, bold upper-case letters (e.g. \(\mathbf{A}\)) to denote matrices and standard lower-case letters (e.g. \(a\)) to denote scalars. For a vector \(\mathbf{a}\), we use \(a_{i}\) to denote its \(i\)-th entry. For a matrix \(\mathbf{A}\), we use \(a_{ij}\) to denote its \((i,j)\)-th entry. \(\left\|\mathbf{a}\right\|_{2}\) denotes the \(\ell_{2}\)-norm of vector \(\mathbf{a}\), and \(\left\|\mathbf{A}\right\|_{F}\) denotes the Frobenius norm of \(\mathbf{A}\). For two vectors \(\mathbf{a}_{1},\mathbf{a}_{2}\), we use \((\mathbf{a}_{1},\mathbf{a}_{2})\) to denote the concatenation of \(\mathbf{a}_{1}\) and \(\mathbf{a}_{2}\). For a matrix \(\mathbf{A}\) with columns \(\mathbf{a}_{1},\ldots,\mathbf{a}_{n}\), we use \(\forall(\mathbf{A})=(\mathbf{a}_{1},\ldots,\mathbf{a}_{n})\) to denote the vectorized form of \(\mathbf{A}\).

### Problem Setup and Assumptions

Optimization literature often focuses on the constraint-free minimization of a function \(\hat{f}:\mathbb{R}^{d}\rightarrow\mathbb{R}\). In this paper, we reformulate this problem using the following definition.

**Definition 1**.: _(Partitioned Equivalence) A function \(f:\mathbb{R}^{d_{1}}\times\mathbb{R}^{d_{2}}\rightarrow\mathbb{R}\) is called a partitioned equivalence of \(\hat{f}:\mathbb{R}^{d}\rightarrow\mathbb{R}\), if \(i)\)\(d_{1}+d_{2}=d\), and \(ii)\) there exists a permutation function \(\pi:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}\) over the parameters of \(\hat{f}\), such that \(\hat{f}(\mathbf{w})=f(\mathbf{x},\mathbf{u})\) if and only if \(\pi(\mathbf{w})=(\mathbf{x},\mathbf{u})\). We say that \((\mathbf{x},\mathbf{u})\) is a partition of \(\mathbf{w}\)._

Despite the difference in the representation of their parameters, \(\hat{f}\) and \(f\) share the same properties, and any algorithm for \(\hat{f}\) would produce the same result for \(f\). Therefore, we turn our focus from the minimization problem of \(\hat{f}\) to the minimization problem in Equation (1). We further assume that \(f\) is a composition of a loss function \(g:\mathbb{R}^{\hat{d}}\rightarrow\mathbb{R}\) and a possibly non-smooth and non-convex model function \(h:\mathbb{R}^{d_{1}}\times\mathbb{R}^{d_{2}}\rightarrow\mathbb{R}^{\hat{d}}\), for some dimension \(\hat{d}\in\mathbb{Z}_{+}\); i.e., \(f(\mathbf{x},\mathbf{u})=g(h(\mathbf{x},\mathbf{u}))\). We obey the following notation with respect to gradients of \(f\):

\[\nabla_{1}f(\mathbf{x},\mathbf{u})=\frac{\partial f(\mathbf{x},\mathbf{u})}{ \partial\mathbf{x}};\quad\nabla_{2}f(\mathbf{x},\mathbf{u})=\frac{\partial f (\mathbf{x},\mathbf{u})}{\partial\mathbf{u}};\;\nabla f(\mathbf{x},\mathbf{u })=\left(\nabla_{1}f(\mathbf{x},\mathbf{u}),\nabla_{2}f(\mathbf{x},\mathbf{u })\right). \tag{2}\]

We consider Nesterov's momentum with constant step size \(\eta\) and momentum parameter \(\beta\):

\[(\mathbf{x}_{k+1},\mathbf{u}_{k+1}) =(\mathbf{y}_{k},\mathbf{v}_{k})-\eta\nabla f(\mathbf{y}_{k}, \mathbf{u}_{k}) \tag{3}\] \[(\mathbf{y}_{k+1},\mathbf{v}_{k+1}) =(\mathbf{x}_{k+1},\mathbf{u}_{k+1})+\beta\left((\mathbf{x}_{k+1},\mathbf{u}_{k+1})-(\mathbf{x}_{k},\mathbf{u}_{k})\right)\]

with \(\mathbf{y}_{0}=\mathbf{x}_{0}\) and \(\mathbf{v}_{0}=\mathbf{u}_{0}\). Let \(\mathcal{B}^{(1)}_{R_{\mathbf{x}}}\subseteq\mathbb{R}^{d_{1}}\) and \(\mathcal{B}^{(2)}_{R_{\mathbf{u}}}\subseteq\mathbb{R}^{d_{2}}\) denote the balls centered as \(\mathbf{x}_{0}\):

\[\mathcal{B}^{(1)}_{R_{\mathbf{x}}}=\{\mathbf{x}\in\mathbb{R}^{d_{1}}:\left\| \mathbf{x}-\mathbf{x}_{0}\right\|_{2}\leq R_{\mathbf{x}}\};\quad\mathcal{B}^{ (2)}_{R_{\mathbf{u}}}=\{\mathbf{u}\in\mathbb{R}^{d_{2}}:\left\|\mathbf{u}- \mathbf{u}_{0}\right\|_{2}\leq R_{\mathbf{u}}\}\]

With the definition of \(\mathcal{B}^{(1)}_{R_{\mathbf{x}}}\) and \(\mathcal{B}^{(2)}_{R_{\mathbf{u}}}\), We make the following assumptions:

**Assumption 1**.: \(f\) _is \(\mu\)-strongly convex with \(\mu>0\) with respect to the first part of its parameters:_

\[f(\mathbf{y},\mathbf{u})\geq f(\mathbf{x},\mathbf{u})+\left\langle\nabla_{1}f( \mathbf{x},\mathbf{u}),\mathbf{y}-\mathbf{x}\right\rangle+\frac{\mu}{2}\left\| \mathbf{y}-\mathbf{x}\right\|_{2}^{2},\quad\forall\mathbf{x},\mathbf{y}\in \mathbb{R}^{d_{1}};\;\mathbf{u}\in\mathcal{B}^{(2)}_{R_{\mathbf{u}}}.\]

**Assumption 2**.: \(f\) _is \(L_{1}\)-smooth with respect to the first part of its parameters:_

\[f(\mathbf{y},\mathbf{u})\leq f(\mathbf{x},\mathbf{u})+\left\langle\nabla_{1}f( \mathbf{x},\mathbf{u}),\mathbf{y}-\mathbf{x}\right\rangle+\frac{L_{1}}{2} \left\|\mathbf{y}-\mathbf{x}\right\|_{2}^{2},\quad\forall\mathbf{x},\mathbf{y }\in\mathbb{R}^{d_{1}};\;\mathbf{u}\in\mathcal{B}^{(2)}_{R_{\mathbf{u}}}.\]
Based on Assumption 1 and 2, we define the condition number of \(f\).

**Definition 2**.: _(Condition Number) The condition number \(\kappa\) of \(f\) is given by \(\kappa=\nicefrac{{L_{1}}}{{\mu}}\)._

**Assumption 3**.: \(g\) _satisfies \(\min_{\mathbf{s}\in\mathbb{R}^{d}}g(\mathbf{s})=\min_{\mathbf{x}\in\mathbb{R}^{ d_{1}},\mathbf{u}\in\mathbb{R}^{d_{2}}}f(\mathbf{x},\mathbf{u})\), and is \(L_{2}\)-smooth:_

\[g(\mathbf{s}_{1})\leq g(\mathbf{s}_{2})+\langle\nabla g(\mathbf{s}_{1}), \mathbf{s}_{2}-\mathbf{s}_{1}\rangle+\frac{L_{2}}{2}\left\|\mathbf{s}_{2}- \mathbf{s}_{1}\right\|_{2}^{2},\quad\forall\mathbf{s}_{1},\mathbf{s}_{2}\in \mathbb{R}^{\hat{d}}.\]

Assumptions 1 and 2 are relaxed versions of the smoothness and strong convexity, accordingly, also made in classical optimization literature: Instead of assuming that the objective is smooth and strongly convex over the whole set of parameters, we only assume such property to hold with respect to a subset of the parameters, while the rest of the parameters are in a ball near initialization. Assumption 3 is standard in prior literature (Liu et al., 2020; Song et al., 2021) on the convergence of neural network training, and holds for a large number of loss functions such as the MSE loss and the logistic loss.

**Assumption 4**.: \(h\) _satisfies \(G_{1}\)-Lipschitzness with respect to the second part of its parameters:_

\[\left\|h(\mathbf{x},\mathbf{u})-h(\mathbf{x},\mathbf{v})\right\|_{2}\leq G_{ 1}\left\|\mathbf{u}-\mathbf{v}\right\|_{2},\quad\forall\mathbf{x}\in\mathcal{ B}_{R_{\mathbf{x}}}^{(1)};\;\mathbf{u},\mathbf{v}\in\mathcal{B}_{R_{\mathbf{u}}}^{(2)}.\]

**Assumption 5**.: _The gradient of \(f\) with respect to the first part of its parameter, \(\nabla_{1}f(\mathbf{x},\mathbf{u})\), satisfies \(G_{2}\)-Lipschitzness with respect to the second part of its parameters:_

\[\left\|\nabla_{1}f(\mathbf{x},\mathbf{u})-\nabla_{1}f(\mathbf{x},\mathbf{v}) \right\|_{2}\leq G_{2}\left\|\mathbf{u}-\mathbf{v}\right\|_{2},\quad\forall \mathbf{x}\in\mathcal{B}_{R_{\mathbf{x}}}^{(1)};\;\mathbf{u},\mathbf{v}\in \mathcal{B}_{R_{\mathbf{u}}}^{(2)}.\]

**Assumption 6**.: _All the minimum values of \(f\), when restricted to the optimization over \(\mathbf{x}\), equal to the global minimum value:_

\[\min_{\mathbf{x}\in\mathbb{R}^{d_{1}}}f(\mathbf{x},\mathbf{u})=f^{\star}:=\min _{\mathbf{x}\in\mathbb{R}^{d_{1}},\mathbf{u}\in\mathbb{R}^{d_{2}}}f(\mathbf{x},\mathbf{u});\quad\forall\mathbf{u}\in\mathcal{B}_{R_{\mathbf{u}}}^{(2)}.\]

Notice that we do not assume \(f\) to be either convex or smooth with respect to \(\mathbf{u}\). As such, there is no reason to believe that the updates in (3) on \(\mathbf{u}\) will make positive progress towards finding the global minimum. We treat the change in the second part of the parameters as errors induced by the updates. Assumptions 4 and 5 are made such that the effect on the change of the model output \(h(\mathbf{x},\mathbf{u})\) and the gradient with respect to \(\mathbf{x}\) caused by the change of \(\mathbf{u}\) can be controlled. Moreover, without Assumption 6, it is possible that the change of \(\mathbf{u}\) will cause the optimization trajectory to get stuck at some point of \(\mathbf{u}\), such that the global minimum value cannot be achieved even when \(\mathbf{x}\) is fully optimized. As a sanity check, we can show that Assumptions 1-6 are satisfied by a smooth and strongly convex function:

**Theorem 1**.: _Let \(\tilde{f}\) be \(\tilde{\mu}\)-strongly convex and \(\tilde{L}\)-smooth. Then \(\tilde{f}\) satisfies Assumptions 1-6 with:_

\[R_{\mathbf{x}}=R_{\mathbf{u}}=\infty;\;\mu=\tilde{\mu};\;L_{1}=L_{2}=\tilde{L} ;\;G_{1}=G_{2}=0.\]

Theorem 1 shows that the combination of Assumptions 1-6 is no stronger than the assumption that the objective is smooth and strongly convex. Therefore, the minimization of the class of functions satisfying Assumptions 1-6 does not lead to better lower complexity bound than the class of smooth and strong convex functions. That is, the best convergence rate we can achieve is \(1-\Theta\left(\nicefrac{{1}}{{\sqrt{\kappa}}}\right)\).

**Remark 1**.: _Notice that Nesterov's momentum in (3) is agnostic to the partition of the parameters, i.e., it executes exactly the same operator on \(\mathbf{x}\) and \(\mathbf{u}\). Therefore, while our result requires the existence of a parameter partition \((\mathbf{x},\mathbf{u})\) that satisfies Assumptions 1-6, it does not need to explicitly identify such a partition. Moreover, in Sections 5.1 and 5.2, we will show that Assumptions 1-6 can be satisfied by two concrete models, where training of deep ReLU networks is one of them._
Accelerated Convergence under Partial Strong Convexity

### Warmup: Convergence of Gradient Descent

We first focus on the convergence of gradient descent. To start, we notice that Assumptions 1 and 6 impliy the PL condition:

**Lemma 1**.: _Suppose that Assumption 1, 6 hold. Then, for all \(\mathbf{x}\in\mathbb{R}^{d}\) and \(\mathbf{u}\in\mathcal{B}^{(2)}_{R_{\mathbf{u}}}\), we have:_

\[\left\|\nabla f(\mathbf{x},\mathbf{u})\right\|_{2}^{2}\geq\left\|\nabla_{1}f( \mathbf{x},\mathbf{u})\right\|_{2}^{2}\geq 2\mu\left(f(\mathbf{x},\mathbf{u})-f^{ \star}\right).\]

Recall that, due to the minimum assumption made on the relationship between \(f(\mathbf{x},\mathbf{u})\) and \(\mathbf{u}\), we treat the change of \(\mathbf{u}\) during the iterates as an error. Thus, we need the following lemma, which bounds how much \(f\) is affected by the change of \(\mathbf{u}\).

**Lemma 2**.: _Suppose that Assumptions 3, 4 hold. For any \(\hat{\mathcal{Q}}>0\) and \(\mathbf{x}\in\mathcal{B}^{(1)}_{R_{\mathbf{x}}},\mathbf{u},\mathbf{v}\in \mathcal{B}^{(2)}_{R_{\mathbf{u}}}\), we have:_

\[f(\mathbf{x},\mathbf{u})-f(\mathbf{x},\mathbf{v})\leq\hat{\mathcal{Q}}^{-1}L_ {2}\left(f(\mathbf{x},\mathbf{v})-f^{\star}\right)+\frac{G_{1}^{2}}{2}\left(L _{2}+\hat{\mathcal{Q}}\right)\left\|\mathbf{u}-\mathbf{v}\right\|_{2}^{2}.\]

With the help of Lemmas 1 and 2, we can show the linear convergence of gradient descent:

**Theorem 2**.: _Suppose that Assumptions 1-4 and 6 hold with \(G_{1}^{4}\leq\frac{\mu^{2}}{8L_{2}^{2}}\) and_

\[R_{\mathbf{x}}\geq 16\eta\kappa\sqrt{L_{1}}\left(f(\mathbf{x}_{0},\mathbf{u}_{ 0})-f^{\star}\right)^{\frac{1}{2}};\ R_{\mathbf{u}}\geq 16\eta\kappa G_{1} \sqrt{L_{2}}\left(f(\mathbf{x}_{0},\mathbf{u}_{0})-f^{\star}\right)^{\frac{1} {2}}.\]

_Then, there exists constant \(c>0\) such that the gradient descent on \(f\) given by:_

\[(\mathbf{x}_{k+1},\mathbf{u}_{k+1})=(\mathbf{x}_{k},\mathbf{u}_{k})-\eta \nabla f(\mathbf{x}_{k},\mathbf{u}_{k}),\]

_with \(\eta=\nicefrac{{c}}{{L_{1}}}\), converges linearly according to:_

\[f(\mathbf{x}_{k},\mathbf{u}_{k})-f^{\star}\leq\left(1-\frac{c}{4\kappa} \right)^{k}\left(f(\mathbf{x}_{0},\mathbf{u}_{0})-f^{\star}\right).\]

I.e., Theorem 2 shows that gradient descent applied to \(f\) converges linearly with a rate of \(1-\Theta(\nicefrac{{1}}{{\kappa}})\) within our settings. The proofs for Lemma 1 and 2, and Theorem 2 are deferred to Appendix B.

### Technical Difficulty

We will now study the convergence property of Nesterov's momentum in equation 3 under only Assumptions 1-6. Our analysis bears two major difficulties.

**Difficulty 1**.: _Most previous analyses of Nesterov's momentum use the global minimum as a reference point to construct the Lyapunov function; see Bansal and Gupta (2019), d'Aspremont et al. (2021). In the original proof of Nesterov, the construction of estimating sequence also assumes the uniqueness of global minimum Nesterov (2018). However, our objective function is non-convex and allows the existence of multiple global minima, which prevents us from directly applying the Lyapunov function or estimating sequence as in previous works._

While the non-convexity of \(f\) introduces the possibility of multiple global minima, Assumption 1 implies that, with a fixed \(\mathbf{u}\), there exists a unique \(\mathbf{x}^{\star}(\mathbf{u})\) that minimizes \(f(\mathbf{x},\mathbf{u})\). Moreover, Assumption 6 implies that, for all \(\mathbf{u}\in\mathcal{B}^{(2)}_{R_{\mathbf{u}}}\), the local minimum \((\mathbf{x}^{\star}(\mathbf{u}),\mathbf{u})\) is also a global minimum. Thus, we can resolve the difficulty of accounting for multiple global minima by tracking the change of \(\mathbf{x}^{\star}(\mathbf{u})\) with the change of \(\mathbf{u}\). The following lemma gives a characterization of this property.
**Lemma 3**.: _Let \(\mathbf{x}^{\star}(\mathbf{u})=\arg\min_{\mathbf{x}\in\mathbb{R}^{d_{1}}}f( \mathbf{x},\mathbf{u})\). Suppose Assumptions 1 and 5 hold. Then, we have:_

\[\|\mathbf{x}^{\star}(\mathbf{u}_{1})-\mathbf{x}^{\star}(\mathbf{u}_{2})\|_{2} \leq\frac{G_{2}}{\mu}\left\|\mathbf{u}_{1}-\mathbf{u}_{2}\right\|_{2},\quad \forall\mathbf{u}_{1},\mathbf{u}_{2}\in\mathcal{B}_{R_{\mathbf{u}}}^{(2)}.\]

Lemma 3 indicates that, if we view \(\mathbf{x}^{\star}(\mathbf{u})\) as a function of \(\mathbf{u}\), then this function is \(\frac{G_{2}}{\mu}\)-Lipschitz. For a fixed \(\mathbf{u}\), given the nice properties on \(\mathbf{x}\), the iterates of Nesterov's momentum will guide \(\mathbf{x}\) to the minimum, based on the current \(\mathbf{u}\). Lemma 3 guarantees that the progress toward the minimum induced by \(\mathbf{u}_{1}\) does not deviate much from the progress toward the minimum induced by \(\mathbf{u}_{2}\).

**Difficulty 2**.: _Lemmas 2 and 3 show the importance of bounding the change of \(\mathbf{u}\) during iterates, namely \(\left\|\mathbf{u}_{k+1}-\mathbf{u}_{k}\right\|_{2}\). In previous works that focus on the convergence of neural network training, it is shown that this change of \(\mathbf{u}\) is closely related to the norm of the gradient. Under the assumption of smoothness, the norm of the gradient is bounded by a factor times the optimality gap at the current point, which is shown to enjoy a linear convergence by induction Du et al. (2019b); Nguyen (2021). However, in the case of Nesterov's momentum, we cannot directly utilize this relationship, since the gradient is evaluated at the intermediate step \((\mathbf{y}_{k},\mathbf{v}_{k})\), and, while we know that the optimality gap at \((\mathbf{x}_{k},\mathbf{u}_{k})\) converges linearly, we have very little knowledge about the optimality gap at \((\mathbf{y}_{k},\mathbf{v}_{k})\)._

To tackle this difficulty, our analysis starts with a careful bound on \(\|\mathbf{x}_{k+1}-\mathbf{x}_{k}\|_{2}^{2}\) by utilizing the convexity with respect to \(\mathbf{x}\) to characterize the inner product \(\langle\nabla_{1}f(\mathbf{y}_{k},\mathbf{v}_{k}),\mathbf{x}_{k}-\mathbf{x}_ {k-1}\rangle\). After that, we derive a bound of \(\|\nabla_{1}f(\mathbf{y}_{k},\mathbf{v}_{k})\|_{2}^{2}\) using a combination of \(\left\|\mathbf{x}_{k+1}-\mathbf{x}_{k}\right\|_{2}^{2}\) and \(\left\|\mathbf{x}_{k}-\mathbf{x}_{k-1}\right\|_{2}^{2}\). Lastly, we relate \(\|\nabla_{2}f(\mathbf{y}_{k},\mathbf{v}_{k})\|_{2}^{2}\) to \(\|\nabla_{1}f(\mathbf{y}_{k},\mathbf{v}_{k})\|_{2}^{2}\) using the following gradient dominance property.

**Lemma 4**.: _Suppose that Assumptions 1, 3, 4, and 6 hold. Then, we have:_

\[\|\nabla_{2}f(\mathbf{x},\mathbf{u})\|_{2}^{2}\leq\frac{G_{1}^{2}L_{2}}{\mu} \left\|\nabla_{1}f(\mathbf{x},\mathbf{u})\right\|_{2}^{2},\quad\forall \mathbf{x}\in\mathcal{B}_{R_{\mathbf{x}}}^{(1)};\;\mathbf{u}\in\mathcal{B}_{R _{\mathbf{u}}}^{(2)}.\]

Lemma 4 establishes the upper bound of \(\|\nabla_{2}f(\mathbf{x},\mathbf{u})\|_{2}^{2}\) using \(\|\nabla_{1}f(\mathbf{x},\mathbf{u})\|_{2}^{2}\). Direct application of this result contributes to the bound on \(\left\|\mathbf{u}_{k+1}-\mathbf{u}_{k}\right\|_{2}\). Intuitively, this lemma implies that the effect of gradient update on \(\mathbf{u}\) is less significant than the effect of gradient update on \(\mathbf{x}\).

### Acceleration of Nesterov's Momentum

After resolving the two technical difficulties stated above, we are ready to derive our main result. We will first state the main result theorem, and provide a sketch of the proof. The detailed proof will be deferred to Appendix D.

**Theorem 3**.: _Let Assumptions (1)-(6) hold. Consider Nesterov's momentum given by (3) with initialization \(\{\mathbf{x}_{0},\mathbf{u}_{0}\}=\{\mathbf{y}_{0},\mathbf{v}_{0}\}\). There exists absolute constants \(c,C_{1},C_{2}>0\), such that, if \(\mu,L_{1},L_{2},G_{1},G_{2}\) and \(R_{\mathbf{x}},R_{\mathbf{u}}\) satisfy:_

\[G_{1}^{4} \leq\frac{C_{1}\mu^{2}}{L_{2}(L_{2}+1)^{2}}\left(\frac{1-\beta}{ 1+\beta}\right)^{3};\quad G_{1}^{2}G_{2}^{2}\leq\frac{C_{2}\mu^{3}}{L_{2}(L_{ 2}+1)\sqrt{\kappa}}\left(\frac{1-\beta}{1+\beta}\right)^{2}; \tag{4}\] \[R_{\mathbf{x}} \geq\frac{36}{c}\sqrt{\kappa}\left(\frac{\eta(L_{2}+1)}{1-\beta} \right)^{\frac{1}{2}}(f(\mathbf{x}_{0},\mathbf{u}_{0})-f^{\star})^{\frac{1}{2}};\] \[R_{\mathbf{u}} \geq\frac{36}{c}\sqrt{\kappa}\left(\frac{\eta G_{1}^{2}L_{2}(L_{2 }+1)(1+\beta)^{3}}{\mu\beta(1-\beta)^{3}}\right)^{\frac{1}{2}}(f(\mathbf{x}_ {0},\mathbf{u}_{0})-f^{\star})^{\frac{1}{2}},\]
_and, if we choose \(\eta=\nicefrac{{c}}{{L_{\star}}}\), \(\beta=\nicefrac{{(4\sqrt{\kappa}-\sqrt{c})}}{{(4\sqrt{\kappa}+\tau\sqrt{c})}}\), then \(\mathbf{x}_{k},\mathbf{y}_{k}\in\mathcal{B}_{R_{\mathbf{x}}}^{(1)}\) and \(\mathbf{u}_{k},\mathbf{v}_{k}\in\mathcal{B}_{R_{\mathbf{u}}}^{(2)}\) for all \(k\in\mathbb{N}\), and Nesterov's recursion converges according to:_

\[f(\mathbf{x}_{k},\mathbf{u}_{k})-f^{\star}\leq 2\left(1-\frac{c}{4\sqrt{ \kappa}}\right)^{k}(f(\mathbf{x}_{0},\mathbf{u}_{0})-f^{\star}). \tag{5}\]

Theorem 3 shows that, under Assumptions 1-6 with a sufficiently small \(G_{1}\) and \(G_{2}\) as in (4), Nesterov's iteration enjoys an accelerated convergence, as in (5). Moreover, the iterates of Nesterov's momentum \(\{(\mathbf{x}_{k},\mathbf{y}_{k})\}_{k=1}^{\infty}\) and \(\{(\mathbf{u}_{k},\mathbf{v}_{k})\}_{k=1}^{\infty}\) stay in a ball around initialization with radius in (4).

To better interpret our result, we first focus on (4). By our choice of \(\beta\), we have that \(1-\beta=\Theta\left(\nicefrac{{1}}{{\sqrt{\kappa}}}\right)\) and \(1+\beta=\Theta\left(1\right)\). Therefore, the requirement of \(G_{1},G_{2}\) in (4) can be simplified to \(G_{1}^{4}\leq O\left(\nicefrac{{\mu^{\gamma/2}}}{{L_{1}^{3/2}L_{2}^{3}}}\right)\) and \(G_{1}^{2}G_{2}^{2}\leq O\left(\nicefrac{{\mu^{\beta/2}}}{{L_{1}^{3/2}L_{2}^{3 }}}\right)\). This simplified condition implies that we need a smaller \(G_{1}\) and \(G_{2}\), if \(\mu\) is small and if \(L_{1}\) and \(L_{2}\) are large. For the requirement on \(R_{\mathbf{x}}\) and \(R_{\mathbf{u}}\) in (4), we can simplify with \(\eta=O\left(\nicefrac{{1}}{{L_{1}}}\right)\) and \(\beta=\Theta\left(1\right)\). In this way, \(R_{\mathbf{x}}\) and \(R_{\mathbf{u}}\) reduce to \(\Omega\left(\nicefrac{{L_{1}^{3/4}L_{2}^{3/2}}}{\mu^{3/4}}\right)\cdot(f( \mathbf{x}_{0},\mathbf{u}_{0})-f^{\star})^{\frac{1}{2}}\) and \(\Omega\left(G_{1}\nicefrac{{L_{1}^{3/4}L_{2}}}{\mu^{7/4}}\right)\cdot(f( \mathbf{x}_{0},\mathbf{u}_{0})-f^{\star})^{\frac{1}{2}}\), respectively. Both quantities grow with a larger \(L_{1}\) and \(L_{2}\) and a smaller \(\mu\). Noticeably \(R_{\mathbf{u}}\) also scales with \(G_{1}\). Focusing on the convergence property in (5), we can conclude that Nesterov's momentum achieves an accelerated convergence rate of \(1-\Theta\left(\nicefrac{{1}}{{\sqrt{\kappa}}}\right)\) compared with the \(1-\Theta\left(\nicefrac{{1}}{{\kappa}}\right)\) rate in Theorem 2. We sketch the proof of Theorem 3 below.

Sketch of Proof.: Setting \(\mathbf{y}_{-1}=\mathbf{y}_{0}\) and \(\mathbf{v}_{-1}=\mathbf{v}_{0}\), our proof utilize the following Lyapunov function:

\[\phi_{k}=f(\mathbf{x}_{k},\mathbf{u}_{k})-f^{\star}+\mathcal{Q}_{1}\left\| \mathbf{z}_{k}-\mathbf{x}_{k-1}^{\star}\right\|_{2}^{2}+\frac{\eta}{8}\left\| \nabla_{1}f(\mathbf{y}_{k-1},\mathbf{v}_{k-1})\right\|_{2}^{2};\quad\mathbf{x }_{k}^{\star}=\operatorname*{arg\,min}_{\mathbf{x}\in\mathbb{R}^{d_{1}}}f( \mathbf{x},\mathbf{v}_{k}),\]

with properly chosen \(\mathcal{Q}_{1}\) and \(\mathbf{z}_{k}\). The proof recursively establishes the following three properties:

\[\left(1-\frac{c}{2\sqrt{\kappa}}\right)^{-1}\phi_{k+1}-\phi_{k}\leq\frac{c}{4 \sqrt{\kappa}}\left(1-\frac{c}{4\sqrt{\kappa}}\right)^{k}\phi_{0}; \tag{6}\]

\[\left\|\mathbf{x}_{k}-\mathbf{x}_{k-1}\right\|_{2}^{2}\leq\mathcal{Q}_{2}\left( 1-\frac{c}{4\sqrt{\kappa}}\right)^{k}\phi_{0};\quad\left\|\mathbf{u}_{k}- \mathbf{u}_{k-1}\right\|_{2}^{2}\leq G_{1}^{2}\mathcal{Q}_{3}\left(1-\frac{c}{4 \sqrt{\kappa}}\right)^{k}\phi_{0}. \tag{7}\]

By our choice of \(\phi_{k}\), Equation (6) implies that:

\[f(\mathbf{x}_{k},\mathbf{u}_{k})-f^{\star}\leq\phi_{k}\leq\left(1-\frac{c}{4 \sqrt{\kappa}}\right)^{k}\phi_{0}, \tag{8}\]

which further implies the convergence in Equation (5). The following lemma shows that, if Equation (5) holds for all \(k\leq\hat{k}\), we can guarantee Equation (7) with \(k=\hat{k}+1\).

**Lemma 5**.: _Let the Assumptions of Theorem 3 hold. If Equation (8) holds for all \(k\leq\hat{k}\), then Equation (7) holds for \(k=\hat{k}+1\) with \(\mathcal{Q}_{2}=\frac{6\eta L_{2}(L_{2}+1)}{1-\beta}\) and \(\mathcal{Q}_{3}=\frac{6\eta L_{2}(L_{2}+1)(1+\beta)^{3}}{\mu\beta(1-\beta)^{3}}\)._

With Lemma 5, we can first guarantee that up to iteration \(\hat{k}+1\), \(\mathbf{x}_{k}\) and \(\mathbf{u}_{k}\) stay in \(\mathcal{B}_{R_{\mathbf{x}}}^{(1)}\) and \(\mathcal{B}_{R_{\mathbf{u}}}^{(2)}\). Thus, we can utilize Assumptions 1-6 in showing Equation (6) for iteration \(\hat{k}+1\). Moreover, the bound on \(\left\|\mathbf{u}_{k}-\mathbf{u}_{k-1}\right\|_{2}\) connects the following lemma to Equation (6).

**Lemma 6**.: _Let the Assumptions of Theorem (3) hold. Then, we have:_

\[\left(1-\frac{c}{2\sqrt{\kappa}}\right)^{-1}\phi_{k+1}-\phi_{k}\leq c\beta^{2} \sqrt{\kappa}\left(G_{1}^{2}L_{2}+\frac{8\mathcal{Q}_{1}G_{1}^{2}G_{2}^{2}}{\mu^{2 }}\right)\left\|\mathbf{u}_{k}-\mathbf{u}_{k-1}\right\|_{2}^{2}\]

Lemmas 5, 6 together with a small enough \(G_{1}\) and \(G_{2}\) imply Equation (6).
Examples

We consider two examples that satisfy Assumptions 1-6. By enforcing the requirement in Theorem 3, we show that, although being nonconvex and possibly non-smooth, the two models enjoy accelerated convergence when trained with Nesterov's momentum.

### Toy example: Additive model

Given matrices \(\mathbf{A}_{1}\in\mathbb{R}^{m\times m},\mathbf{A}_{2}\in\mathbb{R}^{m\times d}\) and a non-linear function \(\sigma:\mathbb{R}^{m}\rightarrow\mathbb{R}^{m}\), we consider \(h(\mathbf{x},\mathbf{u})=\mathbf{A}_{1}\mathbf{x}+\sigma(\mathbf{A}_{2} \mathbf{u})\) as the summation of a linear model and a non-linear model. If we train \(h(\mathbf{x},\mathbf{u})\) over the loss \(g(\mathbf{s})=\frac{1}{2}\left\|\mathbf{s}-\mathbf{b}\right\|_{2}^{2}\) for some label \(\mathbf{b}\in\mathbb{R}^{m}\), the objective can be written as

\[f(\mathbf{x},\mathbf{u})=g(\mathbf{A}_{1}\mathbf{x}+\sigma\left(\mathbf{A}_{2 }\mathbf{u}\right))=\frac{1}{2}\left\|\mathbf{A}_{1}\mathbf{x}+\sigma\left( \mathbf{A}_{2}\mathbf{u}\right)-\mathbf{b}\right\|_{2}^{2}. \tag{9}\]

Due to the non-linearity of \(\sigma\), \(f(\mathbf{x},\mathbf{u})\) is in general non-convex. If we further choose \(\sigma\) to be some non-smooth function such as ReLU, i.e., \(\sigma(\mathbf{x})_{i}=\max\{0,x_{i}\}\), the objective can also be non-smooth. Yet, assuming that \(\sigma\) is Lipschitz, we can show that \(f(\mathbf{x},\mathbf{u})\) satisfies Assumptions 1-6.

**Lemma 7**.: _Suppose that \(\sigma\) is \(B\)-Lipschitz and \(\sigma_{\min}(\mathbf{A}_{1})>0\). Then \(f(\mathbf{x},\mathbf{u})\) satisfies Assumptions 1-6 with_

\[\begin{split} R_{\mathbf{x}}=& R_{\mathbf{u}}= \infty;\;\mu=\sigma_{\min}\left(\mathbf{A}_{1}\right)^{2};\;L_{1}=\sigma_{ \max}\left(\mathbf{A}_{1}\right)^{2};\;L_{2}=1\\ G_{1}=& B\sigma_{\max}\left(\mathbf{A}_{2}\right); \;G_{2}=B\sigma_{\max}\left(\mathbf{A}_{1}\right)\sigma_{\max}\left(\mathbf{A }_{2}\right).\end{split} \tag{10}\]

Notice that while \(\mu\) and \(L_{1}\) depend entirely on the property of \(\mathbf{A}_{1}\), both \(G_{1}\) and \(G_{2}\) can be made smaller by choosing \(\mathbf{A}_{2}\) with a small enough \(\sigma_{\max}(\mathbf{A}_{2})\). Intuitively, this means that \(G_{1}\) and \(G_{2}\) can be controlled, as long as the component that introduces the non-convexity and non-smoothness \(\sigma\left(\mathbf{A}_{2}\mathbf{u}\right)\) is small enough. Therefore, we can apply Theorem 3 to the minimization of Equation (9).

**Theorem 4**.: _Consider the problem in Equation (9) and assume using Nesterov's momentum to find the minimizer of \(f(\mathbf{x},\mathbf{u})\), where \(\sigma\) is a \(B\)-Lipschitz function. Let \(\kappa=\nicefrac{{\sigma_{\max}\left(\mathbf{A}_{1}\right)^{2}}}{{\sigma_{ \min}\left(\mathbf{A}_{1}\right)^{2}}}\), and suppose:_

\[\sigma_{\min}\left(\mathbf{A}_{1}\right)\geq\tilde{C}\sigma_{\max}\left( \mathbf{A}_{2}\right)B\kappa^{0.75}, \tag{11}\]

_for some large enough constant \(\tilde{C}>0\). Then there exists a constant \(c>0\) such that, if we choose \(\eta=\nicefrac{{c}}{{\sigma_{\max}\left(\mathbf{A}_{1}\right)^{2}}}\) and \(\beta=\nicefrac{{(4\sqrt{\kappa}-\sqrt{c})}}{{(4\sqrt{\kappa}+7\sqrt{c})}}\), Nesterov's momentum in Equation (3) converges according to:_

\[f(\mathbf{x}_{k},\mathbf{u}_{k})\leq 2\left(1-\frac{c}{4\sqrt{\kappa}}\right)^{k }f(\mathbf{x}_{0},\mathbf{u}_{0}).\]

Notice that the requirement in Equation (11) favors a larger \(\sigma_{\min}(\mathbf{A}_{1})\) and smaller \(\sigma_{\max}(\mathbf{A}_{2}),B\) and \(\kappa\). Using this example, we empirically verify the theoretical result of Theorem 3, as shown in Figure 1. The three rows correspond to the cases of \(\sigma_{\max}\left(\mathbf{A}_{2}\right)\in\{0.01,0.3,5\}\), respectively. The lines in the plots denote the average loss/distance among 10 trials, while the shaded region denotes the standard deviation. Observing the plots in the left column, we can see that in all three cases, Nesterov's momentum achieves a faster convergence compared with gradient descent, while the case with the largest \(\sigma_{\max}\left(\mathbf{A}_{2}\right)\) introduces the largest variance between results of the ten trials. Recall that \(\sigma_{\max}\left(\mathbf{A}_{2}\right)\) in our example controls the magnitude of \(G_{1}\) and \(G_{2}\). Thus, this phenomenon shows that when \(G_{1}\) and \(G_{2}\) become larger, the theoretical guarantee in Theorem 3 begins to break down and the
result depends more on the initialization. Figures in the right column plots the evolution of \(\|\mathbf{x}_{k}-\mathbf{x}_{k-1}\|_{2}^{2}\) and \(\|\mathbf{u}_{k}-\mathbf{u}_{k-1}\|_{2}^{2}\). All three figures show that the two quantity decrease linearly. This phenomenon supports the linear decay of the two quantities, as shown in Lemma 5. Moreover, as \(\sigma_{\max}\left(\mathbf{A}_{2}\right)\) increases, the relative magnitude of \(\|\mathbf{u}_{k}-\mathbf{u}_{k-1}\|_{2}^{2}\) to \(\|\mathbf{x}_{k}-\mathbf{x}_{k-1}\|_{2}^{2}\) also increase (the line of "\(\|\Delta\mathbf{u}_{k}\|_{2}\)" get closer to "\(\|\Delta\mathbf{x}_{k}\|_{2}\)"). This corresponds to the fact that \(\|\mathbf{u}_{k}-\mathbf{u}_{k-1}\|_{2}^{2}\) scales with \(G_{1}\), as shown in Lemma 5.

### Deep ReLU Neural Networks

Consider the \(\Lambda\)-layer ReLU neural network with layer widths \(\{d_{\ell}\}_{\ell=0}^{\Lambda}\). Denoting the number of training samples with \(n\), we consider the input and label of the training data given by \(\mathbf{X}\in\mathbb{R}^{n\times d_{0}}\) and \(\mathbf{Y}\in\mathbb{R}^{n\times d_{\Lambda}}\). Let the weight matrix in the \(\ell\)-th layer be \(\mathbf{W}_{\ell}\). We use \(\boldsymbol{\theta}=\{\mathbf{W}_{\ell}\}_{\ell=1}^{\Lambda}\) to denote the collection of all weights and \(\sigma(\mathbf{A})_{ij}=\max\{0,a_{ij}\}\) to denote the ReLU function. Then, the output of each layer is given by:

\[\mathbf{F}_{\ell}\left(\boldsymbol{\theta}\right)=\begin{cases}\mathbf{X},& \text{if }\ell=0;\\ \sigma\left(\mathbf{F}_{\ell-1}\left(\boldsymbol{\theta}\right)\mathbf{W}_{ \ell}\right),&\text{if }\ell\in[\Lambda-1];\\ \mathbf{F}_{\Lambda-1}\left(\boldsymbol{\theta}\right)\mathbf{W}_{\Lambda},& \text{if }\ell=\Lambda.\end{cases} \tag{12}\]

We consider the training of \(\mathbf{F}_{\Lambda}\left(\boldsymbol{\theta}\right)\) over the MSE loss, as in \(\mathcal{L}(\boldsymbol{\theta})=\frac{1}{2}\left\|\mathbf{F}_{\Lambda} \left(\boldsymbol{\theta}\right)-\boldsymbol{Y}\right\|_{F}^{2}\). We can interpret the scenario using our partition model in (1). Let \(g(\mathbf{s})=\frac{1}{2}\left\|\mathbf{s}-\mathtt{V}\left(\mathbf{Y}\right) \right\|_{2}^{2}\). If we partition the parameter \(\boldsymbol{\theta}\) into \(\mathbf{x}=\mathtt{V}\left(\mathbf{W}_{\Lambda}\right)\) and \(\mathbf{u}=\left(\mathtt{V}\left(\mathbf{W}_{1}\right),\ldots,\mathtt{V}\left( \mathbf{W}_{\Lambda-1}\right)\right)\), we can write:

\[h(\mathbf{x},\mathbf{u})=\mathtt{V}\left(\mathbf{F}_{\Lambda}\left(\boldsymbol {\theta}\right)\right);\quad f(\mathbf{x},\mathbf{u})=\frac{1}{2}\left\| \mathtt{V}\left(\mathbf{F}_{\Lambda}\left(\boldsymbol{\theta}\right)\right)- \mathtt{V}\left(\mathbf{Y}\right)\right\|_{2}^{2}=\frac{1}{2}\left\|\mathbf{F} _{\Lambda}\left(\boldsymbol{\theta}\right)-\mathbf{Y}\right\|_{F}^{2}. \tag{13}\]

For some given \(\mathbf{x}\) and \(\mathbf{u}\), we let \(\mathbf{W}_{\Lambda}(\mathbf{x})\) be the matrix such that \(\mathbf{x}=\mathtt{V}\left(\mathbf{W}_{\Lambda}(\mathbf{x})\right)\); similarly, we let \(\mathbf{W}_{\ell}(\mathbf{u})\) with \(\ell\in[L-1]\) be the matrices such that \(\mathbf{u}=\left(\mathtt{V}\left(\mathbf{W}_{1}(\mathbf{u})\right),\ldots, \mathtt{V}\left(\mathbf{W}_{\Lambda-1}(\mathbf{u})\right)\right)\). Denote \(\lambda_{\Lambda}=\sup_{\mathbf{x}\in\mathcal{B}_{\mathbf{Rx}}^{(1)}}\sigma_{ \max}\left(\mathbf{W}_{\Lambda}(\mathbf{x})\right)\) and \(\lambda_{\ell}=\sup_{\mathbf{u}\in\mathcal{B}_{\mathbf{Rx}}^{(2)}}\sigma_{ \max}\left(\mathbf{W}_{\ell}(\mathbf{u})\right)\) for \(\ell\in[\Lambda-1]\). Moreover, denote \(\lambda_{i\to j}=\prod_{\ell=i}^{j}\lambda_{\ell}\). Then we can show that \(f(\mathbf{x},\mathbf{u})\) defined in (13) satisfies Assumptions 1-6.

**Lemma 8**.: _Let \(\boldsymbol{\theta}(0)\) be the initialization of the ReLU network in Equation (12) and \(\alpha_{0}=\sigma_{\min}\left(\mathbf{F}_{\Lambda-1}\left(\boldsymbol{\theta}(0) \right)\right)\). Assume that each \(\mathbf{W}_{\ell}\) is initialized such that \(\|\mathbf{W}_{\ell}(0)\|_{2}\leq\frac{\lambda_{\ell}}{2}\) and \(\alpha_{0}>0\). Then, \(f(\mathbf{x},\mathbf{u})\) satisfies Assumptions 1-6 with:_

\[R_{\mathbf{x}}\leq\tfrac{\lambda_{\Lambda}}{2};\;R_{\mathbf{u}}\leq\tfrac{1}{2} \min_{\ell\in[\Lambda-1]}\lambda_{\ell}\cdot\min\left\{1,\tfrac{\alpha_{0}}{2 \sqrt{\Lambda}\|\mathbf{X}\|_{F}\lambda_{1\to\Lambda-1}}\right\};\;\mu=\tfrac{ \alpha_{0}^{2}}{2};\]

\[L_{1}=\left\|\mathbf{X}\right\|_{F}^{2}\lambda_{1\to\Lambda-1}^{2};\;L_{2}=1;\;G_ {1}= \tfrac{\sqrt{\Lambda}\|\mathbf{X}\|_{F}\lambda_{1\to\Lambda-1}}{\min_{ \ell\in[\Lambda-1]}\lambda_{\ell}};\;G_{2}=\left(2\left\|\mathbf{X}\right\|_{F} \lambda_{1\to\Lambda}+\left\|\mathbf{Y}\right\|_{F}\right)G_{1}\]

As shown in Lemma 3.3 by Nguyen Nguyen (2021), with sufficient over-parameterization, we can guarantee that \(\alpha_{0}>0\). In order to show the accelerated convergence of Nesterov's momentum when training 12, we need to \(i)\) guarantee that the condition of \(R_{\mathbf{x}}\) and \(R_{\mathbf{u}}\) in (4) satisfies the upper bound in Lemma 8, and \(ii)\) the quantities \(\mu,L_{1},L_{2},G_{1}\) and \(G_{2}\) defined in Lemma 8 satisfy the requirement in (4). Enforcing the two conditions with sufficient over-parameterization gives us the following theorem.

**Theorem 5**.: _Consider training the ReLU neural network in Equation (12) using the MSE loss, or equivalently, minimizing \(f(\mathbf{x},\mathbf{u})\) defined in Equation (13) with Nesterov's momentum:_

\[\mathbf{W}_{\ell}(k+1) =\hat{\mathbf{W}}_{\ell}(k)-\eta\nabla_{\mathbf{W}_{\ell}}\mathcal{ L}(\hat{\boldsymbol{\theta}}(k))\] \[\hat{\mathbf{W}}_{\ell}(k+1) =\mathbf{W}_{\ell}(k+1)+\beta(\mathbf{W}_{\ell}(k+1)-\mathbf{W}_{ \ell}(k))\]
_with \(\hat{\mathbf{\theta}}(k+1)=\{\hat{\mathbf{W}}_{\ell}(k+1)\}_{\ell=1}^{\Lambda},\eta= \nicefrac{{c}}{{L_{1}}}\) and \(\beta=\frac{4\sqrt{\kappa-\sqrt{c}}}{4\sqrt{\kappa+\gamma}\zeta}\), where \(\kappa=\nicefrac{{2L_{1}}}{{a_{0}^{2}}}\) and \(\alpha_{0},L_{1}\) as defined in Lemma 8. If the width of the neural network satisfies:_

\[d_{\ell}=\Theta\left(m\right)\quad\forall\ell\in[\Lambda-1];\quad d_{L-1}= \Theta\left(n^{4}m^{2}\right), \tag{14}\]

_for some \(m\geq\max\{d_{0},d_{\Lambda}\}\), and we initialize the weights according to:_

\[\left[\mathbf{W}_{\ell}(0)\right]_{ij}\sim\mathcal{N}\left(0,d_{\ell-1}^{-1} \right),\quad\forall\ell\in[\Lambda-1];\quad\left[\mathbf{W}_{\ell}(0)\right] _{ij}\sim\mathcal{N}\left(0,d_{\Lambda-1}^{-\frac{3}{2}}\right).\]

_Then, with high probability over the initialization, there exists an absolute constant \(c>0\) such that:_

\[\mathcal{L}\left(\mathbf{\theta}(k)\right)\leq 2\left(1-\tfrac{c}{4\sqrt{\kappa}} \right)^{k}\mathcal{L}(\mathbf{\theta}(0)). \tag{15}\]

As in prior work Nguyen (2021), we treat the depth of the neural network \(\Lambda\) to be a constant when computing the over-parameterization requirement. Next, we compare our result with Theorem 2.2 and Corollary 3.2 of Nguyen (2021). To start, in order to deal with the additional complexity of Nesterov's momentum as introduced in Section 4.2, our over-parameterization of \(d_{\Lambda-1}=\Theta\left(n^{4}m^{2}\right)\) is slightly larger than the over-parameterization of \(d_{\Lambda-1}=\Theta\left(n^{3}m^{3}\right)\) in Corollary 3.2 of Nguyen (2021). Moreover, in Theorem 2.2 of Nguyen (2021), since their choice of \(\eta\) is also \(O\left(\nicefrac{{1}}{{L_{1}}}\right)\), the convergence rate they achieve for gradient descent is \(1-\Theta\left(\nicefrac{{1}}{{\kappa}}\right)\). Compared with this rate, Theorem 5 achieves a faster convergence of \(1-\Theta\left(\nicefrac{{1}}{{\sqrt{\kappa}}}\right)\). This shows that Nesterov's momentum enjoys acceleration when training deep ReLU neural networks.

## 6 Conclusion and Broader Impact

We consider the minimization of a new class of objective functions, namely the partition model, where the set of parameters can be partitioned into two groups, and the function is smooth and strongly convex with respect to only one group. This class of objectives is more general than the class of smooth and strongly convex functions. We prove the convergence of both gradient descent and Nesterov's momentum on this class of objectives under certain assumptions and show that Nesterov's momentum achieves an accelerated convergence rate of \(1-\Theta\left(\nicefrac{{1}}{{\sqrt{\kappa}}}\right)\) compared with the \(1-\Theta\left(\nicefrac{{1}}{{\kappa}}\right)\) convergence rate of gradient descent. Moreover, we considered the training of the additive model and deep ReLU networks as two concrete examples of the partition model and showed the acceleration of Nesterov's momentum when training these two models.

Our work bears two limitations. First, since we do not make convexity or smoothness assumptions for the subset \(\mathbf{u}\) of parameters, we have to make a strong assumption of the optimality condition as in Assumption 6. It is possible that this assumption can be relaxed if one assumes a better structure on \(\mathbf{u}\). Second, in the examples we provided, we only considered simple parameter partition schemes including partition by submodel as in Section

Figure 1: Experiment of learning additive model with gradient descent and Nesterov’s momentum.
5.1 and partition by layer as in Section 5.2. Considering the complexity of modern neural networks, it remains a valid open question whether more complicated partition schemes, such as selecting a subset of the weights in each layer, will still satisfy the assumptions in our framework. Future works can focus on resolving these two limitations, as well as extending the analysis to different neural network architectures.

This paper focuses mainly on the theoretical understanding of optimization and machine learning methods. Due to the nature of our work, We do not anticipate negative social impacts. Further impact on society depends on the application of machine learning models.
