Fast Convergence of Natural Gradient Descent for Overparameterized Neural Networks

Introduction

Because training large neural networks is costly, there has been much interest in using second-order optimization to speed up training (Becker and LeCun, 1989; Martens, 2010; Martens and Grosse, 2015), and in particlar natural gradient descent (Amari, 1998, 1997). Recently, scalable approximations to natural gradient descent have shown practical success in a variety of tasks and architectures (Martens and Grosse, 2015; Grosse and Martens, 2016; Wu et al., 2017; Zhang et al., 2018a; Martens et al., 2018). Natural gradient descent has an appealing interpretation as optimizing over a Riemannian manifold using an intrinsic distance metric; this implies the updates are invariant to transformations such as whitening (Ollivier, 2015; Luk and Grosse, 2018). It is also closely connected to Gauss-Newton optimization, suggesting it should achieve fast convergence in certain settings (Pascanu and Bengio, 2013; Martens, 2014; Botev et al., 2017).

Does this intuition translate into faster convergence? Amari (1998) provided arguments in the affirmative, as long as the cost function is well approximated by a convex quadratic. However, it remains unknown whether natural gradient descent can optimize neural networks faster than gradient descent — a major gap in our understanding. The problem is that the optimization of neural networks is both nonconvex and non-smooth, making it difficult to prove nontrivial convergence bounds. In general, finding a global minimum of a general non-convex function is an NP-complete problem, and neural network training in particular is NP-complete (Blum and Rivest, 1992).

However, in the past two years, researchers have finally gained substantial traction in understanding the dynamics of gradient-based optimization of neural networks. Theoretically, it has been shown that gradient descent starting from a random initialization is able to find a global minimum if the network is wide enough (Li and Liang, 2018; Du et al., 2018b, a; Zou et al., 2018; Allen-Zhu et al., 2018; Oymak and Soltanolkotabi, 2019). The key technique of those works is to show that neural networks become well-behaved if they are largely overparameterized in the sense that the number of hidden units is polynomially large in the size of the training data. However, most of these works have focused on standard gradient descent, leaving open the question of whether similar statements can be made about other optimizers.

Most convergence analysis of natural gradient descent has focused on simple convex quadratic objectives (e.g. (Martens, 2014)). Very recently, the convergence properties of NGD were studied in the context of linear networks (Bernacchia et al., 2018). While the linearity assumption simplifies the analysis of training dynamics (Saxe et al., 2013), linear networks are severely limited in terms of their expressivity, and it’s not clear which conclusions will generalize from linear to nonlinear networks.

In this work, we analyze natural gradient descent for nonlinear networks. We give two simple and generic conditions on the Jacobian matrix which guarantee efficient convergence to a global minimum. We then apply this analysis to a particular distribution over two-layer ReLU networks which has recently been used to analyze the convergence of gradient descent (Li and Liang, 2018; Du et al., 2018a; Oymak and Soltanolkotabi, 2019). We show that for sufficiently high network width, NGD will converge to the global minimum. We give bounds on the convergence rate of two-layer ReLU networks that are much better than the analogous bounds that have been proven for gradient descent (Du et al., 2018b; Wu et al., 2019; Oymak and Soltanolkotabi, 2019), while allowing for much higher learning rates. Moreover, in the limit of infinite width, and assuming a squared error loss, we show that NGD converges in just one iteration. The main contributions of our work are summarized as follows:

We show that natural gradient enables us to use a much larger step size, resulting in an even faster convergence rate. Specifically, the maximal step size of natural gradient descent is $\mathcal{O}\left(1\right)$ for (polynomially) wide networks.

We show that K-FAC (Martens and Grosse, 2015), an approximate natural gradient descent method, also converges to global minima with linear rate, although this result requires a higher level of overparameterization compared to GD and exact NGD.

We analyze the generalization properties of NGD, showing that the improved convergence rates provably don’t come at the expense of worse generalization.

Related Works

Recently, there have been many works studying the optimization problem in deep learning, i.e., why in practice many neural network architectures reliably converge to global minima (zero training error). One popular way to attack this problem is to analyze the underlying loss surface (Hardt and Ma, 2016; Kawaguchi, 2016; Kawaguchi and Bengio, 2018; Nguyen and Hein, 2017; Soudry and Carmon, 2016). The main argument of those works is that there are no bad local minima. It has been proven that gradient descent can find global minima (Ge et al., 2015; Lee et al., 2016) if the loss surface satisfies: (1) all local minima are global and (2) all saddle points are strict in the sense that there exists at least one negative curvature direction. Unfortunately, most of those works rely on unrealistic assumptions (e.g., linear activations (Hardt and Ma, 2016; Kawaguchi, 2016)) and cannot generalize to practical neural networks. Moreover, Yun et al. (2018) shows that small nonlinearity in shallow networks can create bad local minima.

