ReLU soothes the NTK condition number and accelerates optimization for
  wide neural networks
Rectified linear unit (ReLU), as a non-linear activation function, is well
known to improve the expressivity of neural networks such that any continuous
function can be approximated to arbitrary precision by a sufficiently wide
neural network. In this work, we present another interesting and important
feature of ReLU activation function. We show that ReLU leads to: {\it better
separation} for similar data, and {\it better conditioning} of neural tangent
kernel (NTK), which are closely related. Comparing with linear neural networks,
we show that a ReLU activated wide neural network at random initialization has
a larger angle separation for similar data in the feature space of model
gradient, and has a smaller condition number for NTK. Note that, for a linear
neural network, the data separation and NTK condition number always remain the
same as in the case of a linear model. Furthermore, we show that a deeper ReLU
network (i.e., with more ReLU activation operations), has a smaller NTK
condition number than a shallower one. Our results imply that ReLU activation,
as well as the depth of ReLU network, helps improve the gradient descent
convergence rate, which is closely related to the NTK condition number.
Chaoyue Liu, Like Hui
# ReLU soothes the NTK condition number and accelerates optimization for wide neural networks

###### Abstract

Rectified linear unit (ReLU), as a non-linear activation function, is well known to improve the expressivity of neural networks such that any continuous function can be approximated to arbitrary precision by a sufficiently wide neural network. In this work, we present another interesting and important feature of ReLU activation function. We show that ReLU leads to: _better separation_ for similar data, and _better conditioning_ of neural tangent kernel (NTK), which are closely related. Comparing with linear neural networks, we show that a ReLU activated wide neural network at random initialization has a larger angle separation for similar data in the feature space of model gradient, and has a smaller condition number for NTK. Note that, for a linear neural network, the data separation and NTK condition number always remain the same as in the case of a linear model. Furthermore, we show that a deeper ReLU network (i.e., with more ReLU activation operations), has a smaller NTK condition number than a shallower one. Our results imply that ReLU activation, as well as the depth of ReLU network, helps improve the gradient descent convergence rate, which is closely related to the NTK condition number.

## 1 Introduction

Non-linear activation functions, such as rectified linear unit (ReLU), are well known for their ability to increase the expressivity of neural networks. A non-linearly activated neural network can approximate any continuous function to arbitrary precision, as long as there are enough neurons in the hidden layers [12; 5; 11], while its linear counterpart - linear neural network, which has no non-linear activation functions applied, can only represent linear functions of the input. In addition, deeper neural networks, which have more non-linearly activated layers, have exponentially greater expressivity than shallower ones [32; 27; 29; 22; 33], indicating that the network depth promotes the power of non-linear activation functions.

In this paper, we present another interesting feature of ReLU activation function: ReLU improves data separation in the feature space of model gradient, and helps to decrease the condition number of neural tangent kernel (NTK) [14]. We also show that the depth of ReLU network further promotes this feature, namely, a deeper ReLU activated neural network has a better data separation and a smaller NTK condition number, than a shallower one.

Specifically, we first show the _better separation phenomenon_, i.e., the improved data separation for similar data in the model gradient feature space. We prove that, for an infinitely wide ReLU network \(f\) at its random initialization, a pair of data input vectors \(\mathbf{x}\) and \(\mathbf{z}\) that have similar directions (i.e., small angle \(\theta_{in}\) between \(\mathbf{x}\) and \(\mathbf{z}\)) become more separated in terms of their model gradient directions (i.e., angle \(\phi\) between \(\nabla f(\mathbf{x})\) and \(\nabla f(\mathbf{z})\) is larger than \(\theta_{in}\)). In addition, we show that a linear neural network \(\tilde{f}\), which is the same as the ReLU network but without activation functions, always keeps the model gradient angle \(\bar{\phi}\) the same as \(\theta_{in}\). With this comparison, we see that the ReLU activation
function results in a better data separation in the model gradient space. Furthermore, we show that deeper ReLU networks result in even better data separation.

We further show the _better conditioning_ property of ReLU, i.e., smaller NTK condition number, which is closely related to the better separation phenomenon. Specifically, we experimentally show that the NTK condition number for ReLU network of any depth is smaller than that of the data Gram matrix, which is equal to the NTK condition number for linear models and linear neural networks. Moreover, the NTK condition number monotonically decreases as the depth of ReLU network increases. We theoretically verify these findings for the cases of infinitely wide shallow ReLU network with arbitrary datasets and infinitely wide multi-layer ReLU network with dataset of size 2. The intuition is that, if there exists a pair of similar inputs \(\mathbf{x}\) and \(\mathbf{z}\) in the training set (i.e., the angle between \(\mathbf{x}\) and \(\mathbf{z}\) is small), which is usually the case for large datasets, then the Gram matrix and NTK of linear neural networks must have close-to-zero smallest eigenvalues, resulting in extremely large NTK condition numbers. The ReLU activation function make these similar data more separated (enlarges the small angles between data), hence it helps to increase the smallest eigenvalues of NTK, which in turn leads to a smaller NTK condition number.

