On the Neural Tangent Kernel of Deep Networks with Orthogonal Initialization

Introduction

Deep learning has been responsible for a step-change in performance across machine learning, setting new benchmarks for state-of-the-art performance in many applications, from computer vision , natural language processing , to reinforcement learning , and more. Beyond its fundamental shift in approach, an array of innovative techniques underpin the success of deep learning, such as residual connections , dropout , and batch normalization . The mean field theory recently opened a gate to analyze the principles behind neural networks with random, infinite width, and fully-connected networks as the first subjects. Broadly, what Schoenholz et al. discovered, and then empirically proved, is that there exists a critical initialization called the edge of chaos, allowing the correlation signal to go infinitely far forward and preventing vanishing or exploding gradients. Later, this theory had been extended to a much wider range of architectures, e.g., convolutional networks , recurrent networks , dropout networks , residual networks , and batch normalization .

Critical initialization requires the mean squared singular value of a network’s input-output Jacobian to be $O(1)$. It was already known that the learning process in deep linear networks could be dramatically accelerated by ensuring all singular values of the Jacobian being concentrated near 1, a property known as dynamical isometry . However, what was not known was how to impose dynamical isometry in deep nonlinear networks. Pennington et al. conjectured that they could do so with techniques based on free probability and random matrix theory, giving rise to a new and improved form of initialization in deep nonlinear networks. Since then, dynamical isometry has been introduced to various architectures, such as residual networks , convolutional networks , or recurrent networks with excellent performance on real-world datasets.

In fully connected networks, two key factors help to ensure dynamical isometry. One is orthogonality, and the other is appropriately tuning weights’ and biases’ parameters to establish a linear regime in nonlinear activation . In straightforward scenarios, orthogonal initialization is usually enough to impose dynamical isometry in a linear network. The benefit of the orthogonality in linear networks has been proven recently . However, the dynamics of nonlinear networks with orthogonal initialization has not been investigated. The roadblock is that it has been unclear how to derive a simple analytic expression for the training dynamics.

Hence, to fill this gap, we look to a recent technique called neural tangent kernel (NTK) , developed for studying the evolution of a deep network using gradient descent in the infinite width limit. NTK is a kernel characterized by a derivative of the output of a network to its parameters. It has been shown that the NTK of a network with Gaussian initialization converges to a deterministic kernel and remains unchanged during gradient descent in the infinite-width limit. We extend these results to the orthogonal initialization case and find that orthogonal weights contribute to the same properties for NTK. Given a sufficiently small learning rate and wide width, the network optimized by gradient descent behaves as a model linearized about its initial parameters , where these dynamics are called NTK regime, or lazy training . As the learning rate gets larger or the network becomes deeper, that is, out of the NTK regime, we expect that there will be new phenomena that can differentiate two initialization. To summarize, our contribution is as follows,

We prove that the NTK of an orthogonally-initialized network converges to the NTK of a network initialized by Gaussian weights in the infinite-width limit. Besides, theoretically, during training, the NTK of an orthogonally-initialized infinite-width network stays constant in the infinite-width limit.

We prove that the NTK of an orthogonally-initialized network across architectures, including FCNs and CNNs, varies at a rate of the same order for finite-width as the NTK of a Gaussian-initialized network. Therefore, there are no significant improvements brought by orthogonal initialization for wide and nonlinear networks compared with Gaussian initialization in the NTK regime.

We conduct a thorough empirical investigation of training speed outside the NTK regime to complement theoretical results. We show that orthogonal initialization can speed up training in the large learning rate and depth regime when the hyper-parameters are set to achieve a linear regime in nonlinear activation.

Related Work

Hu et al. ’s investigation of orthogonal initialization in linear networks provided a rigorous proof that drawing the initial weights from the orthogonal group speeds up convergence relative to standard Gaussian initialization. However, deep nonlinear networks are much more complicated, making generating proof the same in these nonlinear settings much more difficult. For example, Sokol and Park attempted to explain why dynamical isometry imposed through orthogonal initialization can significantly increase training speed. They showed a connection between the maximum curvature of the optimization landscape, as measured by a Fisher information matrix (FIM) and the spectral radius of the input-output Jacobian, which partially explains why networks with greater isometric are able to train much faster. Given that NTK and FIM have the same spectrum for the regression problem, our theoretical results can be seen as a complement to each other since our work provide a precise characterization of the dynamics of wide networks in the NTK regime. At the same time, Sokol and Park investigated the advantage of orthogonal initialization outside of the NTK regime.