Another way to understand the optimization of neural networks is to directly analyze the optimization dynamics. Our work also falls within this category. However, most work in this direction focuses on gradient descent. Bartlett et al. ; Arora et al. (2019a) studied the optimization trajectory of deep linear networks and showed that gradient descent can find global minima under some assumptions. Previously, the dynamics of linear networks have also been studied by Saxe et al. (2013); Advani and Saxe (2017). For nonlinear neural networks, a series of papers (Tian, 2017; Brutzkus and Globerson, 2017; Du et al., 2017; Li and Yuan, 2017; Zhang et al., 2018b) studied a specific class of shallow two-layer neural networks together with strong assumptions on input distribution as well as realizability of labels, proving global convergence of gradient descent. Very recently, there are some works proving global convergence of gradient descent (Li and Liang, 2018; Du et al., 2018b, a; Allen-Zhu et al., 2018; Zou et al., 2018; Gao et al., 2019) or adaptive gradient methods (Wu et al., 2019) on overparameterized neural networks. More specifically, Li and Liang (2018); Allen-Zhu et al. (2018); Zou et al. (2018) analyzed the dynamics of weights and showed that the gradient cannot be small if the objective value is large. On the other hand, Du et al. (2018b, a); Wu et al. (2019) studied the dynamics of the outputs of neural networks, where the convergence properties are captured by a Gram matrix. Our work is very similar to Du et al. (2018b); Wu et al. (2019). We note that these papers all require the step size to be sufficiently small to guarantee the global convergence, leading to slow convergence.

To our knowledge, there is only one paper (Bernacchia et al., 2018) studying the global convergence of natural gradient for neural networks. However, Bernacchia et al. (2018) only studied deep linear networks with infinitesimal step size and squared error loss functions. In this sense, our work is the first one proving global convergence of natural gradient descent on nonlinear networks.

There have been many attempts to understand the generalization properties of neural networks since Zhang et al. (2016)’s seminal paper. Researchers have proposed norm-based generalization bounds (Neyshabur et al., 2015, 2017; Bartlett and Mendelson, 2002; Bartlett et al., 2017; Golowich et al., 2017), compression bounds (Arora et al., 2018) and PAC-Bayes bounds (Dziugaite and Roy, 2017, 2018; Zou et al., 2018). Recently, overparameterization of neural networks together with good initialization has been believed to be one key factor of good generalization. Neyshabur et al. (2019) empirically showed that wide neural networks stay close to the initialization, thus leading to good generalization. Theoretically, researchers did prove that overparameterization as well as linear convergence jointly restrict the weights to be close to the initialization (Du et al., 2018b, a; Allen-Zhu et al., 2018; Zou et al., 2018; Arora et al., 2019b). The most closely related paper is Arora et al. (2019b), which shows that the optimization and generalization phenomenon can be explained by a Gram matrix. The main difference is that our analysis is based on natural gradient descent, which converges faster and provably generalizes as well as gradient descent.

Concurrently and independently, Cai et al. (2019) showed that natural gradient descent (they call it Gram-Gauss-Newton) enjoys quadratic convergence rate guarantee for overparameterized networks on regression problems. Additionally, they showed that it is much cheaper to precondition the gradient in the output space when the number of data points is much smaller than the number of parameters.

Convergence Analysis of Natural Gradient Descent

One main focus of this paper is to analyze the following procedure:

where $\eta>0$ is the step size, and $\mathbf{F}$ is the Fisher information matrix associated with the network’s predictive distribution over $y$ (which is implied by its loss function and is $\mathcal{N}(f(\bm{\theta},\mathbf{x}_{i}),1)$ for the squared error loss) and the dataset’s distribution over $\mathbf{x}$.

where $\mathbf{u}=[\mathbf{u}_{1},...,\mathbf{u}_{n}]^{\top}=[f(\bm{\theta},\mathbf{x}_{1}),...,f(\bm{\theta},\mathbf{x}_{n})]^{\top}$ and $\mathbf{y}=[y_{1},...,y_{n}]^{\top}$.

We now introduce two conditions on the network $f_{\bm{\theta}}$ that suffice for proving the global convergence of NGD to a minimizer which achieves zero training loss (and is therefore a global minimizer). To motivate these two conditions we make the following observations. First, the global minimizer is characterized by the condition that the gradient in the output space is zero for each case (i.e. $\nabla_{\mathbf{u}}\mathcal{L}(\bm{\theta})=\mathbf{0}$). Meanwhile, local minima are characterized by the condition that the gradient with respect to the parameters $\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})$ is zero. Thus, one way to avoid finding local minima that aren’t global is to ensure that the parameter gradient is zero if and only if the output space gradient (for each case) is zero. It’s not hard to see that this property holds as long as $\mathbf{G}$ remains non-singular throughout optimization (or equivalently that $\mathbf{J}$ always has full row rank). The following two conditions ensure that this happens, by first requiring that this property hold at initialization time, and second that $\mathbf{J}$ changes slowly enough that it remains true in a big enough neighborhood around $\bm{\theta}(0)$.

The Jacobian matrix $\mathbf{J}(0)$ at the initialization has full row rank, or equivalently, the Gram matrix $\mathbf{G}(0)=\mathbf{J}(0)\mathbf{J}(0)^{\top}$ is positive definite.

Condition 1 implies that $m\leq n$, which means the Fisher information matrix is singular and we have to use the generalized inverse except in the case where $m=n$.

This condition shares the same spirit with the Lipschtiz smoothness assumption in classical optimization theory. It implies (with small $C$) that the network is close to a linearized network (Lee et al., 2019) around the initialization and therefore natural gradient descent update is close to the gradient descent update in the output space. Along with Condition 1, we have the following theorem.

Let Condition 1 and 2 hold. Suppose we optimize with NGD using a step size $\eta\leq\frac{1-2C}{(1+C)^{2}}$. Then for $k=0,1,2,...$ we have

To be noted, $\left\|\mathbf{u}(k)-\mathbf{y}\right\|_{2}^{2}$ is the squared error loss up to a constant. Due to space constraints we only give a short sketch of the proof here. The full proof is given in Appendix B.

