# Efficient NTK using Dimensionality Reduction

Nir Ailon
Technion
<EMAIL_ADDRESS>&&&Supratim Shit
Technion
<EMAIL_ADDRESS>Corresponding author

###### Abstract

Recently, neural tangent kernel (NTK) has been used to explain the dynamics of
learning parameters of neural networks, at the large width limit. Quantitative
analyses of NTK give rise to network widths that are often impractical and
incur high costs in time and energy in both training and deployment. Using a
matrix factorization technique, we show how to obtain similar guarantees to
those obtained by a prior analysis while reducing training and inference
resource costs. The importance of our result further increases when the input
points’ data dimension is in the same order as the number of input points.
More generally, our work suggests how to analyze large width networks in which
dense linear layers are replaced with a low complexity factorization, thus
reducing the heavy dependence on the large width.

_K_ eywords Deep learning $\cdot$ Neural Network $\cdot$ Neural Tangent Kernel
$\cdot$ Over Parameterized Network $\cdot$ Dimensionality Reduction

##  1 Introduction

Deep learning techniques have overwhelmed the world of machine learning in the
last decade. Among them, one of the popular techniques is the deep neural
network. Usually, in practice, these networks are overparameterized, and they
proved to be extremely successful in various applications like image
recognition [11, 2], natural language understanding [21, 7], speech
recognition [10] etc. Due to the non-convexity of the cost function and the
complex dynamics of the training process, it is incredibly challenging to
explain the reason for their success.

In a recent breakthrough, Jacot et. al., [17] discovered that the dynamics of
learning the network parameters at the infinite width limit regime tend to be
that of a kernelized linear regression setting, with respect to a feature map
known as the _Neural Tangent Kernel_ (NTK). Du et. al., [16] improved the
original analysis by providing a sharper non-asymptotic bound for finite
network width. Further, using the NTK idea, Arora et. al., [18] showed why the
training error goes to zero even when adding adversarial noise to the input
data. They also showed why these overparameterized networks have benign
overfitting behaviour. These results follow when the network width is very
high, $O(\mbox{poly}(n))$, where $n$ is the number of input points.
Unfortunately, in a regime where input points are in $\mathbb{R}^{O(n)}$
dimensional space, the number of trainable parameters is infeasible.
Therefore, training such a network would incur high time and energy costs. We
refer the reader to [23] to understand the adverse societal and environmental
effects of such costs.

In another line of work, matrix factorization has been proven to be a useful
tool in reducing costs of both training and deploying deep networks, e.g. [28,
31, 20]. In this work, we provide a connection between matrix factorization
ideas and NTK analysis, suggesting that the resources required to obtain NTK
guarantees may not necessarily be as high as reported in previous work.

Following the work of [16, 18], we study a network containing a linear
transformation of the input, followed by a RELU activation function, followed
by a scalar valued linear function, and follow the dynamic of this network
with respect to an MSE loss function. The NTK matrix given rise to by this
network is well understood following the work of [16, 18]. Here we replace the
first linear layer with a factorization thereof, and show that by carefully
choosing the factorization parameters we can approximately preserve the NTK
structure and hence the network dynamics.

The factorization we choose, following ideas from [31] is a random
dimensionality reduction matrix (Johnson-Lindenstrauss) followed by a
trainable matrix which increases the dimension. The first matrix is constant,
and does not change after the (random initialization). Using well known
properties of JL matrices, we show how the dynamics of the factorized network,
which is much cheaper to execute, mimic those of the original matrix.

##  2 Related Work

Neural tangent kernel (NTK) was first introduced by Jacot et al. [17]. They
showed the way toward a better understanding of the dynamics of the training
phase of a deep network. These techniques are also known as lazy training. It
has been studied for many network architectures such as CNNs, RNNs, and
ResNets [16, 19, 24]. In [16] author shows the required network size and a
learning rate to achieve optimal training loss on a fully connected neural
network using gradient descent. Given a distribution of input points, the [19]
showed the required number of parameters to ensure the desired generalization
loss. Often these networks end up in an overparameterized regime with the
number of parameters being very large than the size of the input data. These
results are interesting in their own right, but they are still impractical in
various applications.

Randomized and deterministic techniques are extensively used to improve
scalability of various learning algorithms such as clustering [27, 33, 36,
37], classification [30, 35, 34], regression [25, 26, 38] and deep learning
[22, 29]. For a fully connected neural network, [15] uses sampling techniques
to prune a learned network by removing less important nodes in the network.
Their sampling technique is based on a sensitivity framework. In [32], the
authors showed how sketching, or sampling techniques can be used to randomly
select features in a fully connected neural network with the Relu activation
function. Thereby reducing the overall network parameters,

Matrix factorization is an elegant way of representing low rank matrices
succinctly. For a fixed and known matrix, there are many deterministic and
randomized techniques for obtaining low rank approximations [5, 12, 13]. The
most well knows method is by the Johnson-Lindenstrauss transformation [1]. It
was originally designed to preserve distances (or inner products) between
vectors. It is also a method for decomposing large matrices as a composition
of (random) dimensionality reduction followed by a transformation in lower
dimensional space. There are various versions of this transformation [6, 4,
14]. A faster version of the Johnson-Lindenstrauss transformation was
introduced by Ailon et. al. [3] which has been extensively used in various
learning algorithms [8, 9].

##  3 Preliminary

A matrix is represented by a bold capital letter, e.g., $\mathbf{M}$. The
$i^{th}$ row and $j^{th}$ column of the matrix are represented by
$\mathbf{M}_{i}$ and $\mathbf{M}^{j}$ respectively. The $(i,j)^{th}$ entry of
the matrix $\mathbf{M}$ is represented as $\mathbf{M}_{i,j}$. The 2-norm of a
vector $\mathbf{v}$ and spectral norm for matrix $\mathbf{M}$ is represented
by $\|\mathbf{v}\|$ and $\|\mathbf{M}\|$ respectively. We represent the ReLU
function by $\sigma(\cdot)$ which is defined as $\sigma(z)=\max\\{0,z\\}$ for
$z\in\mathbb{R}$. We use $\mathbbm{1}\\{\mathcal{E}\\}$ as an indicator
variable, which is $1$ if the event $\mathcal{E}$ is true else $0$.

Consider a two layer fully connected neural network which takes input from
$\mathbb{R}^{d}$ and returns a scalar output in $\mathbb{R}$. The first layer
has $m$ neurons, represented by a weight matrix
$\mathbf{W}\in\mathbb{R}^{m\times d}$. We use the ReLU function as the
activation function in the first layer. The second (latent) layer is denoted
by $\mathbf{v}\in\mathbb{R}^{m}$. The norm of every input points to the
network is assumed to be $1$. Given $n$ samples of input as
$\\{\mathbf{X}_{i},\mathbf{y}_{i}\\}_{i=1}^{n}$. We represent the set
$\\{\mathbf{X}_{i}\\}$ by a matrix $\mathbf{X}\in\mathbb{R}^{n\times d}$ and
$\\{\mathbf{y}_{i}\\}$ by an $n$-dimensional vector
$\mathbf{y}\in\mathbb{R}^{n}$. Based on some network parameters $\mathbf{W}$,
and $\mathbf{v}$ for every $\mathbf{X}_{i}$ we get $\mathbf{u}_{i}$ as
follows,

${}\mathbf{u}_{i}=\frac{1}{\sqrt{m}}\sum_{r\leq
m}\mathbf{v}_{r}\cdot\sigma(\mathbf{W}_{r}\mathbf{X}_{i}^{T}).$ (1)

We represent the set $\\{\mathbf{u}_{i}\\}$ by an $n$-dimensional vector
$\mathbf{u}$. We learn $\mathbf{W}$ and $\mathbf{v}$ such that it minimizes
$\|\mathbf{y}-\mathbf{u}\|^{2}$.

In this paper we use the properties of the Johnson-Lindenstrauss
transformation, which has been stated below for completeness.

######  Definition 3.1 (Johnson-Lindenstrauss [1]).

Let $\mathbf{M}\in\mathbb{R}^{n\times d}$ represents $n$ points in
$\mathbb{R}^{d}$, $\delta\in(0,1)$. Let $\mathbf{S}$ be a $d\times k$ random
matrix such that $k=O(\log(n/\delta))$. For each $i\in[d]$ and $j\in[k]$,
$\mathbf{S}_{i,j}\sim\mathcal{N}(0,1/\sqrt{d})$. Then for every pair of
$i,j\in[n]$ we have
$|\mathbf{M}_{i}\mathbf{M}_{j}^{T}-\mathbf{M}_{i}\mathbf{S}\mathbf{S}^{T}\mathbf{M}_{j}^{T}|\leq\|\mathbf{M}_{i}\|\cdot\|\mathbf{M}_{j}\|$
with probability at least $1-\delta$.

The matrix $\mathbf{S}$ is also known as the JL matrix.

###  3.1 Neural Tangent Kernel

Recent work by Du et. al. [16] showed that in an overparameterized (number of
neurons $\Omega(n^{6})$) neural network, the dynamics of the gradient descent
for a small learning rate $O(\lambda_{0}n^{-2})$ behaves like a linear model.
The linear model is captured by a kernel matrix called the neural tangent
kernel (NTK). Given $\mathbf{X}$ the NTK matrix
$\mathbf{H}^{\infty}\in\mathbb{R}^{n\times n}$ is a gram matrix with the
$(i,j)^{th}$ term defined as follows,

$\mathbf{H}_{i,j}^{\infty}=\mathbb{E}_{\mathbf{w}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{m})}\left[\mathbf{X}_{i}\mathbf{X}_{j}^{T}\mathbbm{1}\\{\mathbf{w}\mathbf{X}_{i}^{T}\geq
0,\mathbf{w}\mathbf{X}_{j}^{T}\geq 0\\}\right].$

The main result in the paper reads as follows:

######  Theorem 3.1.

Assume $\lambda_{0}=\lambda_{\min}(\mathbf{H}^{\infty})>0$. For
$\delta\in(0,1)$, number of neurons
$m=\Omega(\frac{n^{6}}{\lambda_{0}^{4}\delta^{2}})$, set
$\mathbf{W}_{r}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ for every $r\in[m]$ and
we set $\eta=O(\frac{\lambda_{0}}{n^{2}})$ then during the gradient descent,
we have the following with probability at least $1-\delta$ at every
$t=0,1,\ldots$,

$\|\mathbf{u}(t)-\mathbf{y}\|^{2}\leq\left(1-\frac{\eta\lambda_{0}}{2}\right)^{t}\|\mathbf{u}(0)-\mathbf{y}\|^{2}$
(2)

Now, assuming that we fix the second layer parameter, in order to ensure that
$\|\mathbf{u}(t)-\mathbf{y}\|^{2}\leq\varepsilon\|\mathbf{u}(0)-\mathbf{y}\|^{2}$
for some small $\varepsilon\in(0,1)$, the gradient descent has to run for
$t=\Omega(\frac{\log(\varepsilon)}{\log(1-(\lambda_{0}/n)^{2})})$ steps. The
time taken to update $\mathbf{W}(t+1)$ from $\mathbf{W}(t)$ is $\Omega(nmd)$,
i.e., $\Omega\left(\frac{n^{7}d}{\lambda_{0}^{4}\delta^{2}}\right)$. Now, in
the regime where $d=\Theta(n)$, it will take
$\Omega\left(\frac{n^{8}\log(\varepsilon)}{\lambda_{0}^{4}\delta^{2}\log(1-(\lambda_{0}/n)^{2})}\right)$
running time to achieve the above training loss. Notice the dependence on $m$
(number of neurons) and the dimension of the input vectors.

