Training Overparametrized Neural Networks in Sublinear Time

## Abstract

The success of deep learning comes at a tremendous computational and energy cost, and the scalability of training massively overparametrized neural networks is becoming a real barrier to the progress of artificial intelligence (AI). Despite the popularity and low cost-per-iteration of traditional backpropagation via gradient decent, stochastic gradient descent (SGD) has prohibitive convergence rate in non-convex settings, both in theory and practice. To mitigate this cost, recent works have proposed to employ alternative (Newton-type) training methods with much faster convergence rate, albeit with higher cost-per-iteration.
For a typical neural network with m = poly ⁢ ( n ) 𝑚 poly 𝑛 m=\mathrm{poly}(n) italic_m = roman_poly ( italic_n ) parameters and input batch of n 𝑛 n italic_n datapoints in ℝ d superscript ℝ 𝑑 \mathbb{R}^{d} blackboard_R start_POSTSUPERSCRIPT italic_d end_POSTSUPERSCRIPT , the previous work of [ 6 ] requires ∼ m ⁢ n ⁢ d + n 3 similar-to absent 𝑚 𝑛 𝑑 superscript 𝑛 3 \sim mnd+n^{3} ∼ italic_m italic_n italic_d + italic_n start_POSTSUPERSCRIPT 3 end_POSTSUPERSCRIPT time per iteration. In this paper, we present a novel training method that requires only m 1 − α ⁢ n ⁢ d + n 3 superscript 𝑚 1 𝛼 𝑛 𝑑 superscript 𝑛 3 m^{1-\alpha}nd+n^{3} italic_m start_POSTSUPERSCRIPT 1 - italic_α end_POSTSUPERSCRIPT italic_n italic_d + italic_n start_POSTSUPERSCRIPT 3 end_POSTSUPERSCRIPT amortized time in the same overparametrized regime, where α ∈ ( 0.01 , 1 ) 𝛼 0.01 1 \alpha\in(0.01,1) italic_α ∈ ( 0.01 , 1 ) is some fixed constant. This method relies on a new and alternative view of neural networks, as a set of binary search trees, where each iteration corresponds to modifying a small subset of the nodes in the tree. We believe this view would have further applications in the design and analysis of deep neural networks (DNNs).

## 1 Introduction

Deep learning technology achieves unprecedented accuracy across many domains of AI and human-related tasks, from computer vision, natural language processing, and robotics.
This success, however, is approaching its limit and is largely compromised by the computational complexity of these resource-hungry models.
State-of-art neural networks keep growing larger in size, requiring giant matrix operations to train billions of parameters [, , , , , , ]. This barrier is exacerbated by the empirical phenomenon that *overparametrization* in DNNs [] keeps improving model accuracy, despite the danger of overfitting [], motivating the design of complex networks which need to train billions of parameters.
As such, scalable training of deep neural networks is a major challenges of modern AI [, ].

Training a neural network can be broadly viewed a greedy iterative process, starting from an initial set of weight matrices (one per layer of the network). In each iteration, the algorithm chooses a (possibly complicated) rule for updating the value of current weights $W_{i}$ based on the training data, yielding the new weight matrices $W_{i+1}$. The total running time of DNN training is generally composed of two parts: The *number of iterations* (i.e., convergence rate) and the *cost-per-iteration* (i.e., $\mathsf{CPI}$). A long line of research in convex and non-convex optimization has focused on the former question [, , , , , , , , , , , , ]. This paper’s focus is on the latter question.

The most popular iterative method for training DNNs is via *stochastic gradient descent* and its regularized variations [, , , , , ].
The popularity of this method is justified, to a great extent, by the simplicity and fast $\mathsf{CPI}$. Calculating the gradient of the loss function is linear in the dimension of the gradient in each iteration, especially with mini-batch sampling [, ]). Alas, the theoretical convergence rate (number of iterations) of first-order methods
is dauntingly slow in *non-convex* landscapes due to pathological curvatures ($\Omega(\operatorname{poly}(n)\log(1/\epsilon))$ for reducing the training error below $\epsilon$ in overparametrized networks, see e.g., []).

A recent line of work proposed to mitigate this drawback by replacing (S)GD with *second-order* (Newton-type) methods, which exploit information of the Hessian (curvature) of the loss function,
and are proven to converge dramatically faster, at a rate of $O(\log(1/\epsilon))$ iterations, which is *independent* of the input size [, ].
In contrast, Newton methods have a high $\mathsf{CPI}$. since they need to compute the *inverse* of Hessian matrix, which is dense and changes dynamically.
The recent works of [, ] showed that this computational bottleneck can be mitigated for overparametrized DNNs ($m=\operatorname{poly}(n)$) with smooth (resp. $\mathsf{ReLU}$) activations, and presented a Gauss-Newton (resp. NGD) training algorithm with $O(mn^{2})$ training time per iteration. Here $m$ is the number of neurons. We let $n$ be the number of inputs. This runtime was further improved in the work of [], who showed how to implement the Gauss-Newton algorithm in $O(mnd+n^{3})$ time per iteration, which is *linear time* in the network size, assuming $m\gtrsim n^{2}$ (as the dimensions of the Jacobian matrix of the loss is $\Theta(mnd)$ without simplifying assumptions []).

It is tempting to believe that linear-time per iteration [] is unavoidable – For a network with $m$ neurons and a training set of $n$ points in $\mathbb{R}^{d}$, each iteration spends at least $\sim nmd$ time to go through each training datapoint and each neuron. Indeed, this was a common feature of all aforementioned training methods.

Nevertheless, in this paper we present a novel training method with *sublinear* cost per iteration in the network size, while retaining *the same convergence rate* (number of iterations) as the prior state-of-art methods [, , ].
More formally, let $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ be a neural network

$$ $f(x):=\sum_{r=1}^{m}a_{r}\cdot\phi(\langle w_{r},x\rangle-b)$ $$

with bias $b>0$, $a\in\{\pm 1\}^{m}$, each $w_{r}\in\mathbb{R}^{d}$, for all $r\in[m]$. Our main resuilt is as follows.

###### Theorem 1.1 (Main Result, Informal) .

Suppose there are $n$ training data points in $\mathbb{R}^{d}$.
Let $f_{m,n}$ be a sufficiently wide two-layer $\mathsf{ReLU}$ $\mathsf{NN}$ with $m=\operatorname{poly}(n)$ neurons. Let $\alpha\in(0.01,1)$ be some fixed constant.
Let $\epsilon\in(0,0.1)$ be an accuracy parameter. Let ${\cal T}(\epsilon)$ denote the overall time for shrinking loss down to $\epsilon$.
There is a (randomized) algorithm (Algorithm 1) that, with probability $1-1/\operatorname{poly}(n)$, reduces the training error by $1/2$ in each iteration (note that $f_{t}$ is $f_{m,n}$ at time $t$)

$$ $\displaystyle\ell_{2}\mathrm{-loss}(f_{t+1},y)\leq\frac{1}{2}\cdot\ell_{2} \mathrm{-loss}(f_{t},y)$ $$

in amortized cost-per-iteration ($\mathsf{CPI}$)

$$ $\widetilde{O}(m^{1-\alpha}nd+n^{3}).$ $$

The overall running time (including initialization) ${\cal T}(\epsilon)$ is

$$ $\displaystyle O(mnd)+\widetilde{O}(m^{1-\alpha}nd+n^{3}))\cdot\log(1/\epsilon).$ $$

If the algorithm is allowed to use fast matrix multiplication ($\mathsf{FMM}$), then
the $\mathsf{CPI}$ becomes

$$ $\displaystyle\widetilde{O}(m^{1-\alpha}nd+n^{\omega}),$ $$

and the ${\cal T}(\epsilon)$ becomes

$$ $\displaystyle O(mnd)+\widetilde{O}(m^{1-\alpha}nd+n^{\omega})\cdot\log(1/ \epsilon),$ $$

where $\omega$ is the exponent of matrix multiplication, which is currently approximately equal to $2.373$.

The randomness is from two parts: the first part is random initialization weights, and the second part is due to internal randomness of our algorithm.

###### Remark 1.2 .

Notice that the linear cost term $O(mnd)$ for merely computing the network’s loss matrix, is only incurred once at the *initialization* of our training algorithm, whereas in [] and all prior work [, ], this linear term is payed *every iteration* (i.e., $T\cdot mnd$ as opposed to our $T+mnd$).
Our theorem therefore provides a direct improvement over [] when $m=\operatorname{poly}(n)$.

###### Theorem 1.1 (Main Result, Informal) .

###### Remark 1.2 .

#### Key Insight: DNNs as Binary-Search Trees

Our algorithm is based on an alternative view of DNNs, as a *set of binary search trees*, where the relationship between the network’s weights and a training data point is encoded using a binary tree: Each leaf represents the inner product of a neuron and the training data, and each intermediate (non-leaf) node represents the *larger* out of the left and right child. This simple yet new representation of neural networks turns out to enable fast training – The centerpiece of our result is an analysis proving that in each iteration, only a small subset $K$ of paths in this tree collection needs to be updated (amortized worst-case), due to the *sparsity* of activations. Consequently, we only need to update $nK\log m$ tree nodes per iteration. In the Technical Overview Section 3, we elaborate more on its details.

#### Roadmap.

We describe the organization of this work in next a few sentences. We propose our main problem and present the tools we need to use in Section 2. In Section 3, we specifically overview the techniques used in this paper. In Section 4, we analyze the correctness of our algorithm, specifically, we prove the training loss converges. In Section 5, we analyze the running time of our algorithm. In Section 6, we provide a conclusion.

## 2 Preliminaries

### 2.1 Notations.

For any positive integer $n$, we use $[n]$ to denote set $\{1,2,\cdots,n\}$. For any function $f$, we use $\widetilde{O}(f)$ to denote $f\cdot\operatorname{poly}(\log f)$. For two vectors $w$ and $x$, we use $\langle w,x\rangle$ to denote inner product. We use $a^{\top}$ to denote the transpose of $a$. We use $\operatorname*{{\mathbb{E}}}[]$ to denote expectation and $\Pr[]$ for probability. For convenience, we use $\mathsf{FMM}$ to denote fast matrix multiplication. We use $\mathsf{NTK}$ to denote neural tangent kernel. We use $\mathsf{ReLU}$ to denote rectified linear unit. We use $\mathsf{NN}$ to denote neural network. We use $\mathsf{CPI}$ to represent the cost per iteration. We use PSD to denote positive semidefinite.

### 2.2 A Sketching Tool

Sarlós [] firstly introduced the notation of subspace embedding. Many numerical linear algebra applications have used that concept and its variations [, , , , , ].
The formal definition is:

###### Definition 2.1 (Oblivious subspace embedding, OSE [ 41 ] ) .

Given an $N\times k$ matrix $B$, an $(1\pm\epsilon)$ $\ell_{2}$-subspace embedding for the column space of $B$ is a matrix $S$, such that for any $x\in\mathbb{R}^{k}$,

$$ $\displaystyle(1-\epsilon)\|Bx\|_{2}^{2}\leq\|SBx\|_{2}^{2}\leq(1+\epsilon)\|Bx \|_{2}^{2}.$ $$

Equivalently, let $U$ be the matrix whose columns form an orthonormal basis containing the column vectors of $B$, then

$$ $\displaystyle\|I-U^{\top}S^{\top}SU\|_{2}\leq\epsilon.$ $$

It is known that subspace embedding can be given by a Fast-JL sketching matrix [, , , , , ] with a classical $\epsilon$-net argument,

###### Lemma 2.2 .

Assume that $N=\operatorname{poly}(k)$. Assume $\delta\in(0,0.1)$.
For a matrix $B\in\mathbb{R}^{N\times k}$, we can produce an $(1\pm\epsilon)$ $\ell_{2}$-subspace embedding $S\in\mathbb{R}^{k\operatorname{poly}(\log(k/\delta))/\epsilon^{2}\times N}$ for $B$ with probability at least $1-\delta$.
In addition, $SB$ takes $O(Nk\cdot\operatorname{poly}\log k)$ time to be generated.

###### Definition 2.1 (Oblivious subspace embedding, OSE [ 41 ] ) .

###### Lemma 2.2 .

### 2.3 Model formalization

In this section, we formalize the $\mathsf{NN}$ model and the main problem of this paper.
When there is no ambiguity, we will always use the notations in this section throughout the whole paper.

We first define the 2-layer $\mathsf{ReLU}$ activated neural network and its loss function.

###### Definition 2.3 (2-layer $\mathsf{ReLU}$ activated neural network) .

Suppose the dimension of input is $d$, the number of intermediate nodes (or hidden neurons) is $m$, the dimension of output is $1$, the batch size is $n$ and the shifted parameter is $b$ ($b\geq 0$). Then the weight of the first layer can be characterized by $m$ $d$-dimensional vectors $w_{1},w_{2},\cdots,w_{m}$, and the weight of the second layer can be characterized by $m$ scalars $a_{1},a_{2},\cdots,a_{m}$. For convenience, define

$$ $\displaystyle W=[~{}w_{1}^{\top}~{}w_{2}^{\top}~{}\cdots~{}w_{m}^{\top}~{}]^{\top}$ $$

and

$$ $\displaystyle a=[~{}a_{1}~{}a_{2}~{}\cdots~{}a_{m}~{}]^{\top},$ $$

given an input $x\in\mathbb{R}^{d}$, the 2-layer $\mathsf{ReLU}$ activated neural network outputs

$$ $\displaystyle f(W,x,a)=\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}\phi(\langle w_{r} ,x\rangle)$ $$

where $\phi(x)=\max\{x,b\}$ is called shifted $\mathsf{ReLU}$ activation function.

For simplicity, we suppose the data is normalized, that is, $\|x\|_{2}=1$. This is natural in both practical machine learning, and machine learning theory.

We also suppose $a\in\{-1,+1\}^{m}$ is fixed throughout training. This is also natural in the area of theoretical deep learning [, , , , , , ].

###### Definition 2.4 (Loss function) .

Suppose the dimension of input is $d$, the number of intermediate nodes (or hidden neurons) is $m$, the dimension of output is $1$, the batch size is $n$ and the shifted parameter is $b$ ($b\geq 0$).
For a fixed set of $n$
points $x_{1},x_{2},\cdots,x_{n}\in\mathbb{R}^{d}$ and their corresponding labels $y_{1},y_{2},\cdots,y_{n}\in\mathbb{R}$. Consider the following loss function:

$$ $\displaystyle\mathcal{L}(W):=\frac{1}{2}\sum_{i=1}^{n}(y_{i}-f(W,x_{i},a))^{2}.$ $$

By mathematics,

$$ $\displaystyle\frac{\partial f(W,x,a)}{\partial w_{r}}=\frac{1}{\sqrt{m}}a_{r}x {\bf 1}_{w_{r}^{\top}x\geq b},~{}~{}~{}\forall r\in[m].$ (1) $$

Thus one can calculate the gradient of loss function $\mathcal{L}$

$$ $\displaystyle\frac{\partial\mathcal{L}(W)}{\partial w_{r}}=\frac{1}{\sqrt{m}} \sum_{i=1}^{n}(f(W,x_{i},a)-y_{i})\cdot a_{r}\cdot x_{i}\cdot{\bf 1}_{w_{r}^{ \top}x_{i}\geq b}.$ (2) $$

Then we define the prediction function, the Jacobi matrix, and the Gram matrix.

###### Definition 2.5 (prediction function) .