Proof Sketch. Our proof relies on the following insights. First, if the Jacobian matrix has full row rank, this guarantees linear convergence for infinitesimal step size. The linear convergence property restricts the parameters to be close to the initialization, which implies the Jacobian matrix is always full row rank throughout the training, and therefore natural gradient descent with infinitesimal step size converges to global minima. Furthermore, given the network is close to a linearized network (since the Jacobian matrix is stable with respect to small perturbations around the initialization), we are able to extend the proof to discrete time with a large step size.

In summary, we prove that NGD exhibits linear convergence to the global minimizer of the neural network training problem, under Conditions 1 and 2. We believe our arguments in this section are general (i.e., architecture-agnostic), and can serve as a recipe for proving global convergence of natural gradient descent in other settings.

We note that our analysis can be easily extended to more general loss function class. Here, we take the class of functions that are $\mu$-strongly convex with $L$-Lipschitz gradients as an example. Note that strongly convexity is a very mild assumption since we can always add $L_{2}$ regularization to make the convex loss strongly convex. Therefore, this function class includes regularized cross-entropy loss (which is typically used in classification) and squared error (for regression). For this type of loss, we need a strong version of Condition 2.

The key step of proving Theorem 2 is to show if $m$ is large enough, then natural gradient descent is approximately gradient descent in the output space. Thus the results can be easily derived according to standard bounds for convex optimization. Due to the page limit, we defer the proof to the Appendix C.

In Theorem 2, the convergence rate depends on the condition number $\kappa=\frac{L}{\mu}$, which can be removed if we take into the curvature information of the loss function. In other words, we expect that the bound has no dependency on $\kappa$ if we use the Fisher matrix rather than the classical Gauss-Newton (assuming Euclidean metric in the output space (Luk and Grosse, 2018)) in Theorem 2.

Optimizing Overparameterized Neural Networks

In Section 3, we analyzed the convergence properties of natural gradient descent, under the abstract Conditions 1 and 2. In this section, we make our analysis concrete by applying it to a specific type of overparameterized network (with a certain random initialization). We show that Conditions 1 and 2 hold with high probability. We therefore establish that NGD exhibits linear convergence to a global minimizer for such networks.

2 Problem Setup

Formally, we consider a neural network of the following form:

where $\mathbf{x}\in^{d}$ is the input, $\mathbf{w}=\left[\mathbf{w}_{1}^{\top},...,\mathbf{w}_{r}^{\top}\right]^{\top}\in^{md}$ is the weight matrix (formed into a vector) of the first layer, $a_{r}\in$ is the output weight of hidden unit $r$ and $\phi(\cdot)$ is the ReLU activation function (acting entry-wise for vector arguments). For $r\in[m]$, we initialize the weights of first layer $\mathbf{w}_{r}\sim\mathcal{N}(\mathbf{0},\nu^{2}\mathbf{I})$ and output weight $a_{r}\sim\mathbf{unif}\left[\{-1,+1\}\right]$.

Following Du et al. (2018b); Wu et al. (2019), we make the following assumption on the data.

For all i, $\|\mathbf{x}_{i}\|_{2}=1$ and $|y_{i}|=\mathcal{O}\left(1\right)$. For any $i\neq j$, $\mathbf{x}_{i}\nparallel\mathbf{x}_{j}$.

This very mild condition simply requires the inputs and outputs have standardized norms, and that different input vectors are distinguishable from each other. Datasets that do not satisfy this condition can be made to do so via simple pre-processing.

Following Du et al. (2018b); Oymak and Soltanolkotabi (2019); Wu et al. (2019), we only optimize the weights of the first layerWe fix the second layer just for simplicity. Based on the same analysis, one can also prove global convergence for jointly training both layers., i.e., $\bm{\theta}=\mathbf{w}$. Therefore, natural gradient descent can be simplified to

Though this is only a shallow fully connected neural network, the objective is still non-smooth and non-convex (Du et al., 2018b) due to the use of ReLU activation function. We further note that this two-layer network model has been useful in understanding the optimization and generalization of deep neural networks (Xie et al., 2016; Li and Liang, 2018; Du et al., 2018b; Arora et al., 2019b; Wu et al., 2019), and some results have been extended to multi-layer networks (Du et al., 2018a).

Following Du et al. (2018b); Wu et al. (2019), we define the limiting Gram matrix as follows:

The limiting Gram matrix $\mathbf{G}^{\infty}\in^{n\times n}$ is defined as follows. For $(i,j)$- entry, we have

3 Exact Natural Gradient Descent

In this subsection, we present our result for this setting. The main difficulty is to show that Conditions 1 and 2 hold. Here we state our main result.

Even though the objective is non-convex and non-smooth, natural gradient descent with a constant step size enjoys a linear convergence rate. For large enough $m$, we show that the learning rate can be chosen up to $1$, so NGD can provably converge within $\mathcal{O}\left(1\right)$ steps. Compared to analogous bounds for gradient descent (Du et al., 2018a; Oymak and Soltanolkotabi, 2019; Wu et al., 2019), we improve the maximum allowable learning rate from $\mathcal{O}(1/n)$ to $\mathcal{O}(1)$ and also get rid of the dependency on $\lambda_{0}$. Overall, NGD (Theorem 3) gives an $\mathcal{O}(\lambda_{0}/n)$ improvement over gradient descent.