Jacot et al. , who conceived of the neural tangent kernel, shows that NTK both converges to an explicit limiting kernel in the infinite-width networks and remains constant during training with Gaussian initialization. Lee et al. reached the same conclusion from a different angle with a demonstration that the gradient descent dynamics of the original neural network fall into its linearized dynamics regime. We extend these results to orthogonal initialization and show that both NTK of Gaussian and orthogonally-initialized network are the same kernel in the infinite-width limit. While the original work of NTK is groundbreaking in producing an equation to predict the behavior of gradient descent in the NTK regime, its proof is incomplete as it implicitly assumes gradient independence. Yang subsequently removed this assumption and completed the proof. Besides, Huang and Yau have proven the same proprieties of NTK and global convergence of deep networks in a few different ways. However, all of these studies did not focus on the orthogonal initialization as with our work.

Studies involving NTK commonly adopt the ntk-parameterization , since standard parameterization with an infinite-width limit can lead to divergent gradient flow in the infinite limit. To account for this problem for standard parameterization, the authors of sohl2020infinite developed an improved standard parameterization. We tested both techniques using the “Neural Tangents” Python library with some added codes to support orthogonal initialization.

Preliminaries

Convolutional Neural Network (CNN). For notational simplicity, we consider a 1D convolutional networks with periodic boundary conditions. We denote the filter relative spatial location $\beta\in\{-k,\dots,0,\dots,k\}$ and spatial location $\alpha\in\{1,\dots,m\}$, where $m$ is the spatial size. The forward propagation for $l\in\{1,\dots,L-1\}$ is given by,

Standard parameterization requires the parameter set $\theta=\{W^{l}_{ij},b^{l}_{i}\}$ is an ensemble generated by, $W^{l}_{ij}\sim\mathcal{N}(0,\frac{\sigma^{2}_{w}}{n_{l-1}}),~{}~{}~{}b_{i}^{l}\sim\mathcal{N}(0,\sigma^{2}_{b})$, where $\sigma^{2}_{w}$ and $\sigma^{2}_{b}$ are weight and bias variances. The variance of weights is scaled by the width of previous layer $n_{l-1}$ to preserve the order of post-activations layer to be $O(1)$. We denote this parameterizationas standard parameterizaiton. However, this paramterization can lead to a divergence in derivation of neural tangent kernel. To overcome this problem, ntk-parameterization was introduced, $W^{l}_{ij}=\frac{\sigma_{w}}{\sqrt{n_{l-1}}}\omega^{l}_{ij},~{}~{}~{}b_{i}^{l}=\sigma_{b}\beta_{i}^{l}$, where $\omega^{l}_{ij},\beta_{i}^{l}\sim\mathcal{N}(0,1)$.

2 Dynamical Isometry and Orthogonal Initialization

Consider the input-output Jacobian $J=\frac{\partial h^{L}}{\partial x^{0}}=\prod_{l=1}^{L}D^{l}W^{l}$, where $h^{L}$ is output function, $x^{0}$ is input, and $D^{l}$ is a diagonal matrix with elements $D^{l}_{ij}=\phi^{\prime}(h^{l}_{i})\delta_{ij}$. Ensuring all singular values of the Jacobian concentrate near 1 is a property known as dynamical isometry. In particular, It is shown that two conditions regarding singular values of $W^{l}$ and $D^{l}$ contribute crucially to the dynamical isometry in non-linear networks . More precisely, the singular values of $D^{l}$ can be made arbitrarily close to 1 by choosing a linear regime in a nonlinear activation. On the other hand, adopting a random orthogonal initialization can force the singular values of weights into 1. In particular, weights are drawn from a uniform distribution over scaled orthogonal matrices obeying,

This is the standard parameterization for orthogonal weights, and ntk-parameterization of orthogonality follows,

We show a summary of improved standard parameterization and ntk-parameterization across FCN and CNN for Gaussian and orthogonal initialization in Table 1. The factor $s$ in the layer equation of standard parameterization is introduced to prevent divergence of NTK . The core idea is to write the width of the neural network in each layer in terms of an auxiliary parameter, $n_{l}=sN_{l}$. Instead of letting $n_{l}\rightarrow\infty$, we adopt $s$ as the limiting factor.

3 Neural Tangent Kernel

The neural tangent kernel (NTK) is originated from and defined as,

This equation for $f_{t}$ has no substantial insight in studying the behavior of networks because $\Theta_{t}(X,X)$ varies with the time during training. Interestingly, as shown by , the NTK $\Theta_{t}(X,X)$ converges to a deterministic kernel $\Theta_{\infty}(X,X)$ and does not change during training in the infinite-width limit, i.e. $\Theta_{t}(X,X)=\Theta_{\infty}(X,X)$. As a result, the infinite width limit of the training dynamics are given by,

In the case of an MSE loss, $\mathcal{L}(y,f)=\frac{1}{2}\left\lVert y-f_{t}(\theta,x)\right\lVert^{2}_{2}$, the Equation (7) becomes a linear model with a solution,