For a batch of inputs $\{(x_{i},y_{i})\}_{i\in[n]}\in\mathbb{R}^{d}\times\mathbb{R}$, we denote

$$ $\displaystyle\alpha_{r,i}(t):=\phi(\langle w_{r}(t),x_{i}\rangle)$ $$

for every $r\in[m]$ and $i\in[n]$.

Prediction function $f_{t}:\mathbb{R}^{d\times n}\rightarrow\mathbb{R}^{n}$ at time $t$ is defined as follows

$$ $\displaystyle f_{t}=\frac{1}{\sqrt{m}}\sum_{r\in[m]}\begin{bmatrix}a_{r}\cdot \alpha_{r,1}(t))\\ a_{r}\cdot\alpha_{r,2}(t))\\ \vdots\\ a_{r}\cdot\alpha_{r,n}(t))\\ \end{bmatrix}$ $$

Note that $w_{r}(t)$ is the $r$-th weight of the first layer after training of $t$ times.

For convenience, we define weight matrix

$$ $\displaystyle W_{t}=[w_{1}(t)~{}w_{2}(t)~{}\cdots~{}w_{m}(t)]\in\mathbb{R}^{d \times m}$ $$

In addition, we write
data matrix

$$ $\displaystyle X=[x_{1}~{}x_{2}~{}\cdots~{}x_{n}]\in\mathbb{R}^{d\times n}.$ $$

###### Definition 2.6 (Jacobi matrix and related definitions) .

For each $i\in[n]$, $r\in[m]$ and $t\in[T]$, we define

$$ $\displaystyle\beta_{r,i}(t):={\bf 1}_{\langle w_{r}(t),x_{i}\rangle\geq b}.$ $$

For every time step $t$, we use $J_{t}\in\mathbb{R}^{n\times m}$ to denote the Jacobian matrix at $t$. Formally, it can be written as

$$ $\displaystyle~{}J_{t}$ $\displaystyle=$ $\displaystyle~{}\frac{1}{\sqrt{m}}\left[\begin{matrix}a_{1}x_{1}^{\top}\beta_{ 1,1}(t)&a_{2}x_{1}^{\top}\beta_{2,1}(t)&\cdots&a_{m}x_{1}^{\top}\beta_{m,1}(t) \\ a_{1}x_{2}^{\top}\beta_{1,2}(t)&a_{2}x_{2}^{\top}\beta_{2,2}(t)&\cdots&a_{m}x_ {2}^{\top}\beta_{m,2}(t)\\ \vdots&\vdots&\ddots&\vdots\\ a_{1}x_{n}^{\top}\beta_{1,n}(t)&a_{2}x_{n}^{\top}\beta_{2,n}(t)&\ldots&a_{m}x_ {n}^{\top}\beta_{m,n}(t)\\ \end{matrix}\right].$ $$

For each $i\in[n]$, we define $J_{t}(x_{i})$ as the $i$-th row of $J_{t}$.

###### Definition 2.7 (Gram matrix) .

Let $G_{t}\in\mathbb{R}^{n\times n}$ denote the Gram matrix.
Then $G_{t}$ can be formally written as $G_{t}=J_{t}J_{t}^{\top}$. The $(i,j)$-th entry of $G_{t}$ is the inner product between gradient in terms of $x_{i}$ and the gradient in terms of $x_{j}$, i.e.,

$$ $\displaystyle(G_{t})_{i,j}:=\langle\frac{f(W_{t},x_{i})}{\partial W},\frac{f(W _{t},x_{j})}{\partial W}\rangle.$ $$

[, , ] gave a crucial observation that the asymptotic of the Gram matrix $G$ is equal to a PSD matrix $K\in\mathbb{R}^{n\times n}$. The formal definition is

$$ $\displaystyle K(x_{i},x_{j}):=\operatorname*{{\mathbb{E}}}_{w\sim\mathcal{N}(0 ,I)}\left[x_{i}^{\top}x_{j}\textbf{1}_{\langle w,x_{i}\rangle\geq b,\langle w, x_{j}\rangle\geq b}\right].$ (3) $$

[, ] only consider the case where $b=0$ and [] consider the general case $b\geq 0$.

###### Remark 2.8 .

We use $\lambda$ to denote the minimal eigenvalue of the kernel matrix $K$ defined in Eq. (3).

###### Definition 2.3 (2-layer 𝖱𝖾𝖫𝖴 𝖱𝖾𝖫𝖴 \mathsf{ReLU} sansserif_ReLU activated neural network) .

###### Definition 2.4 (Loss function) .

###### Definition 2.5 (prediction function) .

###### Definition 2.6 (Jacobi matrix and related definitions) .

###### Definition 2.7 (Gram matrix) .

###### Remark 2.8 .

### 2.4 Problem definition

We formalize our main problem as follows.

###### Definition 2.9 (Main problem) .

The goal of this paper is to propose a training algorithm such that for an arbitrary 2-layer $\mathsf{ReLU}$ activated neural network defined in Definition 2.3, it converges with high probability, and the running time of each iteration is sublinear in $nmd$ (i.e. $o(nmd)$).

###### Definition 2.9 (Main problem) .

### 2.5 Neural tangent kernel and its relation with data separability

Neural Tangent Kernel ($\mathsf{NTK}$) is a Kernel matrix related to a multi-layer $\mathsf{ReLU}$ activated neural network. It is crucial in the analysis of Jacobi matrix. [] expanded the related concepts and revealed their properties, especially its relation to the data separability of an input batch.

As for data separability, it is a common assumption to the input of a neural network, and it has been used in many over-parameterized neural network literature [, ]. We first define kernels,

###### Definition 2.10 .

Let $b\geq 0$ be the shift parameter. We define continuous version of the shifted $\mathsf{NTK}$ $H^{\operatorname{cts}}$ and discrete version of shifted $\mathsf{NTK}$ $H^{\operatorname{dis}}$ as

$$ $\displaystyle H_{i,j}^{\operatorname{cts}}:=$ $\displaystyle~{}\operatorname*{{\mathbb{E}}}_{w\sim{\cal N}(0,I)}[x_{i}^{\top} x_{j}{\bf 1}_{w^{\top}x_{i}\geq b,w^{\top}x_{j}\geq b}],$ $\displaystyle H_{i,j}^{\operatorname{dis}}:=$ $\displaystyle~{}\frac{1}{m}\sum_{r=1}^{m}[x_{i}^{\top}x_{j}{\bf 1}_{w_{r}^{ \top}x_{i}\geq b,w_{r}^{\top}x_{j}\geq b}].$ $$

Next, we define data separability,

###### Definition 2.11 (Separability of input data) .

Suppose we are given $n$ (normalized) input data points

$$ $\displaystyle\{x_{1},x_{2},\cdots,x_{n}\}\subseteq\mathbb{R}^{d}.$ $$

Assume those points satisfy that $\forall i\in[n],\|x_{i}\|_{2}=1$. For each $i,j$, we define

$$ $\displaystyle\delta_{i,j}^{+}=x_{i}+x_{j}\text{~{}and~{}}\delta_{i,j}^{-}=x_{i }-x_{j}.$ $$

Let $\delta$ be the data separability parameter, formally,

$$ $\displaystyle\delta:=\min_{i\neq j}\{\min\{\|\delta_{i,j}^{+}\|_{2},\|\delta_{ i,j}^{-}\|_{2}\}\}.$ $$

[] has given a property of the minimal eigenvalue of the $\mathsf{NTK}$ of a shifted $\mathsf{ReLU}$ activated neural network.

###### Lemma 2.12 (Lemma C.1 in [ 47 ] ) .

Let $m$ be number of samples of $H^{\operatorname{dis}}$. As long as

$$ $\displaystyle m=\Omega(\lambda^{-1}n\log(n/\rho)),$ $$

then

$$ $\displaystyle\Pr[\lambda_{\min}(H^{\operatorname{dis}})\geq\frac{3}{4}\lambda] \geq 1-\rho.$ $$

Prior work ([]) has shown the relation between the data separability of the input of a neural network and the eigenvalue of the Kernel. But their work focuses on unshifted $\mathsf{ReLU}$ activated neural network. For shifted $\mathsf{ReLU}$ activated neural network, [] provided a further generalization to the shifted Kernel matrix.

###### Theorem 2.13 (Theorem F.1 in [ 47 ] ) .

Consider $n$ points $x_{1},\dots,x_{n}\in\mathbb{R}^{d}$ with $\ell_{2}$-norm all equal to $1$, and consider a random variable $w\sim{\cal N}(0,I_{d})$. Define matrix

$$ $\displaystyle X\in\mathbb{R}^{n\times d}=[x_{1}~{}\dots~{}x_{n}]^{\top}.$ $$

Suppose the data separability of the $n$ points is $\delta$ where $\delta<\sqrt{2}$. Let shift parameter $b\geq 0$. Recall the continuous Hessian matrix $H^{\operatorname{cts}}$ is defined by

$$ $\displaystyle~{}H^{\operatorname{cts}}_{i,j}$ $\displaystyle:=$ $\displaystyle~{}\operatorname*{{\mathbb{E}}}_{w\sim\mathcal{N}(0,I)}[x_{i}^{ \top}x_{j}{\bf 1}_{w^{\top}x_{i}\geq b,w^{\top}x_{j}\geq b}],\forall(i,j)\in[n ]\times[n].$ $$

Let $\lambda:=\lambda_{\min}(H^{\operatorname{cts}})$. Then $\lambda$ has the follow sandwich bound,

$$ $\displaystyle\lambda\in[\exp(-b^{2}/2)\cdot\frac{\delta}{100n^{2}},\exp(-b^{2} /2)].$ $$

###### Definition 2.10 .

###### Definition 2.11 (Separability of input data) .

###### Lemma 2.12 (Lemma C.1 in [ 47 ] ) .

###### Theorem 2.13 (Theorem F.1 in [ 47 ] ) .

## 3 Technical Overview

Here, we describe the outline of the main ideas required to prove Theorem 1.1.

#### Key Ideas

Our algorithm relies on two simple but powerful observations about training 2-layer neural networks: The first observation is that the Jacobian matrix of the loss function is *sparse* – When weights are initialized randomly (with appropriately chosen bias parameter $b$), the fraction of nonzero entries in the Jacobi matrix is small.
Let $c$ be some fixed constant in $[0.1,1]$. We show that there is a choice of the parameter $b$ ensuring simultaneously that(^1^11We refer the readers to Section F for more details.)

- •
For every input $x_{i}$, there are only $O(m^{1-c})$ activated neurons;
- •
The loss of each iteration is still at most a half of the loss of the last iteration.

Our second observation is that the *positions* of the nonzero entries in Jacobian matrix do not change much. This can be seen using the “gradient flow” equation (via Gauss-Newton method)
$W_{t+1}=W_{t}-J_{t}^{\top}g_{t},$
where
$g_{t}:=\arg\min_{g}\|J_{t}J_{t}^{\top}g_{t}-(f_{t}-y)\|_{2}$.
Since the Jacobian matrix is sparse, it is not hard to see that only a little fraction of the weights need to be modified, i.e., the change from $W_{t}$ to $W_{t+1}$ involves updating only a small number of entries.

These two observations suggests a natural “binary-search” type algorithm for updating the weight matrix in *sublinear time* $o(nmd)$ per iteration.

#### Threshold search data structure

We design a dynamic data structure for detecting and maintaining the non-zero entries of the Jacobian matrix $J$ of the network loss, as it evolves over iterations. Notice that whether an entry of $J$ is nonzero is equivalent to whether the inner product of an input $x_{i}$ and a weight $w_{j}$ is larger than $b$ (hence $\phi(w_{j}^{\top}x_{i})>0$).

Accordingly, for every input $x_{i}$ in a batch, our algorithm maintains a binary search tree $\mathcal{T}_{i}$ where each leaf stores the inner product of $x_{i}$ and a weight $w_{j}$, and every non-leaf node stores the the *maximum* of the values of its two children. In this way, non-zero entries can be found by searching, in all the trees $\{\mathcal{T}_{i}\}_{i\in[n]}$, from root to leaf and ignoring the unnecessary branches.

To implement this process efficiently, our data structure
needs to support the following three operations (See Section B for the formal details):
(1) Initialization. Given input vectors $x_{1},\cdots,x_{n}$ and weight vectors $w_{1},w_{2},\cdots,w_{m}$ as input, it constructs $n$ binary trees $\mathcal{T}_{1},\ldots,\mathcal{T}_{n}$ as described above, in $O(mnd)$ time. (2) Updating of weights. Taken an index $j\in[m]$ and a target value $z$, it replaces $w_{j}$ by $z$ in $O(nd+n\log m)$ time, as if initializing it with $w_{1},w_{2},\cdots,w_{j-1},z,w_{j+1},\cdots,w_{m}$ from scratch.
(3) Threshold Search Query. Given an index $i$ and a threshold $\tau$ as input, our data structure rapidly finds all the weights $w_{j}$ which satisfies $\langle x_{i},w_{j}\rangle\geq\tau$ in $O(K_{q}\log m)$ time, where $K_{q}$ is the number of satisfied weights. They can be used to find the nonzero entries of the Jacobian matrix $J$.

#### A Fast DNN Training Algorithm

Using the above dynamic data structure, we design a fast neural network training algorithm (see Algorithm 1) composed of initialization and the (dynamic) training process. At initialization, it initializes the weight vector $W_{0}$ randomly.

The training process consists of maintaining *sparse-recovery* sketches [, , ], online regression, and implicit weight maintenance. The goal of the first two techniques is to efficiently solve the $t$-th iteration regression problem (cf. [])
$g_{t}:=\arg\min_{g}\|J_{t}J_{t}^{\top}g-(f_{t}-y)\|,$.
The idea of implicit-weight-maintenance (via our data structure) is to update weights using the information propagated by the loss function.

The details of these three tools can be summarized as follows:

- •
Sketch maintenance The goal of sketch computing is to eliminate the disastrous influence of the high dimension of $J_{t}^{\top}$ (it has $md$ rows) when solving regression problem in Eq. (3). Roughly speaking, in sketch computing, we find a sketch matrix $S$ with far smaller rows than $J_{t}^{\top}$ such that for any $d$-dimensional vector $x$, $\|SJ_{t}^{\top}x\|_{2}$ is very close to $\|J_{t}^{\top}x\|_{2}$. We show that sketch computing runs in $o(mnd)$ time.
- •
Iterative regression solver To speed-up the solution of the online regression problem (3), we show how to implement the iterative Conjugate-Gradient solver (a-la []) *in sub-linear time*
to find an approximate solution $g_{t}$ in time $o(mnd)+\widetilde{O}(n^{3})$. We then prove that the (accumulated) approximation errors do not harm the convergence rate and precision in our analysis.
- •
Implicit weight maintenance The goal of implicit weight maintenance is to update weights according to the outcome of the iterative regression solver. Updating a single weight can be done by calling Update once. With the result of iterative regression and the fact that only $m^{-c}$ (where $c$ is some fixed constant $c\in[0.1,1]$)
fraction of entries of $J_{t}$ are nonzero, we show that our algorithm finishes the update of weights in $o(mnd)$ time.

The details can be found in the pseudocode of Algorithm 5.

## 4 Convergence Analysis

We focus on the convergence of our training algorithm in this section and leave the proof of running time in Section 5. Specifically, the goal of this section is to prove the following result, which implies that for the neural network randomly initialized at the beginning of our algorithm, the loss function converges linearly with high probability. This section only contains a proof sketch. For more detailed correctness analysis, we refer the readers to section D.
Our main convergence result is the following:

###### Theorem 4.1 (Formal version of Theorem 1.1 , the convergence part) .

Let $m$ be the width of the $\mathsf{NN}$.
If

$$ $\displaystyle m=\Omega(\max\{\lambda^{-4}n^{4},\lambda^{-2}n^{2}d\log(16n/\rho )\}),$ $$

then there is a constant $c^{\prime}>0$ so that our algorithm obtains

$$ $\displaystyle\|f_{t+1}-y\|_{2}\leq 0.5\cdot\|f_{t}-y\|_{2}.$ $$

It holds with probability

$$ $\displaystyle 1-\frac{5}{2}\rho-n^{2}\cdot\exp(-m\cdot\min\{c^{\prime}e^{-b^{2 }/2},\frac{R}{10\sqrt{m}}\})$ $$

The randomness comes from two parts: the initialization of neural network and iterative algorithm itself.

###### Theorem 4.1 (Formal version of Theorem 1.1 , the convergence part) .

#### Bounding the Function Value and Jacobian at the Initialization

We provide a lemma which shows that, with random initialization, as long as the 2-layer $\mathsf{NN}$ is wide enough, the norm of weight matrix, the initial predicted value and the Frobenius norm of the initial Jacobi matrix are all not large with high probability. We defer its proof into Section D.

###### Lemma 4.2 (Informal version of Lemma D.1 ) .

Consider shifted $\mathsf{ReLU}$. Suppose $m$ is the width of neural network. If $m=\Omega(d\log(16n/\rho))$, then we have

- •
$\|W_{0}\|_{2}=O(\sqrt{m})$.
- •
$\max_{i\in[n]}|f(W,x_{i})|=O(1)$.
- •
$\max_{i\in[n]}\|J_{W_{0},x_{i}}\|_{F}=O(1)$.

holds with probability $1-\rho/2$.

###### Lemma 4.2 (Informal version of Lemma D.1 ) .

#### G 𝐺 G italic_G does not move much when W 𝑊 W italic_W does not move much

We provide a lemma which proves that, as long as the 2-layer $\mathsf{NN}$ is wide enough, then with high probability that, for randomly initialized weights $W_{0}$, if $W_{0}$ changes to $W$ after a small change, then the Gram matrix $G_{W}$ will not move much and the minimal eigenvalue of $G_{W}$ will also not move much. And We leave its proof in Section D.

###### Lemma 4.3 (Shifted Perturbation Lemma, informal version of Lemma D.2 ) .

Consider shifted ReLU with $b$. Let $b\geq 0$. Let $R_{0}>0$.
Suppose

$$ $\displaystyle m\geq\Omega(1)\cdot\max\{b^{2}R_{0}^{2},n^{2}R_{0}^{2}\lambda^{- 2},n\lambda^{-1}\log(n/\rho)\},$ $$

then with prob.

$$ $\displaystyle\geq 1-\rho-n^{2}\cdot\exp\big{(}-m\cdot\min\{c^{\prime}e^{-b^{2} /2},\frac{R_{0}}{10\sqrt{m}}\}\big{)},$ $$

for any weight $W\in\mathbb{R}^{d\times m}$ satisfying $\max_{r\in[m]}\|w_{r}-w_{r}(0)\|_{2}\leq R_{0}/\sqrt{m}$, the following holds:

$$ $\displaystyle\|G_{W}-G_{W_{0}}\|_{F}\leq\lambda/2,~{}~{}\text{and}~{}~{} \lambda_{\min}(G_{W})\geq\lambda/2.$ $$

Note that $w_{r}$ is representing the $r$-th column of $W$.

###### Lemma 4.3 (Shifted Perturbation Lemma, informal version of Lemma D.2 ) .

#### Perturbed weights difference under shifted NTK

We give a lemma which proves that, as long as the 2-layer $\mathsf{NN}$ is wide enough, then with high probability that, for randomly initialized weights $W_{0}$, if $W_{0}$ changes to $W$ after a small change, the each row $J_{W,x_{i}}$ of Jacobi matrix $J_{W}$ will not change much, and the Frobenius norm of $J_{W}$ will also not change much. We leave its proof in Section D.

###### Lemma 4.4 (Informal version of Lemma D.3 ) .

Suppose $R_{0}\geq 1$ and $m=\widetilde{\Omega}(n^{2}R_{0}^{2})$. With probability at least $1-\rho$ over the random initialization of $W_{0}$, the following holds for any set of weights $w_{1},\ldots w_{m}\in\mathbb{R}^{d}$ satisfying $\max_{r\in[m]}\|w_{r}-w_{r}(0)\|_{2}\leq R_{0}/\sqrt{m}$,

- •
$\|W-W_{0}\|=O(R_{0})$,
- •
$\|J_{W,x_{i}}-J_{W_{0},x_{i}}\|_{2}=\widetilde{O}({R_{0}^{1/2}}/{m^{1/4}})$ and $\|J_{W}-J_{W_{0}}\|_{F}=\widetilde{O}({n^{1/2}R_{0}^{1/2}}/{m^{1/4}})$,
- •
$\|J_{W}\|_{F}=O(\sqrt{n})$.

###### Lemma 4.4 (Informal version of Lemma D.3 ) .

#### Induction Hypothesis

Finally, we’re ready to prove our major theorem, Theorem 4.1. Note that we only need to prove the induction hypothesis described in definition 4.5, then Theorem 4.1 holds by mathematical induction. We divide the proof of this hypothesis into 2 parts and prove them in section E.1 and section E.2 respectively.

###### Definition 4.5 (Induction hypothesis) .

Define $R_{0}\approx n/\lambda$. For any fixed $t$, if

$$ $\displaystyle~{}\|f_{t}-y\|_{2}\leq\frac{1}{2}\|f_{t-1}-y\|_{2}$ and $\displaystyle~{}\max_{r\in[m]}\|w_{r}(t)-w_{r}(0)\|_{2}\leq R_{0}/\sqrt{m}.$ $$

Then we have

$$ $\displaystyle~{}\|f_{t+1}-y\|_{2}\leq\frac{1}{2}\|f_{t}-y\|_{2}$ and $\displaystyle~{}\max_{r\in[m]}\|w_{r}(t+1)-w_{r}(0)\|_{2}\leq R_{0}/\sqrt{m}.$ $$

Formally, we describe the process of proving this hypothesis by the following Lemma 4.6, and specific proof can be seen in Section E.

###### Lemma 4.6 .

Suppose initial weights $W_{0}$ satisfies the restriction of Lemma 4.2, 4.3 and 4.4, then the induction hypothesis described in Definition 4.5 holds.

Figure: Algorithm 1 Our training algorithm, informal version of Algorithm 5

###### Definition 4.5 (Induction hypothesis) .

###### Lemma 4.6 .

## 5 Running Time Analysis

This section focuses on analyzing the running time of our algorithm. It will show that when $m$ is large enough, the $\mathsf{CPI}$ is $o(nmd)$.
We first present Theorem 5.1, our main running time result of the paper. We then provide three lemmas (Lemma 5.3, Lemma 5.4 and Lemma 5.5) to prove our main theorem.
For more proof details of the running time, we refer the readers to Section F. For simplicity of presentation, we use $o(m)$ and $o(mnd)$ in this section. In Section F, we explicitly compute time by $m^{1-\alpha}$ and $m^{1-\alpha}nd$ where $\alpha\in[0.01,1)$ is some fixed constant. Our main running time result is the following:

###### Theorem 5.1 (The running time part of Theorem 1.1 ) .

The cost per iteration ($\mathsf{CPI}$) of our algorithm is

$$ $\displaystyle o(mnd)+\widetilde{O}(n^{3})$ $$

without using $\mathsf{FMM}$. The $\mathsf{CPI}$ of our algorithm is

$$ $\displaystyle o(mnd)+\widetilde{O}(n^{\omega})$ $$

with using $\mathsf{FMM}$.

###### Proof.

Combining Lemma F.2, Lemma F.3 and Lemma F.4, the computation time of each iteration is

$$ $\displaystyle~{}\widetilde{O}(n^{2}m^{0.76}d)+\widetilde{O}(nm^{0.76}d+n^{3})$ $\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}+O(n^{2}m^{0.76}(d+\log m))$ $\displaystyle=$ $\displaystyle~{}\widetilde{O}(n^{2}m^{0.76}d+n^{3}).$ $$

And if using $\mathsf{FMM}$, similarly the running time is

$$ $\displaystyle\widetilde{O}(n^{2}m^{0.76}d+n^{\omega}).$ $$

By Theorem 4.1, we have:
The time to reduce the training loss to $\epsilon$ is

$$ $\displaystyle\widetilde{O}((n^{2}m^{0.76}d+n^{3})\log(1/\epsilon)).$ $$

Taking advantage of $\mathsf{FMM}$, the time is

$$ $\displaystyle\widetilde{O}((n^{2}m^{0.76}d+n^{\omega})\log(1/\epsilon)).$ $$

Further, for example, if $m=n^{c}$ where $c$ is some large constant, then

$$ $\displaystyle n^{2}m^{0.76}d\leq nm^{1-\alpha}d$ $$

where $\alpha\in[0.1,0.24)$.
Hence the time of each iteration is

$$ $\displaystyle\widetilde{O}(m^{1-\alpha}nd+n^{3}),$ $$

and the time to reduce the training loss to $\epsilon$ is

$$ $\displaystyle\widetilde{O}((m^{1-\alpha}nd+n^{3})\log(1/\epsilon)).$ $$

Taking advantage of $\mathsf{FMM}$, the time is

$$ $\displaystyle\widetilde{O}((m^{1-\alpha}nd+n^{\omega})\log(1/\epsilon)).$ $$

Thus we complete the proof.
∎

###### Theorem 5.1 (The running time part of Theorem 1.1 ) .

###### Proof.

#### Sketch Computing.

We provide the choice of sketching parameters in the following definition and give a lemma that analyzes the running time of the sketch computing process in Algorithm 5 under these parameters. It will imply that sketch computing is sublinear in $m$.

###### Definition 5.2 (sketch parameters) .

We choose sketch parameters in the following ways:

$$ $\displaystyle\epsilon_{\mathrm{sketch}}=$ $\displaystyle~{}0.1,$ $\displaystyle\delta_{\mathrm{sketch}}=$ $\displaystyle~{}1/\operatorname{poly}(n),$ $\displaystyle s_{\mathrm{sketch}}=$ $\displaystyle~{}n\operatorname{poly}(\epsilon_{\mathrm{sketch}}^{-1},\log(n/ \delta_{\mathrm{sketch}}))$ $$

###### Lemma 5.3 (Step 1, sketch computing. Informal version of Lemma F.2 ) .

The sketch computing process of Algorithm 1 (its formal version is Algorithm 5) runs in time
$o(mnd).$

###### Definition 5.2 (sketch parameters) .

###### Lemma 5.3 (Step 1, sketch computing. Informal version of Lemma F.2 ) .

#### Iterative regression.

We present a lemma that analyzes the running time of the iterative regression process in Algorithm 5. It implies that the running time of the iterative regression is sublinear in $m$.

###### Lemma 5.4 (Step 2, running time of iterative regression. Informal version of Lemma F.3 ) .

The iterative regression of Algorithm 1 (its formal version is Algorithm 5) runs in time

$$ $\displaystyle O(o(mnd)\log(n/\delta)+n^{3}).$ $$

Taking advantage of $\mathsf{FMM}$, it takes time

$$ $\displaystyle O(o(mnd)\log(n/\delta)+n^{\omega}),$ $$

where $\omega$ is the exponent of matrix multiplication. Currently $\omega\approx 2.373$ [].

###### Lemma 5.4 (Step 2, running time of iterative regression. Informal version of Lemma F.3 ) .

#### Implicit weight maintenance.

We give a lemma that analyzes the running time of the implicit weight maintenance process in Algorithm 5. It implies that implicit weight maintenance process is sublinear in $m$.

###### Lemma 5.5 (Step 3, implicit weight maintenance. Informal version of Lemma F.4 ) .

The implicit weight maintenance of Algorithm 1(its formal version is Algorithm 5) runs in time

$$ $\displaystyle o(nm)\cdot(d+\log m).$ $$

###### Lemma 5.5 (Step 3, implicit weight maintenance. Informal version of Lemma F.4 ) .

## 6 Conclusion

The computational cost of training massively overparametrized DNNs is posing a major scalability barrier to the progress of AI, and motivates rethinking the traditional SGD-based training algorithms.
For a neural network with $m$ parameters and an input batch of $n$ datapoints in $\mathbb{R}^{d}$, previous state-of-art [, ] show that dramatically fewer iterations (epochs) $T_{\epsilon}$ can be achieved via second-order methods, albeit with $O(mnd+n^{3})$ cost per iteration, i.e.,

$$ $\displaystyle O(T_{\epsilon}\cdot mnd)$ $$

overal time to reduce training error below $\epsilon$.
Our work proposes a simple yet powerful view of the gradient flow process on wide DNNs ($m=\operatorname{poly}(n)$), as a collection of *slowly-changing binary search trees*, enabling the design of a training algorithm for 2-layer overparametrized DNNs in *sublinear* cost-per-iteration, while enjoying the ultra-fast convergence rate of second-order (Gauss-Newton) methods, i.e., in total time

$$ $\displaystyle\widetilde{O}(T_{\epsilon}+mnd)$ $$

instead of the aforementioned

$$ $\displaystyle\widetilde{O}(T_{\epsilon}\cdot mnd).$ $$

While this paper introduces novel contributions to the field of deep neural networks (DNNs) by proposing a sublinear cost-per-iteration (CPI) training method, there are several limitations and potential directions for future research. For example, the core insight of this paper is the representation of DNNs as binary search trees. This alternative view of DNNs is a promising direction, but it may have limitations when dealing with other types of networks (for example, recurrent neural networks).

Future research could explore adapting this approach to other network architectures. Whatsmore, the research currently focuses on two-layer ReLU neural networks. Future research could explore the applicability of the proposed training method to more complex, deeper architectures or different activation functions, thus broadening the scope of this work.

We believe our new representation of DNNs will have further algorithmic implications to scaling training and inference in DNNs.

## Appendix

#### Roadmap.

The appendix of this paper is organized as follows. Section A presents the preliminary tools which are used in the other parts of appendix. Section B presents the complete description and implementation of the threshold search data structure. Section C presents the formal algorithm representation of our fast neural network training algorithm. Section D shows the omitted proofs of some lemmas in the convergence analysis. Section E presents the complete the induction hypothesis proof to show the convergence of our training algorithm. Section F presents the formal proof of the running time of our training algorithm, especially shows a more specific conclusion compared with the main body. Section G presents a formal analysis of the applying conditions of our training algorithm.

## Appendix A Preliminary

This section shows some preliminary tools to be used later.
In Section A.1 we introduce the result for fast matrix multiplication. In Section A.2 we state the probability tools to be used. In Section A.3 we presented a previous result on the relationship of changes of weights and the change of the shifted NTK matrix. In Section A.4 we provide some useful results about fast regression. In Section A.5 we provide results about sparsity-based preserving.

### A.1 Fast matrix multiplication

We state a standard fact for fast matrix multiplication ($\mathsf{FMM}$).

###### Fact A.1 ( $\mathsf{FMM}$ ) .