Our strategy to prove this result will be to show that for the given choice of random initialization, Condition 1 and 2 hold with high probability. For proving Condition 1 hold, we used matrix concentration inequalities. For Condition 2, we show that $\|\mathbf{J}-\mathbf{J}(0)\|_{2}=\mathcal{O}\left(m^{-1/6}\right)$, which implies the Jacobian is stable for wide networks. For detailed proof, we refer the reader to the Appendix D.1.

4 Approximate Natural Gradient Descent with K-FAC

Exact natural gradient descent is quite expensive in terms of computation or memory. In training deep neural networks, K-FAC (Martens and Grosse, 2015) has been a powerful optimizer for leveraging curvature information while retaining tractable computation. The K-FAC update rule for the two-layer ReLU network is given by

where $\mathbf{X}\in^{n\times d}$ denotes the matrix formed from the $n$ input vectors (i.e. $\mathbf{X}=[\mathbf{x}_{1},...,\mathbf{x}_{n}]^{\top}$), and $\mathbf{S}=[\phi^{\prime}(\mathbf{X}\mathbf{w}_{1}),...,\phi^{\prime}(\mathbf{X}\mathbf{w}_{m})]\in^{n\times m}$ is the matrix of pre-activation derivatives. Under the same argument as the Gram matrix $\mathbf{G}^{\infty}$, we get that ${\mathbf{S}^{\infty}\mathbf{S}^{\infty}}^{\top}$ is strictly positive definite with smallest eigenvalue $\lambda_{\mathbf{S}}$ (see Appendix D.3 for detailed proof).

We show that for sufficiently wide networks, K-FAC does converge linearly to a global minimizer. We further show, with a particular transformation on the input data, K-FAC does match the optimization performance of exact natural gradient for two-layer ReLU networks. Here we state the main result.

The key step in proving Theorem 4 is to show

The convergence rate of K-FAC is captured by the condition number of the matrix $\mathbf{X}^{\top}\mathbf{X}$, as opposed to gradient descent (Du et al., 2018b; Oymak and Soltanolkotabi, 2019), for which the convergence rate is determined by the condition number of the Gram matrix $\mathbf{G}$.

The dependence of the convergence rate on $\kappa_{\mathbf{X}^{\top}\mathbf{X}}$ in Theorem 4 may seem paradoxical, as K-FAC is invariant to invertible linear transformations of the data (including those that would change $\kappa_{\mathbf{X}^{\top}\mathbf{X}}$). But we note that said transformations would also make the norms of the input vectors non-uniform, thus violating Assumption 1 in a way that isn’t repairable. Interestingly, there exists an invertible linear transformation which, if applied to the input vectors and followed by normalization, produces vectors that simultaneously satisfy Assumption 1 and the condition $\kappa_{\mathbf{X}^{\top}\mathbf{X}}=1$ (thus improving the bound in Theorem 4 substantially). See Appendix A for details. Notably, K-FAC is not invariant to such pre-processing, as the normalization step is a nonlinear operation.

To quantify the degree of overparameterization (which is a function of the network width $m$) required to achieve global convergence under our analysis, we must estimate $\lambda_{\mathbf{S}}$. To this end, we observe that $\mathbf{G}=\mathbf{X}\mathbf{X}^{\top}\odot\mathbf{S}\mathbf{S}^{\top}$, and then apply the following lemma:

[Schur (1911)] For two positive definite matrices $\mathbf{A}$ and $\mathbf{B}$, we have

The diagonal entries of $\mathbf{X}\mathbf{X}^{\top}$ are all $1$ since the inputs are normalized. Therefore, we have $\lambda_{0}\geq\lambda_{\mathbf{S}}$ according to Lemma 1, and hence K-FAC requires a slightly higher degree of overparameterization than exact NGD under our analysis.

As pointed out by Allen-Zhu et al. (2018), it is unclear if $1/\lambda_{0}$ is small or even polynomial. Here, we bound $\lambda$ using matrix concentration inequalities and harmonic analysis. To leverage harmonic analysis, we have to assume the data $\mathbf{x}_{i}$ are drawn i.i.d. from the unit sphereThis assumption is not too stringent since the inputs are already normalized. Moreover, we can relax the assumption of unit sphere input to separable input, which is used in Li and Liang (2018); Allen-Zhu et al. (2018); Zou et al. (2018). See Oymak and Soltanolkotabi (2019) (Theorem I.1) for more details..

Under this assumption on the training data, with probability $1-n\exp(-n^{\beta}/4)$,

Basically, Theorem 5 says that the Gram matrix $\mathbf{G}^{\infty}$ should have high chance of having large smallest eigenvalue if the training data are uniformly distributed. Intuitively, we would expect the smallest eigenvalue to be very small if all $\mathbf{x}_{i}$ are similar to each other. Therefore, some notion of diversity of the training inputs is needed. We conjecture that the smallest eigenvalue would still be large if the data are $\delta$-separable (i.e., $\left\|\mathbf{x}_{i}-\mathbf{x}_{j}\right\|_{2}\geq\delta$ for any pair $i,j\in[n]$), an assumption adopted by Li and Liang (2018); Allen-Zhu et al. (2018); Zou et al. (2018).

Generalization analysis

It is often speculated that NGD or other preconditioned gradient descent methods (e.g., Adam) perform worse than gradient descent in terms of generalization (Wilson et al., 2017). In this section, we show that NGD achieves the same generalization bounds which have been proved for GD, at least for two-layer ReLU networks.

