On the linearity of large non-linear models: when and why the tangent kernel is constant

Introduction

As the width of certain non-linear neural networks increases, they become linear functions of their parameters. This remarkable property of large models was first identified in jacot2018neural (12) where it was stated in terms of the constancy of the (neural) tangent kernel during the training process. More precisely, consider a neural network or, generally, a machine learning model $f({\mathbf{w}};{\mathbf{x}})$, which takes ${\mathbf{x}}$ as input and has ${\mathbf{w}}$ as its (trainable) parameters. Its tangent kernel $K_{({\mathbf{x}},{\mathbf{z}})}({\mathbf{w}})$ is defined as follows:

The key finding of jacot2018neural (12) was the fact that for some wide neural networks the kernel $K_{({\mathbf{x}},{\mathbf{z}})}({\mathbf{w}})$ is a constant function of the weight ${\mathbf{w}}$ during training. While in the literature, including jacot2018neural (12), this phenomenon is described in terms of the (linear) training dynamics, it is important to note that the tangent kernel is associated to the model itself. As such, it does not depend on the optimization algorithm or the choice of a loss function, which are parts of the training process.

The goal of this work is to clarify a number of issues related to the constancy of the tangent kernel, to provide specific conditions when the kernel is constant, i.e., when non-linear models in the limit, as their width approach infinity, become linear, and also to explicate the regimes when they do not. One important conclusion of our analysis is that the “transition to linearity” phenomenon discussed in this work (equivalent to constancy of tangent kernel) cannot be explained by “lazy training” chizat2019lazy (6) associated to small change of parameters from the initialization point or model rescaling, which is widely held to be the reason for constancy of the tangent kernel, e.g., sun2019optimization (20, 2, 10) (see Section 1.1 for a detailed discussion). The transition to linearity is neither due to a choice of a scaling of the model, nor is a universal property of large models including infinitely wide neural networks. In particular, the models shown to transition to linearity in this paper become linear in a Euclidean ball of an arbitrary fixed radius, not just in a small vicinity of the initialization point, where higher order terms of the Taylor series can be ignored.

Our first observationWhile it is a known mathematical fact (see 868044 (9, 19)), we have not seen it in the neural network literature, possibly due to the discussion typically concerned with the dynamics of optimization. As a special case, note that while $\nabla f\equiv const$ clearly implies that $f$ is linear, it is not a priori obvious that the weaker condition $\|\nabla f\|\equiv const$ is also sufficient. is that a function $f({\mathbf{w}},{\mathbf{x}})$ has a constant tangent kernel if and only if it is linear in ${\mathbf{w}}$, that is

for some “feature map” $\phi$ and function $f_{0}$. Thus the constancy of the tangent kernel is directly linked to the linearity of the underlying model.

So what is the underlying reason that some large models transition to linearity as a function of the parameters and when do we expect it to be the case? As known from the mathematical analysis, the deviation from the linearity is controlled by the second derivative, which is represented, for a multivariate function $f$, by the Hessian matrix $H$. If its spectral norm $\|H\|$ is small compared to the gradient $\nabla_{{\mathbf{w}}}f$ in a ball of a certain radius, the function $f$ will be close to linear and will have near-constant tangent kernel in that ball. Crucially, the spectral norm $\|H\|$ depends not just on the magnitude of its entries, but also on the structure of the matrix $H$. This simple idea underlies the analysis in this paper. Note that throughout this paper we consider the Hessian of the model $f$, not of any related loss function.

Constant tangent kernel for neural networks with linear output layer. In what follows we analyze the class of neural networks with linear output layer, which includes networks that have been found to have constant tangent kernel in jacot2018neural (12, 16, 8) and other works. We show that while the gradient norm $\|\nabla_{{\mathbf{w}}}f\|$ is (omitting log factors) of the order $\Theta(1)$ w.r.t. the network width $m$, the spectral norm of the Hessian matrix $\|H\|$ scales with $m$ as ${1/\sqrt{m}}$. In the infinite width limit, this implies a vanishing Hessian and hence transition to linearity of the model in a ball of an arbitrary fixed radius. A consequence of this analysis is the constancy of the tangent kernel, providing a different perspective on the results in jacot2018neural (12) and the follow-up works.

We proceed to expose the underlying reason why the Hessian matrix scales differently from the gradient and delimit the regimes where this phenomenon exists. As we show, the scaling of the Hessian spectral norm is controlled by both the $\infty$-norms of the vectors $\partial f/\partial\alpha^{(l)},l\in[L]$, where $\alpha^{(l)}$ is the (vector) value of the $l$-th hidden layer, and the norms of layer-wise derivatives (specifically, the $(2,1,1)$-norm of the corresponding order $3$ tensors). On the other hand, the scaling of the gradient and the tangent kernel is controlled by the $2$-norms (i.e., Euclidean norms) of $\partial f/\partial\alpha^{(l)}$. As the network width $m$ (i.e., minimal width of hidden layers) is sufficiently large, the discrepancy between the the $\infty$-norm and $2$-norm increases, while the $(2,1,1)$-norms remain of the same order. Hence we obtain the discrepancy between the scaling behaviors of the Hessian and gradient.

Non-constancy of tangent kernels. We proceed to demonstrate, both theoretically (Section 4) and experimentally (Section 6), that the constancy of tangent kernel is not a general property of large models, including wide networks, even in the “lazy” training regime. In particular, if the output layer of a network is nonlinear, e.g., if there is a non-linear activation on the output, the Hessian norm does not tend to zero as $m\to\infty$, and constancy of tangent kernel will not hold in any fixed neighborhood and along the optimization path, although each individual parameter may undergo only a small change. This demonstrates that the constancy of the tangent kernel relies on specific structural properties of the models. Similarly, we show that inserting a narrow “bottleneck” layer, even if it is linear, will generally result in the loss of near-linearity, as the Hessian norm becomes large compared to the gradient $\nabla_{{\mathbf{w}}}f$ of the model.

Importantly, as we discuss in Section 5, non-constancy of the tangent kernel does not preclude efficient optimization. We construct examples of wide networks which can be provably optimized by gradient descent, yet with tangent kernel provably far from constant along the optimization path and with Hessian norm $\Omega(1)$, same as the gradient.

We proceed to make a number of remarks in the context of some recent work on the subject.

Is the weight change from the initialization to convergence small? In the recent literature (e.g.,sun2019optimization (20, 2, 10)) it is sometimes asserted that the constancy of tangent kernel is explained by small change of weight vector during training, a property related to “lazy training” introduced in chizat2019lazy (6). It is important to point out that the notion “small” depends crucially on the measurement. Indeed, as we discuss below, when measured correctly in relation to the tangent kernel, the change from initialization is not small.

Let ${\mathbf{w}}_{0}$ and ${\mathbf{w}}^{*}$ be the weight vectors at initialization and at convergence respectively. For example, consider a one hidden layer network of width $m$. Each component of the weight vector is updated by $O(1/\sqrt{m})$ under gradient descent, as shown in jacot2018neural (12), and hence for wide networks $\|{\mathbf{w}}^{*}-{\mathbf{w}}_{0}\|_{\infty}=O(1/\sqrt{m})$, a quantity that vanishes with the increasing width. In contrast, the change of the Euclidean norm is not small in training, $\|{\mathbf{w}}^{*}-{\mathbf{w}}_{0}\|^{2}=\sum_{i=1}^{m}(w^{*}_{i}-w_{0,i})^{2}=\Theta(1)$. Thus convergence happens within a Euclidean ball with radius independent of the network width.

In fact, the Euclidean norm of the change of the weight vector cannot be small for Lipschitz continuous models, even in the limit of infinite parameters. This is because

where is $y$ is the label at ${\mathbf{x}}$. Note that the difference $|f({\mathbf{w}}_{0};{\mathbf{x}})-y|$, between the initial prediction $f({\mathbf{w}}_{0};{\mathbf{x}})$ and the ground truth label $y$, is of the same order as $\|\nabla_{{\mathbf{w}}}f\|$. Thus, we see that $\|{\mathbf{w}}^{*}-{\mathbf{w}}_{0}\|=\Omega(1)$, no matter how many parameters the model $f$ has.

