On the Neural Tangent Kernel of Deep Networks with Orthogonal Initialization

1 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 [Nixon and Aguado, 2019], natural language processing [Devlin et al., 2018], to reinforcement learning [Mnih et al., 2015], and more. The mean field theory [Poole et al., 2016; Schoenholz et al., 2016] 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., 2016] discovered, and then empirically proved, is that there exists a critical initialization called the edge of chaos, allowing the correlation signal from the data to go infinitely far forward and preventing vanishing or exploding gradients.

∗Equal contribution.

†Contact Author

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 [Saxe et al., 2013]. However, what was not known was how to impose dynamical isometry in deep nonlinear networks. [Pennington et al., 2017] 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 [Tarnowski et al., 2018; Ling and Qiu, 2019], convolutional networks [Xiao et al., 2018], or recurrent networks [Chen et al., 2018] 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 [Pennington et al., 2017]. In straightforward scenarios, orthogonal initialization is usually enough to impose dynamical isometry in a linear network. The benefit of orthogonality in linear networks has been proven recently [Hu et al., 2020]. 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) [Jacot et al., 2018; Huang and Yau, 2019; Arora et al., 2019; Allen-Zhu et al., 2019; Du et al., 2018; Zou et al., 2020], 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 [Lee et al., 2019]. These dynamics are called NTK regime, or lazy training [Chizat et al., 2019]. As the learning rate gets larger or the network becomes deeper, outside 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.

2 RelatedWork

[Hu et al., 2020]'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, 2018] 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 inputoutput Jacobian, which partially explains why networks with greater isometric are able to train much faster.

[Jacot et al., 2018], 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., 2019] 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. While the original work of NTK is groundbreaking in producing an equation to predict the behavior of gradient descent in the NTK regime, it assumes the with goes to infinity in a sequential order. [Yang, 2019; Yang, 2020] strengthened the proof by taking the limit simultaneously. Besides, [Arora et al., 2019; Allen-Zhu et al., 2019; Du et al., 2018] have proven the same proprieties of NTK and global convergence of deep networks in non-asymptotic ways. However, all of these studies did not treat the orthogonal initialization as with our work.

3 Preliminaries

3.1 Networks and Parameterization

Suppose there are D training points denoted by {(xd, yd)}Dd=1 , where input X = (x1, . . . , xD) ∈ Rn0×D , and label Y = (y1, . . . , yD) ∈ RnL×D . We consider the following architectures:

Fully-Connected Network (FCN). Consider a fullyconnected network of L layers of widths nl, for l = 0, · · · , L, where l = 0 is the input layer and l = L is output layer. Following the typical nomenclature of literature, we denote synaptic weight and bias for the l-th layer by W l ∈ Rnl×nl-1 and bl ∈ Rnl , with a point-wise activations function φ : R → R. For each input x ∈ Rn0 , pre-activations and postactivations are denoted by hl (x) ∈ Rnl and xl (x) ∈ Rnl respectively. The information propagation for l ∈ {1, . . . , L} in this network is govern by,

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

where weight W l ∈ Rnl×nl-1×(2k+1) , and nl is the number of channels in the lth layer. The output layer of a CNN is processed with a fully-connected layer, fi(x) = hLi = ∑nL j=1 ∑α W L ij,αxL-1 j,α .

Standard parameterization requires the parameter set θ = {W l ij , bli } is an ensemble generated by, W l ij ∼ N(0, σ2w nl-1 ), bli ∼ N(0, σ2 b ), where σ2w and σ2 b are weight and bias variances. The variance of weights is scaled by the width of previous layer nl-1 to preserve the order of postactivations 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 = σw √nl-1 ωl ij , bli = σbβl i , where ωl ij , βl i ∼ N(0, 1).

3.2 Dynamical Isometry and Orthogonal Initialization

Consider the input-output Jacobian J = ∂hL ∂x0 = ∏Ll=1 DlW l , where hL is output function, x0 is input, and Dl is a diagonal matrix with elements Dlij = φ′ (hli )δ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 Dl contribute crucially to the dynamical isometry in non-linear networks [Pennington et al., 2017]. More precisely, the singular values of Dl 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 [Sohl-Dickstein et al., 2020]. The core idea is to write the width of the neural network in each layer in terms of an auxiliary parameter, nl = sNl. Instead of letting nl →∞, we adopt s as the limiting factor.

3.3 Neural Tangent Kernel

The neural tangent kernel (NTK) is originated from [Jacot et al., 2018] and defined as,

where function ft are the outputs of the network at training time t, i.e. ft(θ,X) = hLt (θ,X) ∈ RD×nL , and ∇θft(θ,X) = vec([∇θft(θ, x)]x∈X) ∈ RDnL . As such, the neural tangent kernel is formulated as a DnL ×DnL matrix. Let η be the learning rate, and L be the loss function. The ynamics of gradient flow for parameters and output function are given by,