It has been shown (Neyshabur et al., 2019) that the Redemacher complexity (Bartlett and Mendelson, 2002) for two-layer ReLU networks depends on $\left\|\mathbf{w}-\mathbf{w}(0)\right\|_{2}$. By the standard Rademacher complexity generalization bound, we have the following bound (see Appendix E.1 for proof):

which matches the bound for gradient descent in Arora et al. (2019b). For detailed proof, we refer the reader to the Appendix E.1.

Conclusion

We’ve analyzed for the first time the rate of convergence to a global optimum for (both exact and approximate) natural gradient descent on nonlinear neural networks. Particularly, we identified two conditions which guarantee the global convergence, i.e., the Jacobian matrix with respect to the parameters has full row rank and stable for perturbations around the initialization. Based on these insights, we improved the convergence rate of gradient descent by a factor of $\mathcal{O}(\lambda_{0}/n)$ on two-layer ReLU networks by using natural gradient descent. Beyond that, we also showed that the improved convergence rates don’t come at the expense of worse generalization.

Acknowledgements

We thank Jeffrey Z. HaoChen, Shengyang Sun and Mufan Li for helpful discussion.

References

Appendix A The Forster Transform

In a breakthrough paper in the area of communication complexity, Forster used the existence of a certain kind of dataset transformation as the key technical tool in the proof of his main result. The Theorem which establishes the existence of this transformation is paraphrased below.

Suppose $\mathbf{X}\in^{n\times d}$ is a matrix such that all subsets of size at most $d$ of its rows are linearly independent. Then there exists an invertible matrix $\mathbf{A}\in^{d\times d}$ such that if we post-multiply $\mathbf{X}$ by $\mathbf{A}$ (i.e. apply $\mathbf{A}$ to each row), and then normalize each row by its 2-norm, the resulting matrix $\mathbf{Z}\in^{n\times d}$ satisfies $\mathbf{Z}^{\top}\mathbf{Z}=\frac{n}{d}\mathbf{I}_{d}$.

Note that the technical condition about linear independence can be easily be made to hold for an arbitrary $\mathbf{X}$ by adding an infinitesimal random perturbation, assuming it doesn’t hold to begin with.

This result basically says that for any set of vectors, there is a linear transformation of said vectors which makes their normalized versions (given by the rows of $\mathbf{Z}$) satisfy $\mathbf{Z}^{\top}\mathbf{Z}=\frac{n}{d}\mathbf{I}_{d}$. So by combining this linear transformation with normalization we produce a set of vectors that simultaneously satisfy Assumption 1, while also satisfying $\kappa_{\mathbf{Z}^{\top}\mathbf{Z}}=1$.

Forster’s proof of Theorem 7 can be interpreted as defining a transformation function on $\mathbf{Z}$ (initialized at $\mathbf{X}$), and showing that it has a fixed point with the required properties. One can derive an algorithm from this by repeatedly applying the transformation to $\mathbf{Z}$, which consists of "whitening" followed by normalization, until $\mathbf{Z}^{\top}\mathbf{Z}$ is sufficiently close to $\frac{n}{d}\mathbf{I}_{d}$. The $\mathbf{A}$ matrix is then simply the product of the "whitening" transformation matrices, up to a scalar constant. While no explicit finite-time convergence guarantees are given for this algorithm by Forster , we have implemented it and verified that it does indeed converge at a reasonable rate. The algorithm is outlined below.

Appendix B Proof of Theorem 1

We prove the result in two steps: we first provide a convergence analysis for natural gradient flow, i.e., natural gradient descent with infinitesimal step size, and then take into account the error introduced by discretization and show global convergence for natural gradient descent.

To guarantee global convergence for natural gradient flow, we only need to show that the Gram matrix is positive definite throughout the training. Intuitively, for successfully finding global minima, the network must satisfy the following condition, i.e., the gradient with respect to the parameters $\nabla_{\bm{\theta}}\mathcal{L}(\bm{\theta})$ is zero only if the gradient in the output space $\nabla_{\mathbf{u}}\mathcal{L}(\bm{\theta})=\mathbf{0}$ is zero. It suffices to show that the Gram matrix is positive definite, or equivalently, the Jacobian matrix is full row rank.

By Condition 1 and Condition 2, we immediately obtain the following lemma that if the parameters stay close to the initialization, then the Gram matrix is positive definite throughout the training.

Accordingly, we can calculate the dynamics of the network predictions.

Since the Gram matrix $\mathbf{G}(t)$ is positive definite, its inverse does exist. Therefore, we have

By the chain rule, we get the dynamics of the loss in the following form:

By integrating eqn. (23), we find that $\left\|\mathbf{y}-\mathbf{u}(t)\right\|_{2}^{2}=\exp(-2t)\left\|\mathbf{y}-\mathbf{u}(0)\right\|_{2}^{2}$.

That completes the continuous time analysis, under the assumption that the parameters stay close to the initialization. The discrete case follows similarly, except that we need to account for the discretization error. Analogously to eqn. (21), we calculate the difference of predictions between two consecutive iterations.

where we have defined $\bm{\theta}(s)=s\bm{\theta}(k+1)+(1-s)\bm{\theta}(k)=\bm{\theta}(k)-s\eta\mathbf{J}(k)^{\top}\mathbf{G}(k)^{-1}(\mathbf{u}(k)-\mathbf{y})$.

Next we bound the norm of the second term (1) in the RHS of eqn. (24). Using Condition 2 and Lemma 2 we have that

In the first inequality, we used the fact (based on Condition 2) that