We note that the (approximate) linearity of a model in a certain region (and hence the constancy of the tangent kernel) is predicated on the second-order term of the Taylor expansion $({\mathbf{w}}-{\mathbf{w}}_{0})^{T}H({\mathbf{w}}-{\mathbf{w}}_{0})$. That term depends on the Euclidean distance from the initialization $\|{\mathbf{w}}-{\mathbf{w}}_{0}\|$ (and the spectral norm of the Hessian), instead of the $\infty$-norm $\|{\mathbf{w}}-{\mathbf{w}}_{0}\|_{\infty}$. Since, as we discussed above, these norms are different by a factor of $\sqrt{m}$, an argument based on small change of individual parameters from initialization cannot explain the remarkable phenomenon of constant tangent kernel.

In contrast to these interpretations, we show that certain large networks have near constant tangent kernel in a ball of fixed radius due to the vanishing Hessian norm, as their widths approach infinity. Indeed, that is the case for networks analyzed in the NTK literature jacot2018neural (12, 16, 8, 7).

Can the transition to linearity be explained by model rescaling? The work chizat2019lazy (6) introduced the term “lazy training” and proposed a mechanism for the constancy of the tangent kernel based on rescaling the model. While, as shown in chizat2019lazy (6), model rescaling can lead to lazy training, as we discuss below, it does not explain the phenomenon of constant tangent kernel in the setting of the original paper jacot2018neural (12) and consequent works.

Specifically, chizat2019lazy (6) provides the following criterion for the near constancy of the tangent kernel (using their notation):

Here $y$ is the ground truth label, $D^{2}f({\mathbf{w}}_{0})$ is the Hessian of the model $f$ at initialization and $\|Df({\mathbf{w}}_{0})\|^{2}$ is the norm of the gradient, i.e., a diagonal entry of the tangent kernel.

Assuming that $f({\mathbf{w}}_{0})=0$ and choosing a large $\alpha$, forces $\kappa_{\alpha f}\ll 1$, by rescaling the factor $\mathcal{A}$ to be small, while keeping $\mathcal{B}$ unchanged.

While rescaling the model, together with the important assumption of $f({\mathbf{w}}_{0})=0$, leads to a lazy training regime, we point out that it is not the same regime as observed in the original work jacot2018neural (12) and followup papers such as lee2019wide (16, 8) and also different from practical neural network training, since we usually have $\mathcal{A}=\|f({\mathbf{w}}_{0})-y\|=\Theta(1)$ in these settings. Specifically:

The assumption of $f({\mathbf{w}}_{0})=0$ is necessary for the rescaled models in chizat2019lazy (6) to have $\mathcal{A}\ll 1$. Yet, the networks, such as those analyzed in jacot2018neural (12), are initialized so that $f({\mathbf{w}}_{0})=\Theta(1)$.

From Eq.(4), we see that rescaling the model $f$ by $\alpha$ is equivalent to rescaling the ground truth label $y$ by $1/\alpha$ without changing the model (this can also be seen from the loss function, cf. Eq.(2) of chizat2019lazy (6)). When $\alpha$ is large, the rescaled label $y/\alpha$ is close to zero. However, no such rescaling happens in practice or in works, such as jacot2018neural (12, 16, 8). The training dynamics of the model with the label $y/\alpha$ does not generally match the dynamics of the original problem with the label $y$ and will result in a different solution.

Since $\mathcal{A}=\Theta(1)$, in the NTK setting and many practical settings, to satisfy the criterion in Eq.(3), the model needs to have $\mathcal{B}={\|D^{2}f({\mathbf{w}}_{0})\|}/{\|Df({\mathbf{w}}_{0})\|^{2}}\ll 1$. In fact, we note that the analysis of 2-layer networks in chizat2019lazy (6) uses a different argument, not based on model rescaling. Indeed, as we show in this work, $\mathcal{B}$ is small for a broad class of wide neural networks with linear output layer, due to a vanishing norm of the Hessian as the width of the network increases.

In summary, the rescaled models satisfy the criterion, $\kappa\ll 1$, by scaling the factor $\mathcal{A}$ to be small, while the neural networks, such as the ones considered in the original work jacot2018neural (12), satisfy this criterion by having $\mathcal{B}\ll 1$, while $\mathcal{A}=\Theta(1)$.

Is near-linearity necessary for optimization? In this work we concentrate on understanding the phenomenon of constant tangent kernel, when large non-linear systems transition to linearity with increasing number of parameters. The linearity implies convergence of gradient descent assuming that the tangent kernel is non-degenerate at initialization. However, it is important to emphasize that the linearity or near-linearity is not a necessary condition for convergence. Instead, convergence is implied by uniform conditioning of the tangent kernel in a neighborhood of a certain radius, while the linearity is controlled by the norm of the Hessian. These are conceptually and practically different phenomena as we show on an example of a wide shallow network with a non-linear output layer in Section 5. See also our paper liu2020toward (17) for an in-depth discussion of optimization.

Notation and Basic Results on Tangent Kernel and Hessian

We use $\nabla_{{\mathbf{w}}}f$ to represent the derivative of $f({\mathbf{w}};{\mathbf{x}})$ with respect to ${\mathbf{w}}$. For (vector-valued) functions, we use the following definition of its Lipschitz continuity:

For an order $3$ tensor, we define its $(2,2,1)$-norm:

We will later need the following proposition which is essentially a special case of the the Holder inequality.

Consider a matrix $A$ with components $A_{ij}=\sum_{k}T_{ijk}v_{k}$, where $T_{ijk}$ is a component of the order $3$ tensor ${\mathbf{T}}$ and $v_{k}$ is a component of vector ${\mathbf{v}}$. Then the spectral norm of $A$ satisfies

Note that spectral norm is defined as $\|A\|=\sup_{\|{\mathbf{x}}\|=\|{\mathbf{z}}\|=1}{\mathbf{x}}^{T}A{\mathbf{z}}$. Then

2 Tangent kernel and the Hessian

As discovered in jacot2018neural (12) and analyzed in the consequent works lee2019wide (16, 8) the tangent kernel is constant for certain infinitely wide networks during training by gradient descent methods. First, we observe that the constancy of the tangent kernel is equivalent to the linearity of the model. While the mathematical result is not new (see 868044 (9, 19)), we have not seen this stated in the machine learning literature (the proof can be found in Appendix C).

The tangent kernel of a differentiable function $f({\mathbf{w}};{\mathbf{x}})$ is constant if and only if $f({\mathbf{w}};{\mathbf{x}})$ is linear in ${\mathbf{w}}$.

Of course for a model to be linear it is necessary and sufficient for the Hessian to vanish. The following proposition extends this result by showing that small Hessian norm is a sufficient condition for near-constant tangent kernel. The proof can be found in Appendix D.

As we shall see in Section 3, all neural networks that are proven in jacot2018neural (12, 8, 7) to have (near) constant tangent kernel during training, have small (zero, in the limit of $m\to\infty$) spectral norms of the corresponding Hessian matrices.

Transition to linearity: non-linear neural networks with linear output layer

In this section, we analyze the class of neural networks with linear output layer, i.e., there is no non-linear activation on the final output. We show that the spectral norm of the Hessian matrix becomes small, when the width of each hidden layer increases. In the limit of infinite width, these spectral norms vanish and the models become linear, with constant tangent kernels. We point out that the neural networks that are already shown to have constant tangent kernels in jacot2018neural (12, 16, 8) fall in this category.

As a warm-up for the more complex setting of deep networks, we start by considering the simple case of a shallow fully-connected neural network with a fixed output layer, defined as follows:

This definition of a shallow neural network (i.e., with the presence of a factor $1/\sqrt{m}$ and $v_{i}$ and $w_{i}$ of order $O(1)$) is consistent with the NTK parameterization used to show constancy of tangent kernel in jacot2018neural (12, 16).