Theoretical Results

As stated in , the pre-activation $h^{l}_{i}$ of Gaussian initialized network tends to Gaussian processes (GPs) in the infinite-width limit. This is the proposition to construct the NTK in networks with Gaussian weights . We extend this result to the orthogonal initialization across Fully Connected Networks (FCNs) and Convolutional Neural Networks (CNNs):

Consider a FCN of the form (1) at orthogonal initialization, with a Lipschitz nonlinearity $\phi$, and in the limit as $n_{1},...,n_{L-1}\to\infty$, the pre-activations $h^{l}_{i}$, for $i=1,...,n_{l}$ and $l\in\{1,\dots,L\}$, tend to i.i.d centered Gaussian processes of covariance $\Sigma^{l}$ which is defined recursively by:

For a CNN of the form (2) at orthogonal initialization, and in the limit as $n_{1},...,n_{L-1}\to\infty$, the pre-activations $h^{l}_{i,\alpha}$ tend to Gaussian processes of covariance $\Sigma_{\alpha,\alpha^{\prime}}^{l}$, which is defined recursively by:

Different from the independence property of Gaussian initialization, the entries of the orthogonal matrix are correlated. We use the Stein method and exchangeable sequence to overcome this difficulty and leave the detailed proof in the appendix. As shown by Theorem 1, neural networks with Gaussian and orthogonal initialization are in correspondence with an identical class of GPs.

2 Neural Tangent Kernel at Initialization

According to , the NTK of a network with Gaussian weights converges in probability to a deterministic kernel in the infinite-width limit. We show that the NTK of an orthogonally initialized network is identical to the one with Gaussian weights in the infinite-width limit.

Consider a FCN of the form (1) at orthogonal initialization, with a Lipschitz nonlinearity $\phi$, and in the limit as the layers width $n_{1},...,n_{L-1}\to\infty$, the NTK $\Theta^{L}_{0}(x,x^{\prime})$, converges in probability to a deterministic limiting kernel:

The scalar kernel $\Theta^{L}_{\infty}(x,x^{\prime})$ is defined recursively by

For a CNN of the form (2) at orthogonal initialization, and in the limit as $n_{1},...,n_{L-1}\to\infty$, the NTK $\Theta^{L}_{0}(x,x^{\prime})$, converges in probability to a deterministic limiting kernel:

The scalar kernel $\Theta^{L}_{\infty}(x,x^{\prime})$ is given recursively by

Since the Lipschitz function is differentiable besides a measure zero set, then taking the expectation would not destroy the whole statement, which allows for the ReLU activation.

In general, the NTK of CNNs can be computed recursively in a similar manner to the NTK for FCNs. However, the NTK of CNNs propagate differently by averaging over the NTKs regarding the neuron location of the previous layer. According to Theorem 2, the NTK of an orthogonally initialized network converges to an identical kernel as Gaussian initialization. This suggests these two NTKs are equivalent when the network structure (depth of $L$, filter size of $2k+1$, and activation of $\phi$) and choice of hyper-parameters ($\sigma^{2}_{w}$ and $\sigma^{2}_{b}$) are the same in the infinite-width limit.

3 Neural Tangent Kernel During Training

It is shown that the NTK of a network with Gaussian initialization stays asymptotically constant during gradient descent training in the infinite-width limit, providing a guarantee for loss convergence . We find that the NTK of orthogonally initialized networks have the same property, which is demonstrated below in an asymptotic way,

where $\hat{\Theta}_{t}$ are empirical kernels of networks with finite width.

proved the stability of NTK under the assumption of global convergence of neural networks, while provided a self-contained proof of both global convergence and stability of NTK simultaneously. In this work, we refer to the proof strategy from and extend it to the orthogonal case, as shown in the appendix.

To certificate this theorem empirically, we adopt three hidden layers Erf networks trained by gradient descent with learning rate $\eta=1.0$ on a subset of the MNIST dataset of $D=20$. We measure changes of weights, empirical NTK after $T=2^{15}$ steps of gradient descent for varying width at both Gaussian and orthogonal initialization. Figure 1(a,b) show that the relative change in the first and last layer weights scales as $1/\sqrt{n}$ while second and third layer weights scale as $1/n$ with Gaussian and orthogonal weights respectively. In Figure 2(c), we observe the change in NTK is upper bounded by $O(1/\sqrt{n})$ but is closer to $O(1/n)$ for both Gaussian and orthogonal initialization. The discrepancy between theoretical bound ($O(n^{-1/2})$) and experimental observation ($O(n^{-1})$) has been solved in , where they prove that relative change of empirical NTK of Gaussian initialized networks is bounded by $O(1/n)$. Without loss of generality, we infer that the proof framework is suitable for orthogonal weights.