**Condition number and optimization theory.** Recent optimization theories showed that the NTK condition number, or the smallest eigenvalue of NTK, controls the theoretical convergence rate of gradient descent algorithms on wide neural networks [8; 6; 20]. Combined with these theories, our findings imply that: (a), wide ReLU networks have faster convergence rate than linear models and linear neural networks, and (b), deeper wide ReLU networks have faster convergence rate than shallower ones. Experimentally, we indeed find that deeper ReLU networks converges faster than shallower ones.

Contributions.We summarize our contributions below. We find that:

* the ReLU activation function induces better separation between similar data in the feature space of model gradient. A larger depth of the ReLU network enhances this better separation phenomenon.
* ReLU activated neural networks have better NTK conditioning (i.e., smaller NTK condition number), than linear neural networks and linear models. A larger depth of the ReLU network further enhances this better NTK conditioning property.
* This better NTK conditioning property leads to faster convergence rate of gradient descent. We empirically verify this on various real world datasets.

The paper is organized as follow: in Section 2, we describe the setting and define the key quantities and concepts; in Section 3, we analyze linear neural networks as the baseline for comparison; in Section 4 and 5, we show our main results on the better separation and the better conditioning of ReLU neural networks, respectively; in Section 6, we discuss the connection between NTK condition number and convergence rates of gradient descent; in Section 7, we conclude the paper. The proofs of theorems and main corollaries can be found in the appendix.

### Related work

**NTK and its spectrum** have been extensively studied [18; 21; 9; 10; 4], since the discovery of constant NTK for infinitely wide neural networks [14]. For example, the NTK spectrum of an infinitely wide neural network is shown to be similar to that of Laplace kernel [10; 4], and can be computed [9]. Recent optimization theories on wide neural networks find connections between the theoretical convergence rate of gradient descent and the NTK condition number [8; 20].

**ReLU** has emerged to be the dominant choice of activation functions in neural networks used in practice, since [23; 16]. Since then, ReLU activated neural networks have received wide research attention, ranging from optimization [19; 8; 37], expressivity [11; 35; 33], generalization [36; 15; 3], etc. However, to the best of our knowledge, our work is the first to disclose the effect of ReLU on the NTK and its condition number, as well as on the convergence rate of gradient descent. In addition, we show the influence of network depth \(L\) on these features of ReLU.

We are aware of a prior work [2] which has results of similar flavor. It shows that the depth of a linear neural network may help to accelerate optimization via an implicit pre-conditioning of gradient descent. We note that this prior work is in an orthogonal direction, as its analysis is based on the
linear neural network, which is activation-free, while our work focus on the better-conditioning effect of ReLU activation function.

## 2 Setup and Preliminaries

Notations for general purpose.We denote the set \(\{1,2,\cdots,n\}\) by \([n]\). We use bold lowercase letters, e.g., \(\mathbf{v}\), to denote vectors, and capital letters, e.g., \(A\), to denote matrices. Given a vector, \(\|\cdot\|\) denotes its Euclidean norm. Inner product between two vectors is denoted by \(\langle\cdot,\cdot\rangle\). Given a matrix \(A\), we denote its \(i\)-th row by \(A_{i:}\), its \(j\)-th column by \(A_{:j}\), and its entry at \(i\)-th row and \(j\)-th column by \(A_{ij}\). We also denote the expectation (over a distribution) of a variable by \(\mathbb{E}[\cdot]\), and the probability of an event by \(\mathbb{P}[\cdot]\). For a model \(f(\mathbf{w};\mathbf{x})\) which has parameters \(\mathbf{w}\) and takes \(\mathbf{x}\) as input, we use \(\nabla f\) to denote its first derivative w.r.t. the parameters \(\mathbf{w}\), i.e., \(\nabla f:=\partial f/\partial\mathbf{w}\).

(Fully-connected) ReLU neural network.Let \(\mathbf{x}\in\mathbb{R}^{d}\) be the input, \(m_{l}\) be the width (i.e., number of neurons) of the \(l\)-th layer, \(W^{(l)}\in\mathbb{R}^{m_{l}\times m_{l-1}}\), \(l\in[L+1]\), be the matrix of the parameters at layer \(l\), and \(\sigma(z)=\max\{0,z\}\) be the ReLU activation function. A (fully-connected) ReLU neural network \(f\), with \(L\) hidden layers, is defined as:

\[\alpha^{(0)}(\mathbf{x})=\mathbf{x} \tag{1}\] \[\alpha^{(l)}(\mathbf{x})=\frac{\sqrt{2}}{\sqrt{m_{l}}}\sigma \left(W^{(l)}\alpha^{(l-1)}(\mathbf{x})\right),\ \ \forall l\in\{1,2,\cdots,L\},\] \[f(\mathbf{x})=W^{(L+1)}\alpha^{(L)}(\mathbf{x}).\]