In this paper, we apply dimensionality reduction techniques to reduce the
effective number of neurons $m$ and the dimension of input points $d$ and
thereby improving the overall running time. We present our results in two
parts. First, we show how to reduce the dimension of the input points and how
it affects the training process. Next, we show how to reduce the number of
effective neurons in the latent layer.

##  4 Reducing Input Dimension

Let $\mathbf{X}\in\mathbb{R}^{n\times d}$ be a set of $n$ input points in
$\mathbb{R}^{d}$. Recall that we assume each point is normalized to a unit
vector. We apply a Johnson-Lindenstrauss [1] transformation $\mathbf{C}$ on
the input $\mathbf{X}$ before feeding it into the network. Every index of
$\mathbf{C}$ is an i.i.d. sample of $\mathcal{N}(0,1/\sqrt{\ell})$, where
$\ell=O(\log(n/\delta))$ for some $\delta\in(0,1)$. This transformation is
drawn once before training, and not updated in the training process. We use a
low rank weight matrix $\mathbf{W}$ which is defined as
$\mathbf{W}=\mathbf{B}\mathbf{C}$, where
$\mathbf{B}\in\mathbb{R}^{m\times\ell}$ and $\mathbf{C}$ is the $\ell\times d$
JL transformation matrix. Now, based on the above low rank weight matrix
$\mathbf{W}$ for every input $\mathbf{X}_{i}$ and network outputs
$\mathbf{u}_{i}$ as follows,

$\mathbf{u}_{i}=\frac{1}{\sqrt{m}}\sum_{r\leq
m}\mathbf{v}_{r}\cdot\sigma(\mathbf{B}_{r}\mathbf{C}\mathbf{X}_{i}^{T}).$

Notice, how $\mathbf{B}$ and $\mathbf{C}$ replace $\mathbf{W}$ in equation (1)
in a plug and play fashion. Now we discuss the effect of this on the training
phase. Fixing, $\mathbf{C}$ (the dimension reduction layer) and $\mathbf{v}$
the goal is to learn $\mathbf{B}$ that minimizes,

$\min_{\mathbf{B}}\frac{1}{2}\sum_{i\leq
n}\left(\mathbf{y}_{i}-\mathbf{u}_{i}\right)^{2}.$ (3)

The above cost is a function of $\mathbf{B}$ and it can be represented in
vector form as $\Phi(\mathbf{B})=\frac{1}{2}\|\mathbf{y}-\mathbf{u}\|^{2}$.
The matrix $\mathbf{C}$ is first initialized as mentioned above. The parameter
$\mathbf{v}$ is also drawn as a random vector sampled from $\\{-1,+1\\}^{m}$.
Both $\mathbf{C}$ and $\mathbf{v}$ are fixed and do not vary during training.

We now randomly initialize the trainable parameters $\mathbf{B}(0)$ at time
$0$, where every entry is an i.i.d. sample from $\mathcal{N}(0,1)$. Now with a
fixed $\mathbf{C}$ and $\mathbf{v}$, we learn $\mathbf{B}$ using gradient
descent on the above loss function equation (3) based on a small learning rate
$\eta$. Every row of $\mathbf{B}_{r}$, $\forall r\in[m]$ are updated as
follows,

$\displaystyle\mathbf{B}_{r}(t+1)$ $\displaystyle=$
$\displaystyle\mathbf{B}_{r}(t)-\eta\frac{\partial\Phi(\mathbf{B}(t))}{\partial\mathbf{B}_{r}(t)}$
$\displaystyle=$
$\displaystyle\mathbf{B}_{r}(t)-\frac{\eta}{\sqrt{m}}\sum_{j\leq
n}(\mathbf{y}_{j}-\mathbf{u}_{j}(t))\mathbf{v}_{r}\mathbf{X}_{j}\mathbf{C}^{T}\mathbbm{1}\\{\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{j}^{T}\geq
0\\}$

The NTK matrix $\mathbf{H}^{\infty}\in\mathbb{R}^{n\times n}$ is now defined
as follows:

$\mathbf{H}_{i,j}^{\infty}=\mathbb{E}_{\mathbf{b}}\left[\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T}\mathbbm{1}\\{\mathbf{b}\mathbf{C}\mathbf{X}_{i}^{T}\geq
0,\mathbf{b}\mathbf{C}\mathbf{X}_{j}^{T}\geq 0\\}\right]$ (4)

Here $\mathbf{b}$ is a random vector from
$\mathcal{N}(\mathbf{0},\mathbf{I})$. The above term is the expected dot
product between two vectors which corresponds to the change in the output for
the inputs $\mathbf{X}_{i}$ and $\mathbf{X}_{j}$ with respect to $\mathbf{b}$,
i.e.,
$\langle\frac{\partial\mathbf{u}_{i}}{\partial\mathbf{b}},\frac{\partial\mathbf{u}_{j}}{\partial\mathbf{b}}\rangle$
(for details on this derivation, refer to [18]).

Now for a random initialization $\mathbf{B}(0)$, fixed $\mathbf{C}$ and
$\mathbf{v}$ we define $\mathbf{H}(0)\in\mathbb{R}^{n\times n}$ such that
$\forall i,j\in[n]$,

${}\mathbf{H}_{i,j}(0)=\frac{1}{m}\sum_{r\leq
m}\left[\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T}\mathbbm{1}\\{\mathbf{B}_{r}(0)\mathbf{C}\mathbf{X}_{i}^{T}\geq
0,\mathbf{B}_{r}(0)\mathbf{C}\mathbf{X}_{j}^{T}\geq 0\\}\right]$ (5)

######  Lemma 4.1.

For $\mathbf{H}^{\infty}$ and $\mathbf{H}(0)$ as defined above, with at least
$1-\delta$ probability,

$\|\mathbf{H}^{\infty}-\mathbf{H}(0)\|\leq
O\left(\frac{n^{2}\log(n/\delta)}{m}\right)$

###### Proof.

Every single term in the difference matrix
$\mathbf{H}_{i,j}^{\infty}-\mathbf{H}_{i,j}(0)$ is

$\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{i}\left(\mathbb{E}[\mathbbm{1}\\{\mathbf{b}\mathbf{C}\mathbf{X}_{i}^{T}\geq
0,\mathbf{b}\mathbf{C}\mathbf{X}_{j}^{T}\geq 0\\}]-\frac{1}{m}\sum_{r\leq
m}\mathbbm{1}\\{\mathbf{B}_{r}(0)\mathbf{C}\mathbf{X}_{i}^{T}\geq
0,\mathbf{B}_{r}(0)\mathbf{C}\mathbf{X}_{j}^{T}\geq 0\\}\right)$

Here the randomness is over $\mathbf{b}\in\mathbb{R}^{\ell}$, which takes $m$
i.i.d random samples from $\mathcal{N}(\mathbf{0},\mathbf{I}_{\ell})$. For all
$r\in[m]$ the random variable
$\mathbbm{1}\\{\mathbf{B}_{r}(0)\mathbf{C}\mathbf{X}_{i}^{T}\geq
0,\mathbf{B}_{r}(0)\mathbf{C}\mathbf{X}_{j}^{T}\geq 0\\}\in[0,1]$. Now since
the above term is the difference between the expectation of a random variable
and the empirical average of $m$ i.i.d samples of the random variable, so by
Hoeffding’s inequality we get the following with at least $1-\delta^{\prime}$
probability,

$|\mathbf{H}_{i,j}^{\infty}-\mathbf{H}_{i,j}(0)|\leq|\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T}|\frac{\log(2/\delta^{\prime})}{2m}.$

Next, for a random $\mathbf{C}\in\mathbb{R}^{\ell\times d}$ such that each
entry of $\mathbf{C}$ is an i.i.d. sample from $\mathcal{N}(0,1/\sqrt{\ell})$
where $\ell=O(\frac{\log(n/\delta)}{\varepsilon^{2}})$, for all $i,j\in[n]$ we
have
$\mathbb{P}(|\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T}-\mathbf{X}_{i}\mathbf{X}_{j}^{T}|\geq\varepsilon)<\delta/(2n^{2})$.
Now for $\varepsilon=O(1)$, setting $\delta^{\prime}=\delta/(2n^{2})$ and
taking a union bound over all pairs of $(i,j)\in[n]$ we get the following with
at least $1-\delta$,

$\displaystyle\|\mathbf{H}^{\infty}-\mathbf{H}(0)\|\leq\sum_{i,j}|\mathbf{H}_{i,j}^{\infty}-\mathbf{H}_{i,j}(0)|\leq
O\left(n^{2}\frac{\log(2n^{2}/\delta)}{2m}\right)$

∎

Now with the following lemma, we show that $\mathbf{H}(0)$ does not change too
much during the training period.

######  Lemma 4.2.

If $\mathbf{B}_{1},\ldots,\mathbf{B}_{m}$ are random vectors whose entries are
iid sample from $\mathcal{N}(0,1)$ then with probability $1-\delta$, the
following holds.

For any set of vectors $\mathbf{B}_{1}(0),\ldots,\mathbf{B}_{m}(0)$ if
$\|\mathbf{B}_{r}\mathbf{C}-\mathbf{B}_{r}(0)\mathbf{C}\|\leq
O\left(\frac{\delta\lambda_{0}}{n^{2}}\right)\stackrel{{\scriptstyle\Delta}}{{=}}R$,
then with $\mathbf{B}$ the matrix $\mathbf{H}\in\mathbb{R}^{n\times n}$ whose
$(i,j)^{th}$ entries are defined as (5) satisfies
$\|\mathbf{H}-\mathbf{H}(0)\|\leq\frac{\lambda_{0}}{4}$ and
$\lambda_{\min}(\mathbf{H})\geq\frac{\lambda_{0}}{4}$.

###### Proof.

For a fixed $\mathbf{C}$ and a random $\mathbf{B}_{r}$ we define a random
vector $\mathbf{W}_{r}=\mathbf{B}_{r}\mathbf{C}$. Similarly we also define
$\mathbf{W}_{r}(0)=\mathbf{B}_{r}(0)\mathbf{C}$. Now consider the following
event for every input $\mathbf{X}_{i}\in\mathbf{X}$ and neuron $r\in[m]$,

$\mathcal{E}_{r,i}=\\{\exists\mathbf{B}_{r}:\|\mathbf{W}_{r}-\mathbf{W}_{r}(0)\|\leq
R,\mathbbm{1}\\{\mathbf{W}_{r}(0)\mathbf{X}_{i}^{T}\geq
0\\}\neq\mathbbm{1}\\{\mathbf{W}_{r}\mathbf{X}_{i}^{T}\geq 0\\}\\}.$

The above event happens if and only if
$|\mathbf{W}_{r}(0)\mathbf{X}_{i}^{T}|\leq R$. Now notice that for a fixed
$\mathbf{C}$ and a randomly initialized $\mathbf{B}_{r}(0)$ we have
$z\sim\mathcal{N}(0,1)$ where
$z=\mathbf{B}_{r}(0)\mathbf{C}\mathbf{X}_{i}^{T}$. So
$\mathbb{P}(\mathcal{E}_{r,i})=\mathbb{P}(|z|\leq
R)\leq\int_{-R}^{R}\frac{1}{\sqrt{2\pi}}\exp(-x^{2})dx\leq\frac{\sqrt{2}R}{\sqrt{\pi}}$.
Now we bound deviation of every entry $(i,j)\in[n]\times[n]$ as follows,