In the last inequality of eqn. (27), we use the assumption that $\eta\leq\frac{1-2C}{(1+C)^{2}}$.

So far, we have assumed the parameters fall within a certain radius around the initialization. We now justify this assumption.

We use the norm of each update to bound the distance of the parameters to the initialization.

At first glance, the proofs in Lemma 2 and 3 seem to be circular. Here, we prove that their assumptions continue to be jointly satisfied.

Appendix C Proof of Theorem 2

Here, we prove Theorem 2 by induction. Our inductive hypothesis is the following condition.

At the $k$-th iteration, we have $\|\mathbf{y}-\mathbf{u}(k+1)\|_{2}^{2}\leq\left(1-\frac{2\eta\mu L}{\mu+L}\right)\|\mathbf{y}-\mathbf{u}(k)\|_{2}^{2}$.

In analogy to eqn. (24), $\mathbf{P}(k)=\int_{s=0}^{1}\left(\mathbf{J}(k)-\mathbf{J}(\bm{\theta}(s))\right)\mathbf{J}(k)^{\top}\mathbf{G}(k)^{-1}ds$. Next, we introduce a well-known Lemma for $\mu$-strongly convex and $L$-Lipschitz gradient loss.

If the loss function is $\mu$-strongly convex with $L$-Lipschitz gradient, for any $\mathbf{u},\mathbf{y}\in^{n}$, the following inequality holds.

Now, we are ready to bound $\left\|\mathbf{y}-\mathbf{u}(k+1)\right\|_{2}^{2}$:

For the second inequality we used Lemma 5 and the fact that $\nabla_{\mathbf{u}}\mathcal{L}(\mathbf{y})=\mathbf{0}$. For the third inequality we used the result of Condition 3 that $\|\mathbf{P}(\mathbf{u}(k))\|_{2}\leq C$. For the last inequality we used the fact that the loss is $\mu$-strongly convex and $\eta<\frac{2}{\mu+L}\frac{1-(1+\kappa)C}{(1+C)^{2}}$.

For convex function $f$ with Lipschtiz gradient, we have the following co-coercivity property [Boyd and Vandenberghe, 2004]:

Then, for $\mu$-strongly convex loss, we can define $g(\mathbf{x})=f(\mathbf{x})-\frac{\mu}{2}\|\mathbf{x}\|_{2}^{2}$, which is a convex function with $L-\mu$ Lipschitz gradient. By co-coercivity of $g$, we get

After plugging $\nabla g(\mathbf{x})=\nabla f(\mathbf{x})-\mu\mathbf{x}$ and some manipulations, we have

Appendix D Proofs for Section 4

Our strategy to prove this result will be to show that for the given choice of random initialization, Conditions 1 and 2 hold with high probability.

We start with Condition 1, which requires that the inital Gram matrix $\mathbf{G}(0)$ is non-singular, or equivalently that $\mathbf{J}(0)$ has full row rank. Because $\mathbf{G}(0)$ is a sum of random matrices, where the expectation of each random matrix is $\frac{1}{m}\mathbf{G}^{\infty}$, we can bound its smallest eigenvalue using matrix concentration inequalities. Doing so gives us the following lemma.

Next we introduce another lemma, which we will use to show that Condition 2 holds with high probability.

With probability at least $1-\delta$, for all weight vectors $\mathbf{w}$ that satisfy $\left\|\mathbf{w}-\mathbf{w}(0)\right\|_{2}\leq R$, we have the following bound:

Notably, Lemma 7 says that as long as the weights are close to the random initialization, the corresponding Jacobian matrix also stays close to the Jacobian matrix of inital weights. Therefore, we might expect that if the distance to the initialization is sufficiently small, then Condition 2 would hold with high probability.

Thus by Markov’s inequality, we have probability at least $1-\delta$, $\|\mathbf{y}-\mathbf{u}(0)\|_{2}^{2}=\mathcal{O}\left(\frac{n}{\delta}\right)$. Thus the condition on $m$ can be written as $m=\Omega\left(\frac{n^{4}}{\nu^{2}\lambda_{0}^{4}\delta^{3}}\right)$. We note that the larger $m$ is, the smaller the constant $C$ in Condition 2 is ($C\sim\mathcal{O}\left(m^{-1/6}\right)$). We now finish the proof.

Notice that $\mathbf{G}(0)$ can be written as the sum of random symmetric matrices:

Letting the RHS of eqn. (40) be $\delta$, we have $m=\mathcal{O}\left(\frac{n}{\lambda_{0}}\log\frac{n}{\delta}\right)$. ∎

Let $[\mathbf{v}_{r}]_{k-}$ denote the $k$-th smallest entry of $\{\mathbf{v}_{1},...,\mathbf{v}_{m}\}$ after sorting its entries in terms of absolute value. We first state a intermediate lemma we prove later.

By taking $k=\frac{R^{2/3}m^{2/3}}{\nu^{2/3}\delta^{2/3}}$, we have $\left\|\mathbf{J}-\mathbf{J}(0)\right\|_{2}^{2}\leq\frac{2nR^{2/3}}{\nu^{2/3}\delta^{2/3}m^{1/3}}$. To complete the proof, all that remains is to prove $R\leq\sqrt{k}\left[\mathbf{w}_{r}(0)^{\top}\mathbf{x}_{i}\right]_{k-}$.

By applying Markov’s inequality, we obtain