Hessian matrix. We observe that the Hessian matrix $H$ of the neural network $f$ is sparse, specifically, diagonal:

Consequently, if the input $x$ is bounded, say $|x|\leq C$, the spectral norm of the Hessian $H$ is

In the limit of $m\to\infty$, the spectral norm $\|H\|$ converges to 0.

Tangent kernel and gradient. On the other hand, the magnitude of the norm of the tangent kernel of $f$ is of order $\Theta(1)$ in terms of $m$. Specifically, for each diagonal entry we have

Therefore, from Eq. (9) and Eq. (10) we observe that the tangent kernel scales as $\Theta(1)$ while the norm of the Hessian scales as $O(1/\sqrt{m})$ with the size of the neural network $f$. Furthermore, as $m\to\infty$, the norm of the Hessian converges to zero and, by Proposition 2.3, the tangent kernel becomes constant.

So why should there be a discrepancy between the scaling of the Hessian spectral norm and the norm of the gradient? This is not a trivial question. There is no intrinsic reason why second and first order derivatives should scale differently with the size of an arbitrary model. In the rest of this subsection we analyze the source of that phenomenon in wide neural networks, connecting it to disparity of different norms in high dimension.

Specifically, we show that the Hessian spectral norm is controlled by ${\infty}$-norm of the vector $\|{\partial f}/{\partial\alpha}\|_{\infty}$. In contrast, the tangent kernel and the norm of the gradient are controlled by its Euclidean norm $\|{\partial f}/{\partial\alpha}\|$. The disparity between these norms is the underlying reason for the transition to linearity in the limit of infinite width.

$\bullet$ Hessian is controlled by $\|{\partial f}/{\partial\alpha}\|_{\infty}$. Given a model $f$ in Eq. (8), its Hessian matrix $H(f)$ is defined as

where $\frac{\partial^{2}\alpha_{i}}{\partial{\mathbf{w}}^{2}}$ are the components of the order 3 tensor of partial derivatives $\frac{\partial^{2}\alpha}{\partial{\mathbf{w}}^{2}}$. When there is no ambiguity, we suppress the argument and denote the Hessian matrix by $H$. By Proposition 2.1 (essentially the Holder’s inequality: $|{\mathbf{a}}^{T}{\mathbf{b}}|\leq\|{\mathbf{a}}\|_{1}\|{\mathbf{b}}\|_{\infty}$), we have

For this $1$-hidden layer network, the tensor $\frac{\partial^{2}\alpha}{\partial{\mathbf{w}}^{2}}$ is given by

Thus, we conclude that the Hessian spectral norm $\|H\|=O\left(\left\|{\partial f}/{\partial\alpha}\right\|_{\infty}\right)$.

$\bullet$ Tangent kernel and the gradient are controlled by $\|{\partial f}/{\partial\alpha}\|$. Note that the norm of the tangent kernel is lower bounded by the average of diagonal entries: $\|K\|\geq\frac{1}{n}\sum_{i=1}^{n}K_{(x_{i},x_{i})}$, where $n$ is the size of the dataset. Consider an arbitrary diagonal entry $K_{(x,x)}$ of the tangent kernel matrix.

Note that, $\frac{\partial\alpha}{\partial{\mathbf{w}}}$ is a diagonal matrix with $\frac{\partial\alpha_{i}}{\partial w_{i}}=\sigma^{\prime}(w_{i}x)x$. By the Lipschitz continuity of $\sigma(\cdot)$, $\left\|\frac{\partial\alpha}{\partial{\mathbf{w}}}\right\|$ is finite. Therefore, the tangent kernel is of the same order as the $2$-norm $\left\|{\partial f}/{\partial\alpha}\right\|$.

$\bullet$ The discrepancy between the norms. For the network in Eq. (8) we have $\frac{\partial f}{\partial\alpha}=\frac{1}{\sqrt{m}}{\mathbf{v}}$. Hence,

The transition to linearity stems from this observation and the fact discussed above that the Hessian norm scales as $\left\|\frac{\partial f}{\partial\alpha}\right\|_{\infty}$, while the tangent kernel is of the same order as $\left\|\frac{\partial f}{\partial\alpha}\right\|$.

In what follows, we show that this is a general principle applicable to wide neural networks. We start by analyzing two hidden layer neural networks, which are mathematically similar to the general case, but much less complex in terms of the notation.

2 Two hidden layer neural networks

Now, we demonstrate that analogous results hold for $2$-hidden layer neural networks. Consider the $2$-hidden layer neural network:

We denote the output of the first hidden layer by $\alpha^{(1)}(W_{1};{\mathbf{x}})=\sigma(W_{1}{\mathbf{x}})$ and the output of the second hidden layer by $\alpha^{(2)}(W_{1},W_{2};{\mathbf{x}})=\sigma\left(\frac{1}{\sqrt{m}}W_{2}\sigma(W_{1}{\mathbf{x}})\right)$.

Hessian is controlled by $\|{\partial f}/{\partial\alpha}\|_{\infty}$. Similarly to Eq.(13), we can bound the Hessian spectral norm by ${\infty}$-norms of $\partial f/\partial\alpha^{(1)}$ and $\partial f/\partial\alpha^{(2)}$.

Here $\frac{\partial}{\partial W_{l}}$ denotes partial derivatives w.r.t. each element of $W_{l}$, i.e. after flattening the matrix $W_{l}$ .

As this Proposition is a special case of Theorem 3.1, we omit the proof.

When $W_{1},W_{2}$ are initialized as random Gaussians, every term in Eq. (17), except for $\left\|{\partial f}/{\partial\alpha^{(1)}}\right\|_{\infty}$ and $\left\|{\partial f}/{\partial\alpha^{(2)}}\right\|_{\infty}$, is of order $O(1)$, with high probability within a ball of a finite radius (see the discussion in Subsection 3.3 for details).

Hence, just like the one hidden layer case, the magnitude of Hessian spectral norm is controlled by these ${\infty}$-norms:

Tangent kernel and the gradient are controlled by $\|{\partial f}/{\partial\alpha}\|$. A diagonal entry of the kernel matrix can be decomposed into

with each additive term being related to each layer. As the matrix $\partial\alpha^{(l)}/\partial W_{l}$ and the vector $\partial f/\partial\alpha^{(l)}$ are independent from each other and random at initialization, we expect $\left\|\frac{\partial\alpha^{(l)}}{\partial W_{l}}\frac{\partial f}{\partial\alpha^{(l)}}\right\|^{2}$ to be of the same order as $\left\|\frac{\partial\alpha^{(l)}}{\partial W_{l}}\right\|^{2}\left\|\frac{\partial f}{\partial\alpha^{(l)}}\right\|^{2}$, for $l=1,2$.

3 Multilayer neural networks

Now, we extend the analysis to general deep neural networks.

First, we show that, in parallel to one and two hidden layer networks, the Hessian spectral norm and the tangent kernel of a multilayer neural network are controlled by ${\infty}$-norms and $2$-norms of the vectors $\partial f/\partial\alpha^{(l)}$, respectively. Then we show that the magnitudes of the two types of vector norms scales differently with respect to the network width.

We consider a general form of a deep neural network $f$ with a linear output layer:

Initialization and parameterization. In this paper, we consider the NTK initialization/ parameterization jacot2018neural (12), under which the constancy of the tangent kernel had been initially observed. Specifically, the parameters, (weights), ${\mathbf{W}}:=\{{\mathbf{w}}^{(1)},{\mathbf{w}}^{(2)},\cdots,{\mathbf{w}}^{(L)},{\mathbf{w}}^{(L+1)}:={\mathbf{v}}\}$ are drawn i.i.d. from a standard Gaussian, i.e., $w_{i}^{(l)}\sim\mathcal{N}(0,1)$, at initialization, denoted as ${\mathbf{W}}_{0}$. The factor $1/\sqrt{m}$ in the output layer is required by the NTK parameterization in order that the output $f$ is of order $\Theta(1)$. Different parameterizations (e.g., LeCun initialization: $w_{i}^{(l)}\sim\mathcal{N}(0,1/m)$) rescale the tangent kernel and the Hessian by the same factor, and thus do not change our conclusions (see Appendix A).