$\displaystyle\mathbb{E}[|\mathbf{H}_{i,j}-\mathbf{H}_{i,j}(0)|]$
$\displaystyle=$
$\displaystyle\mathbb{E}\bigg{[}O\left(\frac{1}{m}\right)\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T}\sum_{r\leq
m}(\mathbbm{1}\\{\mathbf{B}_{r}\mathbf{C}\mathbf{X}_{i}^{T}\geq
0,\mathbf{B}_{r}\mathbf{C}\mathbf{X}_{j}^{T}\geq 0\\}$ $\displaystyle-$
$\displaystyle\mathbbm{1}\\{\mathbf{B}_{r}(0)\mathbf{C}\mathbf{X}_{i}^{T}\geq
0,\mathbf{B}_{r}(0)\mathbf{C}\mathbf{X}_{j}^{T}\geq 0\\})\bigg{]}$
$\displaystyle\leq$ $\displaystyle O\left(\frac{1}{m}\right)\sum_{r\leq
m}\mathbb{E}[\mathbbm{1}\\{\mathcal{E}_{r,i}\\}\cup\mathbbm{1}\\{\mathcal{E}_{r,j}\\}]\leq
O\left(\frac{2\sqrt{2}R}{\sqrt{\pi}}\right)\leq O(R)$

So, $\mathbb{E}[\sum_{i,j}|\mathbf{H}_{i,j}-\mathbf{H}_{i,j}(0)|]\leq
O(n^{2}R)$ and by Markov inequality we have
$\sum_{i,j}|\mathbf{H}_{i,j}-\mathbf{H}_{i,j}(0)|\leq
O\left(\frac{n^{2}R}{\delta}\right)$ with at least $1-\delta$ probability.
Therefore with at least $1-\delta$ probability
$\|\mathbf{H}-\mathbf{H}(0)\|\leq O\left(\frac{n^{2}R}{\delta}\right)$.
Finally,

$\lambda_{\min}(\mathbf{H})\geq\lambda_{\min}(\mathbf{H}(0))-O\left(\frac{n^{2}R}{\delta}\right)\geq\frac{\lambda_{0}}{2}$

∎

Now we state our main theorem of this section describing the required network
width and the learning rate to get to the desired training loss.

######  Theorem 4.1.

Assume $\lambda_{0}=\lambda_{\min}(\mathbf{H}^{\infty})>0$. Fix
$\delta\in(0,1)$, and take the number of neurons in the latent layer $m$ to be
$\Omega\left(\frac{n^{6}}{\lambda_{0}^{4}\delta^{2}}\right)$. $\mathbf{B}(0)$
is a random matrix whose entries are i.i.d. sample from $\mathcal{N}(0,1)$ and
we let the learning rate $\eta$ be $O\left(\frac{\lambda_{0}}{n^{2}}\right)$
then during the gradient descent we have the following with probability at
least $1-\delta$ for $t=0,1,\ldots$,

$\|\mathbf{u}(t)-\mathbf{y}\|^{2}\leq\left(1-\frac{\eta\lambda_{0}}{2}\right)^{t}\|\mathbf{u}(0)-\mathbf{y}\|^{2}$
(6)

###### Proof.

We prove it by induction. The induction hypothesis is the following,

$\|\mathbf{u}(t)-\mathbf{y}\|^{2}\leq\left(1-\frac{\eta\lambda_{0}}{2}\right)^{t}\|\mathbf{u}(0)-\mathbf{y}\|^{2}$
(7)

The following corollary follows from the hypothesis.

######  Corollary 4.1.

If the above hypothesis is true then we have the following for all $r\in[m]$.

$\|\mathbf{W}_{r}(t+1)-\mathbf{W}_{r}(0)\|\leq\frac{4\sqrt{n}\|\mathbf{u}(0)-\mathbf{y}\|}{\sqrt{m}\lambda_{0}}\stackrel{{\scriptstyle\Delta}}{{=}}R^{\prime}$
(8)

###### Proof.

$\displaystyle\|\mathbf{W}_{r}(t+1)-\mathbf{W}_{r}(0)\|$ $\displaystyle\leq$
$\displaystyle\sum_{t^{\prime}\leq
t}\left\|-\eta\frac{\partial\Phi(\mathbf{B}(t^{\prime}))}{\partial\mathbf{B}_{r}(t^{\prime})}\mathbf{C}\right\|$
(9) $\displaystyle=$ $\displaystyle\eta\sum_{t^{\prime}\leq
t}\left\|\sum_{j\leq
n}\frac{(\mathbf{y}_{j}-\mathbf{u}_{j}(t^{\prime}))\mathbf{v}_{r}\mathbf{X}_{j}\mathbf{C}^{T}\mathbf{C}(\mathbbm{1}\\{\mathbf{B}_{r}(t^{\prime})\mathbf{C}\mathbf{X}_{j}^{T}\geq
0\\})}{\sqrt{m}}\right\|$ $\displaystyle\leq$
$\displaystyle\eta\sum_{t^{\prime}\leq t}\left\|\sum_{j\leq
n}\frac{(\mathbf{y}_{j}-\mathbf{u}_{j}(t^{\prime}))\mathbf{X}_{j}^{T}\mathbf{C}^{T}\mathbf{C}}{\sqrt{m}}\right\|$
$\displaystyle\leq$ $\displaystyle\frac{\eta}{\sqrt{m}}\sum_{t^{\prime}\leq
t}\left(\sum_{j\leq
n}|\mathbf{y}_{j}-\mathbf{u}_{j}(t^{\prime})|\right)\max_{j}\|\mathbf{X}_{j}\mathbf{C}^{T}\mathbf{C}\|$
(10) $\displaystyle\leq$ $\displaystyle
O\left(\frac{\eta}{\sqrt{m}}\right)\sum_{t^{\prime}\leq t}\left(\sum_{j\leq
n}|\mathbf{y}_{j}-\mathbf{u}_{j}(t^{\prime})|\right)$ (11) $\displaystyle\leq$
$\displaystyle O\left(\frac{\eta}{\sqrt{m}}\right)\sum_{t^{\prime}\leq
t}\sqrt{n}\|\mathbf{y}-\mathbf{u}(t^{\prime})\|$ (12) $\displaystyle\leq$
$\displaystyle O\left(\frac{\eta\sqrt{n}}{\sqrt{m}}\right)\sum_{t^{\prime}\leq
t}\left(1-\frac{\eta\lambda_{0}}{2}\right)^{t^{\prime}/2}\|\mathbf{u}(0)-\mathbf{y}\|$
$\displaystyle\leq$ $\displaystyle
O\left(\frac{\eta\sqrt{n}}{\sqrt{m}}\right)\sum_{t^{\prime}=0}^{\infty}\left(1-\frac{\eta\lambda_{0}}{2}\right)^{t^{\prime}/2}\|\mathbf{u}(0)-\mathbf{y}\|=O\left(\frac{4\sqrt{n}}{\sqrt{m}\lambda_{0}}\right)\|\mathbf{u}(0)-\mathbf{y}\|$

In (9) we upper bound
$\mathbbm{1}\\{\mathbf{b}_{r}(t^{\prime})\mathbf{C}\mathbf{x}_{j}^{T}\geq
0\\}$ by $1$. In (10) we used holder’s inequality. In (11) we use the upper
bound $|\mathbf{x}_{i}\mathbf{C}^{T}\mathbf{C}|\leq O(1)$. Using
Cauchy–Schwartz we get (12). In (4) we used our hypothesis. ∎

Notice that the hypothesis is trivially true for the base case $t=0$. Suppose
it holds true for $t^{\prime}=0,1,\ldots,t$ and now we show that it also holds
for $t^{\prime}=t+1$. For this we assume the same event $\mathcal{E}_{r,i}$
for all input $\mathbf{x}_{i}$ and neuron $r$. We know that,
$R=O\left(\frac{\delta\lambda_{0}}{n^{2}}\right)$. Let
$\mathbf{S}_{i}=\\{r\in[m]:\mathbbm{1}\\{\mathcal{E}_{r,i}\\}=0\\}$ and
$\mathbf{S}_{i}^{\perp}=[m]\backslash\mathbf{S}_{i}$. Note that
$\mathcal{E}_{r,i}$ is non empty if and only if
$|\mathbf{W}_{r}(0)\mathbf{X}_{i}^{T}|\leq R$. Now for fixed $\mathbf{C}$ as
defined above and a random $\mathbf{B}_{r}(0)$, $z\sim\mathcal{N}(0,1)$ where
$z=\mathbf{B}_{r}(0)\mathbf{C}\mathbf{X}_{i}^{T}$. So we have,

$\mathbb{P}(\mathcal{E}_{r,i})=\mathbb{P}(|z|\leq
R)\leq\int_{-R}^{R}\frac{1}{\sqrt{2\pi}}exp(-x^{2})dx\leq\frac{2R}{\sqrt{2\pi}}\leq
R.$

So,

$\mathbb{E}[|\mathbf{S}_{i}^{\perp}|]=\sum_{r\leq
m}\mathbb{P}(\mathcal{E}_{r,i})\leq mR.$

Further, $\mathbb{E}[\sum_{i\leq n}|\mathbf{S}_{i}^{\perp}|]\leq mnR$. So by
using Markov inequality we have the following with at least $1-\delta$
probability.

$\sum_{i\leq n}|\mathbf{S}_{i}^{\perp}|\leq\frac{mnR}{\delta}$

Here $C$ is some appropriate constant. Now consider the following term,

$\displaystyle\mathbf{u}_{i}(t+1)-\mathbf{u}_{i}(t)$ $\displaystyle=$
$\displaystyle\frac{1}{\sqrt{m}}\sum_{r\leq
m}v_{r}\left[\sigma(\mathbf{B}_{r}(t+1)\mathbf{C}\mathbf{X}_{i}^{T})-\sigma(\mathbf{W}_{r}(t)\mathbf{C}\mathbf{X}_{i}^{T})\right]$
$\displaystyle=$ $\displaystyle\frac{1}{\sqrt{m}}\sum_{r\leq
m}v_{r}\left[\sigma\left(\left(\mathbf{B}_{r}(t)-\eta\frac{\partial\Phi(\mathbf{B}(t))}{\partial\mathbf{B}_{r}(t)}\right)\mathbf{C}\mathbf{X}_{i}^{T}\right)-\sigma(\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{i}^{T})\right]$
$\displaystyle=$ $\displaystyle I^{i}_{1}+I^{i}_{2}$

Here,

$\displaystyle I^{i}_{1}$ $\displaystyle=$
$\displaystyle\frac{1}{\sqrt{m}}\sum_{r\in\mathbf{S}_{i}}\mathbf{v}_{r}\left[\sigma\left(\left(\mathbf{B}_{r}(t)-\eta\frac{\partial\Phi(\mathbf{B}(t))}{\partial\mathbf{B}_{r}(t)}\right)\mathbf{C}\mathbf{X}_{i}^{T}\right)-\sigma(\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{i}^{T})\right]$
$\displaystyle I^{i}_{2}$ $\displaystyle=$
$\displaystyle\frac{1}{\sqrt{m}}\sum_{r\in\mathbf{S}_{i}^{\perp}}\mathbf{v}_{r}\left[\sigma\left(\left(\mathbf{B}_{r}(t)-\eta\frac{\partial\Phi(\mathbf{B}(t))}{\partial\mathbf{B}_{r}(t)}\right)\mathbf{C}\mathbf{X}_{i}^{T}\right)-\sigma(\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{i}^{T})\right]$

First we upper bound $I^{i}_{2}$.