Given an $n\times n$ matrix $A$ and another $n\times n$ matrix $B$, the time of multiplying $A$ and $B$ is $n^{\omega}$, where $\omega\approx 2.373$ is the exponent of matrix multiplication. Currently, $\omega\approx 2.373$ [].

###### Fact A.1 ( 𝖥𝖬𝖬 𝖥𝖬𝖬 \mathsf{FMM} sansserif_FMM ) .

### A.2 Probability Tools

We list some probability tools which are useful in our analysis.

###### Lemma A.2 (Chernoff bound [ 9 ] ) .

Let $Z=\sum_{i=1}^{n}Z_{i}$, where $Z_{i}=1$ with probability $p_{i}$ and $Z_{i}=0$ with probability $1-p_{i}$, and all $Z_{i}$ are independent. We define $\mu=\operatorname*{{\mathbb{E}}}[Z]=\sum_{i=1}^{n}p_{i}$. Then
1. $\Pr[Z\geq(1+\delta)\mu]\leq\exp(-\delta^{2}\mu/3)$, $\forall\delta>0$ ;
2. $\Pr[Z\leq(1-\delta)\mu]\leq\exp(-\delta^{2}\mu/2)$, $\forall\delta\in(0,1)$.

###### Lemma A.3 (Hoeffding bound [ 19 ] ) .

Let $Z_{1},\cdots,Z_{n}$ denote $n$ independent bounded variables in $[a_{i},b_{i}]$. Let $c_{i}=(b_{i}-a_{i})$ Let $Z=\sum_{i=1}^{n}Z_{i}$, then we have

$$ $\displaystyle\Pr[|Z-\operatorname*{{\mathbb{E}}}[Z]|\geq t]\leq 2\exp\left(- \frac{2t^{2}}{\sum_{i=1}^{n}c_{i}^{2}}\right).$ $$

###### Lemma A.4 (Anti-concentration inequality) .

Let $Z\sim{\cal N}(0,\sigma^{2})$,
that is,
the probability density function of $Z$ is given by $\phi(x)=\frac{1}{\sqrt{2\pi\sigma^{2}}}e^{-\frac{x^{2}}{2\sigma^{2}}}$.
Then

$$ $\displaystyle\Pr[|Z|\leq t]\leq\frac{4}{5}\frac{t}{\sigma}.$ $$

###### Lemma A.2 (Chernoff bound [ 9 ] ) .

###### Lemma A.3 (Hoeffding bound [ 19 ] ) .

###### Lemma A.4 (Anti-concentration inequality) .

### A.3 Perturbed w 𝑤 w italic_w for Shifted NTK

We present a lemma from previous work in []. They show that in general, small changes of weights only lead to small change of the Shifted NTK matrix.

###### Lemma A.5 (Lemma C.2 in [ 47 ] , perturbed $w$ for shifted NTK) .

Suppose $b>0$. Assume $R\leq 1/b$. Suppose $m=\Omega(\lambda^{-1}n\log(n/\rho))$. Define function $H$ which maps $\mathbb{R}^{m\times d}$ to $\mathbb{R}^{n\times n}$ as follows:

$$ $\displaystyle\text{~{}the~{}$(i,j)$-th~{}entry~{}of~{}}H(W)\text{~{}is~{}} \frac{1}{m}x_{i}^{\top}x_{j}\sum_{r=1}^{m}{\bf 1}_{w_{r}^{\top}x_{i}\geq b,w_{ r}^{\top}x_{j}\geq b}.$ $$

Let $m$ vectors $w_{1},w_{2},\cdots,w_{m}$ sampled from ${\mathcal{N}}(0,I_{d})$ and let $\widetilde{W}=[w_{1}~{}w_{2}~{}\cdots~{}w_{m}]$. Then there exist constants $c>0$ and $c^{\prime}>0$ such that, for all $W\in\mathbb{R}^{d\times m}$ with $\|\widetilde{W}-W\|_{\infty,2}\leq R$, the following holds:

- •
Part 1, $\|H(\widetilde{W})-H(W)\|_{F}\leq n\cdot\min\{c\cdot\exp(-b^{2}/2),3R\}$ holds with prob. $\geq 1-n^{2}\cdot\exp(-m\cdot\min\{c^{\prime}\cdot\exp(-b^{2}/2),R/10\})$.
- •
Part 2, $\lambda_{\min}(H(W))\geq\frac{3}{4}\lambda-n\cdot\min\{c\cdot\exp(-b^{2}/2),3R\}$ holds with prob. $\geq 1-n^{2}\cdot\exp(-m\cdot\min\{c^{\prime}\cdot\exp(-b^{2}/2),R/10\})-\rho$.

###### Lemma A.5 (Lemma C.2 in [ 47 ] , perturbed w 𝑤 w italic_w for shifted NTK) .

### A.4 Fast regression solver

We list some useful conclusions about fast regression from [].

###### Lemma A.6 (Lemma B.2 in [ 6 ] ) .

Consider the the regression problem

$$ $\displaystyle\min_{x}\|Bx-y\|_{2}^{2}.$ $$

Suppose $B$ is a PSD matrix with $\frac{3}{4}\leq\|Bx\|_{2}\leq\frac{5}{4}$ holds for all $\|x\|_{2}=1$. Using gradient descent, after $t$ iterations, we obtain

$$ $\displaystyle\|B(x_{t}-x^{\star})\|_{2}\leq c^{t}\cdot\|B(x_{0}-x^{\star})\|_{2}$ $$

for some constant $c\in(0,0.9]$.

###### Lemma A.7 (Lemma B.1 in [ 6 ] ) .

Suppose there is a matrix $Q\in\mathbb{R}^{N\times k}$ ($N\geq k\operatorname{poly}(\log k)$), with condition
number $\kappa$ (i.e., $\kappa=\sigma_{\max}(Q)/\sigma_{\min}(Q)$), consider this minimization problem

$$ $\displaystyle\min_{x\in\mathbb{R}^{k}}\|Q^{\top}Qx-y\|_{2}.$ (4) $$

It is able to find a vector $x^{\prime}$

$$ $\displaystyle\|Q^{\top}Qx^{\prime}-y\|_{2}\leq\|y\|_{2}\cdot\epsilon$ $$

in ${\cal T}_{\mathrm{precond}}+{\cal T}_{\mathrm{iters}}\cdot{\cal T}_{\mathrm{
cost}}$ time
where

- •
${\cal T}_{\mathrm{precond}}=\widetilde{O}(Nk+k^{3})$ without using $\mathsf{FMM}$, $\widetilde{O}(Nk+k^{\omega})$ using $\mathsf{FMM}$.
- •
${\cal T}_{\mathrm{iters}}=O(\log(\kappa/\epsilon))$,
- •
${\cal T}_{\mathrm{cost}}=\widetilde{O}(Nk)$.

The above lemma and preconditioning property implies that the iterative regression will take $\log(\kappa/\epsilon)$ iterations.

###### Corollary A.8 .

Solving regression problem (4) needs $O(\log(\kappa/\epsilon))$ iterations using the above method.

The cost per iteration in the iterative regression is too slow for our application. In Section F, we will show how to improve the cost per iteration while maintaining the same number of iterations.

###### Lemma A.6 (Lemma B.2 in [ 6 ] ) .

###### Lemma A.7 (Lemma B.1 in [ 6 ] ) .

###### Corollary A.8 .

### A.5 Sparsity-based Preserving

We present a tool from the paper []. Firstly, we provide a definition.

###### Definition A.9 .

For every $t\in\{0,1,\cdots,T\}$.
For every $i\in[n]$. We use ${\cal S}_{i,\mathrm{fire}}(t)\subset[m]$ to represent the set of neurons that are “fire” at time $t$, i.e.,

$$ $\displaystyle{\cal S}_{i,\mathrm{fire}}(t):=\{r\in[m]:\langle w_{r}(t),x_{i} \rangle>b\}.$ $$

For all $t\in\{0,1,\cdots,T\}$, define $k_{i,t}:=|{\cal S}_{i,\mathrm{fire}}(t)|$ to express the number of fire neurons for $x_{i}$.

The following lemma (Lemma 3.8 in []) show that with the increase of the shifted paramater, the initial neural network will become sparser.

###### Lemma A.10 (Sparsity preserving) .

Assume $m$ is number of neurons.
For shifted parameter $b>0$, if we use $\phi_{b}$ as the activation function of a 2-layer neural network, then after initialization, with prob. $\geq 1-n\cdot\exp(-\Omega(m\cdot\exp(-b^{2}/2)))$, we have for every $i$, $k_{i,0}$ is not larger than $O(m\cdot\exp(-b^{2}/2))$.

Using the above lemma, we can obtain the following result,

###### Corollary A.11 ( ) .

If we set shifted parameter $b=\sqrt{0.48\log m}$ then $k_{0}=m^{0.76}$.
For $t=m^{0.76}$,

$$ $\displaystyle\Pr\left[|{\cal S}_{i,\mathrm{fire}}(0)|>2m^{0.76}\right]\leq\exp \big{(}-\min\{mR,O(m^{0.76})\}\big{)}.$ $$

###### Definition A.9 .

###### Lemma A.10 (Sparsity preserving) .

###### Corollary A.11 (
) .

## Appendix B Threshold search data structure

This section gives a data structure which can efficiently find all the weights $w_{j}$ such that $\langle w_{j},x_{i}\rangle\geq\tau$ for each given input $x_{i}$ and real number $\tau$. Specifically, Section B.1 formally proposes this data structure. Section B.2 proves the running time of Init satisfies the requirement of Theorem B.1. Section B.3 proves the running time of Update satisfies the requirement of Theorem B.1. Section B.4 proves the running time of Query satisfies the requirement of Theorem B.1. Section B.5 proves the correctness of Query in Theorem B.1.

### B.1 Main result

In this section, we are going to present our key theorem (Theorem B.1).

###### Theorem B.1 (Our tree data structure) .

There exists a data structure which requires $O(mn+nd+md)$ spaces and supports the following procedures:

- •
Init$(\{w_{1},w_{2},\cdots,w_{m}\}\subset\mathbb{R}^{d},\{x_{1},x_{2},\cdots,x_{n}
\}\subset\mathbb{R}^{d}$. Given a series of weights $w_{1},w_{2},\cdots,w_{m}$ and datas $x_{1},x_{2},\cdots,x_{n}$, it preprocesses in time $O(mnd)$.
- •
Update$(z\in\mathbb{R}^{d},j\in[m])$. Given a new weight vector $z\in\mathbb{R}^{d}$ and index $j\in[m]$, it updates weight $w_{j}$ with $z$ in time $O(n(d+\log m))$.
- •
Query$(i\in[n],\tau\in\mathbb{R})$. Given a query index $i\in[n]$ and a threshold $\tau\in\mathbb{R}$, it finds all index $j\in[m]$ such that $\langle w_{j},x_{i}\rangle\geq\tau$ in time $O(K_{q}\cdot\log m)$, where $K_{q}:=|\{j\in[m]~{}|~{}\langle w_{j},x_{i}\rangle\geq\tau\}|$.

###### Proof.

Since $W$ takes $O(md)$ space, $X$ takes $O(nd)$ space, each binary tree $T_{i}$ stores $O(m)$ data, the data structure uses $O(mn+nd+md)$. Then we use the following Lemma B.2, B.3, B.4 and B.5 to prove the correctness and running time of this data structure.
∎

Figure: Algorithm 2 Our tree data structure: members, init

Figure: Algorithm 3 Our dynamic data structure: update

Figure: Algorithm 4 Our dynamic data structure: query

###### Theorem B.1 (Our tree data structure) .

###### Proof.

### B.2 Running Time of Init

We prove Lemma B.2, which presents the running time for the Init operation. The corresponding algorithm is shown in Algorithm 2.

###### Lemma B.2 (Running time of Init ) .

Given a series of weights $\{w_{1},w_{2},\cdots,w_{m}\}\subset\mathbb{R}^{d}$ and datas $\{x_{1},x_{2},\cdots,x_{n}\}\subset\mathbb{R}^{d}$, the procedure Init (Algorithm 2) preprocesses in time $O(nmd)$.

###### Proof.

The Init consists of two independent for loops and two recursive for loopss. The first for loop (start from line 10) has $n$ iterations, which takes $O(nd)$ time. The second for loop (start from line 13) has $m$ iterations, which takes $O(md)$ time.
Now we consider the recursive for loop. The outer loop (line 16) has $n$ iterations. In inner loop has $m$ iterations. In every iteration of the inner loop, line 18 runs in $O(d)$ time. Line 20 takes $O(m)$ time.
Putting it all together, the Init runs in time

$$ $\displaystyle~{}O(nd+md+n(md+m))$ $\displaystyle=$ $\displaystyle~{}O(nmd)$ $$

So far, the proof is finished.
∎

###### Lemma B.2 (Running time of Init ) .

###### Proof.

### B.3 Running Time of Update

We prove Lemma B.3. The corresponding algorithm is shown in Algorithm 3.

###### Lemma B.3 (Running time of Update ) .

Given a weight $z\in\mathbb{R}^{d}$ and index $j\in[m]$, the procedure Update (Algorithm 3) updates weight $w_{j}$ with $z$ in $O(n\cdot(d+\log m))$ time.

###### Proof.

The time of Update mainly comes from the forloop (line 5), which consists of $n$ iterations. In each iteration, line 7 takes $O(d)$ time, and the while loop takes $O(\log m)$ time since it go through a path bottom up.
Putting it together, the running time of Update is $O(n(d+\log m))$.
∎

###### Lemma B.3 (Running time of Update ) .

###### Proof.

### B.4 Running Time of Query

We prove Lemma B.4, which is the running time for the Query operation. The corresponding algorithm is shown in Algorithm 4.

###### Lemma B.4 (Running time of Query ) .

Given a query index $i\in[n]$ and a threshold $\tau>0$, the procedure Query (Algorithm 4) runs in time $O(K_{q}\cdot\log m)$, where $K_{q}:=|\{j\in[m]:\langle w_{j},x_{i}\rangle>\tau\}|$.

###### Proof.

The running time comes from QRecursive with input $\tau$ and $\mathrm{root}(T_{i})$. In QRecursive, we start from the root node $r$ and find indices in a recursive way. The Init guarantees that for a node $r$ satisfying $r.\mathrm{value}>\tau$, the sub-tree with root $r$ must contains a leaf whose value is greater than $\tau$ If not satisfied, all the values of the nodes in the sub-tree with root $r$ is less than $\tau$. This guarantees that all the paths it searches do not have any branch that leads to unnecessary leaves.
Our data structure will report all the indices $i$ satisfying $\langle w_{i},q\rangle>\tau$. Since the depth of $T$ is $O(\log m)$, the running time of Query is $O|K_{q}|\cdot\log m)$.
∎

###### Lemma B.4 (Running time of Query ) .

###### Proof.

### B.5 Correctness of Query

We prove Lemma B.5, which shows the correctness for the Query operation.

###### Lemma B.5 (Correctness of Query ) .

Given a query index $i\in[n]$ and a threshold $\tau>0$, the procedure Query (Algorithm 4) finds all index $j\in[m]$ such that $\langle x_{i},w_{j}\rangle>\tau$.

###### Proof.

Fix $i\in[n]$, for all $j\in[m]$, suppose the $j$-th leaf of $T_{i}$ is $l$, the root of $T_{i}$ is $r$, and the path from $r$ to $l$ is

$$ $\displaystyle r=p_{0}\to p_{1}\to\cdots\to p_{k}=l.$ $$