Numerical Experiments

Our theoretical result indicates that ultra-wide networks with Gaussian and orthogonal initialization should have the same convergence rate during the gradient descent training. This means that two different initializations have similar training dynamics for loss and accuracy function in the NTK regime. Thus, it is now for us to verify our theories in practice. To this end, we perform a series of experiments on MNIST and CIFAR10 dataset. All the experiments are performed with the standard parameterization with TensorFlow.

We compare the train and test loss and accuracy with two different initialization, i.e., Gaussian and orthogonal weights using $D=256$ samples on full CIFAR-10 and MNIST dataset, as summarized in Figure 2. To reduce noise, we averaged the results over 30 different instantiations of the networks. Figures 2(a,b) show the results of the experiments on networks of depth $L=5$, width $n=800$, and activation tanh function, using SGD optimizer with a small learning rate of $\eta=10^{-3}$ for $T=10^{5}$ steps on CIFAR-10 dataset. Figure 2(c)(d) display the results on networks of depth $L=9$, width $n=1600$, and activation ReLU function, using PMSProp optimizer with a small learning rate of $\eta=10^{-5}$ for $T=1.2\times 10^{4}$ steps on MNIST. In all cases, we see an excellent agreement between the training dynamics of the two initialization, which is consistent with our theoretical finding (Theorem 3).

Having confirmed the consistency between training speed of networks with Gaussian and orthogonal initialization in the NTK regime, our primary interest is to find when orthogonal initialization accelerates the training speed for nonlinear networks. We need to go beyond the NTK regime and experiment with an additional requirement for hyper-parameters according to the evidence that orthogonal initialization increases learning speeds when the variance of weights and biases is set to achieve a liner regime in nonlinear activation .

Following , we set $\sigma^{2}_{w}=1.05$, and $\sigma^{2}_{b}=2.01\times 10^{-5}$, and $\phi(x)={\rm tanh}(x)$. We then vary the width of network in one set of experiments as $n=400,800$ and $1600$ when $L=50$, and the depth in another as $L=50,100$ and $200$, when $n=400$. All networks are trained by SGD optimizer on CIFAR-10 dataset. To evaluate the relationship between the learning rate and training speed, we select a threshold accuracy of $p=0.25$ and measure the first step $\tau$ when accuracy exceeds $p$. Figure 3 shows the steps of $\tau$ as a function of the learning rate of $\eta$ for both the training and testing sets.

The results in Figure 3 suggest a more quantitative analysis of the learning process until convergence. We train networks listed in Figure 3 for $5\times 10^{4}$ steps with a certain learning rate. We show the results of a certain network of depth $L=100$ and width $n=400$ trained with a learning rate $\eta=0.01$ as a typical example in Figure 4. The results of other network structures can be found in the appendix. It is shown that the training speed of orthogonally initialized networks is faster than that of Gaussian initialized networks outside the NTK regime. At the same time, orthogonally initialized networks can finally obtain a higher generalization result.

We draw two main conclusions from these experiments. First, orthogonal initialization results in faster training speeds and better generalization than Gaussian initialization in the large learning rate phase. It was shown that the large learning rate phase has many different properties from the small learning rate phase . Our finding can be seen as another effect in the large learning rate phase. Second, given the constant width, the greater the depth of the network, the more significant the difference in performance between orthogonal and Gaussian initialization. This phenomenon is consistent with the theoretical result observed in deep linear networks. It was found that the width needed for efficient convergence to a global minimum with orthogonal initialization is independent of the depth. In contrast, the width needed for efficient convergence with Gaussian initialization scales proportionally in depth .

Conclusion

This study on the neural tangent kernel of wide and nonlinear networks with orthogonal initialization has proven, theoretically and empirically, that the NTK of an orthogonally-initialized network across both FCN and CNN converges to the same deterministic kernel of a network initialized from Gaussian weights in the finite-width limit. We find that with an infinite-width network and a gradient descent (gradient flow) training scheme, the NTK of an orthogonally initialized network does not change during training. Further, it has the same order convergence rate from a finite to an infinite width limit as that of a Gaussian initialized network. Our theoretical results suggest that the dynamics of wide networks with orthogonal initialization behave similarly to that of Gaussian networks with a small learning rate verified by experiments. This observation implies that orthogonal initialization is only effective when not in the lazy (NTK) regime. And it is consistent with the fact that the infinite-width analysis does not explain the practically observed power of deep learning . Last, we find that orthogonal networks can outperform Gaussian networks in the large learning rate and depth on both train and test sets.

Appendix A Appendix