$\displaystyle|I^{i}_{2}|$ $\displaystyle=$
$\displaystyle\left|\frac{1}{\sqrt{m}}\sum_{r\in\mathbf{S}_{i}^{\perp}}\mathbf{v}_{r}\left[\sigma\left(\left(\mathbf{B}_{r}(t)-\eta\frac{\partial\Phi(\mathbf{B}(t))}{\partial\mathbf{B}_{r}(t)}\right)\mathbf{C}\mathbf{X}_{i}^{T}\right)-\sigma(\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{i}^{T})\right]\right|$
$\displaystyle\leq$
$\displaystyle\left|\frac{1}{\sqrt{m}}\sum_{r\in\mathbf{S}_{i}^{\perp}}\left[-\eta\left(\frac{\partial\Phi(\mathbf{B}(t))}{\partial\mathbf{B}_{r}(t)}\right)\mathbf{C}\mathbf{X}_{i}^{T}\right]\right|$
$\displaystyle\leq$
$\displaystyle\frac{\eta}{\sqrt{m}}\left|\sum_{r\in\mathbf{S}_{i}^{\perp}}\sum_{j\leq
n}\frac{(\mathbf{y}_{j}-\mathbf{u}_{j}(t))\mathbf{v}_{r}\mathbf{X}_{j}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{i}^{T}\mathbbm{1}\\{\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{j}\geq
0\\}}{\sqrt{m}}\right|$ $\displaystyle\leq$
$\displaystyle\frac{\eta}{\sqrt{m}}\left|\sum_{r\in\mathbf{S}_{i}^{\perp}}\max_{j}|\mathbf{X}_{j}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{i}^{T}|\sum_{j\leq
n}|\frac{(\mathbf{y}_{j}-\mathbf{u}_{j}(t))\mathbf{v}_{r}\mathbbm{1}\\{\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{j}^{T}\geq
0\\}|}{\sqrt{m}}\right|$ $\displaystyle\leq$
$\displaystyle\frac{2\eta|\mathbf{S}_{i}^{\perp}|}{m}\sum_{j\leq
n}|(\mathbf{y}_{j}-\mathbf{u}_{j}(t))\mathbf{v}_{r}|\leq\frac{2\eta\sqrt{n}|\mathbf{S}_{i}^{\perp}|}{m}\|\mathbf{y}-\mathbf{u}(t)\|$

Since ReLU is $1$-Lipschitz function and $|\mathbf{v}_{r}|=1$, so we get (4).

Now we analyze $I^{i}_{1}$. From Corollary 4.1 we have
$\|\mathbf{W}_{r}(t+1)-\mathbf{W}_{r}(0)\|\leq R^{\prime}$ and
$\|\mathbf{W}_{r}(t+1)-\mathbf{W}_{r}(t)\|\leq R^{\prime}$. Now to always
ensure $R^{\prime}<R$ we have,

$\displaystyle\frac{4\sqrt{n}}{\sqrt{m}\lambda_{0}}\|\mathbf{y}-\mathbf{u}(0)\|$
$\displaystyle\leq$ $\displaystyle
O\left(\frac{\delta\lambda_{0}}{n^{2}}\right)$ (15)
$\displaystyle\frac{4\sqrt{n}}{\sqrt{m}}O(\sqrt{n})$ $\displaystyle\leq$
$\displaystyle O\left(\frac{\delta\lambda_{0}^{2}}{n^{2}}\right)$ (16)
$\displaystyle\frac{O(n)}{\sqrt{m}}$ $\displaystyle\leq$ $\displaystyle
O\left(\frac{\delta\lambda_{0}^{2}}{n^{2}}\right)$ (17)
$\displaystyle\frac{O(n^{3})}{\sqrt{m}}$ $\displaystyle\leq$ $\displaystyle
O\left(\delta\lambda_{0}^{2}\right)$ (18) $\displaystyle m$
$\displaystyle\geq$ $\displaystyle
O\left(\frac{n^{6}}{\lambda_{0}^{4}\delta^{2}}\right)$ (19)

As, $R^{\prime}<R$, so for $I^{i}_{1}$ we have
$\mathbbm{1}\\{\mathbf{W}_{r}(t+1)\mathbf{X}_{i}^{T}\\}=\mathbbm{1}\\{\mathbf{W}_{r}(t)^{T}\mathbf{x}_{i}^{T}\\}$.
Now $I^{i}_{1}=\frac{du_{i}(t)}{dt}$, so we express $I^{i}_{1}$ as follows,

$\displaystyle I^{i}_{1}$ $\displaystyle=$
$\displaystyle-\frac{\eta}{m}\sum_{j\leq
n}\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T}(\mathbf{y}_{j}-\mathbf{u}_{j}(t))(\sum_{r\in\mathbf{S}_{i}}\mathbbm{1}\\{\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{i}^{T}\geq
0,\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{j}^{T}\geq 0\\}$ $\displaystyle=$
$\displaystyle-\eta\sum_{j\leq
n}(\mathbf{u}_{j}(t)-\mathbf{y}_{j})(\mathbf{H}_{i,j}(t)-\mathbf{H}_{i,j}^{\perp}(t))$

Here $\mathbf{H}_{i,j}(t)=\sum_{r\leq
m}\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T}\mathbbm{1}\\{\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{i}^{T}\geq
0,\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{j}^{T}\geq 0\\}$ and
$\mathbf{H}_{i,j}(t)=\sum_{r\in\mathbf{S}_{i}^{\perp}}\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T}\mathbbm{1}\\{\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{i}^{T}\geq
0,\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{j}^{T}\geq 0\\}$.
$\mathbf{H}^{\perp}$ is an $n\times n$ psd matrix whose spectral norm can be
bounded as follows,

$\|\mathbf{H}^{\perp}\|\leq\sum_{i,j}|\mathbf{H}_{i,j}^{\perp}|\leq\frac{n}{m}\sum_{i\leq
n}|\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T}|\cdot|\mathbf{S}_{i}^{\perp}|\leq\frac{Cn^{2}R}{\delta}$

Further,

$\displaystyle\|\mathbf{u}(t+1)-\mathbf{u}(t)\|^{2}$ $\displaystyle\leq$
$\displaystyle\eta^{2}\sum_{i\leq n}\left(\sum_{j\leq
n}\frac{\partial\mathbf{u}_{i}(t)}{dt}\right)^{2}$ $\displaystyle\leq$
$\displaystyle\eta^{2}\sum_{i\leq n}\left(\sum_{j\leq
n}\frac{1}{m}(\mathbf{y}_{j}-\mathbf{u}_{j}(t))(\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T})\sum_{r\leq
m}\mathbbm{1}\\{\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{i}^{T}\geq
0,\mathbf{B}_{r}(t)\mathbf{C}\mathbf{X}_{j}^{T}\geq 0\\}\right)^{2}$
$\displaystyle\leq$
$\displaystyle\eta^{2}n\left(\max_{j}|\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T}|\sum_{j\leq
n}|\mathbf{y}_{j}-\mathbf{u}_{j}(t)|\right)^{2}$ $\displaystyle\leq$
$\displaystyle\eta^{2}n\left(\sum_{j\leq
n}|\mathbf{y}_{j}-\mathbf{u}_{j}(t)|\right)^{2}\leq\eta^{2}n^{2}\|\mathbf{y}-\mathbf{u}(t)\|^{2}$

Finally we bound our desired term,

$\displaystyle\|\mathbf{y}-\mathbf{u}(t+1)\|^{2}$ $\displaystyle=$
$\displaystyle\|\mathbf{y}-\mathbf{u}(t)-(\mathbf{u}(t+1)-\mathbf{u}(t))\|^{2}$
$\displaystyle=$
$\displaystyle\|\mathbf{y}-\mathbf{u}(t)\|^{2}+\|\mathbf{u}(t+1)-\mathbf{u}(t)\|^{2}-2(\mathbf{y}-\mathbf{u}(t))^{T}(\mathbf{u}(t+1)-\mathbf{u}(t))$
$\displaystyle=$
$\displaystyle\|\mathbf{y}-\mathbf{u}(t)\|^{2}+\|\mathbf{u}(t+1)-\mathbf{u}(t)\|^{2}-2(\mathbf{y}-\mathbf{u}(t))^{T}\mathbf{I}_{2}$
$\displaystyle-$ $\displaystyle
2\eta(\mathbf{y}-\mathbf{u}(t))^{T}\mathbf{H}(t)(\mathbf{y}-\mathbf{u}(t))+2\eta(\mathbf{y}-\mathbf{u}(t))^{T}\mathbf{H}(t)^{\perp}(\mathbf{y}-\mathbf{u}(t))$
$\displaystyle\leq$ $\displaystyle\left(1+\eta^{2}n^{2}-\frac{\sum_{i\leq
n}2\eta\sqrt{n}|\mathbf{S}_{i}^{\perp}|}{m}-\eta\lambda_{0}+\frac{2C\eta
n^{2}R}{\delta}\right)\|\mathbf{y}-\mathbf{u}(t)\|^{2}$ $\displaystyle\leq$
$\displaystyle\left(1-\frac{\eta\lambda_{0}}{2}\right)\|\mathbf{y}-\mathbf{u}(t)\|^{2}\leq\left(1-\frac{\eta\lambda_{0}}{2}\right)^{t}\|\mathbf{y}-\mathbf{u}(0)\|^{2}$

∎

Due to our low rank network architecture we only need to update
$O\left(\frac{n^{6}\log(n)}{\lambda_{0}^{4}\delta^{2}}\right)$ parameters. Now
by using $\mathbf{C}$ as Fast Johnson Lindenstrauss transformation [3], the
$\mathbf{C}\mathbf{X}_{i}^{T}$ takes $O(d\log(\ell))$. We compute the
transformation $\mathbf{C}\mathbf{X}^{T}$ in $O(nd\log(\ell))$. So the time
taken to update from $\mathbf{B}(t)$ to $\mathbf{B}(t+1)$ is
$O(nd\log(\ell)+nm\ell)$, i.e.,
$O\left(\frac{n^{7}\log(n)}{\lambda_{0}^{4}\delta^{2}}\right)$. Since, the
guarantees are same upto a constant factor, so to ensure that
$\|\mathbf{u}(t)-\mathbf{y}\|^{2}\leq\varepsilon\|\mathbf{u}(0)-\mathbf{y}\|^{2}$
even when $d=\Omega(n)$ the algorithm takes
$\Omega\left(\frac{n^{7}\log(n)\log(\varepsilon)}{\lambda_{0}^{4}\delta^{2}\log(1-\eta\lambda_{0})}\right)$.

##  5 Reducing Effective Network Width

The matrix $\mathbf{W}$ represents the first layer parameters. We represent
this weight matrix $\mathbf{W}$ as the product
$\mathbf{A}\mathbf{B}\mathbf{C}$. Here $\mathbf{A}\in\mathbb{R}^{m\times k}$,
and as before $\mathbf{B}\in\mathbb{R}^{k\times\ell}$ and
$\mathbf{C}\in\mathbb{R}^{\ell\times d}$. In other words, we add a dimension
reduction gadget as input to the latent layer. It is important to note that
the number of neurons in the first layer remains the same. However, due to the
low rank structure of the weight matrix $\mathbf{W}$, not all neurons are
independent of others. Hence effectively, in the training phase, not every
neuron needs to be learned. Now we study the effect of this addition on the
NTK bound.

Now we discuss how this affects the training phase. First we randomly and
independently initialize $\mathbf{A}\in\mathbb{R}^{m\times k}$ such that all
its entries are iid sampled from $\mathcal{N}(0,1/\sqrt{k})$ for some $k<m$.
We also initialize $\mathbf{C}$ as mentioned in the previous section. Next,
recall that $\mathbf{v}$ is initialized as a random vector sampled from
$\\{-1,+1\\}^{m}$. Now, we learn $\mathbf{B}$ using gradient descent on the
loss function $\Phi(\mathbf{B})=\frac{1}{2}\|\mathbf{y}-\mathbf{u}\|^{2}$
based on a small learning rate $\eta$ by keeping rest of the network
parameters fixed. The matrix $\mathbf{B}$ is updated as follows,