We also denote \(\tilde{\alpha}^{(l)}(\mathbf{x})\triangleq\frac{\sqrt{2}}{\sqrt{m_{l}}}W^{(l )}\alpha^{(l-1)}(\mathbf{x})\). Following the NTK initialization scheme [14], these parameters are randomly initialized i.i.d. according to the normal distribution \(\mathcal{N}(0,1)\). The scaling factor \(\sqrt{2}/\sqrt{m_{l}}\) is introduced to normalize the hidden neurons [7]. We denote the collection of all the parameters by \(\mathbf{w}\).

In this paper, we typically set the layer widths as

\[m_{0}=d,\ \ m_{L+1}=1,\ \ and\ m_{l}=m,\ \ for\ l\in[L]. \tag{2}\]

and call \(m\) as the network width. We focus on the infinite network width limit, \(m\rightarrow\infty\). We also define the network depth \(L\) as the number of hidden layers.

Linear neural network.For a comparison purpose, we also consider a linear neural network \(\bar{f}\), which is the same as the ReLU neural network \(f\) (defined above), except that the activation function is the identity function \(\sigma(z)=z\) and that the scaling factor is \(1/\sqrt{m}\) (we adopt the network width setting in Eq.(2)):

\[\bar{\alpha}^{(0)}(\mathbf{x})=\mathbf{x},\ \ \bar{\alpha}^{(l)}(\mathbf{x})=\frac{1}{ \sqrt{m}}W^{(l)}\bar{\alpha}^{(l-1)}(\mathbf{x}),\ \ \forall l\in\{1,2,\cdots,L\},\ \bar{f}(\mathbf{x})=W^{(L+1)}\bar{\alpha}^{(L)}( \mathbf{x}). \tag{3}\]

Input feature and Gram matrix.Given a dataset \(\mathcal{D}=\{(\mathbf{x}_{i},y_{i})\}_{i=1}^{n}\), we denote its (input) feature matrix by \(X\), where each row \(X_{i}\). \(=\mathbf{x}_{i}^{T}\). The Gram matrix is defined as \(G=XX^{T}\in\mathbb{R}^{d\times d}\), with each \(G_{ij}=\mathbf{x}_{i}^{T}\mathbf{x}_{j}\).

Gradient feature and neural tangent kernel (NTK).Given a model \(f\) (e.g., a neural network) with parameters \(\mathbf{w}\), we consider the vector \(\nabla f(\mathbf{w};\mathbf{x})\) is the gradient feature for the input \(\mathbf{x}\). The NTK \(\mathcal{K}\) is defined as

\[\mathcal{K}(\mathbf{w};\mathbf{x}_{1},\mathbf{x}_{2})=\langle\nabla f(\mathbf{ w};\mathbf{x}_{1}),\nabla f(\mathbf{w};\mathbf{x}_{2})\rangle, \tag{4}\]

where \(\mathbf{x}_{1}\) and \(\mathbf{x}_{2}\) are two arbitrary network inputs. For a given dataset \(\mathcal{D}=\{(\mathbf{x}_{i},y_{i})\}_{i=1}^{n}\), there is a gradient feature matrix \(F\) such that each row \(F_{i}\).\((\mathbf{w})=\nabla f(\mathbf{w};\mathbf{x}_{i})\) for all \(i\in[n]\). The \(n\times n\) NTK matrix \(K(\mathbf{w})\) is defined such that its entry \(K_{ij}(\mathbf{w})\), \(i,j\in[n]\), is \(\mathcal{K}(\mathbf{w};\mathbf{x}_{i},\mathbf{x}_{j})\). It is easy to see that the NTK matrix

\[K(\mathbf{w})=F(\mathbf{w})F(\mathbf{w})^{T}. \tag{5}\]

Note that the NTK for a linear model reduces to the Gram matrix \(G\).

Recent discovery is that, when \(m\) is sufficiently large or infinite, the NTK and gradient feature becomes almost constant during training by gradient descent [14, 21]. Hence, it suffices to analyze these quantities only at the network initialization, which shall extend to all the optimization procedure.
Condition number.The _condition number_\(\kappa\) of a positive definite matrix \(A\) is defined as the ratio between its maximum eigenvalue and minimum eigenvalue:

\[\kappa=\lambda_{max}(A)/\lambda_{min}(A). \tag{6}\]

Embedding angle and model gradient angle.For a specific input \(\mathbf{x}\), we call the vector \(\alpha^{(l)}(\mathbf{x})\) as the \(l\)-embedding of \(\mathbf{x}\). We also call \(\nabla f\), i.e., the derivative of model \(f\) with respect to all its parameters, as the model gradient. In the following analysis, we frequently use the following concepts:

**Definition 2.1** (embedding angle and model gradient angle).: _Given two arbitrary inputs \(\mathbf{x},\mathbf{z}\in\mathbb{R}^{d}\), define the \(l\)-embedding angle, \(\theta^{(l)}(\mathbf{x},\mathbf{z})\triangleq\arccos\left(\frac{\langle\alpha ^{(l)}(\mathbf{x}),\alpha^{(l)}(\mathbf{z})\rangle}{\|\alpha^{(l)}(\mathbf{x })\|\|\alpha^{(l)}(\mathbf{z})\|}\right)\), as the angle between the \(l\)-embedding vectors \(\alpha^{(l)}(\mathbf{x})\) and \(\alpha^{(l)}(\mathbf{z})\), and the model gradient angle, \(\phi(\mathbf{x},\mathbf{z})\triangleq\arccos\left(\frac{\langle\nabla f( \mathbf{x}),\nabla f(\mathbf{x})\rangle}{\|\nabla f(\mathbf{x})\rangle\|\| \nabla f(\mathbf{x})(\mathbf{z})\|}\right)\), as the angle between the model gradient vectors \(\nabla f(\mathbf{x})\) and \(\nabla f(\mathbf{z})\)._

We also denote \(\theta^{(0)}\) by \(\theta_{in}\), as \(\theta^{(0)}\) is just the angle between the original inputs.

Connection between condition number and model gradient angle.The smallest eigenvalue and condition number of NTK are closely related to the smallest model gradient angle \(\min_{i,j\in[n]}\phi(\mathbf{x}_{i},\mathbf{x}_{j})\), through the gradient feature matrix \(F\). Think about the case if \(\phi(\mathbf{x}_{i},\mathbf{x}_{j})=0\) (i.e., \(\nabla f(\mathbf{x}_{i})\) is parallel to \(\nabla f(\mathbf{x}_{j})\)) for some \(i,j\in[n]\), then \(F\), hence NTK \(K\), is not full rank and the smallest eigenvalue \(\lambda_{min}(K)\) is zero, leading to an infinite condition number \(\kappa\). Similarly, if \(\min_{i,j\in[n]}\phi(\mathbf{x}_{i},\mathbf{x}_{j})\) is small, the smallest eigenvalue \(\lambda_{min}(K)\) is also small, and condition number \(\kappa\) is large, as stated in the following proposition (see proof in Appendix B).

**Proposition 2.2**.: _Consider a \(n\times n\) positive definite matrix \(A=BB^{T}\), where matrix \(B\in\mathbb{R}^{n\times d}\), with \(d>n\), is of full row rank. Suppose that there exist \(i,j\in[n]\) such that the angle \(\phi\) between vectors \(B_{i}\). and \(B_{j}\). is small, i.e., \(\phi\ll 1\), and that there exist constant \(C>c>0\) such that \(c\leq\|B_{k}\cdot\|\leq C\) for all \(k\in[n]\). Then, the smallest eigenvalue \(\lambda_{min}(A)=O(\phi^{2})\), and the condition number \(\kappa=\Omega(1/\phi^{2})\)._

Hence, a good data angle separation in the model gradient features, i.e., \(\min_{i,j\in[n]}\phi(\mathbf{x}_{i},\mathbf{x}_{j})\) not too small, is a necessary condition such that the condition number \(\kappa\) is not too large. Given this connection, in the following section we first discuss the separation of gradient features through the model gradient angles.

In the rest of the paper, we specifically refer the NTK matrix, NTK condition number, \(l\)-embedding angle and model gradient angle for the ReLU neural network as \(K\), \(\kappa\), \(\theta^{(l)}\) and \(\phi\), respectively, and refer their linear neural network counterparts as \(\bar{K}\), \(\bar{\kappa}\), \(\bar{\theta}^{(l)}\) and \(\bar{\phi}\), respectively. We also denote the condition number of Gram matrix \(G\) by \(\kappa_{0}\).

## 3 Linear neural network: the baseline for comparison

In this section, we analyze the linear neural network \(\bar{f}\), the linear counterpart of a ReLU network. Note that the only difference between the two architectures is the absence of the ReLU activation function in the linear neural network.

**Theorem 3.1**.: _Consider the linear neural network \(\bar{f}\) as defined in Eq.(3). In the limit of infinite network width \(m\to\infty\) and at network initialization \(\mathbf{w}_{0}\), the following relations hold:_

* _for any input_ \(\mathbf{x}\in\mathbb{R}^{d}\)_:_ \(\|\bar{\alpha}^{(l)}(\mathbf{x})\|=\|\mathbf{x}\|\)_,_ \(\forall l\in[L]\)_; and_ \(\|\nabla f(\mathbf{w}_{0};\mathbf{x})\|=(L+1)\|\mathbf{x}\|\)_._
* _for any inputs_ \(\mathbf{x},\mathbf{z}\in\mathbb{R}^{d}\)_:_ \(\bar{\theta}^{(l)}(\mathbf{x},\mathbf{z})=\theta_{in}(\mathbf{x},\mathbf{z})\)_,_ \(\forall l\in[L]\)_; and_ \(\bar{\phi}(\mathbf{x},\mathbf{z})=\theta_{in}(\mathbf{x},\mathbf{z})\)_._