To simplify the notation, we start by defining the following useful quantities:

It is important to note that the quantity $\mathcal{Q}_{\infty}(f)$ is simply the maximum of the $\infty$-norms $\left\|{\partial{f}}/{\partial\alpha^{(l)}}\right\|_{\infty}$, and that $\mathcal{Q}_{L}(f)$ and $\mathcal{Q}_{2,2,1}(f)$ are independent of the vectors $\partial f/\partial\alpha^{(l)},l\in[L]$.

The Hessian spectral norm is bounded by these quantities via the following theorem (see Appendix E for the proof).

Consider a $L$-layer neural network in the form of Eq.(3.3). For any ${\mathbf{W}}$ in the parameter space, the following inequality holds:

where $C_{1}=L(L^{2}\mathsf{L}_{\phi}^{2L}+L\mathsf{L}_{\phi}^{L}+1)$ and $C_{2}=L\mathsf{L}_{\phi}^{L}$.

The factor $1/\sqrt{m}$ in the second term comes from the definition of the output layer in Eq. (3.3) and is useful to make sure the model output at initialization is of the same order as the ground truth labels.

Tangent kernel and $2$-norms. A diagonal entry of the kernel matrix can be decomposed into

with each additive term being related to each layer. As before, we expect each term $\left\|\frac{\partial\alpha^{(l)}}{\partial{\mathbf{w}}^{(l)}}\frac{\partial f}{\partial\alpha^{(l)}}\right\|^{2}$ has the same order as $\left\|\frac{\partial\alpha^{(l)}}{\partial{\mathbf{w}}^{(l)}}\right\|^{2}\left\|\frac{\partial f}{\partial\alpha^{(l)}}\right\|^{2}$.

3.2 Small Hessian spectral norm and constant tangent kernel

To simplify our analysis, we make the following assumption.

Assumptions. We assume the hidden layer width $m_{l}=m$ for all $l\in[L]$, the number of parameters in each layer $p_{l}\geq m$, and the output is a scalar.The assumption $m_{l}=m$ is to simplify the analysis, as we discuss below we only need $m_{l}\geq m$. We assume that (vector-valued) layer functions $\phi_{l}({\mathbf{w}};\alpha),l\in[L],$ are $\mathsf{L}_{\phi}$-Lipschitz continuous and twice differentiable with respect to input $\alpha$ and parameters ${\mathbf{w}}$.

A fully connected neural network has the form as in Eq.(3.3), with each layer function specified by

where $\sigma(\cdot)$ is a $\mathsf{L}_{\sigma}$-Lipschitz continuous, $\beta_{\sigma}$-smooth activation function, such as $sigmoid$ and $tanh$. The layer parameters $W^{(l)}$ are reshaped into an $m\times m$ matrix. The Euclidean norm of ${\mathbf{W}}$ becomes: $\|{\mathbf{W}}\|=(\sum_{l=1}^{L}\|W^{(l)}\|_{F}^{2})^{1/2}$.

With high probability over the Gaussian random initialization, we have the following lemma to bound the quantities $\mathcal{Q}_{\infty}(f)$, $\mathcal{Q}_{2,2,1}(f)$ and $\mathcal{Q}_{L}(f)$ in a neighborhood of ${\mathbf{W}}_{0}$:

Consider a fully connected neural network $f({\mathbf{W}};{\mathbf{x}})$ with linear output layer and Gaussian random initialization ${\mathbf{W}}_{0}$. Given any fixed $R>0$, at any point ${\mathbf{W}}\in B({\mathbf{W}}_{0},R):=\{{\mathbf{W}}:\|{\mathbf{W}}-{\mathbf{W}}_{0}\|\leq R\}$, with high probability over the initialization, the quantity

See the proof of the lemma in Appendix F. Applying this lemma to Theorem 3.1, we immediately obtain the following theorem:

Consider a fully connected neural network $f({\mathbf{W}};{\mathbf{x}})$ with linear output layer and Gaussian random initialization ${\mathbf{W}}_{0}$. Given any fixed $R>0$, and any ${\mathbf{W}}\in B({\mathbf{W}}_{0},R):=\{{\mathbf{W}}:\|{\mathbf{W}}-{\mathbf{W}}_{0}\|\leq R\}$, with high probability over the initialization, the Hessian spectral norm satisfies the following:

We note that the above theorem also applies to more general networks that have different hidden layer widths, as long as the width of each layer is larger than $m$. See Theorem 3.3below.

In the limit of $m\to\infty$, the spectral norm of the Hessian $\|H({\mathbf{W}})\|$ converges to , for all ${\mathbf{W}}\in B({\mathbf{W}}_{0},R)$. By Proposition 2.3, this immediately implies constancy of tangent kernel and linearity of the model, in the ball $B({\mathbf{W}}_{0},R)$.

On the other hand, the tangent kernel is of order $\Theta(1)$ (see for example du2018gradientdeep (7), where the smallest eigenvalue of the tangent kernel is lower bounded by a width-independent constant). Intuitively, the order of tangent kernel stems from the fact that the $2$-norms $\|\partial f/\partial\alpha^{(l)}\|$ are of order $\Theta(1)$.

By the optimization theory built in our work liu2020toward (17), a finite radius $R$ is enough to include the gradient descent solution, for the square loss. Hence, for very wide networks, the tangent kernel is constant during gradient descent training.

3.3 Neural networks with hidden layers of different width and general architectures

Our analysis above is applicable to other common neural architectures including Convolutional Neural Networks (CNN) and ResNets, as well as networks with a mixed architectural types. Below we briefly highlight the main differences from the fully connected case. Precise statements can be found in Appendix G.

The key observation is that a convolutional layer is “fully connected” in the channel dimension. In contrast, the convolutional operation, which is sparse, is only within the spatial dimensions. Hence, we can apply our analysis to the channel dimension with only minor modifications. As the spatial dimension sizes are independent of the network width, the convolutional operation only contributes constant factors to our analysis. Therefore, our norm analysis extends to the CNN setting.

Architecture with mixed layer types. Neural networks used in practice are often a mixture of different layer types, e.g., a series of convolutional layers followed by fully connected layers. Since our analysis relies on layer-wise quantities, our results extend to such networks.

We have the following general theorem which summarizes our theoretical results.

Consider a general neural network $f({\mathbf{W}};x)$ of the form Eq.(3.3), which can be a fully connected network, CNN, ResNet or a mixture of these types. Let $m$ be the minimum of the hidden layer widths, i.e., $m=\min_{l\in[L]}m_{l}$. Given any fixed $R>0$, and any ${\mathbf{W}}\in B({\mathbf{W}}_{0},R):=\{{\mathbf{W}}:\|{\mathbf{W}}-{\mathbf{W}}_{0}\|\leq R\}$, with high probability over the initialization, the Hessian spectral norm satisfies the following:

Constant tangent kernel is not a general property of wide networks

In this section, we show that a class of infinitely wide neural networks with non-linear output, do not generally have constant tangent kernels. It also demonstrates that a linear output layer is a necessary condition for transition to linearity.

We note that the term $B$ vanishes as $m\to\infty$ due to the constancy of the tangent kernel of $f$. However the term $A$ is generally of the order $\Theta(1)$, when $\phi$ is non-linearIf $\phi$ is linear, the term $A$ is identically zero.. To see that consider any solution ${\mathbf{w}}^{*}$ such that $f({\mathbf{w}}^{*};{\mathbf{x}})=y$ (which exists for over-parameterized networks). Since $f({\mathbf{w}}_{0};{\mathbf{x}})$ is generally not equal to $y$, we obtain the result.

where $H$ is the Hessian matrix of model $f$. Hence, the spectral norm satisfies

This makes the quantity $\mathcal{Q}_{2,2,1}(f)$ to be the order of $O(m)$. Then, Theorem 3.1 indicates that the Hessian spectral norm is no longer arbitrarily small, suggesting a non-constant tangent kernel during training.