$\displaystyle\mathbf{B}(t+1)$ $\displaystyle=$
$\displaystyle\mathbf{B}(t)-\eta\frac{\partial\Phi(\mathbf{B}(t))}{\partial\mathbf{B}(t)}$
$\displaystyle=$ $\displaystyle\mathbf{B}(t)-\frac{\eta}{\sqrt{m}}\sum_{j\leq
n}(\mathbf{y}_{j}-\mathbf{u}_{j}(t))\mathbf{A}^{T}\mathbf{Z}_{j}(t)\mathbf{X}_{j}\mathbf{C}^{T}\
.$

Here $\mathbf{Z}_{j}(t)\in\\{-1,0,1\\}^{m}$ is such that its $r^{th}$ index is
$\mathbf{v}_{r}\cdot\mathbbm{1}\\{\mathbf{A}_{r}\mathbf{B}(t)\mathbf{C}\mathbf{X}_{j}^{T}>0\\}$
for all $r\in[m]$. Now given the input $\mathbf{X}$, the gram matrix
$\mathbf{H}^{\infty}\in\mathbb{R}^{n\times n}$ is defined with its
$(i,j)^{th}$ term as follows:

$\mathbf{H}_{i,j}^{\infty}=\mathbb{E}_{\mathbf{B}}\left[\frac{1}{m}(\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T})(\mathbf{Z}_{i}\mathbf{A}\mathbf{A}^{T}\mathbf{Z}_{j}^{T})\right]\
.$ (20)

Note that $\mathbf{Z}_{i}$ and $\mathbf{Z}_{j}$ are random because they depend
on $B$. Now for a random $\mathbf{B}(0)$ upon initialization, we define
$\mathbf{H}(0)\in\mathbb{R}^{n\times n}$ by $\forall i,j\in[n]$:

${}\mathbf{H}_{i,j}(0)=\frac{1}{m}(\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T})(\mathbf{Z}_{i}(0)\mathbf{A}\mathbf{A}^{T}\mathbf{Z}_{j}^{T}(0))\
.$ (21)

######  Lemma 5.1.

If $\mathbf{B}\in\mathbb{R}^{k\times\ell}$ is a random matrix whose entries
are iid sample from $\mathcal{N}(0,1)$ then with probability $1-\delta$ the
following holds.

For any $\mathbf{B}(0)$, if
$\|\mathbf{A}_{r}\mathbf{B}\mathbf{C}-\mathbf{A}_{r}\mathbf{B}(0)\mathbf{C}\|\leq
O\left(\frac{\delta^{2}\lambda_{0}}{n^{2}}\right)\stackrel{{\scriptstyle\Delta}}{{=}}R$
then with $\mathbf{B}$ the matrix $\mathbf{H}\in\mathbb{R}^{n\times n}$ whose
$(i,j)^{th}$ entries are defined (21), satisfies
$\|\mathbf{H}-\mathbf{H}(0)\|\leq\frac{\lambda_{0}}{4}$ and
$\lambda_{\min}(\mathbf{H})\geq\frac{\lambda_{0}}{4}$.

###### Proof.

For a fixed $\mathbf{A},\mathbf{C}$ and a random $\mathbf{B}$ we get a random
vector as $\mathbf{W}_{r}=\mathbf{A}_{r}\mathbf{B}\mathbf{C}$. Similarly we
also define $\mathbf{W}_{r}(0)$ corresponding to $\mathbf{B}(0)$. Now consider
the following event for every input $\mathbf{X}_{i}\in\mathbf{X}$ and neuron
$r\in[m]$,

$\mathcal{E}_{r,i}=\\{\exists\mathbf{B}:\|\mathbf{W}_{r}-\mathbf{W}_{r}(0)\|\leq
R,\mathbbm{1}\\{\mathbf{W}_{r}(0)\mathbf{X}_{i}^{T}\geq
0\\}\neq\mathbbm{1}\\{\mathbf{W}_{r}\mathbf{X}_{i}^{T}\geq 0\\}\\}.$

The above event happens if and only if
$|\mathbf{W}_{r}(0)\mathbf{X}_{i}^{T}|\leq R$. Now notice that for a fixed
$\mathbf{A}_{r},\mathbf{C}$ and a randomly initialized $\mathbf{B}(0)$ we have
$h\sim\mathcal{N}(0,1)$ where
$h=\mathbf{A}_{r}\mathbf{B}(0)\mathbf{C}\mathbf{X}_{i}^{T}$. So
$\mathbb{P}(\mathcal{E}_{r,i})=\mathbb{P}(|h|\leq
R)\leq\int_{-R}^{R}\frac{1}{\sqrt{2\pi}}\exp(-x^{2})dx\leq R$. Now we bound
deviation of every entry $(i,j)\in[n]\times[n]$ as follows,

$\displaystyle\mathbb{E}[|\mathbf{H}_{i,j}-\mathbf{H}_{i,j}(0)|]$
$\displaystyle=$
$\displaystyle\mathbb{E}\left[\frac{1}{m}\left|(\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T})(\mathbf{Z}_{i}\mathbf{A}\mathbf{A}^{T}\mathbf{Z}_{j}^{T}-\mathbf{Z}_{i}(0)\mathbf{A}\mathbf{A}^{T}\mathbf{Z}_{j}^{T}(0))\right|\right]$

Now notice that,

$\displaystyle(\mathbf{Z}_{i}\mathbf{A}\mathbf{A}^{T}\mathbf{Z}_{j}^{T}-\mathbf{Z}_{i}(0)\mathbf{A}\mathbf{A}^{T}\mathbf{Z}_{j}^{T}(0))$
$\displaystyle=$
$\displaystyle\sum_{p,q}\mathbf{Z}_{i,p}\mathbf{M}_{p,q}\mathbf{Z}_{j,q}-\mathbf{Z}_{i,p}(0)\mathbf{M}_{p,q}\mathbf{Z}_{j,q}(0)$
$\displaystyle=$
$\displaystyle\sum_{p,q}\mathbf{M}_{p,q}\mathbf{v}_{p}\mathbf{v}_{q}(\mathbbm{1}(\mathbf{A}_{p}\mathbf{B}\mathbf{C}\mathbf{X}_{i}^{T}\geq
0)\cdot\mathbbm{1}(\mathbf{A}_{q}\mathbf{B}\mathbf{C}\mathbf{X}_{j}^{T}\geq
0)$ $\displaystyle-$
$\displaystyle\mathbbm{1}(\mathbf{A}_{p}\mathbf{B}(0)\mathbf{C}\mathbf{X}_{i}^{T}\geq
0)\cdot\mathbbm{1}(\mathbf{A}_{q}\mathbf{B}(0)\mathbf{C}\mathbf{X}_{j}^{T}\geq
0))$

Here $\mathbf{M}_{p,q}$ is the inner product between the $p^{th}$ and the
$q^{th}$ row of $\mathbf{A}$, i.e.,
$\mathbf{M}_{p,q}=\mathbf{A}_{p}\mathbf{A}_{q}^{T}$. Next, since
$\mathbb{E}[|\mathbbm{1}(\mathbf{A}_{p}\mathbf{B}\mathbf{C}\mathbf{X}_{i}\geq
0)\cdot\mathbbm{1}(\mathbf{A}_{q}\mathbf{B}\mathbf{C}\mathbf{X}_{j}\geq
0)-\mathbbm{1}(\mathbf{A}_{p}\mathbf{B}(0)\mathbf{C}\mathbf{X}_{i}\geq
0)\cdot\mathbbm{1}(\mathbf{A}_{q}\mathbf{B}(0)\mathbf{C}\mathbf{X}_{j}\geq
0)|]\leq\mathbb{P}(\mathcal{E}_{p,i})+\mathbb{P}(\mathcal{E}_{q,j})$, hence
$\mathbbm{1}(\mathbf{A}_{p}\mathbf{B}\mathbf{C}\mathbf{X}_{i}\geq
0)\cdot\mathbbm{1}(\mathbf{A}_{q}\mathbf{B}\mathbf{C}\mathbf{X}_{j}\geq
0)-\mathbbm{1}(\mathbf{A}_{p}\mathbf{B}(0)\mathbf{C}\mathbf{X}_{i}\geq
0)\cdot\mathbbm{1}(\mathbf{A}_{q}\mathbf{B}(0)\mathbf{C}\mathbf{X}_{j}\geq 0)$
is either $-1$ or $1$ with probability at most $2R$. Here the randomness is
due to the random matrix $\mathbf{B}$. So we can rewrite,
$(\mathbf{Z}_{i}\mathbf{A}\mathbf{A}^{T}\mathbf{Z}_{j}^{T}-\mathbf{Z}_{i}(0)\mathbf{A}\mathbf{A}^{T}\mathbf{Z}_{j}^{T}(0))=\mathbf{x}\mathbf{A}\mathbf{A}^{T}\mathbf{y}$
where $\mathbf{x}$ and $\mathbf{y}$ are $m$ dimensional vectors such that with
at least $1-\delta/2$ probability
$\|\mathbf{x}\|^{2}=\|\mathbf{y}\|^{2}\leq\frac{2mR}{\delta}$. We use the
following claim, to discuss the desired property of $\mathbf{A}$.

######  Claim 5.1.

Let $\mathcal{S}$ be a set such that every element $\mathbf{s}\in\mathcal{S}$
is an $m$ dimensional vector in $\\{-1,0,+1\\}^{m}$ and
$\|\mathbf{s}\|_{1}\leq\frac{2mR}{\delta}$. The size of such set
$|\mathcal{S}|=O(\binom{m}{mR/\delta}\cdot 3^{mR/\delta})$. Let
$k=O(\log(|\mathcal{S}|))=\tilde{O}\left(\frac{mR}{\delta}\right)$ and recall
matrix $\mathbf{A}$ has random i.i.d. samples from
$\mathcal{N}(0,1/\sqrt{k})$. Now for all $\mathbf{s}\in\mathcal{S}$ we get,
$\|\mathbf{s}\mathbf{A}\|=O(\|\mathbf{s}\|)=O(\sqrt{mR/\delta})$ with high
probability.

Now the above claim yields,
$\mathbb{E}[|\mathbf{H}_{i,j}-\mathbf{H}_{i,j}(0)|]\leq\frac{1}{m}(O(1)\sqrt{mR/\delta}\sqrt{mR/\delta})=O(R/\delta)$.
So, $\mathbb{E}[\sum_{i,j}|\mathbf{H}_{i,j}-\mathbf{H}_{i,j}(0)|]\leq
O\left(\frac{n^{2}R}{\delta}\right)$ and by Markov inequality we have
$\sum_{i,j}|\mathbf{H}_{i,j}-\mathbf{H}_{i,j}(0)|\leq
O\left(\frac{n^{2}R}{\delta^{2}}\right)$ with at least $1-\delta$ probability.
Therefore with at least $1-\delta$ probability
$\|\mathbf{H}-\mathbf{H}(0)\|\leq O\left(\frac{n^{2}R}{\delta^{2}}\right)$.
Finally,

$\lambda_{\min}(\mathbf{H})\geq\lambda_{\min}(\mathbf{H}(0))-O\left(\frac{n^{2}R}{\delta^{2}}\right)\geq\frac{\lambda_{0}}{2}$

∎

Note that the above lemma ensures that
$R=O\left(\frac{\lambda_{0}\delta^{2}}{n^{2}}\right)$. Further, it implies
that from the point of initialization of $\mathbf{B}(0)$, out of $m$ neurons
only $O\left(\frac{mR}{\delta}\right)$ of them change their activation
function value from their initial function value. The next lemma states the
required learning rate $\eta$ and the number of neurons $m$ to get the desired
training error after certain iterations.

######  Theorem 5.1.