This theorem states that, without a non-linear activation function, both the feature embedding maps \(\alpha^{(l)}:\mathbf{x}\mapsto\alpha^{(l)}(\mathbf{x})\) and the model gradient map \(\nabla f:\mathbf{x}\mapsto\nabla f(\mathbf{x})\) fail to change the geometrical relationship between any data samples. For any input pairs, the embedding angles \(\bar{\theta}^{(l)}\) and \(\bar{\phi}\) remain the same as the input angle \(\theta_{in}\). Therefore, it is not surprising that the NTK of a linear network is the same as the Gram matrix (up to a constant factor), as formally stated in the following corollary.
**Corollary 3.2** (NTK condition number of linear networks).: _Consider a linear neural network \(\bar{f}\) as defined in Eq.(3). In the limit of infinite network width \(m\to\infty\) and at network initialization, the NTK matrix \(\bar{K}=(L+1)^{2}G\). Moreover, \(\bar{\kappa}=\kappa_{0}\)._

This corollary tells that, for a linear neural network, regardless of its depth \(L\), the NTK condition number \(\bar{\kappa}\) is always equal to the condition number \(\kappa_{0}\) of the Gram matrix \(G\).

## 4 ReLU network has better data separation in model gradient space

In this section, we show that the ReLU non-linearity helps data separation in the model gradient space. Specifically, for two arbitrary inputs \(\mathbf{x}\) and \(\mathbf{z}\) with small \(\theta_{in}(\mathbf{x},\mathbf{z})\), we show that the model gradient angle \(\phi(\mathbf{x},\mathbf{z})\) is strictly larger than \(\theta_{in}(\mathbf{x},\mathbf{z})\), implying a better angle separation of the two data points in the model gradient space. Moreover, we show that the model gradient angle \(\phi(\mathbf{x},\mathbf{z})\) monotonically increases with the number of layers \(L\), indicating that deeper network (more ReLU non-linearity) has better angle separation.

Embedding vectors and embedding angles.We start with investigating the relations among the \(l\)-embedding vectors \(\alpha^{(l)}\) and the embedding angles \(\theta^{(l)}\).

**Theorem 4.1**.: _Consider the ReLU network \(f\) defined in Eq.(1) at its initialization. In the infinite network width limit \(m\to\infty\), for all \(l\in[L]\), the following relations hold:_

* _for any input_ \(\mathbf{x}\in\mathbb{R}^{d}\)_,_ \(\|\alpha^{(l)}(\mathbf{x})\|=\|\alpha^{(l-1)}(\mathbf{x})\|\)_;_
* _for any two inputs_ \(\mathbf{x},\mathbf{z}\in\mathbb{R}^{d}\)_,_ \[\cos\theta^{(l)}(\mathbf{x},\mathbf{z})=\frac{\pi-\theta^{(l-1)}(\mathbf{x}, \mathbf{z})}{\pi}\cos\theta^{(l-1)}(\mathbf{x},\mathbf{z})+\frac{1}{\pi}\sin \theta^{(l-1)}(\mathbf{x},\mathbf{z}).\] (7)

The theorem states that, during forward propagation, the \(l\)-embedding vectors for each input keeps unchanged in magnitude, and the embedding angles \(\theta^{(l)}\) between any two inputs obey a closed form relation from layer to layer. Therefore, we have the following corollary.

**Corollary 4.2**.: _Consider the same ReLU neural network \(f\) as in Theorem 4.1. For all \(l\in[L]\), the \(l\)-embedding \(\alpha^{(l)}\) satisfies:_

* _For any input_ \(\mathbf{x}\in\mathbb{R}^{d}\)_,_ \(\|\alpha^{(l)}(\mathbf{x})\|=\|\mathbf{x}\|\)_;_
* _Define function_ \(g:[0,\pi)\to[0,\pi)\) _as_ \[g(z)=\arccos\left(\frac{\pi-z}{\pi}\cos z+\frac{1}{\pi}\sin z\right),\] (8) _and let_ \(g^{l}(\cdot)\) _be the_ \(l\)_-fold composition of_ \(g(\cdot)\)_. Given any two inputs_ \(\mathbf{x},\mathbf{z}\in\mathbb{R}^{d}\)_,_ \[\theta^{(l)}(\mathbf{x},\mathbf{z})=g^{l}\left(\theta^{(0)}(\mathbf{x}, \mathbf{z})\right).\] (9)

We see that the the embedding angles between inputs are governed by the function \(g\) defined in Eq.(8). Please see Appendix A for the plot of the function and detailed discussion about its properties. As a highlight, \(g\) has the following property: \(g\) is approximately the identity function \(g(z)\approx z\) for small \(z\), i.e., \(z\ll 1\). This property directly implies the following corollary.

**Corollary 4.3**.: _Given any inputs \(\mathbf{x},\mathbf{z}\) such that \(\theta_{in}(\mathbf{x},\mathbf{z})=o(1/L)\), for each \(l\in[L]\), the \(l\)-embedding angle \(\theta^{(l)}(\mathbf{x},\mathbf{z})\) can be expressed as_

\[\theta^{(l)}(\mathbf{x},\mathbf{z})=\theta_{in}(\mathbf{x},\mathbf{z})-\frac{l }{3\pi}(\theta_{in}(\mathbf{x},\mathbf{z}))^{2}+o\left((\theta_{in}(\mathbf{x },\mathbf{z}))^{2}\right).\]