Indeed, as we prove below, the Hessian spectral norm is lower bounded by a positive constant, which in turn implies that the linearity does not hold for this kind of neural networks.

Specifically, consider a bottleneck network with of the following form:

This network has three hidden layers, where the first and third hidden layer have an arbitrarily large width $m$, and the second hidden layer, as the bottleneck layer, has a width $m_{b}=1$. Each individual parameter is initialized by the standard norm distribution. For simplicity of the analysis, the activation function is identity for the bottleneck layer is identity, and is quadratic for the first and third layers, i.e. $\sigma(z)=\frac{1}{2}z^{2}$.

The following theorem gives a lower bound for the Hessian spectral norm $\|H\|$ in a ball around the initialization ${\mathbf{W}}_{0}$.

Consider the bottleneck network $f({\mathbf{W}};x)$ defined in Eq. (31). Given an arbitrary radius $R>0$, for any ${\mathbf{W}}\in B({\mathbf{W}}_{0},R)$ and any $\delta\in(0,1)$, the Hessian matrix $H({\mathbf{W}})$ of the model satisfies

for some constant $C_{1}>0$, with probability at least $1-2\delta-e^{-m/16}$.

In particularly, in the limit of $m\to\infty$,

See the proof in Appendix H. With this lower bounded Hessian, Proposition 2.2 directly implies that the linearity of the model does not hold for this network. As Eq. (2) shows that $\|{\mathbf{W}}^{*}-{\mathbf{W}}_{0}\|=\Omega(1)$, our analysis implies the model is not linear, hence tangent kernel is not constant, along the optimization path. In Section 6, we empirically verify this finding.

In table 1, we summarize the key findings of this section and compare them with the case of neural networks with linear output layer.

Optimization of wide neural networks

A number of recent analyses show convergence of gradient descent for wide neural networks du2018gradientshallow (8, 7, 1, 23, 3, 13, 4). While an extended discussion of optimization is beyond the scope of this work, we refer the interested reader to our separate paper liu2020toward (17). The goal of this section is to clarify the important difference between the (near-)linearity of large models and convergence of optimization by gradient descent. It is easy to see that a wide model undergoing the transition to linearity can be optimized by gradient descent if its tangent kernel is well-conditioned at the initialization point. The dynamics of such a model will be essentially the same as for a linear model, an observation originally made in jacot2018neural (12).

However near-linearity or, equivalently, near-constancy of the tangent kernel is not necessary for successful optimization. What is needed is that the tangent kernel is well-conditioned along the optimization path, a far weaker condition.

The technical result is a consequence of Corollary 8.1 in liu2020toward (17).

Numerical Verification

We conduct experiments to verify the non-constancy of tangent kernels for certain types of wide neural networks, as theoretically observed in Section 4.

Specifically, we use gradient descent to train each neural network described below on a synthetic data until convergence. We compute the following quantity to measure the max (relative) change of tangent kernel from initialization to convergence: $\Delta K:=\sup_{t>0}\|K({\mathbf{w}}_{t})-K({\mathbf{w}}_{0})\|_{F}/\|K({\mathbf{w}}_{0})\|_{F}.$ For a network that has a nearly constant tangent kernel during training, $\Delta K$ is expected to be close to , while a network with a non-constant tangent kernel, $\Delta K$ should be $\Omega(1)$. Detailed experimental setup and data description are given in Appendix B.

In Figure 1, right panel, we demonstrate the evolution of tangent kernel with respect to the training time for a very wide neural network (width $m=10^{4}$). We see that, for the neural network with a non-linear output layer, tangent kernel changes significantly from initialization, while tangent kernel of the linear output network is nearly unchanged during training.

Wide neural networks with a bottleneck. We consider a fully connected neural network with $3$ hidden layers and a linear output layer. The second hidden layer, i.e., the bottleneck layer, has a width $m_{b}$ which is typically small, while the width $m$ of the other hidden layers are typically very large, $m=10^{4}$ in our experiment. For different bottleneck width $m_{b}=\{3$, $5$, $10$, $50$, $100$, $500$,$1000\}$, we train the network on a synthetic dataset using gradient descent until convergence, and compute $\Delta K$.

The change of tangent kernels for different bottleneck width is shown in Figure 2. We can see that a narrow bottleneck layer in a wide neural network prevent the neural tangent kernel from being constant during training. As expected, increasing the width of the bottleneck layer, makes the change of the tangent kernel smaller. We observe that the scaling of the tangent kernel change with width follows close to $\Theta\left(1/\sqrt{m}\right)$ (dashed line in Figure 2) in alignment with our theoretical results (Theorem 3.3).

Acknowledgements

The authors acknowledge support from NSF, the Simons Foundation and a Google Faculty Research Award. We thank James Lucas for correcting the proof of Prop. 2.3. The GPU used for the experiments was donated by Nvidia.

References

Appendix A Other Parameterization Strategies

Throughout the paper, our analysis is based on the NTK prameterization jacot2018neural (12), under which the constancy of tangent kernel is originally observed. In this section, we show that different parameterization strategies (e.g., LeCun initialization lecun2012efficient (15) : $w_{0;i}^{(l)}\sim\mathcal{N}(0,1/m)$) do not change our conclusions.

Specifically, we show that, compared to the NTK prameterization, a different parameterization strategy only rescales the tangent kernel $K$ and the spectral norm of the Hessian $\|H\|$ by the same factor, hence the ratio between tangent kernel $K$ and Hessian spectral norm keeps the same and $\|H\|=o(\|K\|)$ still holds. This still implies that the tangent kernel is almost constant during training.

Recall that we initialize the parameters ${\mathbf{W}}=\{{\mathbf{w}}^{(1)},{\mathbf{w}}^{(2)},\cdots,{\mathbf{w}}^{(L)},{\mathbf{w}}^{(L+1)}:={\mathbf{v}}\}$ of the general form of a deep neural network $f$, Eq.(3.3) by a standard Gaussian, i.e. $w_{i}^{(l)}\sim\mathcal{N}(0,1)$. If we apply another parameterization strategy $\bar{{\mathbf{W}}}$ here, for example, $\bar{w}_{i}^{(l)}\sim\mathcal{N}(0,\sigma_{m}^{2})$, where $\sigma_{m}$ can be a function of $m$, we can see every $\bar{w}_{i}^{(l)}=\sigma_{m}w_{i}^{(l)}$ where $w_{i}^{(l)}\sim\mathcal{N}(0,1)$.

In this case, the gradient of the model $f$ w.r.t. the weights of layer $l$ is

And by the same reason, the Hessian of the model f w.r.t. the weights of layer $l_{1}$ and $l_{2}$ is

Therefore, it’s easy to see the ratio of the norm of the tangent kernel to the norm of the Hessian keeps the same:

In many practical machine learning tasks, it is popular to use the LeCun initialization/parameterization: each individual parameter $(W^{(l)}_{0})_{ij}\sim\mathcal{N}(0,\frac{1}{m})$, while there is no factor $1/\sqrt{m}$ in the definition of the layer function, e.g., for fully connected layers

In this setting, the factor $\sigma_{m}=1/\sqrt{m}$. Then, by the analysis above, we see that

It is also interesting to note that, the Euclidean norm of the parameter change ${\mathbf{w}}^{*}-{\mathbf{w}}_{0}$ also scales:

Appendix B Experimental Setup

We use a synthetic dataset of size $N=60$ which contains $C=3$ classes. Each data point $(x,y)$ is sampled as follows: label $y$ is randomly sampled from $\{0,1,2\}$ with equal probability; given $y$, $x$ is drawn from the following distribution:

We encode each $y_{i}\in\{0,1,2\}$ in $\{(x_{i},y_{i})\}_{i=1}^{N}$ by a one-hot vector ${\mathbf{y}}_{i}\in\{0,1\}^{3}$. And ${\mathbf{y}}_{i,j}$ means the $j$-th component of ${\mathbf{y}}_{i}$. We use this dataset for all the optimization tasks mentioned below.

B.1 Wide neural networks with non-linear output layers

In the experiments, we train three different neural networks:

Neural network with a linear output layer

Neural network with a softmax-activated (non-linear) output layer

Neural network with a swish-activated (non-linear) output layer

We use gradient descent to minimize the loss functions until convergence is achieved (i.e. loss less than $10^{-4}$). To measure the change of tangent kernels, we compute the max (relative) change of tangent kernel from initialization to convergence: $\Delta K:=\sup_{t>0}\|K({\mathbf{w}}_{t})-K({\mathbf{w}}_{0})\|_{F}/\|K({\mathbf{w}}_{0})\|_{F}.$ For each training, we take $10$ independent runs and report the average $\Delta K$.

B.2 Wide neural networks with a bottleneck

In the experiment, we use a fully connected neural network with $3$ hidden layers and a linear output layer. Its second hidden layer, i.e., the bottleneck layer has a width $m_{b}$, while the other hidden layers has a width $m$. Specifically, it is defined as:

For each bottleneck width, we use gradient descent to minimize the loss functions until convergence is achieved (i.e. loss less than $10^{-4}$) and compute the max (relative) change of tangent kernel from initialization to convergence: $\Delta K:=\sup_{t>0}\|K({\mathbf{w}}_{t})-K({\mathbf{w}}_{0})\|_{F}/\|K({\mathbf{w}}_{0})\|_{F}.$ For each training, take $10$ independent runs and report the average $\Delta K$.

Appendix C Proof for Proposition 2.2

Recall that the tangent kernel is defined as

For a constant tangent kernel, each element $K_{ii}({\mathbf{w}})$ is constant. Noting that $K_{ii}({\mathbf{w}})=\|\nabla_{\mathbf{w}}f({\mathbf{w}},{\mathbf{x}}_{i})\|^{2}$, we have $\|\nabla f({\mathbf{w}},{\mathbf{x}})\|$ is constant in ${\mathbf{w}}$, for all input ${\mathbf{x}}$.

The following arguments basically follow the idea from 868044 (9) (a more general result was shown in sakai1996riemannian (19)).

To simplify the notation, in the rest of the proof, we hide the argument ${\mathbf{x}}$, and we use $f({\mathbf{w}})$ to denote $f({\mathbf{w}};{\mathbf{x}})$.

Let $\|\nabla f({\mathbf{w}})\|=c$. Consider the ordinary differential equation (ODE)

For any $t_{1},t_{2}$, since $\|\nabla f({\mathbf{w}})\|=c$, we have

but $|{\mathbf{w}}(t_{1})-{\mathbf{w}}(t_{2})|=|\int_{t_{2}}^{t_{1}}\|d{\mathbf{w}}(t)/dt\|dt|=c|t_{1}-t_{2}|$, which indicates

Dividing by $t$ and taking $t$ to $\pm\infty$ allows us to have $\langle\nabla f({\mathbf{w}}_{0}),{\mathbf{v}}-{\mathbf{w}}_{0}\rangle=0$. Then we construct the level set

where $g^{\prime}(t)=c$ for all $t$ which shows $f$ is linear.

Appendix D Proof of Proposition 2.3

Since the function $f$ is twice differentiable w.r.t. ${\mathbf{w}}$, according to Taylor’s theorem, we have the following expression for the gradient:

Then the Euclidean norm of the gradient change is bounded by

Since $t\in$ and the ball $B({\mathbf{w}}_{0},R)$ is convex, the point ${\mathbf{w}}_{0}+t({\mathbf{w}}-{\mathbf{w}}_{0})$ is within $B({\mathbf{w}}_{0},R)$. Hence,

Since $f$ is smooth, the gradients $\nabla_{{\mathbf{w}}}f({\mathbf{w}}_{0})$ and $\nabla_{{\mathbf{w}}}f({\mathbf{w}})$ are bounded. Therefore, $|K_{({\mathbf{x}},{\mathbf{z}})}({\mathbf{w}})-K_{({\mathbf{x}},{\mathbf{z}})}({\mathbf{w}}_{0})|=O(\epsilon R)$. ∎

Appendix E Proof of Theorem 3.1

The Hessian matrix $H$ of the neural network can be written as the following structure:

Here, each Hessian block $H^{(l_{1},l_{2})}:=\frac{\partial^{2}f}{\partial{\mathbf{w}}^{(l_{1})}\partial{\mathbf{w}}^{(l_{2})}}$ is the second derivative of $f$ w.r.t. its weights of $l_{1}$-th and $l_{2}$-th layers, where we treat the final layer parameters ${\mathbf{v}}$ as ${\mathbf{w}}^{(L+1)}$.

The following lemma allows us to bound the Hessian spectral norm by the norms of its blocks (see proof in Appendix I.1).

Spectral norm of a matrix $H$ (58) is upper bounded by the sum of the spectral norm of its blocks, i.e. $\|H\|\leq\sum_{l_{1},l_{2}}\|H^{(l_{1},l_{2})}\|$, $l_{1},l_{2}\in[L+1]$.

Now, we analyze the Hessian blocks case by case. Since the Hessian matrix is symmetry, without loss of generosity, we assume $1\leq l_{1}\leq l_{2}\leq L+1$.

By the chain rule, the gradient of the model $f$ w.r.t. the weights of layer $l$, can be written as

Then, the Hessian block has the following expression:

Hence, the spectral norm of Hessian block $H^{(l_{1},l_{2})}$ is bounded by

with $C^{\prime}_{1}=L^{2}\mathsf{L}_{\phi}^{2L}+L\mathsf{L}_{\phi}^{L}+1$.

𝐿11\leq l_{1}

𝐿1l_{1}=l_{2}=L+1. In this case, the Hessian block $H^{(L+1,L+1)}$ is simply zero. Hence, the spectral norm is zero.

Applying Lemma E.1, we immediately obtain the desired result. ∎

Appendix F Proof for Lemma 3.1

According to the definitions of the quantities $\mathcal{Q}_{\infty}(f)$, $\mathcal{Q}_{2,2,1}(f)$ and $\mathcal{Q}_{L}(f)$ in Eq.(21), it suffices to show that the followings layer-wise properties hold everywhere in the ball $B({\mathbf{W}}_{0},R)$ with high probability over the initialization:

The matrix spectral norm $\left\|\frac{\partial\alpha^{(l)}}{\partial{\mathbf{w}}^{(l)}}\right\|=O(1)$ w.r.t. $m$, for all $l\in[L]$;

The $(2,2,1)$-norms of order $3$ tensors, $\left\|\frac{\partial^{2}\alpha^{(l)}}{\partial{\mathbf{w}}^{(l)2}}\right\|_{2,2,1}$, $\left\|\frac{\partial^{2}\alpha^{(l)}}{\partial\alpha^{(l-1)}\partial{\mathbf{w}}^{(l)}}\right\|_{2,2,1}$ and $\left\|\frac{\partial^{2}\alpha^{(l)}}{(\partial\alpha^{(l-1)})^{2}}\right\|_{2,2,1}$ are all of the order $O(1)$ w.r.t. $m$, for all $l\in[L]$.

We start the proof with some preliminary results, and then prove the above statements one by one.

The fully connected neural network is defined in the following way:

where $m_{0}=d$ which is the dimension of the input ${\mathbf{x}}$, and $m_{l}=m$ for all $l\in[L]$. The trainable parameters of this network are ${\mathbf{W}}:=\{W^{(1)},W^{(2)},\cdots,W^{(L)},W^{(L+1)}:={\mathbf{v}}\}$, and are initialized by the random Gaussian initialization, i.e., each parameter $(W^{(l)}_{0})_{ij}\sim\mathcal{N}(0,1),\forall l\in[L]$, and $v_{0,i}\sim\mathcal{N}(0,1)$, $i,j\in[m]$. As the parameters $W^{(l)}$ of each layer are reshaped into matrices, the Euclidean norm of parameters becomes $\|{\mathbf{W}}\|:=(\sum_{l=1}^{L+1}\|W^{(l)}\|_{F}^{2})^{1/2}$, where $\|\cdot\|_{F}$ is the Frobenius norm of a matrix.