This equation for ft has no substantial insight in studying the behavior of networks because Θt(X,X) varies with the time during training. Interestingly, as shown by [Jacot et al., 2018], the NTK Θt(X,X) converges to a deterministic kernel Θ∞(X,X) and does not change during training in the infinite-width limit, i.e. Θt(X,X) = Θ∞(X,X). As a result, the infinite width limit of the training dynamics are given by,

In the case of an MSE loss, L(y, f) = 12 ‖y - ft(θ, x)‖22 , the Equation (7) becomes a linear model with a solution,

4 Theoretical Results

4.1 An Orthogonally Initialized Network is a Gaussian Process in the Infinite Width Limit

As stated in [Lee et al., 2017; Matthews et al., 2018], the preactivation hli 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 [Jacot et al., 2018]. We extend this result to the orthogonal initialization across Fully Connected Networks (FCNs) and Convolutional Neural Networks (CNNs):

Theorem 1. Consider a FCN of the form (1) at orthogonal initialization, with a Lipschitz nonlinearity φ, and in the limit as n1, ..., nL-1 →∞, the pre-activations hli , for i = 1, ..., nl and l ∈ {1, . . . , L}, tend to i.i.d centered Gaussian processes of covariance Σl which is defined recursively by:

For a CNN of the form (2) at orthogonal initialization, and in the limit as n1, ..., nL-1 → ∞, the pre-activations hli,α tend to Gaussian processes of covariance Σ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.

4.2 Neural Tangent Kernel at Initialization

According to [Jacot et al., 2018], 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.

Theorem 2. Consider a FCN of the form (1) at orthogonal initialization, with a Lipschitz nonlinearity φ, and in the limit as the layers width n1, ..., nL-1 → ∞, the NTK ΘL0 (x, x′ ), converges in probability to a deterministic limiting kernel:

The scalar kernel ΘL∞(x, x′ ) is defined recursively by

For a CNN of the form (2) at orthogonal initialization, and in the limit as n1, ..., nL-1 → ∞, the NTK ΘL0 (x, x′ ), converges in probability to a deterministic limiting kernel:

The scalar kernel ΘL∞(x, x′ ) is given recursively by

Remark 1. 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 φ) and choice of hyper-parameters (σ2w and σ2 b ) are the same in the infinite-width limit.

4.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 [Jacot et al., 2018]. We find that the NTK of orthogonally initialized networks have the same property, which is demonstrated below in an asymptotic way,

Theorem 3. Assume that λmin(Θ∞) > 0 and ηcritical = λmin(Θ∞)+λmax(Θ∞) 2 . Let n = n1, ..., nL-1 be the width of hidden layers. Consider a FCN of the form (1) at orthogonal initialization, trained by gradient descent with learning rate η < ηcritical (or gradient flow). For every input x ∈ Rn0 with ‖x‖2 ≤ 1, with probability arbitrarily close to 1,

where ˆΘt are empirical kernels of networks with finite width.

For a CNN of the form (2) at orthogonal initialization, trained by gradient descent with learning rate η < ηcritical (or gradient flow), for every input x ∈ Rn0 with ‖x‖2 ≤ 1, and filter relative spatial location β ∈ {-k, . . . , 0, . . . , k}, with probability arbitrarily close to 1,

[Jacot et al., 2018] proved the stability of NTK under the assumption of global convergence of neural networks, while [Lee et al., 2019] provided a self-contained proof of both global convergence and stability of NTK simultaneously. In this work, we refer to the proof strategy from [Lee et al., 2019; Liu et al., 2020] 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 η = 1.0 on a subset of the MNIST dataset of D = 20. We measure changes of weights, empirical NTK after T = 215 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/√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/√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 [Huang and Yau, 2019], 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.

5 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 η = 10-3 for T = 105 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 [Hinton et al., 2012] optimizer with a small learning rate of η = 10-5 for T = 1.2×104 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 [Pennington et al., 2017].

Following [Pennington et al., 2017], we set σ2w = 1.05, and σ2 b = 2.01 × 10-5 , and φ(x) = 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

t

t

Figure 4: Learning dynamics measured by the optimization and generalization accuracy on train set and test set, for networks of depth L = 100 and width n = 400. We additionally average our results over 30 different instantiations of the network to reduce noise. Black curves are the results of orthogonal initialization, and red curves are performances of Gaussian initialization. (a) The training speed of an orthogonally initialized network is faster than that of a Gaussian initialized network. (b) On the test set, compared to the network with Gaussian initialization, the orthogonally initialized network not only learns faster but ultimately converges to a higher generalization performance.

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 τ when accuracy exceeds p. Figure 3 shows the steps of τ as a function of the learning rate of η 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×104 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 η = 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 [Lewkowycz et al., 2020; Li et al., 2019]. 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 [Hu et al., 2020].

6 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 orthogonallyinitialized 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 [Arora et al., 2019; Chizat et al., 2019]. Last, we find that orthogonal networks can outperform Gaussian networks in the large learning rate and depth on both train and test sets.