This appendix is dedicated to proving the key results of this paper, namely Theorem 1, Theorem 2, and Theorem 3 based on a series of lemmas, which describe the asymptotics of neural networks with orthogonal weights at initialization and during training. We prove the Gaussian process behavior of the output function, the limiting deterministic kernel for the orthogonal initialization, and its stability during training in the first three sections. Besides, we provide more experiments in the NTK regime and outside the NTK regime in the last section.

Throughout this section, we assume that $n_{1}=n_{2}=\cdots=n_{L}=n$. That is, $\{W_{ij}\}_{n\times n}$ is an orthogonal mapping. We first cite the following lemma from , which describes how the post-activations of one-layer transform by multiplying a random orthogonal matrix.

Let $(W_{ij})_{n\times n}$ be an orthogonal matrix randomly sampled by the Haar measure of orthogonal matirx. Let $B$ be a $n\times n$ matrix s.t $Tr(BB^{T})=n$. Then $Tr(BW)$ converges to a standard Gaussian distribution as the size of the matrix $n$ tends to infinity.

If we condition on the previous layer’s output, the next layer’s pre-activations are Gaussian when the width tends to infinity. Therefore if we take the limit of previous layers $n_{1},\dots,n_{l-1}\rightarrow\infty$ sequentially, the pre-activation $\{h_{i}^{l}\}$ of the l-th layer tends to a centered Gaussian process with respect to the input vector x. However, our goal is to take all the previous layers’ width simultaneously. The main technical difficulty is that we lose the independence between different index $i$ in the finite width when we implement orthogonal ensemble. Hence analysis based on the central limit theorem for i.i.d random sequences would not work. Instead, we follow the strategy in to apply a modified version of the exchangeable random sequence central limit theorem. Note that the Haar probability measure is invariant under row and column permutations. This implies that permuting the output index $i$ won’t change the joint law of $\{h_{i}^{l}(x)\}_{1\leq i\leq n_{l}}$.

We will use the following adapted version of the central limit theorem for exchangeable sequences in .

For each positive integer $n$ let $\left(X_{n,i};i=1,2,...\right)$ be an infinitely exchangeable process with zero mean, finite variance $\sigma^{2}_{n}$, and finite absolute third moment. Suppose also that the variance has a limit $\lim_{n\to\infty}\sigma^{2}_{n}=\sigma^{2}_{*}$. Define

Then $S_{n}$ converges in distribution to $\mathcal{N}(0,\sigma^{2}_{*})$.

As it was pointed out in , convergence with respect to arbitary finite-dimensional marginals is equivalent to convergence of all possible linear projections to real-valued random variable. We restate Definition 7 of as follows

The projections are defined in terms of a finite linear projection of the l-th layer’s input values without the biases:

Since we impose Haar ensemble to $\{W_{ij}\}_{1\leq i,j\leq n_{l-1}}$, it’s invariant under the column permutation. This implies that $\{\gamma^{l}_{j}(\mathcal{L},\alpha)[n]\}_{1\leq j\leq n_{l-1}}$ form an exchangeable sequence with respect to the column index $j$.

To fit the three conditions of Lemma 2, we need the following lemma on the moment calculation of orthogonal random matrices:

For all $i,j,r,s,\alpha,\beta,\lambda,\mu$,

In our scaling setting, the second moment should multiply by $n$ and the fourth moment should multiply by $n^{2}$.

Let $X_{n,i}:=\gamma^{l}_{j}(\mathcal{L},\alpha)[n]$, then

Thus, condition a) of Lemma 2 is satisfied.

The rest of the proof goes by induction. We assume that for the $l-1$-th layer, the pre-activations tend to independent Gaussian processes as the previous layers tend to infinite width simultaneously. Then we check that the moment conditions b) and c) of Lemma 2 for the l-th layer under the inductive hypothesis and Lemma 3. We first calculate the covariance formula non-rigorously where we take the limit sequentially. Since the recursion formula of the covariance doesn’t depend on how we take the limit.

We impose an assumption on the nonlinear activation function $\phi(x)$:

For a network of depth $L$ at initialization, with a Lipschitz nonlinearity $\phi$, and as $n_{1},...,n_{l-1}\to\infty$ sequentially, the pre-activations $h^{l}_{i}$, for $i=1,...,n_{l}$, tend to i.i.d centered Gaussian processes specified by the covariance $\{\Sigma^{l}\}_{1\leq l\leq L}$, where $\Sigma^{l}$ is defined recursively by:

taking the expectation with respect to a centered Gaussian process of covariance $\Sigma^{l-1}$ denoted by $f$. In the CNN case, $\Sigma_{\alpha,\alpha^{\prime}}^{l}$ is defined recursively by:

We show the result in the framework of NTK parameterization, while the argument for standard parameterization can be derived in the same way. In the FCN case, When $L=1$, there are no hidden layers and $h^{1}_{i}$ has the form:

Then we check the variance $\Sigma^{1}$ of output layer $h^{l}_{i}$. By Lemma 3,

We note that the fraction $\frac{1}{n_{0}}$ is from the scaling of the orthogonal distribution, while the term $n_{0}$ comes from the initialization. Thus we can compute the covariance of the first layer explicitly:

The next step is to use the induction method. Consider an $l$-network as the function mapping the input to the pre-activations $h^{l}_{i}$. The induction hypothesis gives us that as $n_{1},...,n_{l-2}\to\infty$ sequentially, the pre-activations $h^{l-1}_{i}$ tend to i.i.d Gaussian processes with covariance $\Sigma^{l-1}$. Then the inputs of the l-th layer are governed by:

where $x_{j}^{l-1}:=\phi(h^{l-1}_{j})$. When the input $x=x^{\prime}$, let

then $tr(BB^{T})=n_{l-1}$. By Lemma 1, we have that $lim_{n\rightarrow\infty}\sum^{n_{l-1}}_{j=1}\frac{1}{\sqrt{n_{l-1}}}W^{l}_{ij}x^{l-1}_{j}$ tends to $\mathcal{N}(0,lim_{n\rightarrow\infty}\frac{\sum_{i=1}^{n_{l-1}}(x^{l-1}_{i})^{2}}{n_{l-1}})$. As a result, $\{h_{i}^{l}\}$ are centered Gaussian variables. With the help of Lemma 3, we get the following covariance expression between two input data $x$ and $x^{\prime}$:

Since $h^{l-1}_{j}$, $h^{l-1}_{k}$ are independent for $j\neq k$ by the inductive hypothesis, we have

By the symmetry of the underlying index j, we can choose index $1$ as a representative:

We still need to verify the independence of $h_{i}^{l}$, $h_{j}^{l}$ for $i\neq j$. This follows from computing the covariance between $h_{i}^{l}$, $h_{j}^{l}$:

Note that we ignore the bias $b_{i}^{l}$, since they are independent with the weight parameters $\{W^{l}_{ij}\}$. By Lemma 3,

Since we know that as $n_{l-1}\rightarrow\infty$, $h_{i}^{l}(x)$, $h_{j}^{l}(x^{\prime})$ are Gaussian variables, zero covariance means they are independent.

With the explicit covariance formula in hand, we can prove by induction that the limiting covariance of the finite width output functions match with the covariance of the infinite Gaussian processes obtained above.

Let $\sigma^{2}(l,\mathcal{L},\alpha)[n]$ be the variance of the random variable $\gamma^{(l)}_{j}(\mathcal{L},\alpha)[n]$, then

By Lemma 3 and independence between l-th layer and (l-1)-th layer, we have

Once we prove that $\phi(h^{l-1}_{1}(x_{i})[n])\phi(h^{l-1}_{1}(x_{j})[n])$ is uniformly integrable with respect to $n$, the proof goes exactly the same as the proof of Lemma 17 in . To build the uniform integrability, we need the Lipschitz nonlinearity property of $\phi$.

is uniformly bounded as $n\rightarrow\infty$. This can be down by induction.

Since $\{W_{ij}\}$ is a scaled orthogonal matrix,

Applying the lipschitz nonlinearity property, we get

To verify condition (2) and (3) of Lemma 2 under the inductive hypothesis, we go through the same procedure as the proof of lemmas 15 and 16 in . We neglect the detail and conclude that by Lemma 3 and the inductive hypothesis,

Since it holds for arbitary linear functional $\alpha$, we have shown that $\{h_{i}^{l}(x)\}$ converge to independent Gaussian processes, where the covariance is specified by the recursion formula in Proposition 1. In the CNN case, adding the convolution index $\alpha$ won’t change the symmetry of index $j$ and Lemma 2 can be extended to CNN in the same way without any nontrivial change. Therefore we have simultaneously proved the following proposition for FNN and CNN:

Consider a FCN at orthogonal initialization, with a Lipschitz nonlinearity $\phi$, and in the limit as $n_{1},...,n_{L-1}\to\infty$, the pre-activations $h^{l}_{i}$, for $i=1,...,n_{l}$ for $l\in\{1,\dots,L\}$, tend to i.i.d centered Gaussian processes of covariance $\Sigma^{l}$ which is defined recursively by:

For a CNN at orthogonal initialization, and in the limit as $n_{1},...,n_{L-1}\to\infty$, the pre-activations $h^{l}_{i,\alpha}$ tend to Gaussian processes of covariance $\Sigma_{\alpha,\alpha^{\prime}}^{l}$ for $l\in\{1,\dots,L-1\}$, which is defined recursively by:

The covariance of last layer is a trace over $\alpha,\alpha^{\prime}\in m\times m$, because the last layer is fully connected:

A.2 NTK at Initialization

In the infinite-width limit, the neural tangent kernel which is random at networks with orthogonal initialization, converges in probability to a deterministic limit.

Consider a FCN at orthogonal initialization, with a Lipschitz nonlinearity $\phi$, and in the limit as the layers width $n_{1},...,n_{L-1}\to\infty$, the NTK $\Theta^{L}_{0}(x,x^{\prime})$, converges in probability to a deterministic limiting kernel:

The scalar kernel $\Theta^{L}_{\infty}(x,x^{\prime})$ is defined recursively by

For a CNN at orthogonal initialization, and in the infinite-channel limit, the NTK $\Theta^{L}_{0}(x,x^{\prime})$, converges in probability to a deterministic limiting kernel:

The scalar kernel $\Theta^{L}_{\infty}(x,x^{\prime})$ is given recursively by

We show the result in the framework of NTK parameterization, while the argument for standard parameterization can be derived in the same way with a similar result, see the details for the Gaussian initialization in . For the input layer $L=1$, taking the derivative with respect to $W^{1}_{ij}$, $b^{1}_{j}$, we have

From $(l-1)$-th layer to $l$-th layer, by the inductive hypothesis, as $n_{1},\dots,n_{l-2}\rightarrow\infty$, the pre-activations $h^{l-1}_{i}$ are i.i.d centered Gaussian with covariance $\Sigma^{l-1}$ and $\Theta^{l-1}_{ii^{\prime}}(x,x^{\prime})$ converges to:

Now we calculate the NTK for $l$ layer network and divide the parameters into two parts. The first part only involves the parameters of the $l$-th layer, and the other is the collection of parameters from previous $1,\dots,l-1$ layers. The first part of the NTK is given by

Note that in the Gaussian initialization,

for $i\neq j$ before taking the $n_{l-1}\rightarrow\infty$ limit. Thus for the $\frac{1}{n_{l-1}}$ scaling, we already observe that $\sum_{i}\frac{\sigma^{2}_{w}}{n_{l-1}}\phi(h^{l-1}_{i}(x))\phi(h^{l-1}_{i}(x^{\prime}))$ tends to its mean by the classical law of large numbers. If we take the limit sequentially, since $h_{i}^{l-1}(x)$ and $h_{j}^{l-1}(x)$ are independent Gaussian by the previous section, we know that it tends to $\Sigma^{l}(x,x^{\prime})$ for the same reason.

If we want to know how $\alpha_{l}$ behaves when the width tends to infinity simultaneously, the nonlinear activation would break the non-asymptotic independence. So we first assume that we work in the linear network category. Then

The right-hand side is independent of index $i$ as expected.

Let $\mu_{n}$ denote the mean of $\alpha_{l}[n]$, then by Chebyshev’s inequality, for $\forall\epsilon>0$,

Since $\lim_{n\rightarrow\infty}\mu_{n}=\Sigma_{l}$, we have

for n large enough. In conclusion, we have $\alpha_{l}[n]$ tends to $\Sigma_{l}(x,x^{\prime})$.

In the non-linear activation case, if the activation satisfies definition 2, we still have the asymptotic result

Therefore, by Chebyshev’s inequality, we get

By the induction hypothesis, the NTK of $(l-1)$-layer networks $\Theta^{l-1}_{kk^{\prime}}(x,x^{\prime})$ converges to a diagonal kernel as $n_{1},\dots,n_{l-2}\rightarrow\infty$. we denote the second part of NTK by,

Note that $\{W_{ik}^{l}W_{i^{\prime}k^{\prime}}^{l}\}$ is independent of $\Theta^{l-1}_{ii^{\prime}}(x,x^{\prime})\dot{\phi}(h^{l-1}_{i}(x))\dot{\phi}(h^{l-1}_{i^{\prime}}(x^{\prime}))$, this allows us to prove that it converges to a deterministic kernel by induction. By Lemma 3,

if $i\neq j\neq i^{\prime}\neq j^{\prime}$. This implies that $n_{l}(n_{l-1}-1)(n_{l-1}-2)(n_{l-1}-3)$ number of terms in the expansion are zero. Now we reorganize the left terms into three groups:

$i\neq i^{\prime}$: The only terms in the expansion that survive after taking the expectation are of the form: $W_{ik}W_{i^{\prime}k^{\prime}}W_{ik^{\prime}}W_{i^{\prime}k}.$ The expectation is separated into

There are $O(n_{l-1}\cdot(n_{l-1}-1))$ such terms in the expansion, by Lemma 3,

$(i=i^{\prime})\neq(j=j^{\prime})$: We need to estimate a fourth order moment