We see that, at the small angle regime \(\theta_{in}=o(1/L)\), the embedding angles \(\theta^{(l)}\) at any layer \(l\) is the same as the input angle \(\theta_{in}\) at the lowest order. In addition, the higher order corrections are always negative making \(\theta^{(l)}<\theta_{in}\). We also note that the correction term \(\Delta\theta^{(l)}\triangleq\theta^{(l)}-\theta_{in}\) is linearly dependent on layer \(l\) at its lowest order.
Model gradient angle.Now, we investigate the model gradient angle \(\phi\) and its relation with the embedding angles \(\theta^{(l)}\) and input angle \(\theta_{in}\), for the ReLU network.

**Theorem 4.4**.: _Consider the ReLU network defined in Eq.(1) with \(L\) hidden layers and infinite network width \(m\). Given two arbitrary inputs \(\mathbf{x}\) and \(\mathbf{z}\), the angle \(\phi(\mathbf{x},\mathbf{z})\) between the model gradients \(\nabla f(\mathbf{x})\) and \(\nabla f(\mathbf{z})\) satisfies_

\[\cos\phi(\mathbf{x},\mathbf{z})=\frac{1}{L+1}\sum_{l=0}^{L}\left[\cos\theta^{ (l)}(\mathbf{x},\mathbf{z})\prod_{l^{\prime}=l}^{L-1}(1-\theta^{(l^{\prime})}( \mathbf{x},\mathbf{z})/\pi)\right]. \tag{10}\]

_Moreover, \(\|\nabla f(\mathbf{x})\|=(L+1)\|\mathbf{x}\|\), for any \(\mathbf{x}\)._

Better data separation with ReLU.Comparing with Theorem 3.1 for linear neural networks, we see that the non-linear ReLU activation only affects the relative direction, but not the the magnitude, of the model gradient. Combining Theorem 4.4 with Corollary 4.2, we get the relation between \(\phi\) and the input angle \(\theta_{in}\). Figure 1 plots \(\phi\) as a function of \(\theta_{in}\) for different network depth \(L\).

The **key observation** is that: for relatively small input angles (say, \(\theta_{in}<60^{\circ}\)), the model gradient angle \(\phi\) is always greater than the input angle \(\theta_{in}\). This suggests that, after the mapping \(\nabla f:\mathbf{x}\mapsto\nabla f(\mathbf{x})\) from the input space to model gradient space, data inputs becomes more (directionally) separated, if they are similar in the input space (i.e., with small \(\theta_{in}\)). Comparing to the linear neural network case, where \(\bar{\phi}(\mathbf{x},\mathbf{z})=\theta_{in}(\mathbf{x},\mathbf{z})\) as in Theorem 3.1, we see that the ReLU non-linearity results in a better angle separation \(\phi(\mathbf{x},\mathbf{z})>\bar{\phi}(\mathbf{x},\mathbf{z})\) for similar data.

Another important observation is that: deeper ReLU networks lead to larger model gradient angles, when \(\theta_{in}<60^{\circ}\). This indicates that deeper ReLU networks, which has more layers of ReLU non-linear activation, makes the model gradient more separated between inputs. Note that, in the linear network case, the depth does not affect the gradient angle \(\bar{\phi}\).

We theoretically confirm these two observations in the regime of small input angle \(\theta_{in}=o(1/L)\), by the following theorem.

**Theorem 4.5** (Better separation with ReLU).: _Consider two network inputs \(\mathbf{x},\mathbf{z}\in\mathbb{R}^{d}\), with small input angle \(\theta_{in}(\mathbf{x},\mathbf{z})=o(1/L)\), and the ReLU network defined in Eq.(1) with \(L\) hidden layers and infinite network width \(m\). At the network initialization, the angle \(\phi(\mathbf{x},\mathbf{z})\) between the model gradients \(\nabla f(\mathbf{x})\) and \(\nabla f(\mathbf{z})\) satisfies_

\[\cos\phi(\mathbf{x},\mathbf{z})=\left(1-\frac{L}{2\pi}\theta_{in}+o(\theta_{ in})\right)\cos\theta_{in}. \tag{11}\]

Figure 1: **Model gradient angles \(\phi\) vs. input angle \(\theta_{in}\) (according to Theorem 4.4).** Linear neural networks, of any depth \(L\), always have \(\bar{\phi}=\theta_{in}\), as the black dash line showed. ReLU neural networks with various depths have better data separation \(\phi>\theta_{in}\) for similar data (i.e., small \(\theta_{in}\)). Moreover, deeper ReLU networks have better separation than shallow ones for similar data. All neural networks are infinitely wide.
Noticing the negative sign within the factor \(\left(1-\frac{1}{2\pi}\theta_{in}+o(\theta_{in})\right)\), we know that the factor is less than \(1\) and we obtain that: \(\phi(\mathbf{x},\mathbf{z})>\theta_{in}(\mathbf{x},\mathbf{z})=\phi(\mathbf{x}, \mathbf{z})\). Noticing the depth \(L\) dependence of this factor, we also get that: the deeper the ReLU network (i.e., larger \(L\)) is, the larger \(\phi\) is, in the regime \(\theta_{in}=o(1/L)\).