Let $\lambda_{0}$ is the smallest non zero eigenvalue of
$\mathbf{H}^{\infty}$. For $\delta\in(0,1)$, number of neurons
$m=\Omega\left(\frac{n^{6}}{\lambda_{0}^{4}\delta^{7}}\right)$,
$\mathbf{B}(0)$ is randomly initialized and we set
$\eta=O\left(\frac{\lambda_{0}\delta^{4}}{n^{2}}\right)$. Then during gradient
descent we have the following with probability at least $1-\delta$ for
$t=0,1,\ldots$,

$\|\mathbf{u}(t)-\mathbf{y}\|^{2}\leq\left(1-\frac{\eta\lambda_{0}}{2}\right)^{t}\|\mathbf{u}(0)-\mathbf{y}\|^{2}$
(22)

###### Proof.

We prove this by induction. The induction hypothesis is the following,

$\|\mathbf{u}(t)-\mathbf{y}\|^{2}\leq\left(1-\frac{\eta\lambda_{0}}{2}\right)^{t}\|\mathbf{u}(0)-\mathbf{y}\|^{2}$
(23)

Notice that the hypothesis is trivially true for the base case $t=0$. Suppose
it holds true for $t^{\prime}=0,1,\ldots,t$ and now we show that it also holds
for $t^{\prime}=t+1$. We have,

$\displaystyle\|\mathbf{u}(t+1)-\mathbf{u}(t)\|^{2}$ $\displaystyle\leq$
$\displaystyle\eta^{2}\sum_{i\leq
n}\left(\frac{\partial\mathbf{u}_{i}(t)}{dt}\right)^{2}$ $\displaystyle\leq$
$\displaystyle\eta^{2}\sum_{i\leq n}\left(\sum_{j\leq
n}\frac{1}{m}(\mathbf{y}_{j}-\mathbf{u}_{j}(t))(\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T})(\mathbf{Z}_{i}(t)\mathbf{A}\mathbf{A}^{T}\mathbf{Z}_{j}^{T}(t))\right)^{2}$
$\displaystyle\leq$
$\displaystyle\eta^{2}n\left(\frac{1}{m}\max_{j}|\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T}|\sum_{j\leq
n}\left|(\mathbf{Z}_{i}(t)\mathbf{A}^{T}\mathbf{A}\mathbf{Z}_{j}^{T}(t))\left(\mathbf{y}_{j}-\mathbf{u}_{j}(t)\right)\right|\right)^{2}$
$\displaystyle\leq$
$\displaystyle\eta^{2}n\left(\frac{O(1)}{m}\left(\sum_{j\leq
n}(\mathbf{Z}_{i}(t)\mathbf{A}\mathbf{A}^{T}\mathbf{Z}_{j}^{T}(t))^{2}\right)^{1/2}\|\mathbf{y}-\mathbf{u}(t)\|\right)^{2}$
$\displaystyle\leq$
$\displaystyle\eta^{2}n\left(\frac{O(1)}{m}\left(\sum_{j\leq
n}(\|\mathbf{Z}_{i}(t)\mathbf{A}\|\|\mathbf{A}^{T}\mathbf{Z}_{j}^{T}(t)\|)^{2}\right)^{1/2}\|\mathbf{y}-\mathbf{u}(t)\|\right)^{2}$

Notice that,

$\displaystyle\|\mathbf{Z}_{i}(t)\mathbf{A}\|^{2}=\sum_{r_{1}\leq
k}(\mathbf{Z}_{i}(t)\mathbf{A}_{r_{1}}^{T})^{2}=\sum_{r_{1}\leq
k}\left(\sum_{r_{2}\leq
m}\mathbf{v}_{r_{2}}\mathbbm{1}(\mathbf{A}_{r_{2}}\mathbf{B}(t)\mathbf{C}\mathbf{X}_{i}^{T}\geq
0)\mathbf{A}_{r_{2},r_{1}}\right)^{2}=\sum_{r_{1}\leq
k}(\mathbf{v}\tilde{\mathbf{A}}_{r_{1}}^{T})^{2}=\|\mathbf{v}\tilde{\mathbf{A}}\|^{2}$

Here $\tilde{\mathbf{A}}_{r_{1}}$ represents the entrywise vector product
between $\mathbf{A}_{r_{1}}$ and corresponding indicator variable. Notice that
$\|\tilde{\mathbf{A}}\|_{F}^{2}\leq\|\mathbf{A}\|_{F}^{2}$ and with at least
$1-\delta/2$ probability $\|\mathbf{A}\|_{F}^{2}\leq 2m/\delta$. So, for all
$i\in[n]$,
$\mathbb{E}_{\mathbf{v}}[\|\mathbf{Z}_{i}(t)\mathbf{A}\|^{2}]=\sum_{r_{1}}\|\tilde{\mathbf{A}}_{r_{1}}\|^{2}\leq\|\mathbf{A}\|_{F}^{2}\leq
2m/\delta$. Now by Markov we get,
$\sum_{i}\|\mathbf{Z}_{i}(t)\mathbf{A}\|^{2}\leq O(nm/\delta^{2})$ with at
least $1-\delta$ probability. Now,

$\displaystyle\|\mathbf{u}(t+1)-\mathbf{u}(t)\|^{2}$ $\displaystyle\leq$
$\displaystyle\eta^{2}n\left(\frac{O(1)}{m}\left(\sum_{j\leq
n}(\|\mathbf{Z}_{i}(t)\mathbf{A}^{T}\|\|\mathbf{A}\mathbf{Z}_{j}^{T}(t)\|)^{2}\right)^{1/2}\|\mathbf{y}-\mathbf{u}(t)\|\right)^{2}$
$\displaystyle\leq$
$\displaystyle\eta^{2}n\left(\frac{O(1)}{m}\left(\frac{nm^{2}}{\delta^{4}}\right)^{1/2}\|\mathbf{y}-\mathbf{u}(t)\|\right)^{2}\leq
O\left(\frac{\eta^{2}n^{2}}{\delta^{4}}\right)\|\mathbf{y}-\mathbf{u}(t)\|^{2}$

From the above analysis notice that,

$\displaystyle|\mathbf{u}_{i}(t+1)-\mathbf{u}_{i}(0)|$ $\displaystyle\leq$
$\displaystyle\sum_{t^{\prime}<t}\left|\frac{\partial\mathbf{u}_{i}(t^{\prime})}{dt^{\prime}}\right|\leq
O\left(\frac{\eta}{\delta^{2}}\right)\sum_{t^{\prime}<t}\sqrt{n}\|\mathbf{y}-\mathbf{u}(t)\|$
$\displaystyle\leq$ $\displaystyle
O\left(\frac{\eta\sqrt{n}}{\delta^{2}}\right)\sum_{t^{\prime}<t}\left(1-\frac{\eta\lambda_{0}}{2}\right)^{t^{\prime}/2}\|\mathbf{u}(0)-\mathbf{y}\|$
$\displaystyle\leq$ $\displaystyle
O\left(\frac{\eta\sqrt{n}}{\delta^{2}}\right)\sum_{t^{\prime}=1}^{\infty}\left(1-\frac{\eta\lambda_{0}}{2}\right)^{t^{\prime}/2}\|\mathbf{u}(0)-\mathbf{y}\|\leq\frac{\sqrt{n}}{\delta^{2}\lambda_{0}}\|\mathbf{u}(0)-\mathbf{y}\|\leq\frac{n}{\delta^{5/2}\lambda_{0}}$

Our network parameters at time $t=0$, ensures that the output
$\mathbf{u}_{i}(0)$ is $O(1)$, for all $i\in[n]$. Further we know that
$\mathbf{y}_{i}=O(1)$. Hence by Markov’s inequality we have
$\|\mathbf{u}(0)-\mathbf{y}\|^{2}\leq n/\delta$ with at least probability
$1-\delta$.

SO Now, based on the events $\mathcal{E}_{r,i}$ for all input $\mathbf{X}_{i}$
and neuron $r$, let $\mathbf{u}_{i}$ be the network output for
$\mathbf{X}_{i}$. Next, for all $r\in[m]$, since the change from
$\mathbf{W}_{r}(0)$ to $\mathbf{W}_{r}$ is bounded, so we have
$|\mathbf{u}_{i}-\mathbf{u}_{i}(0)|=\frac{1}{\sqrt{m}}\mathbf{v}(\sigma(\mathbf{W}\mathbf{X}_{i}^{T})-\sigma(\mathbf{W}(0)\mathbf{X}_{i}^{T}))\leq\frac{1}{\sqrt{m}}\|\mathbf{v}\|\|\mathbf{W}-\mathbf{W}(0)\|\|\mathbf{X}_{i}\|\leq\frac{1}{\sqrt{m}}\sqrt{m}\sqrt{m}R=\sqrt{m}R$.
Now at each step $t$ of the gradient descent the parameter $\mathbf{W}_{r}(t)$
corresponding to every neuron $r\in[m]$ are affected due to the change in
$\mathbf{B}(t)$. So in order to insure that the condition in lemma 5.1 holds,
we need ensure
$|\mathbf{u}_{i}(t+1)-\mathbf{u}_{i}(0)|\leq|\mathbf{u}_{i}-\mathbf{u}_{i}(0)|$.
So we get, $m=\Omega\left(\frac{n^{6}}{\lambda_{0}^{4}\delta^{7}}\right)$.
Further, we know that, $R=O\left(\frac{\delta^{2}\lambda_{0}}{n^{2}}\right)$
and $\mathbb{P}(\mathcal{E}_{r,i})\leq R$. Let
$\mathbf{S}_{i}=\\{r\in[m]:\mathbbm{1}\\{\mathcal{E}_{r,i}\\}=0\\}$ and
$\mathbf{S}_{i}^{\perp}=[m]\backslash\mathbf{S}_{i}$. Let
$\mathbf{s}_{i}\in\mathbb{R}^{m}$ be such that $\mathbf{s}_{i,r}=1,\forall
r\in[\mathbf{S}_{i}^{\perp}]$. So, $\mathbb{E}[\|\mathbf{s}_{i}\|^{2}]\leq
mR$. Further, $\mathbb{E}[\sum_{i\leq n}|\mathbf{S}_{i}^{\perp}|]\leq mnR$. So
by using Markov inequality we have $\sum_{i\leq
n}|\mathbf{S}_{i}^{\perp}|\leq\frac{mnR}{\delta}$ with at least $1-\delta$
probability. Now consider the following term,

$\displaystyle\mathbf{u}_{i}(t+1)-\mathbf{u}_{i}(t)$ $\displaystyle=$
$\displaystyle\frac{1}{\sqrt{m}}\sum_{r\leq
m}\mathbf{v}_{r}\left[\sigma(\mathbf{A}_{r}\mathbf{B}(t+1)\mathbf{C}\mathbf{X}_{i}^{T})-\sigma(\mathbf{A}_{r}\mathbf{B}(t)\mathbf{C}\mathbf{X}_{i}^{T})\right]$
$\displaystyle=$ $\displaystyle\frac{1}{\sqrt{m}}\sum_{r\leq
m}\mathbf{v}_{r}\left[\sigma\left(\mathbf{A}_{r}\left(\mathbf{B}(t)-\eta\frac{\partial\Phi(\mathbf{B}(t))}{\partial\mathbf{B}(t)}\right)\mathbf{C}\mathbf{X}_{i}^{T}\right)-\sigma(\mathbf{A}_{r}\mathbf{B}(t)\mathbf{C}\mathbf{X}_{i}^{T})\right]$
$\displaystyle=$ $\displaystyle I^{i}_{1}+I^{i}_{2}$

We analyze the term based on two sets of neurons. We have $I^{i}_{1}$
corresponding to the neurons which do not change their activation function
value based for the network parameter $\mathbf{B}(t)$ and $\mathbf{B}(t+1)$.
The term $I^{i}_{2}$ corresponds to rest of the neurons (that changes).