In each case, we get that $\gamma\sim O(\frac{1}{n_{l-1}})$. Thus we have

Note that the weights are drawn from the NTK parameterization, and by Lemma 3,

Therefore the number of terms with the same index i are of order $O(\frac{1}{n_{l-1}})$ and we get

By the inductive hypothesis and the integrability result from the last section,

Thus we obtain that second part of the NTK tends to $\sigma^{2}_{w}\Theta^{l-1}_{\infty}(x,x^{\prime})\dot{\Sigma}^{l}(x,x^{\prime})$ if we let the width go to infinity simultaneously. Combining the two parts, we have

by the same moment argument as before. For the second part, we still have

For the last (fully-connected) layer, we have,

Thus, we believe that the previous theorems also hold without assuming that the weight matrix is a square matrix.

A.3 NTK during Training

The NTK of networks with orthogonal initialization stays constant during training in the infinite-width limit. More precisely, we can give an upper bound for the change of parameters and NTK at wide width.

where $\hat{\Theta}_{t}$ are empirical kernels of network with finite width.

$\exists K>0$ s.t for every $c>0$, $\exists n_{c}$ s.t for $n\geq n_{c}$, we have the following bound:

We prove the result by induction for the standard parameterization, while the proof for ntk parameterization can be derived in the same way. For $l\geq 1$, let

As in the original proof for Gaussian initialization , there is a constant $K_{1}$, depending on $D$, $\sigma^{2}_{w}$, $\sigma_{b}^{2}$, and number of layers $L$ s.t with high probability over orthogonal initialization,

Note: there is a scaling factor $\frac{1}{\sqrt{n}}$ along with $||x^{l}(\theta,X)||_{2}$ in the Gaussian case, since from the input layer to the first layer, for the Gaussian initialization, we have

Decomposing the $J(\theta)$ into two parts, we have

Note: In our setting, since we want the orthogonal and Gaussian initialization have the same NTK, we need to scale the initialization for the input layer:

So we recover the same lipschitz constant as in the Gaussian initialization condition:

(local Lipschitzness of $J(\theta)$ with scale condition) Under the above scaling condition on the input layer, $\exists K>0$ s.t for every $c>0$, $\exists n_{c}$ s.t for $n\geq n_{c}$, we have the following bound: $\exists K>0$ s.t for every $c>0$, $\exists n_{c}$ s.t for $n\geq n_{c}$, we have the following bound:

Where $B(\theta_{0},R):=\{\theta:||\theta-\theta_{0}||_{2}<R\}$.

With this Corollary, the left proof for Theorem 3 is exact same as that of Theorem G.1 and Theorem G.2 for Gaussian initialization in .

For the CNN case, we refer the reader to , especially $P_{26}\sim P_{29}$. The constancy of NTK can be reduced to control of the Hessian norm. Most of the estimates are independent of the initialization, we only need to replace Lemma F.4 of into the following lemma:

From or , by Stein’s method, we have a total variation bound:

where $Z\sim N(0,\|x\|^{2})$. By the definition of total variation distance, we conclude that under $1-\frac{2\sqrt{3}}{n_{0}-1}$ probability, $\sum_{k=1}^{n_{0}}W_{ik}^{1}x_{k}\sim N(0,\|x\|^{2})$. Hence by the concentration inequality for Gaussian random variable, we have

Since $\frac{d_{0}}{m^{2}\|x\|^{2}}$ is of order $O(1)$, we know that

When $l\geq 2$, by induction, we can prove the lemma along the same line as $l=1$. ∎

Combing the lemma with the argument in , we have proved theorem 3 in the main text. We remark that to further extend other results of , we also need a Chi-Squared Stein method for nonlinear functional of the orthogonal matrix ensemble, which can be achieved by the exchangeable pair method in .

A.4 Additional Numerical Experiments

In the first experiment, summarized in Figure 5, we compare the train and test loss and accuracy with two different initialization, i.e., Gaussian and orthogonal weights using $D=256$ samples from the CIFAR10 dataset. To reduce noise, we averaged the results over 30 different instantiations of the networks. Figures 5(a,b) show the results of the experiments on the $L=3$, $n=400$ network with tanh activation, while figures 5(c,d) display the results for the $L=7$, $n=800$ network with tanh activation. All networks are optimized using gradient descent with a learning rate of $\eta=10^{-4}$ for $T=10^{4}$ steps. Consistent with our theoretical findings, the loss and accuracy of both networks are almost the same in the NTK regime.

As for learning dynamics outside the NTK regime, we preform more experiments on the convergence properties of networks in the large depth and large learning rate phase. Figure 6 not only confirms that the orthogonally initialized networks train faster and learn better than the orthogonally initialized network, but the performance gap between the two initializations can be very large as the depth increases.

References