Therefore, we have $\sqrt{k}\left[\mathbf{w}_{r}(0)^{\top}\mathbf{x}_{i}\right]_{k-}\geq\frac{k^{3/2}\nu\delta}{m}=R$. ∎

We prove the lemma by contradiction. First, we define the event

We can then bound the Jacobian perturbation.

D.2 Proof of Theorem 4

Based on the update in the weight space, we can get the update in output space accordingly.

The first equality we used properties of Khatri-Rao, Hadamard and Kronecker products while the second equality used the generalized inverse.

The key of bounding 2 is to analyze $(\mathbf{X}^{\top}\mathbf{X}^{\top})^{-1}\mathbf{X}^{\top}\star\mathbf{S}^{\top}(\mathbf{S}\mathbf{S}^{\top})^{-1}$. For convenience, we denote this term as 3. By the identity of $(\mathbf{A}\ast\mathbf{B})(\mathbf{A}^{\top}\star\mathbf{B}^{\top})=\mathbf{A}\mathbf{A}^{\top}\odot\mathbf{B}\mathbf{B}^{\top}$, we have

Next, we move on to show the weights of the network remain close to the initialization point.

Now to prove that $\mathbf{S}^{\infty}{\mathbf{S}^{\infty}}^{\top}$ is strictly positive definite, it is equivalent to show $\phi^{\prime}(\mathbf{x}_{1}),...,\phi^{\prime}(\mathbf{x}_{n})\in\mathcal{H}$ are linearly independent. Suppose that there are $\alpha_{1},...,\alpha_{n}\in$ such that

We now prove that $\alpha_{i}=0$ for all $i$.

We define $D_{i}=\left\{\mathbf{w}\in^{d}:\mathbf{w}^{\top}\mathbf{x}_{i}=0\right\}$. As shown by Du et al. [2018b], $D_{i}\not\subset\cup_{j\neq i}D_{j}$. For a fixed $i\in[n]$, we can choose $\mathbf{z}\in D_{i}\setminus\cup_{j\neq i}D_{j}$. We can pick a small enough radius $r_{0}>0$ such that $B(\mathbf{z},r)\cap D_{j}=\emptyset,\forall j\neq i,r\leq r_{0}$. Let $B(\mathbf{z},r)=B_{r}^{+}\cup B_{r}^{-}$, where $B_{r}^{+}=B(\mathbf{z},r)\cap D_{i}$.

For $j\neq i$, $\phi^{\prime}(\mathbf{x}_{j})$ is continuous in the neighborhood of $\mathbf{z}$, therefore for any $\epsilon>0$ there is a small enough $r$ such that

Let $\mu$ be Lebesgue measure on d, we have

Now recall that $\sum_{i}\alpha_{i}\phi^{\prime}(x_{i})\equiv 0$, we have

where $\delta_{ij}$ is the Kronecker delta. We complete the proof.

D.4 Proof of Theorem 5

It has been shown that the Gram matrix for infinite width networks has the following form [Xie et al., 2016, Arora et al., 2019b]:

We note that any function defined on the unit sphere has a spherical harmonic decomposition:

Due to the fact that $\mathbf{G}^{\infty}-\left[\mathbf{G}^{\infty}\right]^{n}$ is PSD, we can then bound the smallest eigenvalue of $\left[\mathbf{G}^{\infty}\right]^{n}$. Define a matrix $\mathbf{K}\in^{n\times n}$ whose rows are

Therefore, the matrix Chernoff bound gives

According to Xie et al. , $\gamma_{n}$ decays slower than $\mathcal{O}\left(1/n\right)$. This completes the proof.

Appendix E Proofs for Section 5

Similar to Neyshabur et al. , Arora et al. [2019b], we analyze the generalization error based on Rademacher Complexity [Bartlett and Mendelson, 2002]. Based on Rademacher complexity, we have the following generalization bound.

Therefore, to get the generalization bound, we only need to calculate the Rademacher complexity of a certain function class. Lemma 3 suggests that the learned function $f(\mathbf{w},\mathbf{a})$ from NGD is in a restricted class of neural networks whose weights are close to the initialization $\mathbf{w}(0)$. The following lemma bounds the Rademacher complexity of this function class.

For given $A,B>0$, with probability at least $1-\delta$ over the random initialization, the following function class

has empirical Rademacher complexity bounded as

With Lemma 9 and 10 at hand, we are only left to bound the distance of the weights to their initialization. Recall the update rule for $\mathbf{w}$:

The following lemma upper bounds the distance for each hidden unit.

If two conditions hold for $s=0,...,k$, then we have

To analyze the whole weight vector, we start with ideal case – infinite width network. In that case, both $\mathbf{J}$ and $\mathbf{G}$ are constant matrix throughout the training, and the function error decay exponentially. It is easy to show that the distance is given by

In the case of finite wide networks, $\mathbf{J}$ and $\mathbf{G}$ would change along the weights, but we can bound the changes and show the norms are small if the network are wide enough. Therefore, the distance is dominated by eqn. (72). From Lemma 3, we know $\|\mathbf{w}(k+1)-\mathbf{w}(0)\|_{2}=\mathcal{O}\left(\sqrt{\frac{n}{\lambda_{0}\delta}}\right)$. According to Lemma 7, it is easy to show that $\|\mathbf{J}(k)-\mathbf{J}(0)\|_{2}=\mathcal{O}\left(\frac{n^{2/3}}{\nu^{1/3}\lambda_{0}^{1/6}m^{1/6}\delta^{1/2}}\right)$ and $\|\mathbf{G}(k)-\mathbf{G}(0)\|_{2}=\mathcal{O}\left(\frac{n^{4/3}}{\nu^{2/3}\lambda_{0}^{1/3}m^{1/3}\delta}\right)$. With these bounds at hand, we are ready to bound $\|\mathbf{w}(\infty)-\mathbf{w}(0)\|_{2}$ for finite wide networks.