To make the presentation of the proof as simple as possible, we first make the following assumption about the initial parameters ${\mathbf{W}}_{0}$. Then we prove it in Lemma F.1 that the assumption is satisfied with high probability over the random Gaussian initialization.

We assume that there exists a constant $c_{0}>0$ such that, for all initial weight matrices/vector $W_{0}^{(l)}$, $\|W_{0}^{(l)}\|\leq c_{0}\sqrt{m}$, where $l\in[L+1]$.

If the parameters are initialized as $(W^{(l)}_{0})_{ij}\sim\mathcal{N}(0,1)$ for all $l\in[L+1]$ and $m>d$, then, for each layer $l\in[L+1]$, we have with probability at least $1-2\exp(-\frac{m}{2})$,

We prove the following lemma which states that the norm of the matrix $W^{(l)}$ keeps its order in a finite ball around the $W_{0}^{(l)}$ .

If ${\mathbf{W}}_{0}$ satisfies Assumption F.1, then for any ${\mathbf{W}}$ such that $\|{\mathbf{W}}-{\mathbf{W}}_{0}\|\leq R$, we have

See the proof in Appendix I.3. The following lemma gives bounds on the Euclidean norm of the vector of hidden neurons for each layer.

If ${\mathbf{W}}_{0}$ satisfies Assumption F.1, then, for any ${\mathbf{W}}$ such that $\|{\mathbf{W}}-{\mathbf{W}}_{0}\|\leq R$, we have, at all hidden layers

When $l=2,3,\cdots,L$. Recall from Eq.(22) that, a fully connected layer $\alpha^{(l)}$ is defined as, for $l=2,3,\cdots,L$:

Note that, in this case, the parameter vector ${\mathbf{w}}^{(l)}$ is reshaped to an $m\times m$ matrix $W^{(l)}$. The first derivatives of $\alpha^{(l)}$ are

By the definition of spectral norm, $\|A\|=\sup_{\|{\mathbf{v}}\|=1}\|A{\mathbf{v}}\|$, we have, for all $2\leq l\leq L$,

In the last inequality, we used Lemma F.3 and the Lipschitz continuity of the activation $\sigma(\cdot)$.

In this layer, the input ${\mathbf{x}}$ is fixed (independent of trainable parameters) and not a dynamical variable. Hence, $\partial\alpha^{(1)}/\partial{\mathbf{x}}$ is not an interesting object in our Hessian analysisIndeed, it does not show up in the Hessian analysis (c.f. the proof of Theorem 3.1 in Section E)..

For $\partial\alpha^{(1)}/\partial W^{(1)}$, we have (with a similar analysis as in Eq.(71)),

F.3 (2,2,1)221(2,2,1)-norms of order 333 tensors are O​(1)𝑂1O(1)

We consider the first layer i.e. $l=1$ and the rest of the layers i.e. $l=2,3,\cdots,L$ separately.

When $l=2,3,\cdots,L$. The second derivatives of the vector-valued layer function $\alpha^{(l)}$, which are order $3$ tensors, have the following expressions:

By the definition of the $(2,2,1)$-norm for order $3$ tensors, and Lemma F.2, we get

Similarly, by using Lemma F.2 and Lemma F.3, we have,

When $l=1$. As discussed in Section F.2, the input $\alpha^{(0)}={\mathbf{x}}$ is constant, we only need to analyze the tensor $\frac{\partial^{2}\alpha^{(l)}}{(\partial W^{(l)})^{2}}$ in this case. With a similar analysis in Eq.(77), we have

First of all, we present a few useful facts, Lemma F.4-F.6 that will be used during the proof. The proofs of the following lemmas are in Appendix I.5-I.7.

We first show that each activation of the hidden layers is bounded at initialization, with high probability.

The following lemma gives an upper bound to Euclidean norms of ${\mathbf{b}}^{(l)}$ in the ball $B({\mathbf{W}}_{0},R)$.

If the initial parameters ${\mathbf{W}}_{0}$ of the multi-layer neural network $f({\mathbf{W}})$ satisfies Assumption F.1, then, for any ${\mathbf{W}}$ such that $\|{\mathbf{W}}-{\mathbf{W}}_{0}\|\leq R$, we have, at all hidden layers, i.e., $\forall l\in[L]$,

First of all, we prove, by induction, the following claim: for all $l\in[L]$,

In the base case, we consider $l=L$. We have

To bound the second additive term above, we need the following inequality:

where the last equality is the result of Lemma F.3 that ${\|\alpha^{(l-1)}}\|=O(\sqrt{m})$.

Also, note that $\Sigma^{\prime}$ is a diagonal matrix, then, we have

where we used Lemma F.5 and Eq.(85) in the last equality. Now, insert Eq.(86) into Eq.(84), and apply Lemma F.5 and the induction hypothesis, then we have

Appendix G Generalization to other architectures

In this section, we apply Theorem 3.1 to both convolutional neural networks (CNN) and residual networks (ResNets), and show that they both have small Hessian spectral norms when the network width $m$ is sufficiently large and last layer is of linear form.

A convolutional neural network (CNN) is a network of the type in Eq.(3.3), with each convolutional layer function $\phi_{l}$ defined as

where $\ast$ is the convolution operator (see the definition below), and the layer width $m_{l}=m$ for all $l=2,3,\cdots,L$, and $m_{1}=d$ with $d$ as the number of channels of the input.

To simplify the notation, we consider a one-dimensional CNN, i.e., a “image” is an $1$-D array of “pixels”, and one will find that the analysis in this section also applies to higher dimensional CNNs. We also drop the layer indices $l$, wherever there is no ambiguity.

Reformulation of convolutional layer. Now, we reformulate the convolutional layer function in Eq.(90) into a fully-connected-like function. Then, we can use the techniques developed in Section F to prove for the CNN. Specifically, for all $k\in[K]$, define matrices $W^{[k]}$ and $\alpha^{[k]}$ such that each entry $(W^{[k]})_{ij}=W_{k,i,j}$ and $(\alpha^{[k]})_{jq}=\alpha_{j,q+k-\frac{K+1}{2}}$. Then, the convolution operator in Eq.(91) can be rewritten as

Here in the summation, it is matrix multiplication. Note that, while $W^{[k]}$ are independent from each other for different $k\in[K]$, the inputs $\alpha^{[k]}$ are not independent from each other; instead, they share pixels: $(\alpha^{[k]})_{j,q}=(\alpha^{[k^{\prime}]})_{j,q+k-k^{\prime}}$, i.e., each $\alpha^{[k]}$ is a pixel-shifted version of $\alpha$ (newly generated pixels after shift is filled with zeros).

Therefore, the convolutional layer function can also be written as (for $l>1$)

Here, we can see we will use this expression of convolutional layer function for analysis in this section.

Before proceeding to the proof for CNN, we first point out a few useful facts, as summarized in the following lemmas.

Given matrices $A,B$ and $C$ such that $A=BC$, we have $\|A\|_{F}\leq\|B\|\|C\|_{F}$, where $\|B\|$ is the spectral norm of matrix $B$.

See the proof in Appendix I.8. The following two lemmas provide bounds on the spectral norm of weights and Frobenius norm of hidden layers. These two lemmas (and the proofs) are analogous to Lemma F.2 and F.3, and we omit the proof.

Suppose the parameters are initialized as $(W^{[k]}_{0})_{i,j}\sim\mathcal{N}(0,1)$, for all $k\in[K],i,j\in[m]$. Then, with high probability of the random initialization, we have for any ${\mathbf{W}}\in B({\mathbf{W}}_{0},R)$ the following holds

Suppose the parameters are initialized as $(W^{[k]}_{0})_{i,j}\sim\mathcal{N}(0,1)$, for all $k\in[K],i,j\in[m]$ and for all layers. Then, with high probability of the random initialization, we have for any ${\mathbf{W}}\in B({\mathbf{W}}_{0},R)$ the following holds at all hidden layers