If $\langle x_{i},w_{j}\rangle>\tau$, first $j\in$ QRecursive($p_{k}$), then, suppose $j\in$ QRecursive($p_{t+1}$), then $p_{t+1}$.value $\geq\langle w_{j},x_{i}\rangle>\tau$, thus $j\in$ QRecursive($p_{t+1}$) $\subseteq$ QRecursive($p_{t}$). Hence by induction, $j\in$ QRecursive($p_{0}$)=Query($i,\tau$).
If $\langle x_{i},w_{j}\rangle\leq\tau$, since $l$.value $\geq\tau$, $j$ will not be returned.
Thus Query finds exactly all the index $j\in[m]$ such that $\langle x_{i},w_{j}\rangle>\tau$.

∎

###### Lemma B.5 (Correctness of Query ) .

###### Proof.

## Appendix C Formal Algorithm Representation

We have given a concise representation of our training algorithm (Algorithm 1) in previous sections, for facilitating understanding. For the sake of completeness and convenient implementation, this section gives a formal algorithm representation of our fast neural network training algorithm. (See Algorithm 5.)

This algorithm starts with initializing weights $W_{0}$ and setting shifted parameter $b$. After that, it repeatedly executes sketch computing, iterative regression and implicit weight maintenance until enough times. Specifically, sketch computing computes a sketch matrix $S$ for $J_{t}^{\top}$ with property $\|SJ_{t}^{\top}x\|$ is closed to $\|J_{t}^{\top}x\|$ for every $x$ with large probability. Iterative regression makes use of a fast regression solver to find an approximate solution of

$$ $\displaystyle g_{t}:=\arg\min_{g}\|J_{t}J_{t}^{\top}g-(f_{t}-y)\|$ $$

with the help of the sketch matrix $S$.

Implicit weight maintenance utilizes the threshold search data structure to update weights using the information propagated by the iterative regression.

Figure: Algorithm 5 Our training algorithm, Formal version of Algorithm 1

## Appendix D More Details about Convergence Analysis

The convergence analysis is shown in Section 4. It uses Lemma 4.2, Lemma 4.3 and Lemma 4.4 without proofs. In this section, we formally present the proofs of the three lemmas. In Section D.1, we provide the proof of Lemma 4.2. In Section D.2, we provide the proof of Lemma 4.3. In Section D.3, we provide the proof of Lemma 4.4.

### D.1 Proof of Lemma 4.2

###### Lemma D.1 (Formal version of Lemma 4.2 ) .

For 2-layer ReLU activated neural network, suppose $m=\Omega(d\log(16n/\rho))$, then the following

- •
$\|W_{0}\|_{2}=O(\sqrt{m})$.
- •
$|f(W,x_{i})|=O(1)$, for $i\in[n]$.
- •
$\|J_{W_{0},x_{i}}\|_{F}=O(1)$, for $i\in[n]$.

holds with prob. $\geq 1-\rho/2$.

###### Proof.

(a) The first term can be seen in Corollary 5.35 of []. Notice that $W_{0}\in\mathbb{R}^{m\times d}$ is a
Gaussian random matrix, the Corollary gives

$$ $\displaystyle\Pr[\|W_{0}\|_{2}\leq\sqrt{m}+\sqrt{d}+t]\geq 1-2e^{-\frac{t^{2}} {2}}.$ $$

Let us set $m=\max\{d,\sqrt{2\log(8/\rho)}\}$, it gives $\|W_{0}\|_{2}\leq 3\sqrt{m}$ with probability $1-\rho/4$.

(b) For the second term, first, $a_{r},r\in[m]$ are Rademacher variables, thereby 1-sub-Gaussian, so with probability $1-2e^{-mt^{2}/2}$ we have $\frac{1}{m}|\sum_{r=1}^{m}a_{r}|\leq t$. This means if we take $m=\Omega(\log(16/\rho))$,

$$ $\displaystyle\Pr[\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}=O(1)]\geq 1-\frac{\rho} {8}.$ (5) $$

Next, the vector $v_{i}=W_{0}^{\top}x_{i}\in\mathbb{R}^{m}$ is standard Gaussion vector. Write $a=\left[\begin{array}[]{cccc}a_{1}&a_{2}&\cdots&a_{m}\end{array}\right]^{\top}$, since activation function $\phi_{b}$ is $1$-Lipschitz, with a vector $a$ fixed, the function

$$ $\displaystyle\Phi:\mathbb{R}^{m}\to\mathbb{R},v_{i}\mapsto\frac{1}{\sqrt{m}}a^ {\top}\phi_{b}(v_{i})=f(W_{0},x_{i})$ $$

has a Lipschitz parameter of $1/\sqrt{m}$.

Due to the concentration of a Lipschitz function under Gaussian variables (Theorem 2.26 in []),

$$ $\displaystyle\Pr[|\Phi(v_{i})-\mathbb{E}_{W_{0}}(\Phi(v_{i}))|\geq t]\leq 2e^{ -\frac{mt^{2}}{2}},$ $$

which means if $m=\Omega(\log(16n/\rho))$,

$$ $\displaystyle\left|\frac{1}{\sqrt{m}}a^{\top}\phi_{b}(W_{0}x_{i})-\frac{1}{ \sqrt{m}}\left(\sum_{r\in[m]}a_{r}\right)\operatorname*{{\mathbb{E}}}_{w\sim{ \cal N}(0,I_{d})}[\phi_{b}(\langle w,x_{i}\rangle)]\right|=O(1)$ (6) $$

holds at the same time for all $i\in[n]$ with probability $1-\frac{\rho}{8}$.

We know

$$ $\displaystyle|\mathbb{E}_{w\sim{\cal N}(0,I_{d})}[\phi_{b}(wx_{i})]|\leq|\phi_ {b}(0)|+\mathbb{E}_{\xi\sim{\cal N}(0,1)}[|\xi|]=O(1).$ (7) $$

Plugging in Eq. (5), (7) into Eq. (6), we see that once $m=\Omega(\log(16n/\rho))$, then with probability $1-\rho/4$, for all $i\in[n]$,

$$ $\displaystyle|f(W_{0},x_{i})|=|\frac{1}{\sqrt{m}}a^{\top}\phi_{b}(W_{0}x_{i})| =O(1).$ $$

(c) Let $d_{W,x}=\phi_{b}^{\prime}(Wx)$ denote the element-wise derivative of the activation function, since $\phi_{b}$ is $1$-Lipschitz, we have $\|d_{W,x}\|_{\infty}=O(1)$. Note that $J_{W,x}=\frac{1}{\sqrt{m}}((d_{W,x}\circ a)x^{\top})$ where $\circ$ denotes the element-wise product, we can easily know

$$ $\displaystyle\|J_{W,x_{i}}\|_{F}\leq\frac{1}{\sqrt{m}}\cdot\|\text{Diag}(d)\|_ {2}\cdot\|a\|_{2}\cdot\|x\|_{2}=O(1).$ $$

∎

###### Lemma D.1 (Formal version of Lemma 4.2 ) .

###### Proof.

### D.2 Proof of Lemma 4.3

###### Lemma D.2 (Shifted Perturbation Lemma, formal version of Lemma 4.3 ) .

For 2-layer ReLU activated neural network. Suppose the shifted parameter is $b$ ($b\geq 0$). Let $R_{0}>0$ be a parameter.
Suppose

$$ $\displaystyle m\geq\Omega(1)\cdot\max\{b^{2}R_{0}^{2},n^{2}R_{0}^{2}\lambda^{- 2},n\lambda^{-1}\log(n/\rho)\},$ $$

then with prob. $\geq$
$1-\rho-n^{2}\cdot\exp\big{(}-m\cdot\min\{c^{\prime}e^{-b^{2}/2},\frac{R_{0}}{1
0\sqrt{m}}\}\big{)},$
for every $W\in\mathbb{R}^{d\times m}$ satisfying $\max_{r\in[m]}\|w_{r}-w_{r}(0)\|_{2}\leq R_{0}/\sqrt{m}$, the following holds

$$ $\displaystyle\|G_{W}-G_{W_{0}}\|_{F}\leq\lambda/2,~{}~{}~{}~{}~{}~{}\lambda_{ \min}(G_{W})\geq\lambda/2.$ $$

###### Proof.

We use Lemma A.5 by setting $R=R_{0}/\sqrt{m}$ (that lemma require that $R\leq 1/b$) and letting $W=\left[\begin{array}[]{cccc}w_{1}&w_{2}&\cdots&w_{m}\\
\end{array}\right]$.

Since $R_{0}/\sqrt{m}\leq 1/b$, then we have $m\geq R_{0}^{2}b^{2}$ (this is the corresponding to the first term of $m$ lower bound in lemma statement).

Note $H(W)$ is essentially $G_{W}$, and $\|w_{r}(t)-w_{r}(0)\|_{2}\leq R$ for any $r$, thus by Lemma A.5, we have

- •
$\|G_{W}-G_{0}\|_{F}\leq n\cdot\min\{ce^{-b^{2}/2},3R\}=n\cdot\min\{ce^{-b^{2}/
2},3R_{0}/\sqrt{m}\}$ with prob.
$\displaystyle 1-n^{2}\exp(-m\cdot\min\{c^{\prime}e^{-b^{2}/2},R/10\})=1-n^{2}
\exp(-m\cdot\min\{c^{\prime}e^{-b^{2}/2},\frac{R_{0}}{10\sqrt{m}}\}),$
- •
$\lambda_{\min}(G_{W})\geq\frac{3}{4}\lambda-n\min\{ce^{-b^{2}/2},3R\}=\frac{3}
{4}\lambda-n\min\{ce^{-b^{2}/2},3R_{0}/\sqrt{m}\}$ with prob.
$\displaystyle 1-\rho-n^{2}\exp(-m\cdot\min\{c^{\prime}e^{-b^{2}/2},R/10\})=1-
\rho-n^{2}\exp(-m\cdot\min\{c^{\prime}e^{-b^{2}/2},\frac{R_{0}}{10\sqrt{m}}\}).$

Then it remains to prove

$$ $\displaystyle n\cdot\min\{ce^{-b^{2}/2},3R_{0}/\sqrt{m}\}\leq\frac{\lambda}{2}.$ $$

Since $m\geq\Omega(n^{2}R_{0}^{2}\lambda^{-2})$, we have $3nR_{0}/\sqrt{m}\leq\frac{\lambda}{2}$, which finishes the proof.
∎

###### Lemma D.2 (Shifted Perturbation Lemma, formal version of Lemma 4.3 ) .

###### Proof.

### D.3 Proof of Lemma 4.4

###### Lemma D.3 (The shifted $\mathsf{NTK}$ version of Lemma C.4 in [ 6 ] , formal version of Lemma 4.4 ) .

Suppose $R_{0}\geq 1$ and $m=\widetilde{\Omega}(n^{2}R_{0}^{2})$. Then for every $w\in\mathbb{R}^{d\times m}$ satisfying $\max_{r\in[m]}\|w_{r}-w_{r}(0)\|_{2}\leq R_{0}/\sqrt{m}$, the following holds

- •
$\|W-W_{0}\|=O(R_{0})$,
- •
$\|J_{W,x_{i}}-J_{W_{0},x_{i}}\|_{2}=\widetilde{O}({R_{0}^{1/2}}/{m^{1/4}})$ and $\|J_{W}-J_{W_{0}}\|_{F}=\widetilde{O}({n^{1/2}R_{0}^{1/2}}/{m^{1/4}})$,
- •
$\|J_{W}\|_{F}=O(\sqrt{n})$,

with prob. $\geq 1-\rho$. The randomness comes from the initialization of $W_{0}$.

###### Proof.

(1) The first claim follows from

$$ $\displaystyle\|W-W_{0}\|\leq$ $\displaystyle~{}\|W-W_{0}\|_{F}$ $\displaystyle=$ $\displaystyle~{}\Big{(}\sum_{r=1}^{m}\|w_{r}-w_{r}(0)\|_{2}^{2}\Big{)}^{1/2}$ $\displaystyle\leq$ $\displaystyle~{}\sqrt{m}\cdot R_{0}/\sqrt{m}$ $\displaystyle=$ $\displaystyle~{}R_{0}.$ $$

where the first step comes from $\|\cdot\|\leq\|\cdot\|_{F}$, the second step comes from definition of Frobenius norm, the third step comes from $\|w_{r}-w_{r}(0)\|_{2}\leq R_{0}/\sqrt{m}$, and the last step comes from canceling $\sqrt{m}$.
(2) For the second claim, we have for any $i\in[n]$

$$ $\displaystyle\|J_{W,x_{i}}-J_{W_{0},x_{i}}\|^{2}=$ $\displaystyle~{}\frac{1}{m}\sum_{r=1}^{m}a_{r}^{2}\cdot\|x_{r}\|_{2}^{2}\cdot| \textbf{1}_{\langle w_{r},x_{i}\rangle\geq b}-\textbf{1}_{\langle w_{r}(0),x_{ i}\rangle\geq b}|^{2}$ $\displaystyle=$ $\displaystyle~{}\frac{1}{m}\sum_{r=1}^{m}|\textbf{1}_{\langle w_{r},x_{i} \rangle\geq b}-\textbf{1}_{\langle w_{r}(0),x_{i}\rangle\geq b}|.$ (8) $$

The second equality follows from $a_{r}\in\{-1,1\}$, $\|x_{i}\|_{2}=1$ and

$$ $\displaystyle s_{i,r}:=|\textbf{1}_{\langle w_{r},x_{i}\rangle\geq b}-\textbf{ 1}_{\langle w_{r}(0),x_{i}\rangle\geq b}|\in\{0,1\}.$ (9) $$

We define the event $A_{i,r}$ as

$$ $\displaystyle A_{i,r}=\left\{\exists\widetilde{w}~{}:~{}\|\widetilde{w}-w_{r}( 0)\|\leq R_{0}/\sqrt{m},~{}~{}~{}\textbf{1}_{\langle\widetilde{w},x_{i}\rangle \geq b}\neq\textbf{1}_{\langle w_{r}(0),x_{i}\rangle\geq b}\right\}.$ $$

It is not hard to see $A_{i,r}$ holds if and only if $\langle w_{r}(0),x_{i}\rangle\in[b-R_{0}/\sqrt{m},b+R_{0}/\sqrt{m}]$. Since $w_{r}(0)$ is sampled from Gaussian ${\cal N}(0,I_{d})$ and $\|x_{i}\|=1$, we have $\langle w_{r}(0),x_{i}\rangle$ is sampled from Gaussian ${\cal N}(0,1)$, thus by the anti-concentration of Gaussian (see Lemma A.4), we have

$$ $\displaystyle\operatorname*{{\mathbb{E}}}[s_{i,r}]=\Pr[A_{i,r}]$ $\displaystyle=\Pr[~{}\langle w_{r}(0),x_{i}\rangle\in[b-R_{0}/\sqrt{m},b+R_{0} /\sqrt{m}]~{}]$ $\displaystyle\leq\Pr[~{}\langle w_{r}(0),x_{i}\rangle\in[-R_{0}/\sqrt{m},R_{0} /\sqrt{m}]~{}]$ $\displaystyle\leq\frac{4}{5}R_{0}/\sqrt{m}.$ $$

Thus we have