$\displaystyle I^{i}_{1}$ $\displaystyle=$
$\displaystyle\frac{1}{\sqrt{m}}\sum_{r\in\mathbf{S}_{i}}\mathbf{v}_{r}\left[\sigma\left(\mathbf{A}_{r}\left(\mathbf{B}(t)-\eta\frac{\partial\Phi(\mathbf{B}(t))}{\partial\mathbf{B}(t)}\right)\mathbf{C}\mathbf{X}_{i}^{T}\right)-\sigma(\mathbf{A}_{r}\mathbf{B}(t)\mathbf{C}\mathbf{X}_{i}^{T})\right]$
$\displaystyle I^{i}_{2}$ $\displaystyle=$
$\displaystyle\frac{1}{\sqrt{m}}\sum_{r\in\mathbf{S}_{i}^{\perp}}\mathbf{v}_{r}\left[\sigma\left(\mathbf{A}_{r}\left(\mathbf{B}(t)-\eta\frac{\partial\Phi(\mathbf{B}(t))}{\partial\mathbf{B}(t)}\right)\mathbf{C}\mathbf{X}_{i}^{T}\right)-\sigma(\mathbf{A}_{r}\mathbf{B}(t)\mathbf{C}\mathbf{X}_{i}^{T})\right]$

We first upper bound $|I^{i}_{2}|$.

$\displaystyle|I^{i}_{2}|$ $\displaystyle=$
$\displaystyle\left|\frac{1}{\sqrt{m}}\sum_{r\in\mathbf{S}_{i}^{\perp}}\mathbf{v}_{r}\left[\sigma\left(\mathbf{A}_{r}\left(\mathbf{B}(t)-\eta\frac{\partial\Phi(\mathbf{B}(t))}{\partial\mathbf{B}(t)}\right)\mathbf{C}\mathbf{X}_{i}^{T}\right)-\sigma(\mathbf{A}_{r}\mathbf{B}(t)\mathbf{C}\mathbf{X}_{i}^{T})\right]\right|$
(25) $\displaystyle\leq$
$\displaystyle\left|\frac{1}{\sqrt{m}}\sum_{r\in\mathbf{S}_{i}^{\perp}}\left[-\eta\mathbf{A}_{r}\left(\frac{\partial\Phi(\mathbf{B}(t))}{\partial\mathbf{B}(t)}\right)\mathbf{C}\mathbf{X}_{i}^{T}\right]\right|$
$\displaystyle\leq$
$\displaystyle\frac{\eta}{\sqrt{m}}\|\mathbf{s}_{i}\|\cdot\left\|\sum_{j\leq
n}\frac{(\mathbf{y}_{j}-\mathbf{u}_{j}(t))(\mathbf{X}_{j}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{i}^{T})(\mathbf{A}_{\mathbf{S}_{i}^{\perp}}\mathbf{A}^{T}\mathbf{Z}_{j}^{T}(t))}{\sqrt{m}}\right\|$
$\displaystyle\leq$
$\displaystyle\eta\sqrt{\frac{mR}{m\delta}}\left\|\sum_{j\leq
n}\frac{(\mathbf{y}_{j}-\mathbf{u}_{j}(t))(\mathbf{X}_{j}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{i}^{T})(\mathbf{A}_{\mathbf{S}_{i}^{\perp}}\mathbf{A}^{T}\mathbf{Z}_{j}^{T}(t))}{\sqrt{m}}\right\|$
$\displaystyle\leq$
$\displaystyle\eta\sqrt{\frac{R}{m\delta}}\left\|\max_{j}(\mathbf{X}_{j}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{i}^{T})(\mathbf{A}_{\mathbf{S}_{i}^{\perp}}\mathbf{A}^{T}\mathbf{Z}_{j}^{T}(t))\right\|\sum_{j\leq
n}|\mathbf{y}_{j}-\mathbf{u}_{j}(t)|$ $\displaystyle\leq$
$\displaystyle\eta\sqrt{\frac{R}{m\delta}}\left\|\max_{j}\mathbf{A}_{\mathbf{S}_{i}^{\perp}}\mathbf{A}^{T}\mathbf{Z}_{j}^{T}(t)\right\|\sqrt{n}\|\mathbf{y}-\mathbf{u}(t)\|$
$\displaystyle\leq$
$\displaystyle\eta\sqrt{\frac{nR}{m\delta}}\|\mathbf{A}_{\mathbf{S}_{i}^{\perp}}\|\|\mathbf{A}\|\max_{j}\|\mathbf{Z}_{j}(t)\|\|\mathbf{y}-\mathbf{u}(t)\|$
$\displaystyle\leq$
$\displaystyle\eta\sqrt{\frac{nR}{m\delta}}\sqrt{\frac{m}{k}}\sqrt{m}\|\mathbf{y}-\mathbf{u}(t)\|$
$\displaystyle\leq$
$\displaystyle\eta\sqrt{\frac{n}{\delta}}\|\mathbf{y}-\mathbf{u}(t)\|$

Since ReLU is a $1$-Lipschitz function and $|\mathbf{v}_{r}|=1$, we get (25).
We have (5) by applying Cauchy-Schwartz on the sum over
$|\mathbf{S}_{i}^{\perp}|$ terms. We get (5) due to Holder inequality. The
$\mathbf{A}_{\mathbf{S}_{i}^{\perp}}$ is a $m\times d$ dimensional matrix.
Here every row of this matrix is either all zero vector or a row from
$\mathbf{A}$. For all $j\in[m]$, if $j\in\mathbf{S}_{i}^{\perp}$, then the
$j^{th}$ row of $\mathbf{A}_{\mathbf{S}_{i}^{\perp}}$ is equal to
$\mathbf{A}_{j}$, else it is $\\{0\\}^{k}$. So in (5) we upper bound
$\|\mathbf{A}_{\mathbf{S}_{i}^{\perp}}\|\leq 1$. Further,
$\|\mathbf{A}\|^{2}\leq\frac{m}{k}$ and $\forall
j\in[n],\|\mathbf{Z}_{j}(t)\|\leq\sqrt{m}$.

Now we express $I^{i}_{1}$ as follows,

$\displaystyle I^{i}_{1}$ $\displaystyle=$
$\displaystyle-\frac{\eta}{m}\sum_{j\leq
n}(\mathbf{y}_{j}-\mathbf{u}_{j}(t))(\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T})\left(\sum_{r\in\mathbf{S}_{i}}(\mathbf{Z}_{i}(t)\mathbf{A}_{r}^{T}\mathbf{A}_{r}\mathbf{Z}_{j}^{T}(t))\right)$
$\displaystyle=$ $\displaystyle-\eta\sum_{j\leq
n}(\mathbf{y}_{j}-\mathbf{u}_{j}(t))(\mathbf{H}_{i,j}(t)-\mathbf{H}_{i,j}^{\perp}(t))$

Here
$\mathbf{H}_{i,j}(t)=\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T}(\mathbf{Z}_{i}(t)\mathbf{A}\mathbf{A}^{T}\mathbf{Z}_{j}^{T}(t))$
and the term
$\mathbf{H}_{i,j}^{\perp}(t)=\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T}(\mathbf{Z}_{i}(t)\mathbf{A}_{\mathbf{S}_{i}^{\perp}}\mathbf{A}_{\mathbf{S}_{i}^{\perp}}^{T}\mathbf{Z}_{j}^{T}(t))$.
$\mathbf{H}^{\perp}$ is an $n\times n$ PSD matrix whose spectral norm can be
bounded as follows,

$\displaystyle\|\mathbf{H}^{\perp}\|$ $\displaystyle\leq$
$\displaystyle\sum_{i,j}|\mathbf{H}_{i,j}^{\perp}|=\frac{1}{m}\sum_{i\leq
n}\left(\sum_{j\leq
n}\left|\mathbf{X}_{i}\mathbf{C}^{T}\mathbf{C}\mathbf{X}_{j}^{T}(\mathbf{Z}_{i}(t)\mathbf{A}_{\mathbf{S}_{i}^{\perp}}\mathbf{A}_{\mathbf{S}_{i}^{\perp}}^{T}\mathbf{Z}_{j}^{T}(t))\right|\right)$
$\displaystyle\leq$ $\displaystyle\frac{1}{m}\sum_{i\leq n}\left(\sum_{j\leq
n}\left|O(1)\cdot(\mathbf{Z}_{i}(t)\mathbf{A}_{\mathbf{S}_{i}^{\perp}}\mathbf{A}_{\mathbf{S}_{i}^{\perp}}^{T}\mathbf{Z}_{j}^{T}(t))\right|\right)$
$\displaystyle=$ $\displaystyle\frac{1}{m}\sum_{i\leq n}\left(\sum_{j\leq
n}\left|O(1)\cdot(\tilde{\mathbf{Z}}_{i}(t)\mathbf{A}_{\mathbf{S}_{i}^{\perp}}\mathbf{A}_{\mathbf{S}_{i}^{\perp}}^{T}\tilde{\mathbf{Z}}_{j}^{T}(t))\right|\right)$
$\displaystyle\leq$ $\displaystyle\frac{1}{m}\sum_{i\leq n}\left(\sum_{j\leq
n}\left|O(1)\cdot\|\mathbf{A}_{\mathbf{S}_{i}^{\perp}}\tilde{\mathbf{Z}}_{i}^{T}(t)\|\|\mathbf{A}_{\mathbf{S}_{i}^{\perp}}\tilde{\mathbf{Z}}_{j}^{T}(t)\|\right|\right)$
$\displaystyle\leq$ $\displaystyle\frac{1}{m}\sum_{i\leq n}\left(\sum_{j\leq
n}\left|O(1)\cdot\|\mathbf{A}\tilde{\mathbf{Z}}_{i}^{T}(t)\|\|\mathbf{A}\tilde{\mathbf{Z}}_{j}^{T}(t)\|\right|\right)$
$\displaystyle\leq$ $\displaystyle\frac{1}{m}\sum_{i\leq n}\left(\sum_{j\leq
n}O(1)\frac{mR}{\delta}\right)\leq O\left(\frac{n^{2}R}{\delta}\right)$

In (5) the vectors $\tilde{\mathbf{Z}}_{i}(t)$ and $\tilde{\mathbf{Z}}_{j}(t)$
are $m$ dimensional sparse vectors. For every index of
$r\in\mathbf{S}_{i}^{\perp}$,
$\tilde{\mathbf{Z}}_{i,r}(t)=\mathbf{Z}_{i,r}(t)$. Similarly the vector
$\tilde{\mathbf{Z}}_{j}(t)$ is also defined. So their sparsity is bounded by
$|\mathbf{S}_{i}^{\perp}|$.

$\tilde{\mathbf{Z}}_{i}(t)$ and $\tilde{\mathbf{Z}}_{j}(t)$. In (5) we use the
upper bound
$\|\mathbf{A}_{\mathbf{S}_{i}^{\perp}}\mathbf{Y}^{T}\|\leq\|\mathbf{A}\mathbf{Y}^{T}\|$,
where $\mathbf{Y}\in\mathbb{R}^{m}$.

Finally we bound our desired term,