We note that the proof for CNNs is basically analogous to that for fully connected neural networks (FCNs). Here, we refer readers to follow the proof idea for FCNs and only discuss the main differences below. In the following, we focus on analyzing the layers for $l>1$. For the case of $l=1$, we omit the proof, and refer the readers to the discussion in Section F, which also applies here.

Here, in the second inequality, we used Lemma G.1, and in the last equality, we used Lemma G.2. Similarly, using Lemma G.1 and G.3, we also have

Similarly, by using Lemma G.1, G.2 and G.3, we also have

G.2 Residual Networks (ResNet)

The parameters ${\mathbf{W}}:=\{W^{(1)},W^{(2)},\cdots,W^{(L)},W^{(L+1)}:={\mathbf{v}}\}$ are initialized following the random Gaussian initialization strategy, i.e., $(W^{(l)}_{0})_{ij}\sim\mathcal{N}(0,1),\forall l\in[L]$, and $v_{0,i}\sim\mathcal{N}(0,1)$, $i,j\in[m]$.

This definition of ResNet differs from the standard ResNet architecture in he2016deep (11) that the skip connections are at every layer, instead of every two layers. One will find that the same analysis can be easily generalized to cases where skip connections are at every two or more layer. The same definition, up to a scaling factor, was also theoretically studied in du2018gradientdeep (7).

We see that the ResNet is the same as a fully connected neural network, Eq. (F.1), except that the activations $\alpha^{(l)}$ has an extra additive term $\alpha^{(l-1)}$ from the previous layer, interpreted as skip connection. Because of this similarity, the proof for ResNet is almost identical to that for fully connected networks. In the following, we sketch the proof for ResNet. Specifically, we focus on the arguments that are new to ResNet, and omit those identical to the fully connected case.

Parallel to Lemma F.2 and F.3 for fully connected case, we have the following lemmas for the ResNet.

Suppose the parameters are initialized as $(W^{(l)}_{0})_{i,j}\sim\mathcal{N}(0,1)$, for all $l\in[L]$, and $v_{0,i}\sim\mathcal{N}(0,1)$, $i,j\in[m]$. Then, with high probability of the random initialization, we have for any ${\mathbf{W}}\in B({\mathbf{W}}_{0},R)$ the following holds

Suppose the parameters are initialized as $(W^{(l)}_{0})_{i,j}\sim\mathcal{N}(0,1)$, for all $l\in[L],$ and $v_{0,i}\sim\mathcal{N}(0,1)$, $i,j\in[m]$. Then, with high probability of the random initialization, we have for any ${\mathbf{W}}\in B({\mathbf{W}}_{0},R)$ the following holds at all hidden layers

The proofs of the above two lemmas are almost identical to those of Lemma F.2 and F.3. We omit the proofs here, and refer interested readers to proofs of Lemma F.2 and F.3.

We note that $\|\partial\alpha^{(l)}/\partial{\mathbf{w}}^{(l)}\|$ has the same expression as the one of the fully connected networks. By the same argument in Section F.2, as well as Lemma G.5, we have $\|\partial\alpha^{(l)}/\partial{\mathbf{w}}^{(l)}\|=O(1)$.

When $l=1$, the layer function is defined by

In this layer, the input ${\mathbf{x}}$ is fixed (independent of trainable parameters) and not a dynamical variable. Hence, $\partial\alpha^{(1)}/\partial\alpha^{(0)}$ is not an interesting object in our Hessian analysis.

We see that both $\|{\nabla_{\alpha}\phi_{l}}\|$ and $\|\nabla_{\mathbf{w}}\phi_{l}\|$ are bounded, hence, the (vector-valued) layer function of ResNet is Lipschitz continuous.

Note that the skip connection term $\alpha^{(l-1)}$ in Eq.(101) is linear in $\alpha^{(l-1)}$ and independent from $W^{(l)}$. Hence, the order $3$ tensors are exactly the same as in the case of fully connected networks. Applying the same argument as in Section F.3 gives the following:

For a ResNet, define vector ${\mathbf{b}}^{(l)}_{res}:=\partial f/\partial{\alpha^{(l)}}$ for $l\in[L]$. Specifically, ${\mathbf{b}}^{(l)}_{res}$ takes the following form:

G.3 Architecture with mixed layer types

Appendix H Proof of Theorem 4.1

The proof of the lemma is deferred to Appendix I.9.

Consider an arbitrary parameter setting ${\mathbf{W}}\in B({\mathbf{W}}_{0},R)$.

Note that spectral norm of a matrix is lower bounded by the norm of its blocks, then

Hence, it’s sufficient to lower bound the norm of the Hessian block ${\partial^{2}f}/{\left(\partial{\mathbf{w}}^{(2)}\right)^{2}}$.

With simple computation, the gradient of $f$ w.r.t. ${\mathbf{w}}^{(2)}_{i}$ is:

and each entry of the Hessian matrix takes the form:

In the second equality above, we have used $\sigma(x)=\frac{1}{2}x^{2}$.

Then, the spectral norm of this Hessian block is

For the last factor $\|\alpha^{(1)}\|^{2}$, using the tail bound for $\chi^{2}(m)$ laurent2000adaptive (14), we have, with probability at least $1-e^{-m/16}$,

Applying Lemma H.1 to the factor $\left|\sum_{j=1}^{m}{\mathbf{w}}_{j}^{(4)}({\mathbf{w}}_{j}^{(3)})^{2}\right|$ and using union bound, for an arbitrary $\delta\in(0,1)$, we have, with probability at least $1-2\delta-e^{-m/16}$,

Hence, we get the lower bound of the Hessian spectral norm at ${\mathbf{W}}\in B({\mathbf{W}}_{0},R)$

Appendix I Proofs of Technical Lemmas

I.2 Proofs for Gaussian Random Initialization

In particular, for the initial parameter setting ${\mathbf{W}}_{0}$, we have

Letting $t=\sqrt{m}$ and noting that $m>d$, we finish the proof. ∎

I.3 Proof of Lemma F.2

By triangle inequality and the definition $\|{\mathbf{W}}\|=\sum_{l=1}^{L+1}\|W^{(l)}\|_{F}$, we have for all layers, i.e., $l\in[L+1]$,

Note that, at the output layer, $W^{(L+1)}$ i.e. ${\mathbf{v}}$ is a vector, and the Frobenius norm $\|\cdot\|_{F}$ reduces to the Euclidean norm $\|\cdot\|$. ∎

I.4 Proof of Lemma F.3

To analyze $\|\alpha^{(l)}({\mathbf{W}})\|$, let’s first consider the input layer, i.e., $l=0$: $\|{\alpha}^{(0)}\|=\|{\mathbf{x}}\|\leq\sqrt{d}\|{\mathbf{x}}\|_{\infty}\leq\sqrt{d}C_{{\mathbf{x}}}$, where $d$ is the dimension of the input ${\mathbf{x}}$. Then we prove Eq.(66) by induction. For the first hidden layer $l=1$,

Above, we used the $\mathsf{L}_{\sigma}$-Lipschitz continuity and applied Lemma F.2 in the second inequality. Now, suppose for $l$-th layer we have

Then, by a similar argument as in Eq.(114), we can get

I.5 Proof of Lemma F.4

When $2\leq l\leq L$, $|\alpha_{i}^{(l)}|$ takes the following form:

where we can see $\sum_{k=1}^{m}\ W^{(l)}_{ik}\alpha_{k}^{(l-1)}\sim\mathcal{N}(0,\|\alpha^{(l-1)}\|^{2})$ since $W^{(l)}_{ik}\sim\mathcal{N}(0,1)$ at initialization. By the concentration inequality for Gaussian random variable, we have

for $c^{(l)}_{\alpha}=\frac{m}{2\mathsf{L}_{\sigma}^{2}\|\alpha^{(l-1)}\|^{2}}=\Omega(1)$ by Lemma F.3. When $l=1$, we have

Similarly, at initialization, $\sum_{k=1}^{d}W^{(1)}_{ik}{\mathbf{x}}_{k}\sim\mathcal{N}(0,\|{\mathbf{x}}\|^{2})$. Hence