$$ $\displaystyle\Pr\left[\sum_{i=1}^{m}s_{i,r}\geq(t+4/5)R_{0}\sqrt{m}\right]\leq$ $\displaystyle~{}\Pr\left[\sum_{i=1}^{m}(s_{i,r}-\operatorname*{{\mathbb{E}}}[s _{i,r}])\geq tR_{0}\sqrt{m}\right]$ $\displaystyle\leq$ $\displaystyle~{}2\exp\left(-\frac{2t^{2}R_{0}^{2}m}{m}\right)$ $\displaystyle=$ $\displaystyle~{}2\exp(-t^{2}R_{0}^{2})$ $\displaystyle\leq$ $\displaystyle~{}2\exp(-t^{2}).$ (10) $$

holds for any $t>0$.
The second inequality is due to the Hoeffding bound (see Lemma A.3), the last inequality is because $R_{0}>1$. Taking $t=2\log(n/\rho)$ and using union bound over $i$, with prob. $\geq 1-\rho$,

$$ $\displaystyle\|J_{W,x_{i}}-J_{W_{0},x_{i}}\|_{2}^{2}=\frac{1}{m}\sum_{r=1}^{m} s_{i,r}\leq\frac{1}{m}\cdot 2\log(n/\rho)R_{0}\sqrt{m}=\widetilde{O}(R_{0}/ \sqrt{m})$ $$

holds for all $i\in[n]$. The first equality comes from Eq. (8) and Eq. (9), the second inequality comes from Eq. (10). Thus we conclude with

$$ $\displaystyle\|J_{W,x_{i}}-J_{W_{0},x_{i}}\|_{2}=\widetilde{O}(R_{0}^{1/2}/m^{ 1/4})\text{~{}~{}~{}and~{}~{}~{}}\|J_{W}-J_{W_{0}}\|_{F}=\widetilde{O}(n^{1/2} R_{0}^{1/2}/m^{1/4}).$ (11) $$

(3) The thrid claim follows from

$$ $\displaystyle\|J_{W}\|_{F}\leq$ $\displaystyle~{}\|J_{W_{0}}\|_{F}+\|J_{W}-J_{W_{0}}\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}O(\sqrt{n})+\|J_{W}-J_{W_{0}}\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}O(\sqrt{n})+\widetilde{O}(n^{1/2}R_{0}^{1/2}/m^{1/4})$ $\displaystyle=$ $\displaystyle~{}O(\sqrt{n}).$ $$

where the 1st step is due to triangle inequality, the 2nd step is due to the third claim in Lemma 4.2, the 3rd step is due to Eq. (11), and the last step is due to $m=\widetilde{\Omega}(R_{0}^{2}n^{2})$.
∎

###### Lemma D.3 (The shifted 𝖭𝖳𝖪 𝖭𝖳𝖪 \mathsf{NTK} sansserif_NTK version of Lemma C.4 in [ 6 ] , formal version of Lemma 4.4 ) .

###### Proof.

## Appendix E Induction

Section 4 has defined the induction hypothesis (see Definition 4.5) and given a lemma (see Lemma 4.6) to prove that induction hypothesis holds for all time with high probability, but left its proof to this section. Here, we present and prove the following Lemma E.1, the formal version of Lemma 4.6, and then the crucial Theorem 4.1 holds straightforwardly. We divided the proof of each part of the lemma in Section E.1 and Section E.2, and combine them in Section E.3.

###### Lemma E.1 (Formal version of Lemma 4.6 ) .

Define $R_{0}\approx n/\lambda$. With probability at least $1-\frac{5}{2}\rho-n^{2}\cdot\exp\big{(}-m\cdot\min\{c^{\prime}e^{-b^{2}/2},
\frac{R_{0}}{10\sqrt{m}}\}\big{)}$ of the initial weights $W_{0}$, for every $t>0$, if

- •
$\|f_{t}-y\|_{2}\leq\frac{1}{2}\|f_{t-1}-y\|_{2}$
- •
$\max_{r\in[m]}\|w_{r}(t)-w_{r}(0)\|_{2}\leq R_{0}/\sqrt{m}$

then

- •
$\|f_{t+1}-y\|_{2}\leq\frac{1}{2}\|f_{t}-y\|_{2}$
- •
$\max_{r\in[m]}\|w_{r}(t+1)-w_{r}(0)\|_{2}\leq R_{0}/\sqrt{m}$

also holds.

###### Lemma E.1 (Formal version of Lemma 4.6 ) .

### E.1 Proof of Lemma E.1 : the first lemma

As stated in the previous subsection, we use induction. Here we need to break the induction step (Lemma E.1) into two separate steps, Lemma E.2 and Lemma E.3.
Each separated induction step corresponds to prove one part in the Lemma E.1.
We first prove the first part of Lemma E.1.

###### Lemma E.2 (Part 1 of Lemma E.1 ) .

Suppose initial weights $W_{0}$ satisfies the restriction of Lemma 4.2, 4.3 and 4.4, then for any fixed $t$, if

- •
$\|f_{t}-y\|_{2}\leq\frac{1}{2}\|f_{t-1}-y\|_{2}$ holds
- •
$\max_{r\in[m]}\|w_{r}(t)-w_{r}(0)\|_{2}\leq R_{0}/\sqrt{m}$ holds

Then we have

- •
$\|f_{t+1}-y\|_{2}\leq\frac{1}{2}\|f_{t}-y\|_{2}$ holds.

This proof is similar to [], for the completeness, we still provide the details here.

###### Proof.

We prove the first claim holds for time $t+1$. Define

$$ $\displaystyle J_{t,t+1}=\int_{0}^{1}J\Big{(}(1-s)W_{t}+sW_{t+1}\Big{)}\mathrm{ d}s,$ $$

and denote $g^{\star}=(J_{t}J_{t}^{\top})^{-1}(f_{t}-y)$ to be the optimal solution to Eq. (3), then we have

$$ $\displaystyle~{}\|f_{t+1}-y\|_{2}$ $\displaystyle=$ $\displaystyle~{}\|f_{t}-y+(f_{t+1}-f_{t})\|_{2}$ $\displaystyle=$ $\displaystyle~{}\|f_{t}-y+J_{t,t+1}(W_{t+1}-W_{t})\|_{2}$ $\displaystyle=$ $\displaystyle~{}\|f_{t}-y-J_{t,t+1}J_{t}^{\top}g_{t}\|_{2}$ $\displaystyle=$ $\displaystyle~{}\|f_{t}-y-J_{t}J_{t}^{\top}g_{t}+J_{t}J_{t}^{\top}g_{t}-J_{t,t +1}J_{t}^{\top}g_{t}\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\|f_{t}-y-J_{t}J_{t}^{\top}g_{t}\|_{2}+\|(J_{t}-J_{t,t+1})J_{t }^{\top}g_{t}\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\|f_{t}-y-J_{t}J_{t}^{\top}g_{t}\|_{2}+\|(J_{t}-J_{t,t+1})J_{t }^{\top}g^{\star}\|_{2}+\|(J_{t}-J_{t,t+1})J_{t}^{\top}(g_{t}-g^{\star})\|_{2},$ (12) $$

where the 2nd step is from the definiton of $J_{t,t+1}$ and simple calculus, the 3rd step is from the updating rule of the algorithm, the 5th step is due to triangle inequality, and the sixth step is because triangle inequality.
For the first quantity in Eq. (12), we have

$$ $\displaystyle\|J_{t}J_{t}^{\top}g_{t}-(f_{t}-y)\|_{2}\leq\frac{1}{6}\|f_{t}-y \|_{2},$ (13) $$

since $g_{t}$ is an $\epsilon_{0}(\epsilon_{0}\leq\frac{1}{6})$ approximate solution to regression problem (3).
For the second quantity in Eq. (3), we have

$$ $\displaystyle\|(J_{t}-J_{t,t+1})J_{t}^{\top}g^{\star}\|_{2}\leq$ $\displaystyle~{}\|(J_{t}-J_{t,t+1})\|\cdot\|J_{t}^{\top}g^{\star}\|_{2}$ $\displaystyle=$ $\displaystyle~{}\|(J_{t}-J_{t,t+1})\|\cdot\|J_{t}^{\top}(J_{t}J_{t}^{\top})^{- 1}(f_{t}-y)\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\|(J_{t}-J_{t,t+1})\|\cdot\|J_{t}^{\top}(J_{t}J_{t}^{\top})^{- 1}\|\cdot\|(f_{t}-y)\|_{2}$ (14) $$

where the 1st step is due to matrix spectral norm, the 2nd step is because the definition of $g^{*}$, and the 3rd step relies on matrix spectral norm.
We bound these term separately. First,

$$ $\displaystyle\|J_{t}-J_{t,t+1}\|\leq$ $\displaystyle~{}\int_{0}^{1}\|J((1-s)W_{t}+sW_{t+1})-J(W_{t})\|\mathrm{d}s$ $\displaystyle\leq$ $\displaystyle~{}\int_{0}^{1}\left(\|J((1-s)W_{t}+sW_{t+1})-J(W_{0})\|+\|J(W_{0 })-J(W_{t})\|\right)\mathrm{d}s$ $\displaystyle\leq$ $\displaystyle~{}\widetilde{O}(R_{0}^{1/2}n^{1/2}/m^{1/4}),$ (15) $$

where the 1st step comes from simple calculus, the 2nd step comes from triangle inequality, and the 3rd step comes from the second claim in Lemma 4.4 and the fact that

$$ $\displaystyle\|(1-s)w_{r}(t)+sw_{r}(t+1)-w_{0}\|_{2}\leq$ $\displaystyle~{}(1-s)\|w_{r}(t)-w_{r}(0)\|_{2}+s\|w_{r}(t+1)-w_{r}(0)\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}R_{0}/\sqrt{m}.$ $$

Then, we have

$$ $\displaystyle\|J_{t}^{\top}(J_{t}J_{t}^{\top})^{-1}\|=\frac{1}{\sigma_{\min}(J _{t}^{\top})}\leq\sqrt{2/\lambda}$ (16) $$

where the 2nd step comes from $\sigma_{\min}(J_{t})=\sqrt{\lambda_{\min}(J_{t}^{\top}J_{t})}\geq\sqrt{\lambda
/2}$ (see Lemma 4.3).

Combining Eq. (14), (15) and (16), we have

$$ $\displaystyle\|(J_{t}-J_{t,t+1})J_{t}^{\top}g^{\star}\|_{2}$ $\displaystyle\leq\widetilde{O}({R_{0}^{1/2}\lambda^{-1/2}n^{1/2}}/{m^{1/4}})\| f_{t}-y\|_{2}$ $\displaystyle=\widetilde{O}({\lambda^{-1}nm^{-1/4}})\|f_{t}-y\|_{2}$ $\displaystyle\leq\|f_{t}-y\|/6,$ (17) $$

since $m=\widetilde{\Omega}(\lambda^{-4}n^{4})$.
Let us consider the third term in Eq. (12),

$$ $\displaystyle\|(J_{t}-J_{t,t+1})J_{t}^{\top}(g_{t}-g^{\star})\|_{2}\leq\|J_{t} -J_{t,t+1}\|\cdot\|J_{t}^{\top}\|\cdot\|g_{t}-g^{\star}\|_{2}$ (18) $$

by matrix norm. Moreover, one has

$$ $\displaystyle\frac{\lambda}{2}\|g_{t}-g^{\star}\|_{2}\leq$ $\displaystyle~{}\lambda_{\min}(J_{t}J_{t}^{\top})\|g_{t}-g^{\star}\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\|J_{t}J_{t}^{\top}g_{t}-J_{t}J^{\top}_{t}g^{\star}\|_{2}$ $\displaystyle=$ $\displaystyle~{}\|J_{t}J_{t}^{\top}g_{t}-(f_{t}-y)\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\sqrt{\lambda/n}\cdot\|f_{t}-y\|_{2},$ (19) $$

where 1st step comes from $\lambda_{\min}(J_{t}J_{t}^{\top})=\lambda_{\min}(G_{t})\geq\lambda/2$ (see Lemma 4.3), the 2nd step is because simple linear algebra, the 3rd step is because the definition of $g^{*}$, and the last step is because $g_{t}$ is an $\epsilon_{0}$-approximate solution to $\min_{g_{t}}\|J_{t}J_{t}^{\top}g_{t}-(f_{t}-y)\|$ and $\epsilon_{0}\leq\sqrt{\lambda/n}$.
Consequently, we have

$$ $\displaystyle\|(J_{t}-J_{t,t+1})J_{t}^{\top}(g_{t}-g^{\star})\|_{2}\leq$ $\displaystyle~{}\|J_{t}-J_{t,t+1}\|\cdot\|J_{t}^{\top}\|\cdot\|g_{t}-g^{\star} \|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\widetilde{O}(R_{0}^{1/2}n^{1/2}m^{-1/4})\cdot\sqrt{n}\cdot \frac{2}{\sqrt{n\lambda}}\cdot\|f_{t}-y\|_{2}$ $\displaystyle=$ $\displaystyle~{}\widetilde{O}(n\lambda^{-1}m^{-1/4})\cdot\|f_{t}-y\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\frac{1}{6}\|f_{t}-y\|_{2},$ (20) $$

where the 1st step is because of matrix spectral norm, the 2nd step comes from Eq. (15), (E.1) and the fact that $\|J_{t}\|\leq O(\sqrt{n})$ (see Lemma 4.4), and the last step comes from the $m=\Omega(n^{4}\lambda^{-4})$.
Combining Eq. (12), (13), (E.1), and (E.1), we have proved the first claim, i.e.,

$$ $\displaystyle\|f_{t+1}-y\|_{2}\leq\frac{1}{2}\|f_{t}-y\|_{2}.$ (21) $$

Thus, we complete the proof.
∎

###### Lemma E.2 (Part 1 of Lemma E.1 ) .

###### Proof.

### E.2 Proof of Lemma E.1 : the second lemma

We now move to the second part for Lemma E.1. We show it in Lemma E.3.

###### Lemma E.3 (Part 2 of Lemma E.1 ) .

Suppose initial weights $W_{0}$ satisfies the restriction of Lemma 4.2, 4.3 and 4.4,
then for any fixed $t$, if

- •
$\|f_{t}-y\|_{2}\leq\frac{1}{2}\|f_{t-1}-y\|_{2}$ holds
- •
$\max_{r\in[m]}\|w_{r}(t)-w_{r}(0)\|_{2}\leq R_{0}/\sqrt{m}$ holds

Then we have

- •
$\max_{r\in[m]}\|w_{r}(t+1)-w_{r}(0)\|_{2}\leq R_{0}/\sqrt{m}$ holds

This proof is similar to [], for the completeness, we still provide the details here.

###### Proof.

First, we have

$$ $\displaystyle\|g_{t}\|_{2}\leq$ $\displaystyle~{}\|g^{\star}\|_{2}+\|g_{t}-g^{\star}\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\|(J_{t}J_{t}^{\top})^{-1}(f_{t}-y)\|_{2}+\|g_{t}-g^{\star}\|_ {2}$ $\displaystyle\leq$ $\displaystyle~{}\|(J_{t}J_{t}^{\top})^{-1}\|\cdot\|(f_{t}-y)\|_{2}+\|g_{t}-g^{ \star}\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\frac{2}{\lambda}\cdot\|f_{t}-y\|_{2}+\frac{2}{\sqrt{n\lambda} }\cdot\|f_{t}-y\|_{2}$ $\displaystyle\lesssim$ $\displaystyle~{}\frac{1}{\lambda}\cdot\|f_{t}-y\|_{2},$ (22) $$