**Remark 4.6** (Separation in distance).: _Indeed, the better angle separation discussed above implies a better separation in Euclidean distance as well. This can be easily seen by recalling from Theorem 4.4 that the model gradient mapping \(\nabla f\) preserves the norm (up to a universal factor \(L+1\))._

We also point out that, Figure 1 indicates that for large input angles (say \(\theta_{in}>60^{\circ}\)) the model gradient angle \(\phi\) is always large (greater than \(60^{\circ}\)). Hence, non-similar data never become similar in the model gradient feature space.

## 5 Smaller condition number of NTK for ReLU network

In this section, we show both theoretically and experimentally that, with the ReLU non-linear activation function, a ReLU neural network has a smaller condition number \(\kappa\) of NTK than its linear counterpart - linear neural network. Moreover, for larger network depth \(L\), the NTK condition number \(\kappa\) of a ReLU neural network is generically smaller.

### Theoretical analysis

Theoretically, we show these findings for the following cases: two-layer ReLU network with any dataset; deep ReLU network with dataset of size \(2\). In both cases, we take the infinite width limit.

Shallow ReLU neural network.Now, let's look at the NTK of a shallow ReLU neural network

\[f(W;\mathbf{x})=\frac{\sqrt{2}}{\sqrt{m}}\mathbf{v}^{T}\sigma(W\mathbf{x}). \tag{12}\]

Here, we fix \(\mathbf{v}\) to random initialization and let \(W\) trainable.

With the presence of the ReLU activation function, the better data separation theorem (Theorem 4.5) together with the connection between condition number and angle separation ( Proposition 2.2) suggest that the NTK of this ReLU network has a smaller condition number \(\kappa\) than the Gram matrix, as well as the NTK of the linear neural network. The following theorem confirms this expectation.

**Theorem 5.1**.: _Consider the ReLU network in Eq.(12) in the limit \(m\to\infty\) and at initialization. The smallest eigenvalue \(\lambda_{min}(K)\) of its NTK is larger than that of the Gram matrix: \(\lambda_{min}(K)>\lambda_{min}(G)\). Moreover, the NTK condition number is less than that of the Gram matrix: \(\kappa<\kappa_{0}\)._

Deep ReLU neural network.For deep ReLU network, we consider the special case where the dataset \(\mathcal{D}\) consists two samples, with a small angle between their input vectors.

**Theorem 5.2**.: _Consider the ReLU neural network \(f\) as defined in Eq.(1) in the infinite width limit \(m\to\infty\) and at initialization. Consider the dataset \(\mathcal{D}=\{(\mathbf{x}_{1},y_{1}),(\mathbf{x}_{2},y_{2})\}\) with the input angle \(\theta_{in}\) between \(\mathbf{x}_{1}\) and \(\mathbf{x}_{2}\) small, \(\theta_{in}=o(1/L)\). Then, the NTK condition number \(\kappa<\kappa_{0}\). Moreover, for two ReLU neural networks \(f_{1}\) of depth \(L_{1}\) and \(f_{2}\) of depth \(L_{2}\) with \(L_{1}>L_{2}\), we have \(\kappa_{f_{1}}<\kappa_{f_{2}}\)._

This theorem also shows that a ReLU network has a better conditioned NTK than a linear model or a linear neural network. Moreover, the depth of ReLU network enhances this better conditioning. These results also confirm that the ReLU activation helps to decrease the NTK condition number.

### Experimental evidence

In this subsection, we experimentally show that the phenomena of better data separation and better NTK conditioning widely happen in practice.

Dataset.We use the following datasets: synthetic dataset, MNIST [17], FashionMNIST (f-MNIST) [34], SVHN [24] and Librispeech [26]. The synthetic data consists of \(2000\) samples which are randomly drawn from a 5-dimensional Gaussian distribution with zero-mean and unit variance.
The MNIST, f-MNIST and SVHN datasets are image datasets where each input is an image. The Librispeech is a speech dataset including \(100\) hours of clean speeches. In the experiments, we use a subset of Librispeech with \(50,000\) samples, and each input is a 768-dimensional vector representing a frame of speech audio and we follow [13] for the feature extraction.

Models.For each of the datasets, we use a ReLU activated fully-connected neural network architecture to process. The ReLU network has \(L\) hidden layers, and has \(512\) neurons in each of its hidden layers. The ReLU network uses the NTK parameterization and initialization strategy (see [14]). For each dataset, we vary the network depth \(L\) from \(0\) to \(10\). Note that \(L=0\) corresponding to the linear model case. In addition, for comparison, we use a linear neural network, which has the same architecture with the ReLU network except the absence of activation function.