$\displaystyle\|\mathbf{y}-\mathbf{u}(t+1)\|^{2}$ $\displaystyle=$
$\displaystyle\|\mathbf{y}-\mathbf{u}(t)-(\mathbf{u}(t+1)-\mathbf{u}(t))\|^{2}$
$\displaystyle=$
$\displaystyle\|\mathbf{y}-\mathbf{u}(t)\|^{2}+\|\mathbf{u}(t+1)-\mathbf{u}(t)\|^{2}-2(\mathbf{y}-\mathbf{u}(t))^{T}(\mathbf{u}(t+1)-\mathbf{u}(t))$
$\displaystyle=$
$\displaystyle\|\mathbf{y}-\mathbf{u}(t)\|^{2}+\|\mathbf{u}(t+1)-\mathbf{u}(t)\|^{2}-2(\mathbf{y}-\mathbf{u}(t))^{T}\mathbf{I}_{2}$
$\displaystyle-$ $\displaystyle
2\eta(\mathbf{y}-\mathbf{u}(t))^{T}\mathbf{H}(t)(\mathbf{y}-\mathbf{u}(t))+2\eta(\mathbf{y}-\mathbf{u}(t))^{T}\mathbf{H}(t)^{\perp}(\mathbf{y}-\mathbf{u}(t))$
$\displaystyle\leq$
$\displaystyle\left(1+\frac{\eta^{2}n^{2}}{\delta^{4}}-\sum_{i\leq
n}2\eta\sqrt{\frac{n}{\delta}}-\eta\lambda_{0}+\frac{2\eta
n^{2}R}{\delta}\right)\|\mathbf{y}-\mathbf{u}(t)\|^{2}$ $\displaystyle\leq$
$\displaystyle\left(1-\frac{\eta\lambda_{0}}{2}\right)\|\mathbf{y}-\mathbf{u}(t)\|^{2}$
$\displaystyle\leq$
$\displaystyle\left(1-\frac{\eta\lambda_{0}}{2}\right)^{t}\|\mathbf{y}-\mathbf{u}(0)\|^{2}$

By setting $\eta=\frac{\lambda_{0}\delta^{4}}{2n^{2}}$ we get (5). ∎

#### Running Time:

The previous theorem essentially implies that in an overparameterized network
with network width of $O\left(\frac{n^{6}}{\lambda_{0}^{4}\delta^{7}}\right)$
one can find global minimum by gradient descent with a learning rate
$\eta=O\left(\frac{\lambda_{0}\delta^{4}}{n^{2}}\right)$. However, notice that
due to our dimensionality reduction, our network can achieve it much faster
even for network width of $O(n^{6})$. In each iteration we update a
$k\times\ell$ dimensional matrix matrix $\mathbf{B}$. For all $i\in[n]$ we do
a matrix vector product, $\mathbf{Z}_{i}(t)\mathbf{A}$ and
$\mathbf{X}_{i}\mathbf{C}^{T}$ which takes $O(mk+\ell d)=O(mk)$. Now instead
of using $\mathbf{A}$ as naive dimensionality reduction matrix (i.e.,
Johnson–Lindenstrauss), if replace it with Fast Johnson–Lindenstrauss
transformation [3], then our running time improves to
$O(m\log(k))=O\left(n^{6}\log\left(\frac{n}{\lambda_{0}\delta}\right)\right)$.
As in each iteration of the gradient descent process uses all the $n$ inputs,
hence updating $\mathbf{B}(t)$ takes
$O(nm\log(k))=O\left(n^{7}\log\left(\frac{n}{\lambda_{0}\delta}\right)\right)$.
Now for $t\geq\frac{\log(\varepsilon)}{\log(1-\eta\lambda_{0})}$ our learnt
parameter ensures,

$\|\mathbf{y}-\mathbf{u}(t+1)\|^{2}\leq\varepsilon\cdot\|\mathbf{y}-\mathbf{u}(0)\|^{2}.$

So the total running time to achieve this is
$O\left(\frac{n^{7}\log\left(\frac{n}{\lambda_{0}\delta}\right)\log(\varepsilon)}{\log(1-\eta\lambda_{0})}\right)$.
Notice that with existing technique the running time would be $O(nmdt)$ which
is
$O\left(\frac{n^{7}d\log(\varepsilon)}{\lambda_{0}^{4}\delta^{2}\log(1-(\lambda_{0}/n)^{2})}\right)$.

##  6 Conclusion

In this work we presented a simple yet powerful way to reduce the required
number of unknown parameters in an overparameterized network. Our result uses
a dimentionality reduction approach for input points as well as for the vector
fed into the neurons of the latent layer. We get a significant improvement in
the running time while retaining the original theoretical NTK based
guarantees. In particular, when $d=O(n)$ our network achieves a similar
training loss with a faster running time by factor $O(n)$ compared to other
related results.

##  7 Acknowledgements

This project has received funding from the European Union’s Horizon 2020
research and innovation programmed under grant agreement No 682203 -ERC-[ Inf-
Speed-Tradeoff].

## References

  * [1] William B Johnson and Joram Lindenstrauss.  Extensions of lipschitz mappings into a hilbert space 26.  Contemporary mathematics, 26:28, 1984.
  * [2] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner.  Gradient-based learning applied to document recognition.  Proceedings of the IEEE, 86(11):2278–2324, 1998.
  * [3] Nir Ailon and Bernard Chazelle.  Approximate nearest neighbors and the fast johnson-lindenstrauss transform.  In Proceedings of the thirty-eighth annual ACM symposium on Theory of computing, pages 557–563, 2006.
  * [4] Jiří Matoušek.  On variants of the johnson–lindenstrauss lemma.  Random Structures & Algorithms, 33(2):142–156, 2008.
  * [5] Michael W Mahoney and Petros Drineas.  Cur matrix decompositions for improved data analysis.  Proceedings of the National Academy of Sciences, 106(3):697–702, 2009.
  * [6] Nir Ailon and Bernard Chazelle.  Faster dimension reduction.  Communications of the ACM, 53(2):97–104, 2010.
  * [7] Ronan Collobert, Jason Weston, Léon Bottou, Michael Karlen, Koray Kavukcuoglu, and Pavel Kuksa.  Natural language processing (almost) from scratch.  Journal of machine learning research, 12(ARTICLE):2493–2537, 2011\.
  * [8] Petros Drineas, Michael W Mahoney, Shan Muthukrishnan, and Tamás Sarlós.  Faster least squares approximation.  Numerische mathematik, 117(2):219–249, 2011.
  * [9] Petros Drineas, Malik Magdon-Ismail, Michael W Mahoney, and David P Woodruff.  Fast approximation of matrix coherence and statistical leverage.  The Journal of Machine Learning Research, 13(1):3475–3506, 2012\.
  * [10] Geoffrey Hinton, Li Deng, Dong Yu, George E Dahl, Abdel-rahman Mohamed, Navdeep Jaitly, Andrew Senior, Vincent Vanhoucke, Patrick Nguyen, Tara N Sainath, et al.  Deep neural networks for acoustic modeling in speech recognition: The shared views of four research groups.  IEEE Signal processing magazine, 29(6):82–97, 2012.
  * [11] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.  Imagenet classification with deep convolutional neural networks.  Advances in neural information processing systems, 25, 2012.
  * [12] Edo Liberty.  Simple and deterministic matrix sketching.  In Proceedings of the 19th ACM SIGKDD international conference on Knowledge discovery and data mining, pages 581–588, 2013.
  * [13] David P Woodruff et al.  Sketching as a tool for numerical linear algebra.  Foundations and Trends® in Theoretical Computer Science, 10(1–2):1–157, 2014.
  * [14] Kasper Green Larsen and Jelani Nelson.  Optimality of the johnson-lindenstrauss lemma.  In 2017 IEEE 58th Annual Symposium on Foundations of Computer Science (FOCS), pages 633–638. IEEE, 2017.
  * [15] Cenk Baykal, Lucas Liebenwein, Igor Gilitschenski, Dan Feldman, and Daniela Rus.  Data-dependent coresets for compressing neural networks with applications to generalization bounds.  In International Conference on Learning Representations, 2018.
  * [16] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh.  Gradient descent provably optimizes over-parameterized neural networks.  arXiv preprint arXiv:1810.02054, 2018.
  * [17] Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  Advances in neural information processing systems, 31, 2018.
  * [18] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang.  Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks.  In International Conference on Machine Learning, pages 322–332. PMLR, 2019.
  * [19] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang.  On exact computation with an infinitely wide neural net.  Advances in Neural Information Processing Systems, 32, 2019.
  * [20] Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré.  Learning fast algorithms for linear transforms using butterfly factorizations.  In International conference on machine learning, pages 1517–1527. PMLR, 2019.
  * [21] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171–4186, 2019.
  * [22] Ben Mussay, Margarita Osadchy, Vladimir Braverman, Samson Zhou, and Dan Feldman.  Data-independent neural pruning via coresets.  In International Conference on Learning Representations, 2019.
  * [23] Emma Strubell, Ananya Ganesh, and Andrew McCallum.  Energy and policy considerations for deep learning in nlp.  In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 3645–3650, 2019.
  * [24] Sina Alemohammad, Zichao Wang, Randall Balestriero, and Richard Baraniuk.  The recurrent neural tangent kernel.  arXiv preprint arXiv:2006.10246, 2020.
  * [25] Rachit Chhaya, Jayesh Choudhari, Anirban Dasgupta, and Supratim Shit.  Streaming coresets for symmetric tensor factorization.  In International Conference on Machine Learning, pages 1855–1865. PMLR, 2020.
  * [26] Rachit Chhaya, Anirban Dasgupta, and Supratim Shit.  On coresets for regularized regression.  In International Conference on Machine Learning, pages 1866–1876. PMLR, 2020.
  * [27] Lingxiao Huang and Nisheeth K Vishnoi.  Coresets for clustering in euclidean spaces: importance sampling is nearly optimal.  In Proceedings of the 52nd Annual ACM SIGACT Symposium on Theory of Computing, pages 1416–1429, 2020.
  * [28] Mikhail Khodak, Neil A Tenenholtz, Lester Mackey, and Nicolo Fusi.  Initialization and regularization of factorized neural layers.  In International Conference on Learning Representations, 2020.
  * [29] Baharan Mirzasoleiman, Kaidi Cao, and Jure Leskovec.  Coresets for robust training of neural networks against noisy labels.  Neural Information Processing Systems (NeurIPS), 2020.
  * [30] Morad Tukan, Alaa Maalouf, and Dan Feldman.  Coresets for near-convex functions.  Advances in Neural Information Processing Systems, 33, 2020.
  * [31] Nir Ailon, Omer Leibovitch, and Vineet Nair.  Sparse linear networks with a fixed butterfly structure: theory and practice.  In Uncertainty in Artificial Intelligence, pages 1174–1184. PMLR, 2021.
  * [32] Insu Han, Haim Avron, Neta Shoham, Chaewon Kim, and Jinwoo Shin.  Random features for the neural tangent kernel.  arXiv preprint arXiv:2104.01351, 2021.
  * [33] Shaofeng Jiang, Robert Krauthgamer, Xuan Wu, et al.  Coresets for clustering with missing values.  Advances in Neural Information Processing Systems, 34, 2021.
  * [34] Tung Mai, Cameron Musco, and Anup Rao.  Coresets for classification–simplified and strengthened.  Advances in Neural Information Processing Systems, 34:11643–11654, 2021.
  * [35] Alexander Munteanu, Simon Omlor, and David Woodruff.  Oblivious sketching for logistic regression.  In International Conference on Machine Learning, pages 7861–7871. PMLR, 2021.
  * [36] Rachit Chhaya, Anirban Dasgupta, Jayesh Choudhari, and Supratim Shit.  On coresets for fair regression and individually fair clustering.  In International Conference on Artificial Intelligence and Statistics, pages 9603–9625. PMLR, 2022.
  * [37] Supratim Shit, Anirban Dasgupta, Rachit Chhaya, and Jayesh Choudhari.  Online coresets for parameteric and non-parametric bregman clustering.  TMLR, 2022.
  * [38] Murad Tukan, Xuan Wu, Samson Zhou, Vladimir Braverman, and Dan Feldman.  New coresets for projective clustering and applications.  In International Conference on Artificial Intelligence and Statistics, pages 5391–5415. PMLR, 2022.