where the 1st step relies on triangle inequality, the 2nd step replies on the definition of $g^{*}$, the 3rd step uses matrix norm, the 4th step comes from Eq. (E.1) and the last step uses the obvious fact that $1/\sqrt{n\lambda}\leq 1/\lambda$.
Hence, for any $0\leq k\leq t$ and $r\in[m]$, if we use $g_{k,i}$ to denote the $i^{th}$ indice of $g_{k}$, then we have

$$ $\displaystyle\|w_{r}(k+1)-w_{r}(k)\|_{2}=$ $\displaystyle~{}\|(J_{k}^{\top}g_{k})_{r}\|_{2}$ $\displaystyle=$ $\displaystyle~{}\left\|\sum_{i=1}^{n}\frac{1}{\sqrt{m}}a_{r}x_{i}^{\top}{\bf 1 }_{\langle w_{r}(t),x_{i}\rangle\geq b}g_{k,i}\right\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\frac{1}{\sqrt{m}}\sum_{i=1}^{n}|g_{k,i}|$ $\displaystyle\leq$ $\displaystyle~{}\frac{\sqrt{n}}{\sqrt{m}}\|g_{k}\|_{2}$ $\displaystyle\lesssim$ $\displaystyle~{}\frac{\sqrt{n}}{\sqrt{m}}\cdot\frac{1}{2^{k}\lambda}\|f_{0}-y \|_{2}$ $\displaystyle\lesssim$ $\displaystyle~{}\frac{n}{\sqrt{m}\lambda}\cdot\frac{1}{2^{k}},$ (23) $$