Results.For each of the experimental setting, we evaluate both the smallest pairwise model gradient angle \(\min_{i,j\in[n]}\phi(\mathbf{x}_{i},\mathbf{x}_{j})\) and the NTK condition number \(\kappa\), at the network initialization. We take \(5\) independent runs over \(5\) random initialization seeds, and report the average. The results are shown in Figure 2. As one can easily see from the plots, a ReLU network (depth \(L=1,2,\cdots,10\)) always have a better separation of data (i.e., larger smallest pairwise model gradient angle), and a better NTK conditioning (i.e., smaller NTK condition number), than its corresponding linear network (compare the solid line and dash line of the same color). Furthermore, the monotonically decreasing NTK condition number shows that a deeper ReLU network have a better conditioning of NTK.

## 6 Optimization acceleration

Recently studies showed strong connections between the NTK condition number and the theoretical convergence rate of gradient descent algorithms on wide neural networks [8; 6; 31; 1; 37; 25; 20]. In [8; 6], the authors derived the worse case convergence rates explicitly in terms of the smallest eigenvalue of NTK \(\lambda_{min}(K)\), \(L(\mathbf{w}_{t})\leq(1-\eta\lambda_{min}(K)/2)^{t}L(\mathbf{w}_{0})\), where \(L\) is the square loss function and \(t\) is the time stamp of the algorithm. Later on, in [20], the NTK condition number is explicitly involved in the theoretical convergence rate:

\[L(\mathbf{w}_{t})\leq(1-\kappa^{-1})^{t}L(\mathbf{w}_{0}). \tag{13}\]

Although \(\kappa\) is evaluated on the whole optimization path, all these theories used the fact that NTK is almost constant for wide neural networks and an evaluation at the initialization \(\mathbf{w}_{0}\) is enough.

Figure 2: **Better separation (left) and better NTK conditioning (right) of ReLU network.** Solid lines are for ReLU networks, and dash lines are for linear networks. **Left:** ReLU network works better in separating similar data, while linear network remains similar to a linear model. **Right:** ReLU network has better conditioning of NTK than linear network and linear model. Note that \(L=0\) (without any hidden layer) corresponds to the case of a linear model, and the NTK in this case is the Gram matrix.
As a smaller NTK condition number (or larger smallest eigenvalue of NTK) implies a faster theoretical convergence rate, our findings suggest that: (a), wide ReLU networks have faster convergence rate than linear models and linear neural networks, and (b), deeper wide ReLU networks have faster convergence rate than shallower ones.

We experimentally verify this implication on the convergence speed. Specifically, we train the ReLU networks, with depth \(L\) ranging from \(1\) to \(10\), for the datasets MNIST, f-MNIST and Librispeech. For all the training tasks, we use cross entropy loss as the objective function and use mini-batch stochastic gradient descent (SGD) of batch size \(500\) to optimize. For each task, we find its optimal learning rate by grid search. On MNIST and f-MNIST, we train \(500\) epochs, and on Librispeech, we training \(2000\) epochs.

The curves of training loss against epochs are shown in Figure 3. We observe that, for all these datasets, a deeper ReLU network always converges faster than shallower ones. This is consistent with the theoretical prediction that the deeper ReLU network, which has smaller NTK condition number, has faster theoretical convergence rate.

## 7 Conclusion and discussions

In this work, we showed the key role that ReLU (as a non-linear activation function) played in data separation: for any two inputs that are "directionally" close (small \(\theta_{in}\)) in the original input space, a ReLU activated neural network makes them better separated in the model gradient feature space (\(\phi>\theta_{in}\)), while a linear neural network, which has no non-linear activation, keeps the same separation (\(\bar{\phi}=\theta_{in}\)). As a consequence, the NTK condition number \(\kappa\) for the ReLU neural network is smaller than those for the linear model and linear neural networks. Moreover, we showed that the NTK condition number \(\kappa\) is further decreased if the ReLU network is going deeper. The smaller NTK condition number suggests a faster gradient descent convergence rate for deep ReLU neural network.

Finite network width.Our theoretical analysis is based the setting of infinite network width. We note that, for finite network width and at network initialization, the analysis only differs by a zero-mean noisy term, wherever we took the limit \(m\rightarrow\infty\). This noisy term scales down to \(0\) as the network width increases. We believe our main results still hold for finite but large network width.

Infinite depth.In this work, we focused on the finite depth scenario which is the more interesting case from a practical point of view. Our small angle regime analysis (Corollary 4.3, Theorem 4.5 and 5.2) do not directly extend to the infinite depth case. But, as Theorem 4.4 and Figure 1 indicate, the \(\phi(\theta_{in})\) function seems to converge to a step function when \(L\rightarrow\infty\), which implies orthogonality between each pair of model gradient vectors, hence a NTK condition number being \(1\). This is consistent with the prior knowledge that NTK converges to \(1\) in the infinite depth limit [28].

Other activation functions.In this work, we focused on ReLU, which is the mostly used activation function in practice. We believe similar results also hold for other non-linear activation functions. We leave this direction as a future work.

Figure 3: **Training curve of ReLU networks with different depths.** On each of these datasets, we see that deeper ReLU network always converges faster than shallower ones.