Under the same setting as Theorem 3, with probability at least $1-\delta$ over the random initialization, we have

Finally, we know that for any sample $\mathcal{S}$ drawn from data distribution $\mathcal{D}$, with probability at least $1-\delta/3$ over the random initialization, the following hold simutaneously:

This implies an upper bound on the training error:

The learned function $f(\mathbf{w},\mathbf{a})$ belongs to the restricted function class (68).

The function class $\mathcal{F}_{A,B}$ has Rademacher complexity bounded as

Also, with the probability at least $1-\delta/3$ over the sample $\mathcal{S}$, we have

Taking a union bound, we have know that with probability at least $1-\frac{2}{3}\delta$ over the sample $\mathcal{S}$ and the random initialization, we have

E.2 Technical Proofs for Generalization Analysis

To bound the norm of distance, we first decompose the total distance into the sum of each weight update.

We then analyze the term $\mathbf{u}(k)-\mathbf{y}$, which evolves as follows.

Under the same setting as Theorem 3, with probability at least $1-\delta$ over the random initialization, we have

where $\|\mathbf{\zeta}(k)\|_{2}=\mathcal{O}\left((1-\eta)^{k}\sqrt{\frac{n}{\delta}}\nu+k\left(1-\frac{\eta}{2}\right)^{k-1}\eta\frac{n^{7/6}}{\nu^{1/3}\lambda_{0}^{2/3}m^{1/6}\delta}\right)$

Plugging eqn. (85) into eqn. (84), we have

The RHS term in above equation is considered perturbation and we can upper bound their norm easily. By Lemma 13, we have

Plugging eqn. (87) into eqn. (84), we have

For $\mathbf{e}_{1}$, we have $\|\mathbf{e}_{1}\|_{2}=\mathcal{O}\left(\sqrt{\frac{n}{\lambda_{0}\delta}}\nu+\frac{n^{7/6}}{\nu^{1/3}\lambda_{0}^{7/6}m^{1/6}\delta}\right)$ by using the following inequality:

In eqn (90), we need first bound $\|\mathbf{G}(s)^{-1}-(\mathbf{G}^{\infty})^{-1}\|_{2}\|\mathbf{y}\|_{2}$. The following lemma bounds this norm by expanding the inverse with an infinite series.

Under the same setting as Theorem 3, with probability at least $1-\delta$ over the random initialization, we have

It is easy to show that both $\|\mathbf{J}\|_{2}$ and $\|\mathbf{y}\|_{2}$ are $\mathcal{O}\left(\sqrt{n}\right)$. Plugging them back into eqn. (90), we have

Similarly, we can bound $\mathbf{e}_{3}$ as follows,

We also bound the first term in eqn. (88):

Combining bounds (87), (90), (93) and (94), we have

where $\left\|\mathbf{\xi}(k)\right\|_{2}=\mathcal{O}\left(\eta\left(1-\frac{\eta}{2}\right)^{k}\frac{n^{7/6}}{\nu^{1/3}\lambda_{0}^{2/3}m^{1/6}\delta}\right)$. Applying eqn. (96) recursively, we get

For the second term, we have $\|(1-\eta)^{k}\mathbf{u}(0)\|_{2}=\mathcal{O}\left((1-\eta)^{k}\sqrt{\frac{n}{\delta}}\nu\right)$. For the last term, we have

Notice that $\mathbf{G}^{-1}=\alpha\sum_{s=0}^{\infty}\left(\mathbf{I}-\alpha\mathbf{G}\right)^{s}$, as long as $\alpha$ is small enough so that $\mathbf{I}-\alpha\mathbf{G}$ is positive definite. Therefore, instead of bounding $\|\mathbf{G}(s)^{-1}-(\mathbf{G}^{\infty})^{-1}\|_{2}$ directly, we can upper bound the following quantity:

Let $e(s)$ denote $\left\|\left(\mathbf{I}-\alpha\mathbf{G}(k)\right)^{s}-\left(\mathbf{I}-\alpha\mathbf{G}^{\infty}\right)^{s}\right\|_{2}$, we then have the following recursion:

Also we can easily bound the deviation of $\mathbf{G}(k)$ from $\mathbf{G}^{\infty}$ as follows,

Plugging eqn. (102) into eqn. (101), we have

With basic techniques of series theory and the fact $e(1)=\alpha E$, we have the following result:

Plugging eqn. (104) back into eqn. (99), we get

Appendix F Asymptotic analysis

As shown by Lee et al. , infinitely wide neural networks are linearized networks in the sense that the first-order Taylor expansion is accurate and the training dynamics of wide neural networks are well captured by linearized models in practice. Assume linearity (i.e., the Jacobian matrix $\mathbf{J}$ is constant over $\mathbf{w}$), we have the following result:

which means exact natural gradient descent can converge with one iteration if we take $\eta=1$, demonstrating the effectiveness of natural gradient descent.

Moreover, under the linearized network, we can conveniently analyze the trajectories of GD and NGD. Notably, we analyze the paths taken by GD and NGD in both output space and weight space. The dynamics with infinitesimal step size are summarized as follow.

By standard matrix differential equation theory, we have