where the 1st step is because of the updating rule, the 2nd step is because of the definition of $J_{k}$, the 3rd step is because of triangle inequalities and the fact that $a_{r}=\pm 1$, $\|x_{r}\|_{2}=1$, the 4th step comes is because of Cauchy-Schwartz inequality, the 5th step is because of Eq. (21) and Eq. (E.2), and the last step is because of the fact that $\|f_{0}-y\|_{2}\leq O(\sqrt{n})$ (see Lemma 4.2, $f_{0}(x_{i})=O(1)$ for any $i\in[n]$, thus $\|f_{0}-y\|_{2}=\sqrt{\sum_{i=1}^{n}(f(x_{i})-y_{i})^{2}}=O(\sqrt{n}$).
Consequently, we have

$$ $\displaystyle\|w_{r}(t+1)-w_{r}(0)\|_{2}\leq$ $\displaystyle~{}\sum_{k=0}^{t}\|w_{r}(k+1)-w_{r}(k)\|_{2}\lesssim\sum_{k=0}^{t }\frac{n}{\sqrt{m}\lambda}\cdot\frac{1}{2^{k}}\lesssim\frac{R_{0}}{\sqrt{m}},$ $$

where the 1st step is because of triangle inequality, the 2nd step is because of Eq. (23), and the last step is because of simple summation.
Thus we also finish the proof of the second claim.
∎

###### Lemma E.3 (Part 2 of Lemma E.1 ) .

###### Proof.

### E.3 Proof of Lemma E.1 : combination

We use Lemma E.2 and Lemma E.3 to prove Lemma E.1.

###### Proof.

Since the probability of initial weight $W_{0}$ satisfies the restriction of Lemma 4.2, Lemma 4.3 and Lemma 4.4 is $1-\rho/2$, $1-\rho-n^{2}\cdot\exp\big{(}-m\cdot\min\{c^{\prime}e^{-b^{2}/2},\frac{R_{0}}{1
0\sqrt{m}}\}\big{)}$, $1-\rho$ respectively, by union bound, the probability of they all happen is at least

$$ $\displaystyle 1-\frac{5}{2}\rho-n^{2}\cdot\exp\big{(}-m\cdot\min\{c^{\prime}e^ {-b^{2}/2},\frac{R_{0}}{10\sqrt{m}}\}\big{)}$ $$

In this case, for any fixed $t$, combining Lemma E.2 and Lemma E.3, if

- •
$\|f_{t}-y\|_{2}\leq\frac{1}{2}\|f_{t-1}-y\|_{2}$ holds,
- •
$\max_{r\in[m]}\|w_{r}(t)-w_{r}(0)\|_{2}\leq R_{0}/\sqrt{m}$ holds

then we have

- •
$\|f_{t+1}-y\|_{2}\leq\frac{1}{2}\|f_{t}-y\|_{2}$ holds.
- •
$\max_{r\in[m]}\|w_{r}(t+1)-w_{r}(0)\|_{2}\leq R_{0}/\sqrt{m}$ holds

Thus by induction, with prob. $\geq 1-\frac{5}{2}\rho-n^{2}\cdot\exp\big{(}-m\cdot\min\{c^{\prime}e^{-b^{2}/2
},\frac{R_{0}}{10\sqrt{m}}\}\big{)}$,

$$ $\displaystyle\|f_{t}-y\|_{2}\leq\frac{1}{2}\|f_{t-1}-y\|_{2}$ $$

holds for all $t$, hence finished the proof of Lemma E.1.
∎

###### Proof.

### E.4 Number of Iterations for Iterative Regression

###### Lemma E.4 .

The iterative regression in our fast training algorithm requires $O(\log(n/\lambda))$ iterations.

###### Proof.

By Lemma D.1, $\|J_{t}J_{t}^{\top}\|=\|G_{t}\|=O(n)$ and $\lambda_{\min}(J_{t}J_{t}^{\top})=\lambda_{\min}(G_{t})\geq O(\lambda)$. Let $\epsilon_{\mathrm{reg}}$ be chosen as Algorithm 5.

Thus by Corollary A.8, the number of iterations needed by the iterative regression is

$$ $\displaystyle O(\log(\kappa(J_{t}^{\top})/\epsilon_{\mathrm{reg}}))=$ $\displaystyle~{}O(\log(\sqrt{n/\lambda}/\sqrt{\lambda/n}))$ $\displaystyle=$ $\displaystyle~{}O(\log(n/\lambda)).$ $$

∎

###### Lemma E.4 .

###### Proof.

## Appendix F More Running Time Details

Section 5 analyzes the running time of our algorithm. It shows that when $m$ is large enough, the running time of $\mathsf{CPI}$ is $o(mnd)+\widetilde{O}(n^{3})$, and with $\mathsf{FMM}$, the $\mathsf{CPI}$ can be reduced to $o(mnd)+\widetilde{O}(n^{\omega})$.

In this section, we give the specific time complexity hidden by $o(mnd)$, and also give the complete algorithm representation of our training algorithm. It will show that when $m$ is large enough, the $\mathsf{CPI}$ is $\widetilde{O}(m^{1-\alpha}nd)$.
Similar with Section 5, we first present Theorem F.1, the running time result. We then provide three lemmas (Lemma F.2, Lemma F.3 and Lemma F.4) to prove our main theorem.
Our main running time result is the following:

###### Theorem F.1 (Running time part of Theorem 1.1 , formal version of Theorem 5.1 ) .

The $\mathsf{CPI}$ is $\widetilde{O}(m^{1-\alpha}nd+n^{3})$, and the running time for shrinking the training loss to $\epsilon$ is $\widetilde{O}((m^{1-\alpha}nd+n^{3})\log(1/\epsilon))$.

Using $\mathsf{FMM}$, the $\mathsf{CPI}$ is $\widetilde{O}(m^{1-\alpha}nd+n^{\omega})$, the running time is $\widetilde{O}((m^{1-\alpha}nd+n^{\omega})\log(1/\epsilon))$. Note that $\omega$ is the exponent of matrix multiplication. Currently, $\omega\approx 2.373$.

###### Proof.

Combining Lemma F.2, Lemma F.3 and Lemma F.4, the computation time of each iteration is

$$ $\displaystyle~{}\widetilde{O}(n^{2}m^{0.76}d)+\widetilde{O}(nm^{0.76}d+n^{3})+ O(n^{2}m^{0.76}(d+\log m))$ $\displaystyle=$ $\displaystyle~{}\widetilde{O}(n^{2}m^{0.76}d+n^{3}+n^{2}m^{0.76}d)$ $\displaystyle=$ $\displaystyle~{}\widetilde{O}(n^{2}m^{0.76}d+n^{3}),$ $$

where the first step comes from hiding $\log m$ on $\widetilde{O}$, the second step comes from simple merging.
And if using $\mathsf{FMM}$, similarly the running time is $\widetilde{O}(n^{2}m^{0.76}d+n^{\omega})$.

By Theorem 4.1, we have:
The time to reduce the training loss to $\epsilon$ is $\widetilde{O}((n^{2}m^{0.76}d+n^{3})\log(1/\epsilon))$. Taking advantage of $\mathsf{FMM}$, the time is $\widetilde{O}((n^{2}m^{0.76}d+n^{\omega})\log(1/\epsilon))$.

Further, for example, if $m=n^{c}$ where $c$ is some large constant, then $n^{2}m^{0.76}d\leq nm^{1-\alpha}d$ where $\alpha\in[0.1,0.24)$.
Hence the time of each iteration is $\widetilde{O}(m^{1-\alpha}nd+n^{3})$, and the time to reduce the training loss to $\epsilon$ is $\widetilde{O}((m^{1-\alpha}nd+n^{3})\log(1/\epsilon))$. Taking advantage of $\mathsf{FMM}$, the time is $\widetilde{O}((m^{1-\alpha}nd+n^{\omega})\log(1/\epsilon))$. Thus we complete the proof.
∎

For the rest of this section, we provide detailed analysis for the steps. In Section F.1 we analyse the sketch computing step. In Section F.2 we analyse the iterative regression step. In Section F.3 we analyse the implicit weight maintenance step.

###### Theorem F.1 (Running time part of Theorem 1.1 , formal version of Theorem 5.1 ) .

###### Proof.

### F.1 Sketch computing

We delicate to prove the lemma that formally analyzes the running time of the sketch computing process in Algorithm 5 to show its time complexity.

###### Lemma F.2 (Sketch computing, formal version of Lemma 5.3 ) .

The sketch computing process of Algorithm 5 (from line 10 to line 23) runs in time
$\widetilde{O}(m^{0.76}n^{2}d).$

###### Proof.

In the sketch computing process, by Corollary A.11, only $O(m^{0.76}d)$
entries of each column of $A$ is nonzero, thus calculating each column of $B$ takes $O(m^{0.76}dt)$
time, where $t$ is the number of rows of $B$. And according to Lemma 2.2,

$$ $\displaystyle t=$ $\displaystyle~{}n\operatorname{poly}(\log(n/\delta_{\mathrm{sketch}}))/ \epsilon_{\mathrm{sketch}}^{2}$ $\displaystyle=$ $\displaystyle~{}O(n\operatorname{poly}(\log(n/\delta_{\mathrm{sketch}}))).$ $$

Since

$$ $\displaystyle\epsilon_{\mathrm{sketch}}=0.1\mathrm{~{}~{}~{}and~{}~{}~{}} \delta_{\mathrm{sketch}}=\frac{1}{\operatorname{poly}(n)},$ $$

the whole for-loop runs in time
$O(n^{2}m^{0.76}d\operatorname{poly}(\log(n))).$
∎

###### Lemma F.2 (Sketch computing, formal version of Lemma 5.3 ) .

###### Proof.

### F.2 Iterative regression

We delicate to prove a lemma that formally analyzes the running time of the iterative regression process in Algorithm 5 to show its time complexity.

###### Lemma F.3 (Iterative regression, formal version of Lemma 5.4 ) .

The iterative regression of Algorithm 5 (from line 26 to line 35) runs in time

$$ $\displaystyle\widetilde{O}(nm^{0.76}d+n^{3}).$ $$

Taking advantage of $\mathsf{FMM}$, the running time is

$$ $\displaystyle\widetilde{O}(nm^{0.76}d+n^{\omega}),$ $$

where $\omega$ is the exponent of matrix multiplication. Currently $\omega\approx 2.373$ [].

###### Proof.

The algorithm calculate $R$ using $QR$ decomposition in line 26 (Algorithm 5). This step will take $O(n^{3})$ time. Taking advantage of $\mathsf{FMM}$, it will take $O(n^{\omega})$ time [].

For the while-loop from line 31 (Algorithm 5), define $p$ as the number of iterations of the while-loop from line 31 (Algorithm 5), then

$$ $\displaystyle p$ $\displaystyle=O(\log(n/\lambda))$ $\displaystyle=O(\log(\frac{n}{(\exp(-b^{2}/2)\cdot\frac{\delta}{100n^{2}})}))$ $\displaystyle=O(\log(n/\delta)+b^{2})$ $\displaystyle=O(\log(n/\delta)+\log m)$ $\displaystyle=O(\log(mn/\delta)),$ $$

where the 1st step comes from Lemma E.4, the 2nd step comes from Theorem 2.13, the 3rd step comes from identical transformation, and the 4th step comes from $b=\Theta(\sqrt{\log m})$.

And in each iteration, note that $R$ is $n\times n$, $A$ is $md\times n$, $S$ is $t\times md$,

- •
we have calculating $v=R^{\top}A^{\top}(D\otimes I_{d})ARz_{t}-R^{\top}y_{\mathrm{reg}}$ takes
$\displaystyle O(n^{2}+m^{0.76}dn+m^{0.76}dn+n^{2}+n^{2})=O(m^{0.76}dn+n^{2})$
time,
- •
and calculating $(R^{\top}A^{\top}(D\otimes I_{d})AR)^{\top}v$ takes
$\displaystyle O(n^{2}+m^{0.76}dn+m^{0.76}dn+n^{2})=O(m^{0.76}dn+n^{2})$
time.

Thus each iteration in the while-loop from line 31 (Algorithm 5) takes
$O(m^{0.76}dn+n^{2})$
time, the total process of the iterative regression takes
$O((m^{0.76}dn+n^{2})\log(mn/\delta)+n^{3})$
time.

Using $\mathsf{FMM}$, the running time is
$O((m^{0.76}dn+n^{2})\log(mn/\delta)+n^{\omega})$.

In our regime, $O(\log(n/\delta))=O(\log m)$ since $m=\operatorname{poly}(n/\delta)$. Thus, we can hide the log factors in $\widetilde{O}$.
∎

###### Lemma F.3 (Iterative regression, formal version of Lemma 5.4 ) .

###### Proof.

### F.3 Implicit weight maintenance

We give a lemma that formally analyzes the running time of the implicit weight maintenance process in Algorithm 5 to show its time complexity.

###### Lemma F.4 (Implicit weight maintenance, formal version of Lemma 5.5 ) .

The implicit weight maintenance of Algorithm 5 (from line 38 to line 43) runs in time
$O(n^{2}m^{0.76}(d+\log m)).$

###### Proof.

Let us consider every iteration of the for loop starting at line 40 (Algorithm 5), since $(J_{t})_{r}$ is $d\times n$, computing $W_{t+1}$ takes $O(nd)$ time. And by Lemma B.3, updating $W_{t+1}$ takes $O(n(d+\log m))$ time, thus each iteration takes $O(n(d+\log m))$ time. By Theorem G.2,
$|K|=O(nm^{0.76})$,
thus the whole implicit weight maintenance takes
$O(n^{2}m^{0.76}(d+\log m))$
time.
∎

###### Lemma F.4 (Implicit weight maintenance, formal version of Lemma 5.5 ) .

###### Proof.

## Appendix G Combination

Theorem 1.1 shows that as long as the 2-layer neural network is broad enough, then there exists a training algorithm with sublinear running time and large converge probability. Theorem 4.1 gives an analysis about how large $m$ should be, but its result is based on $\lambda$, the minimal eigenvalue of $K$ (^2^22See Section (3) for the definition of $K$.), which is not straightforward.

In this section, we convert the bound of Theorem 4.1 into a bound only related to batch number $n$, data separability $\delta$ and tolerable probability of failure $\rho$.

###### Definition G.1 (Two sparsity definitions) .

We define sparsity of the 2-layer neural network to the number of activated neurons.

We define sparsity of a Jacobi matrix of 2-layer neural network as the maximal number of non-zero entries of a row in the Jacobi matrix $J$ ($J\in\mathbb{R}^{n\times md}$) of weights.

We present the following theorem.

###### Theorem G.2 .

For a 2-layer ReLU activated neural network. Suppose $m$ is the number of neurons, $d$ is the dimension of points, $n$ is represent the number of points, $\rho\in(0,1/10)$ is the failure probability, and $\delta$ is the separability of data points.

For any real number $\overline{\alpha}\in(0,1]$, let $b=\sqrt{0.5(1-\overline{\alpha})\log m}$, if

$$ $\displaystyle m=\Omega((\delta^{-4}n^{12}\log^{4}(n/\rho))^{1/\overline{\alpha }})$ $$

then the training algorithm in Algorithm 5 converges with prob. $\geq 1-\frac{5}{2}\rho-n^{2}\cdot\exp(-m\cdot\min\{c^{\prime}e^{-b^{2}/2},
\frac{R}{10\sqrt{m}}\})$, and the sparsity of the neural network is

$$ $\displaystyle O(m^{\frac{3+\overline{\alpha}}{4}})$ $$

with probability $1-n\cdot\exp(-\Omega(m\cdot\exp(-b^{2}/2)))$.
Especially, for any given parameter $\epsilon_{0}\in(0,1/4]$, if we choose $\overline{\alpha}=0.04$, the sparsity is $O(m^{0.76})$.

###### Proof.

From Theorem 2.13,
we know

$$ $\displaystyle\lambda\geq\exp(-b^{2}/2)\cdot\frac{\delta}{100n^{2}}.$ $$

Since by Theorem 4.1,
we need

$$ $\displaystyle m=\Omega(\lambda^{-4}n^{4}b^{2}\log^{2}(n/\rho))$ $$

to make our algorithm converges, we need to choose

$$ $\displaystyle m=$ $\displaystyle~{}\Omega((\exp(b^{2}/2)\cdot{100n^{2}}\cdot{\delta}^{-1})^{4} \cdot n^{4}b^{2}\log^{2}(n/\rho))$ $\displaystyle=$ $\displaystyle~{}\Omega(\exp(4\cdot b^{2}/2)\cdot\delta^{-4}\cdot n^{12}b^{2} \log^{2}(n/\rho))$ $\displaystyle=$ $\displaystyle~{}\Omega(m^{1-\overline{\alpha}}\cdot\delta^{-4}\cdot n^{12} \cdot(\log m)\cdot\log^{2}(n/\rho))$ $$

where the final step is because $b=\sqrt{0.5(1-\overline{\alpha})\log m}$.

Suppose the constant hidden by $\Omega$ is $C$, then the above equation is equivalent to

$$ $\displaystyle m^{\overline{\alpha}}\geq C\cdot\delta^{-4}\cdot n^{12}\cdot( \log m)\cdot\log^{2}(n/\rho),$ $$

and since $m=\operatorname{poly}(n)$, $\log m\leq\log^{2}n$, thus as long as

$$ $\displaystyle m\geq(C\delta^{-4}n^{12}\log^{4}(n/\rho))^{1/\overline{\alpha}},$ $$

we have $m=\Omega(\lambda^{-4}n^{4}b^{2}\log^{2}(n/\rho))$, then by Theorem 4.1, our algorithm converges.

Then, according to Lemma A.10, the sparsity of this neural network
is equal to

$$ $\displaystyle=$ $\displaystyle~{}O(m\cdot\exp(-b^{2}/2))$ $\displaystyle=$ $\displaystyle~{}m\cdot m^{-(1-\overline{\alpha})/4}$ $\displaystyle=$ $\displaystyle~{}m^{\frac{3+\overline{\alpha}}{4}}$ $$

where the second step is because $b=\sqrt{0.5(1-\overline{\alpha})\log m}$, for any $\overline{\alpha}\in(0,1]$.
∎

###### Definition G.1 (Two sparsity definitions) .

###### Theorem G.2 .

###### Proof.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Can the authors speak to the comments in the weaknesses section above?

If at least a toy experiment can be added I think it would really help the case of the paper, at least in my eyes.
soundness: 2
strengths: The paper makes a theoretical good case for using the implicit sparsity of the Jacobian for overparameterized NNs. The method proposes only updating active neurons, therefor reducing computation/gradient calcs.

The paper proposed a binary search tree structure with sublinear time complexity per iteration enables fast updates and queries on non-zero Jacobian entries.

Paper provides theoretical basis, although I did not check everything in the appendix. 

While the method requires the Gauss-Newton method for its sparse Jacobian updates, the paper proposed using conjugate gradient methods to do the computations faster. They also rely on matrix sketching which makes is fair.
summary: The paper proposed treating the training of a NN as a binary tree search in sublinear time by taking advantage of the sparsity in the Jacobian. The paper's analysis focuses on two-layer ReLU networks.
weaknesses: This paper proposes a very new way (at least to me) to train a NN. It seems to have some nice theoretical benefits but since it's so new there might be a lot of unforeseen issues with the method. Overall, I like the paper but have some concerns on the future empirical side of it. 

One that comes to mind is the complexity (not theoretical complexity -- the authors take care of that in the paper). Adam for example has learning rate learning rate and (beta1,beta2). The proposed method would require conjugate gradient which can be expensive and comes with hyper-params of its own, it has sketching which again has more hyper-params like sketch size, ​sketch tolerance and sketch ​failure probability, also things like tree update frequency need to be taken into account. 

Furthermore, while it would be great to take advantage of sparsity, it is not exactly clear to me how sparse the Jacobian would be, and current hardware isn't really setup to 1 calculate sparse gradients, 2 update weights in a much cheaper way due to the sparsity. That is, harvesting the benefit of sparsity is more difficult than one might think in practice.   

Finally, there are no experiments, it would be really nice to see even some toy experiment to see how this method works in practice.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
questions: My review is as follows:

- It's a bit confusing the Algorithm 1 has two lines where a sketch matrix S is constructed.
- The wording in Algorithm 1 "Maintain the weight implicitly" could perhaps be made clearer?
- I have realized that Algorithm 5 has the full details of Algorithm 1. Still the point stands that the wording of Algorithm 1 is ambiguous.
- Why the references on Locality Sensitive Hashing?
- In my opinion, some simple numerical results that verify the lower cost-per-iteration claim would improve the paper.
- It is not clear why it's "almost" impossible to improve the algorithm, just by considering the way Theorem 1.4 and Theorem 7.3 are written. "This result shows that it is almost impossible to truly improve our algorithm"
- I think the organization of the paper could be improved. For instance, sketching seems important to the approach, but it's mentioned for the first time in page 6.

Minor:
- phi is not defined in line 86 when it first shows up
- line 114, typo payed --> paid?
- typo line 129, "that" supports
- line 152, typo: "next a few"

There are a few other typos.
soundness: 2
strengths: Please see the "Questions" section.
summary: This paper focuses on an alternative method for training shallow neural relu networks. The paper provides an analysis for the computational complexity of the proposed algorithm which employs efficient techniques to reduce the complexity per iteration.
weaknesses: Please see the "Questions" section.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: - Can the authors comment on how an algorithm that approximates the NTK matrix and trains only the second layer would compare to their algorithm?
- Can the authors comment on the implicit bias of their algorithm, compared to previous state-of-the art algorithms?
- Does the analysis extend to other activations or deeper networks?

Minor things and typos:
- Line 129: missing `which/that' after `data structure'.
- Line 133: missing closing bracket `)'.
- The informal Theorem 1.4 is hard to understand, if SETH and OVC are not defined.
- Theorem 5.1: it would be good to define lambda.
soundness: 3
strengths: - The problem of reducing the computational cost of running big machine learning models is an important and relevant line of research.
- The paper contains a rigorous analysis of the algorithm presented.
- The connection between neural nets training and binary search trees might be relevant for future work.
summary: This paper provides a novel algorithm for training overparametrized neural networks. It is proven that for 2-layers ReLU nets, this algorithm can be run in sublinear (in the network's width) per-iteration cost, and it preserves the same convergence rate as previous works.
weaknesses: - The analysis seems to be limited to shifted-ReLU activation, relying heavily on the sparsity of the Jacobian. To the best of my understanding, the analysis does not extend to other activations. 
- I am unsure about the relevance of this result: can an approximation of the NTK matrix together with kernel ridge regression achieve sublinear computation time? The authors should comment on this.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: 1. In Line 281 and Line 285, the online regression problem is referencing Eq. (4), $\min_x \lVert Q^\top Q x - y \rVert_2$. Is this correct? I believe the correct reference should be $g_t:=\arg\min_g\lVert J_tJ_t^\top g - (f_t-y)\rVert$.
2. In Definition 5.5, what does $R_0 \approx n/\lambda$ precisely mean?
soundness: 2
strengths: The proposed algorithm achieves a CPI that is sublinear in the network width, which is an improvement over existing second-order methods in overparameterized settings. It is great that the paper provides formal proof of both the sublinear CPI and the ultra-fast convergence rate typical of second-order methods.
summary: This paper proposes a second-order optimization algorithm based on the Gauss-Newton method to efficiently train an overparameterized 2-layer shifted ReLU network with a fixed second layer, using squared loss in the NTK regime. The challenge with second-order methods is their high cost-per-iteration (CPI), despite requiring fewer iterations for convergence compared to first-order methods. The authors aim to reduce CPI while maintaining fast convergence. Under the overparameterized setting, where the width $m$ is polynomial in the number of training points $n$ ($m = \mathrm{poly}(n)$) with input dimension $d$, the proposed method reduces CPI to $O(m^{1-\alpha}nd + n^3)$ for some constant $\alpha \in (0.01, 1)$, improving upon existing second-order methods with a CPI of $O(mnd + n^3)$. The method also achieves $O(\log(1/\epsilon))$ iterations to reach a training error below $\epsilon$, matching the convergence rate of existing second-order methods.

The key idea behind the algorithm is using a tree-based data structure to efficiently identify activated neurons, leveraging the sparsity of shifted ReLU activations. Combined with sketching techniques for solving the online regression problem from the Gauss-Newton method, this leads to a sublinear CPI.
weaknesses: This paper has several critical weaknesses, including issues with the novelty of the tree-based data structure, poor writing, and a lack of practical implications.

1. **The proposed dynamic tree data structure lacks novelty and proper attribution to prior work.**
-  The dynamic tree data structure is not novel and lacks proper citation of prior work, such as [Alman et al., 2023]. The correlation tree proposed in [Alman et al., 2023] (Algorithms 1 and 2 therein) is nearly identical to the threshold search data structure used in this paper (Algorithms 2, 3, and 4). Notations and text are similar, but the authors do not properly cite this relevant work while claiming novelty.

2. **The writing is unpolished and the presentation of key concepts is imprecise.**
- The paper does not sufficiently describe the model setup or key notations in the main text, making it hard for readers to follow. For example, fundamental elements like the loss function and variables such as $J_t$, $f_t$, and $y$ (which appear on Line 247) are only defined in the appendix. Such definitions should be in the main text or referenced before their use to improve clarity. 
- In Theorem 5.1, important parameters such as $\lambda$ and $R$ appear without definition in the main text. While $\lambda$ is defined in the appendix (Remark A.5), it should be introduced earlier. The definition of $R$ is unclear and absent.
- The assumptions regarding network width $m = \Omega(\max(\lambda^{-4}n^4, \lambda^{-2}n^2d\log(16n/\rho)))$ should be better explained. Since $\lambda$ (minimal eigenvector of kernel matrix) is instance-dependent, its relationship to the training data and parameters (e.g., network width and number of training points) should be discussed. Additionally, $\lambda$ is not always guaranteed to be strictly positive.

3. **The assumptions made by the paper are impractical.**
- The paper relies on using "shifted" ReLU activation $\phi(x) = \max (x, b)$ with shifted parameter $b\ge 0$. The shifted parameter should be appropriately chosen (typically, $b=\sqrt{0.48 \log m}$) to ensure sparsity in the activations, and this plays a key role in the analysis. This raises concerns about the generalizability of the results to practical setups, including standard ReLU networks.
- The assumption that the second layer of the network is fixed during training is unrealistic, as in practical settings both layers are typically trained jointly. Prior works, such as [Du et al., 2019], have analyzed overparameterized networks where both layers are trained together, and it would be important to consider whether this paper's analysis can be extended to such cases.

4. **The lack of numerical experiments weakens the practical implications.**
- While I understand that experiments are not strictly necessary for a theoretical paper, empirical results comparing the total training time and CPI of the proposed method with other second-order methods and SGD would strengthen the paper’s practical relevance.

5. Minor comments:
- Typo in Line 87:  Our main reuilt $\to$ Our main result
- Typo in Line 101: Mismatched parenthesis
- Typo in Line 133: Mismatched parenthesis
- Typo in Line 247: $g_t:=\arg\min_g\lVert J_tJ_t^\top g_t - (f_t-y)\rVert$ $\to$ $g_t:=\arg\min_g\lVert J_tJ_t^\top g - (f_t-y)\rVert$

---

**References**

[Alman et al., 2023] Bypass Exponential Time Preprocessing: Fast Neural Network Training via Weight-Data Correlation Preprocessing, NeurIPS 2023.

[Du et al., 2019] Gradient Descent Provably Optimizes Over-parameterized Neural Networks, ICLR 2019

Meta-review: This paper proposes an algorithm that exploits the sparsity structure in the Jacobian matrix to train overparameterized neural networks in a sublinear (in the number of neurons) time per iteration. The authors prove that the proposed algorithm exhibits fast convergence under sufficient overparameterization.

The reviewers acknowledged the theoretical improvement the paper offers to the community, and the importance of the problem being tackled—making NN training faster is a very important research problem.

However, it is deemed that the paper has some areas of improvement. One is presentation and clarity. Multiple reviewers pointed out that the writing is not fully polished, and I agree with this viewpoint. The description of the main method is only high-level, hence does not provide concrete ideas of how things are actually implemented. Some issues in the presentation of mathematical results were pointed out too; e.g., symbols used without definition and validity of assuming $\lambda > 0$. Moreover, it was pointed out that the idea of tree data structure bears similarity with an existing paper, but this similarity was not properly discussed in the paper.

Another shortcoming of this paper pointed out by the reviewers is that the whole analysis relies on an uncommon activation function shifted ReLU. Given that the proposed algorithm heavily exploits the sparsity induced by the shifted ReLU, I am worried if the proposed algorithm will prove useful in neural networks using other activation functions, such as “unshifted” ReLU.

One additional limitation of this paper that was not particularly highlighted by the reviews is that the paper relies on very heavily overparameterized networks, requiring width of the network polynomially large in term of the number of data points. This level of overparameterization makes a “NTK” style of convergence analysis possible, which exploits the fact that the parameter does not change much during training and hence the training dynamics become close to that of a kernel machine. Given that the convergence analysis is reliant on unrealistically very large width, I’m concerned if the proposed algorithm, if implemented, will be useful for practical NN training.

Given these shortcomings, I believe that the paper does not meet the bar for acceptance at ICLR at the moment.

Decision: Reject