Training Multi-Layer Over-Parametrized Neural Network in Subquadratic Time
Authors: Zhao Song
Lichen Zhang
Ruizhe Zhang
Venue: ICLR 2022 Submitted (2022)
Keywords: Deep learning
optimization
over-parametrization

Abstract: In the recent years of development of theoretical machine learning, over-parametrization has been shown to be a powerful tool to resolve many fundamental problems, such as the convergence analysis of deep neural network. While many works have been focusing on designing various algorithms for over-parametrized network with one-hidden layer, multiple-hidden layers framework has received much less attention due to the complexity of the analysis, and even fewer algorithms have been proposed. In this work, we initiate the study of the performance of second-order algorithm on multiple-hidden layers over-parametrized neural network. We propose a novel algorithm to train such network, in time subquadratic in the width of the neural network. Our algorithm combines the Gram-Gauss-Newton method, tensor-based sketching techniques and preconditioning.

Under review as a conference paper at ICLR 2022
TRAINING MULTI-LAYER OVER-PARAMETRIZEED NEURAL NETWORK IN SUBQUADRATIC TIME
Anonymous authors
Paper under double-blind review
ABSTRACT
In the recent years of development of theoretical machine learning, over- parametrization has been shown to be a powerful tool to resolve many fundamental problems, such as the convergence analysis of deep neural network. While many works have been focusing on designing various algorithms for over-parametrized network with one-hidden layer, multiple-hidden layers framework has received much less attention due to the complexity of the analysis, and even fewer al- gorithms have been proposed. In this work, we initiate the study of the per- formance of second-order algorithm on multiple-hidden layers aver-parametrized neural network. We propose a novel algorithm to train such network, in time sub- quadratic in the width of the neural network. Our algorithm combines the Gram- method, tensor-based based sketching techniques and preconditioning.
INTRODUCTION
Deep neural networks have been playing central role in both practical (such as computer vision (LeCun et al., 1998; Krizhevsky et al., 2012; Szegedy et al., 2015; He et al., 2016), natural language processing (Collobert et al., 2011; Devlin et al., 2018), automatic driving system, game playing (Silver et al., 2016; 2017) ) and theoretical machine learning community (Li & Liang (2018); Jacot al. (2018); Du et al. (2019b); Allen-Zhu et al. (2019a;b); Du et al. (2019a): Song & Yang (2019); Brand et al. (2021); Zou et al. (2018); Cao & Gu (2019); Lee et al. (2019a); Liu et al. (2020; 2021); Chen et al. (2021)) In order analyze the dynamics of neural networks and obtain provable guarantees, using vee-pparametrization has been a growing trend.
In terms of understanding the convergence behavior of peee-parametrized networks, most of the attentions have been directed to the study of first-order method such as gradient descent or stochastic gradient descent. The widespread use of first-order method is explained, computational efficiency, since computing the gradient of the loss function at each iteration is usually cheap and simple, let alone with its compatibility with random sampling- bbased method such as mini- batch. One of the major drawbacks of first-order methods is their convergence rate is typically slow in many non-convex settings (poly(n, L, log(1/e)) iterations, where n is the number of training samples, T. is the number of layers and is the precision of training), e.g., deep neural network with ReI.U activation, as shown in Allen-Zhu et al. 2a),, which is often the case of deep werr-parameterized neural network.
large degree, by its
Second-order method (which employs the information of the Hessian matrix) on the other hand enjoys a much faster convergence rate (only log(1/e) iterations (Zhang al., 2019), but not poly(n) log(1/e) iterations) and exploits the local geometry of the loss function to overcome the pathological curvature issues that are critical in first-order method. Another clear advantage of recond-order method over first-order method is it does not require the tuning of learning rate. The expense of using second-order method its prohibitive cost per iteration, as it is imperative to invert synamically-changii Hessian matrix or equivalently, solving regression task involving Hessian. Given any weight matrix size m X m (m is the width of network), its Hessian matrix has size m² x m². which makes any naive implementation of second-order algorithm takes at least O(m²) time since one needs to write down the Hessian. This explains the scarcity of deploying large-scale second-order method in non-convex setting, such training deep neural networks, in contrast to their popular presence in convex optimization setting (Vaidya (1989); Daitch & Spielman (2008); Lee et al. (2015); Cohen et al. (2019); Lee al. (2019b); Jiang et al. (2020b;a); Song & Yu (2021)).

---

Under review as a conference paper at ICLR 2022
Recent works (Cai et al. (2019); Zhang et al. (2019)) improved the practicality of second-order method on training deep networks and presented algorithms to train one-hidden layer over- parametrized networks with smooth activation functions. Specifically, they achieve running time of (mm²) per iteration. Their methods are essentially variants of Gauss-Newton method combined with using neural tangent kernels (Jacot et al. (2018)) prove the convergence. By using the idea of randomized linear algebra (Clarkson & Woodruff (2013); Woodruff (2014)) and a clever sketching matrix as preconditioner, Brand et al. (2021) further improves the running time to (1m) per iteration.
However, all of these algorithms are for training a shallow network with one-hidden layer and fall short on deep networks- First, it is not clear that their algorithms can be generalized to multi-layer setting, due to the presence of gradient vanishing or exploding. In the seminal work of Allen-Zhu et al. (2019a), they showed that as long as the networks are over-parametrized, first-order methods such gradient descent and stochastic gradient descent won't 't encounter such problem. But does this still hold for second-order method? Can we provably show that second-order method has a good performance when training deep over-parametrized networks? Second, even the fastest of them (Brand al. (2021)) would incur running time of Om)²n per iteration, which seems unavoidable due to the size of intermediate weight matrices is m m.
In this work, we take the first step to tame the beast We propose a second-order method that achieves subquadratic cost per iteration with respect to m. and show that it has linear convergence rate training deep over-paramen over-parametrized neural networks. We emphasize the importance of obtaining subquadratic algorithm, since in multi-layer settings, the network width is typically much larger than in one-hidden layer setting (m 2 n L2, (Zou & Gu, 2019)).
Our work can be decomposed into two parts: algorithmically and analytically. From algorith- mie perspective, our method builds upon variant of Gauss-Newton method (Björck (1996) called Slam-GGass-Nare
method (Cai et al. (2019); Brand et al. (2021)). In order achieve feasible running time, we exploit two features of the gradient. which is the key ingredient to form the Jaco- bian matrix: 1). The gradient is low rank (rank n). 2). The gradient can be formulated as the outer product of two vectors. From an analytical perspective, our work is inspired by Allen-Zhu et al. (2019a). In contrast their proof which is straightforward analysis of the gradient, make of the multi-layer neural tangent kernels (Du et al. (2019a)) and establish a connection between a Gram matrix we compute each iteration and the NTK matrix.
Our Contributions. We summarize our technical contributions below.
We develop analytical framework for the convergence behavior of second-order method on training multi-layer over peerppaaaetrized neural network. To facilitate the analysis, we exploit the equivalence between neural tangent kernels and our over- over-paramen parametrized network.
We design a second-order algorithm to train such networks, and achieve a cost per iteration of o(m²). Our algorithm makes use of Gram-Gauss Gauss-Newton method, tensor-based sketch- ing techniques, and data structures that maintains a low rank representation efficiently.
By combining fast tensor algebra techniques and sketching-based preconditioning, we de- vise an algorithm to efficiently solve regression problem where the involved matrix has its rows being tensor product of vectors.
1.1
OUR RESULT
Our main result can be summarized in the following three theorems, with one analyzing the conver- gence behavior general Gram-based optimization framework, one designing an efficient algo- rithm to realize this second-order optimization scheme, and the other is a novel algorithm to solve tensor-based regression high precision and fast, which is a key step in our second-order method.
Throughout this paper, we will use n to denote the number of training data points, d to denote the dimension of input data points, m to denote the width of the network and T. to denote the number of layers of the network. We use ft E IR to denote the prediction of neural network at time t.
Our first theorem demonstrates the fast convergence rate of our algorithm.
Theorem 1.1 (Convergence, informal version of Theorem F.19). Suppose the width of the neural network satisfies m 2 poly(n, L), then there exists algorithm (Algorithm 1) such that, over
2

---

Under review as a conference paper at ICLR 2022
the randomness of initialization of the network and the algorithm, with probability least 1 e-iloog* m), have
1 |e -yll2 S 2 vll2.
where ft R" is the the prediction produced by neural network at time
The above theorem establishes the linear convergence behavior of our second-order method, which is standard convergence result for second-order method, as well the same behavior as in one- hidden layer over-paraine paraamaaetree networks (Brand et al. (2021)). However, compared to one-hidden layer case, our analysis is much more sophisticated since have to carefully control the probability so that it does not blow up exponentially with respect to the number of layers.
The next theorem concerns the cost per iteration of second-order algorithm.
Theorem 1.2 (Runtime, informal version of Theorem B.1). There exists randomized algorithm
(Algorithm 1) that trains multi-layer neural network of width m. with the cost per training iteration being
O(m2~9(U).
We improve the overall training time of multi-layer over-parametrized networks with second-order method from Tiuit + O(m²) to Tini +TToommr o(m²), where Tint is the initialization time of training, typically takes (m²)) ). As we have argued before, multi-layer over-parametrized networks require m to be in the order of n², hence improving the cost per iteration from quadratic to subquadratic is an important gain in speeding up training. advantage even more evident when one seeks a high precision solution, and hence the number of iterations T is large.
We highlight that non-trivial to obtain subquadratic running time per iteration: not handled properly, computing the matrix-vector product with weight matrices will take O(m²) time! This means that even for first-order methods such as gradient descent, is not clear how to achieve a subquadratic running time, since one has multiply the weight matrix with vector in both forward evaluation and backpropagation. In our case, we have also a Jacobian matrix of size n X m², so forming naively will cost (nm²² time, which is prohibitively large. Finally, note that the update matrix is also an m x m matrix. In order to eircumvent these problems, we exploit the fact that the gradient is of low rank (rank n), hence fast matrix-vector product. We also observe that each row of the Jacobian matrix can be formulated tensor product of two vectors, therefore we can make use of fast randomized linear algebra to approximate the tensor product efficiently. As a byproduct, we have the following technical theorem:
compute rank-n factorization and use it support
Theorem 1.3 (Fast Tensor Regression, informal version of Theorem D.14), Given two X m matri- ces and V with m > n and target vector e E R.. Let J = [vec(ure]] ..... vec v "T] E mnxmm where us is the i-th row of matrix U Vi E [n]. There is randomized algorithm that takes + + Onnm+n2 (log(n/e) + log(m/0)) + n²) time and outputs vector To E RR" such that
|J - S
holds with probability at least 5, and k is the condition number of J.
From a high level, the algorithm proceeds follows: given matrices U and V, it forms an approxi- mation J guxa log (m/6), where each row is generated by applying fast tensor sketching technique to u, and 0; (Ahle al. (2020)). Then, it uses another sketching matrix for J to obtain a good preconditioner for J. Subsequently, runs a gradient descent to solve the regression.
To understand this runtime better, we note that um term is the size of matrices U and V. hence reading the entries from these matrices will take at least O(nm) time. The algorithm then uses tensor-based sketching techniques (Ahle et al. (2020)) to squash length m² tensors length O(n log(m/66)). All subsequent operations are performed on these much smaller vectors. Fi- nally, computing the preconditioner takes Ö(n²) time, and running the gradient descent takes (n² log(s/e)) time.
3

---

Under review as a conference paper at ICLR 2022
1.2
RELATED WORK
Second-order Method in Optimization. Though not prevalent as first-order method in deep learning, second-order methods are one of the most popular in convex setting, such as linear pro- gramming (Vaidya (1989); Daitch & Spielman (2008); Lee et al. (2015); Cohen et al. (2019)), em- pirical risk minimization (Lee et al. (2019b)), cutting plane method (Jiang et al. (2020b)) and semi- definite programming (Jiang al. (2020a)). Due the prohibitive high cost of implementing one step of second-orde method, most of these works focus on improving the cost per iteration.
In non-convex setting, there's a vast body of ongoing works (Martens & Grosse (2015); Botev al. (2017); Pilanci & Wainwright (2017); Agarwal et al. (2017); Bernacchia et al. (2)18); Cai et al. (2019); Zhang al. (2019); Brand al. (2021); Yao et al. (2021)) that try to improve the practicality of second-ordd method and adapt them to train deep neural networks. As shown in Cai et al. (2019), possible to exploit the equivalence between ever-parametrized networks and neural tangent kernel to optimize an n x matrix instead of an m² x m2 matrix, which is an important breakthrough in gaining speedup for second-order method. Sketching and sampling- -baend methods can also be used to accelerate the computation of inverses of the Hessian matrix (Pilanci & Wainwright (2017)). In spirit, our work resembles most with Cai et al. (2019) and Brand et al. (2021), the sense that our optimization also works on an and sampling, as in Pilanci & Wainwright (2017); Brand et al. (2021).
Gram matrix. Our algorithm also makes use of sketching
wee-paaametteriee Neural Networks. In recent deep learning literature, understanding the ge- ometry and convergence behavior of various optimization algorithms on over- werr-parameteiieed neural networks has received lot of attention (Li & Liang (2018); Du et al. (2019b); Allen-Zhu et al. (2019a;b); Du et al. (2019a); Song & Yang (2019); Ji & Telgarsky (2020); Zou et al. (2018); Cao & Gu (2019); Liu al. (2020; 2021)). The seminal work of Jacot et al. (2018) initiates the study of neural tangent kernel (NTK), which is a powerful analytical tool in this area, since as long as the neural network is wide enough (m 2 ((n")), then the optimization dynamic equivalent to that on NTK.
neural network is
Sketching. Using randomized linear algebra to reduce the dimension of the problem and speedup the algorithms for various problems has been a growing trend in machine learning community (Sarlos (2006); Clarkson & Woodruff (2013); Woodruff (2014)) due to its wide range of applications to various tasks, especially the efficient approximation of kernel matrices (Avron et al. (2014); Ahle et al. (2020); Woodruff Zandieh (2020)). The standard "Sketch-and-Solee (2013)) paradigm involves using sketching to reduce the dimension of the problem and then using a blackbox for the original problem to gain an edge on computational efficiency. Another line of work is to use sketching preconditioner (Woodruff (2014); Brand et al. (2021)) obtain a high precision solution.
(Clarkson & Woodruff
Roadmap. In Section 2, we give preliminary view of the training setup consider in this paper. In Section 2.1, we introduce the notations that will be used throughout this paper. In Section 2.2, consider the training setup. In Section 3, overview the techniques employed in this paper. In Section 3.2, we demonstrate various techniques to prove the convergence of our second-order method. In Section 3.1, we examine the algorithmic tools utilized in this paper achieve sub- quadratic cost per iteration. In Section 4, we summarize the results in this paper and point out some future directions.
2 PRELIMINARIES
2.1 NOTATIONS
For any positive integer n, we use [n] to denote the set 2.. We use E[:) to denote expec- tation and Pr[:] for probability. We use ||æl|2 to denote the & norm of a vector x. We use ||A]] || to denote the spectral norm of matrix A. We use |||| to denote the Frobenius norm of A. We use A to denote the transpose of matrix A. We use Im to denote the identity matrix of size m X For matrix A or vector x, we use | ||=||u to denote the number of nonzero entries of A and x respectively. Note that || llo is semi-norm since it satisfies triangle inequality. Given real square matrix A, we use Amaxx (A) and Amin (A) to denote its largest and smallest eigenvalues respectively.
Given a real matrix A, we use m and m (A) denote its largest and smallest singular
4

---

Under review as a conference paper at ICLR 2022
values respectively. We use ((,,, to denote the Gaussian distribution with mean j and variance o². We use (n)(n) denote O(f(n) poly log(f(n)). We use (.-) to denote the inner product, when applying to two vectors, this denotes the standard dot product between two vectors, and when applying two matrices, this means (A, B) tr|A'' B] where tr[A]] denote the trace of matrix A.
2.2 PROBLEM SETUP
Let X E R denote the data matrix with n. data points and ma features. Without loss of generality, we assume | l|2 1, E [n]. Consider an L layer neural network with one vector W2 E Rmaxmi and W1 E " Xma, We will use
E ! and L matrices Wr, _m,,,_,,,,, We(t) denote the weight matrix at layer at time 6, and (W)(t) to denote its gradient. We also (tt)= (Wi(t)..... (t) to denote the collection of weight matrices at time t.
Architecture. We first describe our network architecture. The network consists of hidden layers, each represented by weight matrix We Rm for any & e [L]. The output layer consists of a vector a E mm We define the neural network prediction function f mR - RR as follows:
f(W, x) = a @WW (((]x)))
where à R > R the shifted ReLU activation function (ax(x) = max (x b, 0)) applied coordinate- wise to vector.
We measure the loss via squared-loss function:
L(W)
(31 - J(W,
This is also the objective function for training.
We define the prediction function fe R xn
- R" as
fx) (t) x(W(t),
x(W(),
Initialization. Our neural networks are initialized as follows:
For each e [L]. the initial weight matrix We(0) p is initialized such that each entry is sampled from ((0, 2.).
Each entry of a is an i.i.d. sample from {- 1 - + vmn
uniformly at random.
Gradient. In order to write gradient in an elegant way, we define some artificial variables:
9t,1 = Wiz4,
h41 = @WW1x,)
Vi E [n]
91. Wehae-11
h4.e @(Wehi,e-1)
Vi E [n], AA E [L]M1]
(I)
D 1 diag (((1w:)),
Vi E in]
Di. diag (o' (WWehh,e-1)),
Vi [n], Ve (L]V
Using the definitions of and h, we have f(W, = Th,,, E ER, Vi E [m]
We can compute the gradient of L terms of We pm for all L >
OL(W) aW.
(((W, vi) DEE
II
W
DER
mrXme
k=t+1
my 1 1xm, my Xmm
my m
1xms-1
Note that the gradient for W1 E Xmo (recall that mo d) is slightly different and can not be written by general form. By the chain rule, we can compute the gradient with respect to WT.
ac(w) OW,
I.
E((W,x.)) vi) DE
II
W
D
(3)
my Xm
mi-1 X.ng my X
mL mIX1 ixmo
5

---

Under review as a conference paper at ICLR 2022
It is worth noting that the gradient matrix is of rank n, since it's sum of n rank-1 matrices.
Jacobian. For each layer & [L] and time E [T]. we define the Jacobian matrix Je,e E Puxmame-1 via the following formulation:
Je = veel (1
(f OW7())
(f OW,(e)
vec
vecl
OW?())
The Gram matrix at layer and time is then defined as Ge = Je E Rux whose (i, j-th entry is (O.)) (f(
OW,
oW,
TECHNIQUE OVERVIEW
In this section, we give overview of the techniques employed in this paper. In Section 3.1, we showcase our algorithm and explain various techniques being used to obtain subquadratic cost per iteration. In Section 3.2, we give an overview of the proof to show the convergence of our algorithm. To give simpler and cleaner presentation, we assume my m for all E [L].
3.1 SUBQUADRATIC TIME
In this section, overview the techniques deployed in our implementation of the second-order method. Our main focus is to achieve subquadratic cost per iteration. Instead of using Hessian matrix of size m² m², we use n Gram matrix derived from the neural tangent kernel. However this would still incur a cost of O(mm²) per iteration, since each gradient is an m X m matrix and the Jacobian consists of such gradients.
We start by demonstrating our algorithm:
Algorithm Informal version of our algorithm.
1: procedure OURAL.GORTHMMC Bikee))
b Theorem 1.1,1.2
2: /#Initialization/
3: Initialize (O), Vć E [L]
4: Store Wi (0)0)a; in memory, Vi [n]
Takes (nm²) time Takes O(nIm²) time
5: Store (0) < (W) ve E [Z] i [m] in memory
6: for = 0 T do
7:
Forwwrrdd computation*/ for = 1 do
8:
9:
04e + h4 Vi E [n]
10:
h4.e - <((W(()) + (W.)(h,-)))) Vi E
Takes O(n²m) + o(nm²) time h sparse b Takes (mm) time
12:
Di.e t (OW + AW, (m)
13:
Dr. sparse
14:
end for
15:
ft < [a T hi.L..... a hn.L
b Takes (nm) time
16:
Baackaard computation for T. do
17:
18:
04e a Dr, Form Je, that approximates Je,, using {ure); 1
(eee)
Takes o(nIm²) time
19:
20:
b. Taaes (mm) time, Je, E Rxx where s ((n)
Compute ge that approximates (Je,
Takes Ö(nm) time
22:
Form Je via low rank factorization En De.uui.u.
23:
Implicitly update AW, AW, E M and store it in memory
24:
end for
25:
end for
26:
end procedure

---

Under review as a conference paper at ICLR 2022
Step 1: Invert Gram by solving regression. Recall the update rule of generic algorithm is given by
We(t 1) (E
where cs y after proper scaling. Naively forming the Gram matrix J,, will take m(n²m²² time and inverting will take O(n") time. To avoid the quadratic cost at this step, we instead solve a regression, or a linear system since the Gram matrix has full rank: find the vector ge, E IR" such that
|| |J,e: - 2
is minimized. This enables to utilize the power of sketching to solve the regression efficiently.
Step 2: Solve Gram regression via preconditioning. In order to solve the above regression, we adapt the idea of obtaining good preconditioner via sketching then apply iterative method to solve (Brand al. (2021)). Roughly speaking, we first use random matrix S E Rexm² that has the subspace embedding property (Sarlos (2006) to reduce the number of rows of J then we run QR decomposition on matrix SJJ". This gives us matrix R such that SJ R has orthonormal columns. We then use gradient descent to optimize the objective ||JJ Rz, y11.. Since S is a subspace embedding for J" we can make sure that the condition number of the matrix J R is small (0(1), hence the gradient descent converges after log(k//c) iterations, where R. is the condition number of J. However, in order implement the gradient descent, we still need to multiply an m n matrix with a length vector, in the worst case this will incur time of O(mm²). In order to bypass this barrier, we need to exploit extra structural properties of the Jocobian, which will be demonstrated in the following steps.
Step 3: Low rank structure of the gradient. Instead of studying Jacobian directly, we first try to understand the low rank structure of the gradient. Consider af(W...)) E Rmmm, it can be written as (for simplicity, we use hs,0 to denote x.)):
OW.
(W, hi,e-1 a D. ... Waen aw.
v.ERmx1
uu_ xm
This means the gradient is essentially an outer product of two vectors, and hence has rank one. This has several interesting consequences: for paer-paraeerrized networks, the gradient is merely of rank n instead of m. When using first-order method such as gradient descent or stochastic gradient de- scent, the weight is updated via low rank matrix. To some extent, this explains why the weight does not move too far from initialization in over-parametrized networks when using first-order method to train. Also, as we will illustrate below, this enables the efficient approximation of Jacobian matrices and maintenance of the change.
Step 4: Fast approximation of the Jacobian matrix. We now turn attention to design a fast approximation algorithm to the Jacobian matrix. Recall that Jacobian matrix Jr, E R=mmm is an m² matrix, therefore writing down the matrix will take O(nm²) time. However, it is worth noticing that each row of Jr, is vec(u;o; T, Vi E [m)]. or equivalently, u; o U2 where o denotes the tensor product between two vectors. Suppose and {ui,..., 2n E (IR then we can compute the tensor product sketching techniques, such as TensorSketch (Avron et al. (2014); Diao et al. (2017: 2019)) or TensorSRHT (Ahle et al. (2020); Woodruff & Zandieh (2020)) in time nearly linear in m and the targeted sketching dimension s, in contrast to the naive ((m)) time. Since it suffices to pre- serve the length of all vectors in the column space of JE.- the target dimension s can be chosen as
are given the collection of (u1,-.., Un. (RT")"
o 0 via tensor-based
O - poly (log(m/m)))) Use Jr,, E Raax to denote this approximation of Jr,,, we perform the preconditioned gradient descent we described above on this smaller matrix. This enables to lower the overall cost of the regression step to be subquadratic in m.
Step Efficient update via low rank factorization. The low rank structure of the gradient can fur- ther be utilized to represent the change on weight matrices AW in a way such that any matrix-vector product involving AW can be performed fast. Let ge E RD denote the solution to the regression problem posed in Step 1. Note that by the update rule of our method, we shall use J E R
to update the weight matrix, but writing down the matrix will already take O(m?) time. Therefore,
7

---

Under review as a conference paper at ICLR 2022
it is instructive to find succinct representation for the update. The key observation is that each column of J/, is a tensor product of two vectors: u, o U/ or equivalently, u.y The update matrix can be rewritten as Σi= ge.uu.y." and we can use this representation for the update on the weight, instead of adding it directly. Let
U²= = g1,1u1 22
ge.n Un
E RRmx
Ve 21
22
Un
E RR m
then the update can be represented as UVV.". Consider multiplying a vector E R" with this representation, we first multiply y with V,T E _=xmm which takes (mm) time. Then we multi- ply V. y E Rn with U2 E Rmxx which takes O(mn) time. This drastically reduces the cost of multiplying the weight matrix with a vector from O(m²) to ((m)).
Inspired by this idea, it is tempting to store all intermediate low rank representations across all iterations and use them to facilitate matrix-vector product, which incurs runtime of O(Tmn). This is fine when T' is relatively small, however, if one looks for a high precision solution which requires a large number of iterations, then T' might be too large and (T mn) might be in the order of ((m²). To circumvent this problem, we design the data structure so that it will exactly compute the m x m change matrix and update the weight and clean up the cumulative changes. This can be viewed as a "restart" of the data structure. To choose the correct number of updates before restarting, we utilize the dual exponent of matrix multiplication, a (Gall & Urrutia, 2018), which means it takes (m²) X m by an X m² matrix. Hence, we restart the data structure after around
time to multiply m /n updates. Therefore, we achieve an amortized o(m²) time, which is invariant even though the number of iterations T grows larger and larger.
3.2 CONVERGENCE ANALYSIS
In this section, we demonstrate the strategy to prove that our second-order method achieves a linear convergence rate the training loss.
Step 1: Initialization. Let (O) be the random initialization. We first show that for any data point IG, we have ((W(0), 2;) = O(1). The analysis draws inspiration from Allen-Zhu et al. (2019a). The general idea is, given a fixed unit length vector æ, multiplying it with a random Gaussian matrix W will make sure that || 2. Since W is a random Gaussian matrix, applying shifted ReLU activation gives a random vector with a truncated Gaussian distribution conditioned on a binomial random variable indicating which neurons are activated. We will end up with (((Wx)) well as <(Wx) being sparse. Inductively applying this idea to each layer and carefully controlling the error occurring each layer, we can show that with good probability, f(W (0), 2x.) is constant.
1 as
We also bound the spectral norm of 2((W10)... by Ō( (V/LL//) One of the key part of this matrix is OW, the consecutive product D2 PLWL... DE+1 We+1 De. By studying the distribution of its product with a fixed vector, one can show that the spectral norm of this consecutive product is bounded by O( T.). Finally, we make use the fact that each entry of is Rademacher random variable scaled by 1 1/Vm, hence the norm of DT.W,.... D W4+1 Dr is bounded by ((VL/m) with good probability.
Furthermore, we show that the Gram matrix for the multiple-layer over-parametrized neural net- work, which is defined as Je has nontrivial minimum eigenvalue after the initialization. In particular, we adapt the neural tangent kemee (NTK) for multiple-layer neural networks defined by Du et al. (2019a) into our setting by analyzing the corresponding Gaussian process with shifted ReLU activation function. Then, we can prove that with high probability, the least eigenvalue of the initial Gram matrix is lower bounded by the least eigenvalue of the neural tangent kernel matrix.
Step 2: Small perturbation. The next step is to show that if all weight matrices undergo small perturbation from initialization (in terms of spectral norm), then the corresponding Jacobian matrix has not changed too much. As long as the perturbation is small enough, it is possible to show that the change of the h vector (in terms of l, norm) and the consecutive product (in terms of spectral norm) is also small. Finally, using the fact that a is Rademacher vector with scaling 1/ l/mm, we can show that the change
a (S + AW, a(W: +WW)
as(W(0),
OW.
8

---

Under review as a conference paper at ICLR 2022
has its spectral norm being bounded by D(yLL/m) for any layer & E [L] and input data [n]. Consequently, the Frobenious norm of the Jacobian matrix is bounded by O( ((/L//))
Step 3: Connect everything via double induction. Put things together, we use double induction argument, where we assume the perturbation of weight matrix is small and the gap between ft and y is at most 1/2 of the gap between fi-1 and y. By carefully bounding various terms and exploiting the fact the Jacobian matrix always has relative small spectral norm (Ō (Ö(ynL/m)) mm), we first show that the weights are not moving too far from the initialization, then use this fact to derive a final convergence bound for f] y122.
4 DISCUSSION AND FUTURE DIRECTIONS
In this work, we propose and analyze a second-orde method to train multi-layer over- per-parametrized neural networks. Our algorithm achieves linear convergence rate in terms of training loss, and achieves subquadratic (o(m²) cost per training iteration. From an analytical perspective, we greatly extend the analysis of (Allen-Zhu et al. (2019a)) to second-order method, coupled with the usage of the equivalence between multi-layer over-parametrized networks and neural tangent kernels (Du et al. (2019a)). From an algorithmic perspective, we achieve a subquadratic cost per iteration, which is significant improvement from O(m²n) time per per iteration due to the prohibitively large network width m. Our algorithm combines various techniques, such as training with the Gram matrix, solve the Gram regression via sketching-base bbased preconditioning. fast tensor computation and dimensionality reduction and low rank decomposition of weight updates. Our algorithm is especially valuable when one requires a high precision solution on training loss, and hence the number of iterations is large.
One of the interesting questions from our work is: is it possible to obtain an algorithm that has nearly linear cost per iteration on m as in the case of training one-hidder layer over-pawee werrparametrized networks (Brand al. (2021))? In particular, can this runtime be achieved under the current best width of multi-layer werrppaaaeetrized networks (m 2 n n² We note that the major limitation in our method is the sparsity of the change of the diagonal matrices (AD) is directly related to the magnitude of the change of weights ([]W)).. In our analysis of convergence, we go through a careful double induction argument, which fact imposes on a lower bound on || [[. It seems to us that, in order to achieve nearly linear runtime, one has to adapt a different analytical framework or approach the problem from a different perspective.
A related question is, how can maintain the changes of weight more efficiently? In work, achieve speedup in the neural network training process by observing that the change of the weights small in each iteration. Similar phenomenon also appears in some classical optimization problem (e.g., solving linear program (Cohen et al., 2019; Jiang et al., 2021) and solving semidefinite pro- gram (Jiang et al., 2020a)) and they achieve further speedup by using lazy update and amortization techniques to compute the weight changes, or using more complicated data structure to maintain the changes of the weight changes. Can we adapt their techniques to neural network training? An or- thogonal direction to maintain the change is to design an initialization setup such that while we still have enough randomness obtain provable guarantees, the matrix-vectorr product with the initial weight matrix can be performed faster than O(m²) by sparsifying the Gaussian matrix as in Derez- iński et al. (2021) or imposing extra structural assumption such as using circulant Gaussian (Rauhut et al., 2012; Nelson & NguyÅn, 2013; Krahmer et al., 2014).
Another question concerns activation functions. In this paper, we consider the shifted ReLU acti- vation and design our algorithm and analysis around its properties. Is it possible to generalize our algorithm and analysis to various other activation functions, such as sigmoid, tanh or leaky ReLU? If one chooses a smooth activation, can we get a better result in terms of convergence rate? Can we leverage this structure to design faster algorithms?
Finally, the network architecture considered in this paper is the standard feedforward network. Is it possible to extend our analysis and algorithm to other architectures, such as recurrent neural networks (RNN)? For RNN, the weight matrices for each layer are the same, hence it is trickier to analyze the training dynamics such networks. Though the convergence of first-order method on peerpparmetrrieed multi-layer RNN has been established, it is unclear whether such analysis can be extended to second-order method.

---

Under review as a conference paper at ICLR 2022
Ethics Statement. This is theory paper that proposed efficient algorithm to train multi- -ulti-layer over- parametrized neural networks, hence does not have direct ethics implications.
Reproducibility Statement. All results in this paper can be directly verified and reproduced via reading into the proofs. For complete algorithm description and its runtime analysis, see Section B. For specific data structure used in the algorithm, Section C. For the result regarding the fast tensor regression, see Section D. For convergence analysis of the algorithm, see Section E, F and G.
REFERENCES
Naman Agarwal, Brian Bullins, and Elad Hazan. Second-order stochastic optimization for machine learning in linear time. 2017.
Thomas D. Ahle, Michael Kapralov, Jakob Bæk Tejs Knudsen, Rasmus Pagh, Ameya Velingker, David P. Woodruff, and Amir Zandieh. Oblivious sketching of high-degree polynomial kernels. In Proceedings of the 2020 ACM-SIAM Symposium on Discrete Algorithms (SODA), pp. 141-160, 2020.
Zeyuan Allen-Zhu, Yuanzhi Li. and Zhao Song. A convergence theory for deep learning via over- parameterization. In ICML, 2019a.
Zeyuan Allen- Yuanzhi Li, and Zhao Song. On the convergence rate of training recurrent neural networks. In NeurlPS, 2019b.
Haim Avron, Huy L. Nguyen, and David P. Woodruff. Subspace embeddings for the polynomial kernel. In NeurIPS, 2014.
Alberto Bernacchia, Mate Lengyel, and Guillaume Hennequin. Exact natural gradient in deep linear networks and its application the nonlinear case. In S. Bengio, H. Wallach, H. Larochelle, and R. Garnett (eds.), Advances Neural Information Processing Systems. Curran Associates, Inc., 2018
K. Grauman, N.
Sergei Bernstein. On modification of chebyshev's inequality and of the error formula of laplace. Ann. Sci. Inst. Sav. Ukraine, Sect. Math, 144):3849, 1924.
Åke Björck. Numerical Methods for Least Squares Problems. Society for Industrial and Applied Mathematics, 1996.
Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical Gauss GewssnNewton optimisation for deep learning. In Proceedings of the 34th International Conference on Machine Learning, pp. 557-565, 2017.
Jan van den Brand, Binghui Peng, Zhao Song, and Omri Weinstein. Training (overparametrizeed) neural networks in near-linear time. In ITCS, 2021.
Tianle Cai, Ruiqi Gao, Jikai Hou, Siyu Chen, Dong Wang, Di He, Zhihua Zhang, and Liwei Wang. Gram- trram-gaus--eet gauss-newton method: Learning overparameterized neural networks for regression prob- lems. arXiv preprint arXiv: /675, 2019.
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. NeurlPS, pp. 10835-10845, 2019.
Zixiang Chen, Yuan Cao, Difan Zou, and Quanquan Gu. How much wer-paraaeeterization suf- ficient to learn deep ReLU networks? In International Conference Leurning Representations (ICLR), 2021.
Herman Chernoff. A measure of asymptotic efficiency for tests of a hypothesis based on the sum of observations. The Annals of Mathematical Statistics, pp. 493-507, 1952.
Kenneth L. Clarkson and David P. Woodruff. Low rank approximation and regression in input sparsity time. In Symposium on Theory of Computing Conference (STOC), pp. 81-90, 2013.
Michael B Cohen, Yin Tat Lee, and Zhao Song. Solving linear programs in the current matrix multiplication time. In STOC, 2019.
10

---

Under review as a conference paper at ICLR 2022
Ronan Collobert, Jason Weston, Léon Bottou, Michael Karlen, Koray Kavukcuoglu, and Pavel Kuksa. Natural language processing (almost) from scratch. Journal of machine learning research, 12:2493-2533 2011.
Samuel I Daitch and Daniel A Spielman. Faster approximate lossy generalized flow via interior point algorithms. In Proceedings of the fortieth annual ACM symposium on Theory of computing (STOC), pp. 451-460, 2008
Michat Dereziński, Jonathan Lacotte, Mert Pilanci, and Michael W. Mahoney. Newton-less: Spar- sification without trade-offs for the sketched newton update, 2021.
Jacob Devlin, Ming-Woi Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 181.04000 2018.
Huaian Diao, Zhao Song, Wen Sun, and David Woodruff. Sketching for kronecker product regres- sion and p-splines. In AISTATS, 2017.
Huaian Diao, Rajesh Jayaram, Zhao Song, Wen Sun, and David Woodruff. Optimal sketching for kronecker product regression and low rank approximation. In Advances in Neural Information Processing Systems (NewrIPS), 2019.
Simon Du, Jason D Lee, Haochuan Li, Liwei Wang. and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference Machine Learning (ICML), 2019a.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In ICLR, 2019b.
François Le Gall and Florent Urrutia. Improved rectangular matrix multiplication using powers of the srngersmith-wnnogeee tensor. Proceedings of the Twenty-Ninth Annual ACM-SIAM Symposium on Discrete Algorithms, SODA '18, pp. 1029-1046, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR), pp. 770-778, 2016.
Wassily Hoeffding. Probability inequalities for sums of bounded random variables. Journal of the American Statistical Association, 58(301):13-33) 1963.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: convergence and gen- eralization in neural networks. In Proceedings of the 32nd International Conference on Neural Information Processing Systems (NeurIPS), pp. 8580-8589, 2018.
Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbi- trarily small test error with shallow relu networks. In ICLR, 2020.
Haotian Jiang, Tarun Kathuria, Yin Tat Lee, Swati Padmanabhan, and Zhao Song. A faster interior point method for semidefinite programming. FOCS, 2020a.
Haotian Jiang, Yin Tat Lee, Zhao Song, and Sam Chiu-wai Wong. An improved cutting plane method for convex optimization, convex-concave games and its applications. In STOC, 2020b.
Shunhua Jiang, Zhao Song, Omri Weinstein, and Hengjie Zhang. Faster dynamic matrix inverse for faster Ips. In STOC, 2021.
Felix Krahmer, Shahar Mendelson, and Holger Rauhut. Suprema of chaos processes and the re- stricted isometry property. Communications on Pure and Applied Mathematics, 67(1)):1877- 1904, 2014.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep con- volutional neural networks. Advances in neural information processing systems, 25:1097 2012.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEFF, (6(11):2278-2324 1998.

---

Under review as a conference paper at ICLR 2022
Jachoon Lee, Lechao Xiao, Samuel S. Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide Neural Networks of Any Depth Evolve as Linear Models under Gradient Descent. 2019a.
Yin Tat Lee, Aaron Sidford, and Sam Chiu-wai Wong. A faster cutting plane method and its impli- cations for combinatorial and convex optimization. In Foundations of Computer Science (FOCS), 2015 IEEE 56th Annual Symposium on, pp- 1049-1065. IEHE, 2015.
Yin Tat Lee, Zhao Song, and Qiuyi Zhang. Solving empirical risk minimization in the current matrix multiplication time. In Conference on Learning Theory (COLT)) pp. 2140-2157. PMLR, 2019b.
Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. In NeurlPS, 2018.
Chaoyue Liu, Libin Zhu, and Misha Belkin. On the linearity of large non-linear models: when and why the tangent kernel is constant. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 15954-15964. Curran Associates, Inc., 2020.
Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in over- parameterized non-linear systems and neural networks, 2021.
Yichao Lu, Paramveer Dhillon, Dean P Foster, and Lyle Ungar. Faster ridge regression via the sub- sampled randomized hadamard transform. In Advances in neural information processing systems (NIPS), pp. (677-77,220 2013.
James Martens and Roger Grosse Optimizing neural networks with kronecker-factoree approximate curvature. In Proceedings of the 32nd International Conference on International Conference on Machine Learning Volume 37, ICML' 15, pp. 2408-2417. JMLR.org, 2015.
Jelani Nelson and Huy NguyÅn. Sparsity lower bounds for dimensionality reducing maps. In Proceedings of the forty-fifth annual ACM symposium Theory of computing, pp. 101-110, 2013.
Mert Pilanci and Martin J. Wainwright. Newton sketch: A near linear-time optimization algorithm with quaara-q convergence. SIAM J. Optim., 27:205-22 245, 2017.
Holger Rauhut, Justin Romberg, and Joel A Tropp. Restricted isometries for partial random circulant matrices. Applied and Computational Harmonic Analysis, 32(2):242 25442,42, 2012.
Mark Rudelson and Roman Vershynin. Non-asympiotic theory of random matrices: extreme singu- lar values. In Proceedings of the International Congress of Mathematicians 2010 (ICM 2010) (In 4 Volumes) Vol. I: Plenary Lectures and Ceremonies Vols. I1-IV: Invited Lectures, pp. 1576-1602. World Scientific, 2010.
Tamas Sarlos. Improved approximation algorithms for large matrices via random projections. In 2006 47th Annual IEEE Symposium on Foundations of Computer Science (FOCS), pp. 143-152. IEEE, 2006.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484448) 2016.
David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Mathhw Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676):354-355 2017.
Zhao Song and Xin Yang. Quadratic suffices for over-parametrization via matrix chernoff bound. arXiv preprint arXiv: 906600993, 2019.
Zhao Song and Zheng Yu. Oblivious sketching-bart based central path method for linear programming.
In International Conference on Machine Learning (ICML), pp. 9835-9847 PMLR, 2021.
12

---

Under review as a conference paper at ICLR 2022
Zhao Song, David P. Woodruff, Zheng Yu, and Lichen Zhang. Fast sketching of polynomial kernels of polynomial degree. ICML, 2021.
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Du- mitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and paltern recognition, pp. 1-9, 2015.
Pravin M Vaidya. Speeding-up linear programming using fast matrix multiplication. In 30th Anmual Symposium on Foundations of Computer Science, pp. 332-337. IEEE, 1989.
Virginia Vassilevska Williams. Multiplying matrices faster than coppersmith winthradddddd....o In Proceed- ings of the forty-fourth annual ACM symposium Theory of computing (STOC), pp. 887-898. ACM, 2012.
David P. Woodruff.
Sketching as a tool for numerical linear algebra. Foundations and Trends in Theoretical Computer Science, 10(1-2):1-157, 2014.
David Woodruff and Amir Zandieh. Near input sparsity time kernel embeddings via adaptive sampling. In ICML,, 2020.
Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney. Adahessian: An adaptive second order optimizer for machine learning. Proceedings of the AAAI Conference on Artificial Intelligence, 35(12):10665-1067 10673,
May 2021.
Guodong Zhang, James Martens, and Roger B Grosse. Fast convergence of natural gradient descent for over-parameterizee neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Difan Zou and Quanquan Gu. An improved analysis training over-parameterized deep neural networks. In NeurlPS, pp. 2053-2062, 2019.
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Stochastic gradient descent optimizes over-parameterized deep relu networks, 2018.
13

---

Under review as a conference paper at ICLR 2022
Roadmap. In Section A, we remind readers with the notations and some probability tools. In Section B, we illustrate the complete version of our algorithm and give a runtime analysis of it. In Section C, we design a simple low rank maintenance data structure and show how to use it to efficiently implement matrix- product. In Section D, we introduce an efficient regression solver handling
Jacobian and Gram regression. In Section E, study the spectral property of the Gram matrix at each layer and connects it with multi-layer neural tangent kernels. In Section F, analyze the convergence of our algorithm by using some heavy machinery such as structural analysis of the gradient and careful double induction. In Section G, we give a detailed proof of one technical lemma.
A PRELIMINARIES AND PROBABILITY TOOLS
In this section, we introduce notations that will be used throughout the rest of the paper and several useful probability tools that will be heavily exploited in the later proofs.
For any positive integer n, we use [n] to denote the set {1,2....,,)) We use E[-] to denote expectation and Pr[-] for probability. We use || to denote the l norm of a vector x. We use ||A|]| denote the spectral norm of matrix A. We use ||A|] F denote the Frobenius norm of A. We use A to denote the transpose of matrix A. We use Im to denote the identity matrix of size m m. For matrix or vector z, we use ||.Allo. ||=/|l to denote the number of nonzero entries of A and respectively. Note that - is semi-norm since it satisfies triangle inequality. Given real square matrix A, we use Amax (A) and (A) to denote its largest and smallest eigenvalues respectively. Given real matrix A, we use Tmax (A) and min (A) to denote its largest and smallest singular values respectively. We use (N,, to denote the Gaussian distribution with mean μ and variance o². We use (f)))) to denote (() poly log(f(n)). We use (..r) to denote the inner product, when applying to two vectors, this denotes the standard dot product between two vectors, and when applying to two matrices, this means (A, B) = tr(AA B] where [r[A] denote the trace of matrix A.
Notations.
Lemma A.1 (Chernoff bound Chernoff (1952)). Let X X1, where X; with probability Pa and X, = 0) with probability 1 - Pt, and all X; are independent. Let μ = [[X] > Ps- Then
1. Pr X 2 (1+ 8)pl < exp( -82p-)8)33) 48 > 0;
Pr X (1- S cxp(-62p//2) YO <
Lemma A.2 (Hoeffding bound Hoeffding (1963)). Let X1,
Xn denote n independent bounded
variables in [as, bi]. Let X X X,, then we have
212
Pr[[X - E[X] 2 $ & exx
Ei=ı (b. a,)
Lemma A.3 (Bernstein inequality Bernstein (1924)). Let XI,... Xn be independent zero-mean random variables. Suppose that [Xi] S M almost surely, for all i. Then, for all posilive I,
n
12/2
Pr
Σ
x, >
exp (TE
B[X + Mt/33
Lemma A.4 (Anti-concentraation of Gaussian distribution). Let X NO, o2), then
2 4t ( 3 5
Pr[[X| S
Lemma A.5 (Concentration of subgaussian random variables). Let E B be vector where each entry of a is sampled from a subgaussian distribution of parameter o², then for any vector x E IR"",
12
PrlKa, 2 [|11]11]] 2 eexp(
20
Lemma A.6 (Small ball probability). Let R be vector such that la,| 8 for all [n], Ii...., Fn are n i.i.d. Rademacher random variables. Then, there exist absolute constants C₁ C2 such that for any >
Pr(l(a,)) S $]} S min { Cit Cet [a]]2 6/n
14

---

Under review as a conference paper at ICLR 2022
B COMPLETE ALGORITHM AND ITS RUNTIME ANALYSIS
In this section, we first present our complete algorithm, then give a runtime analysis it.
Algorithm 2 Complete version of our algorithm.
1: procedure COMPLETEALGGRRITHMZ. e Rdx" & E I)
Theorem B.1
2:
/#Iitialization -
3:
Initialize (O), VE E [L]
4: Store W1 (O) in memory, Vi E [n]
Takes (nm²²) time D Algorithm
5:
NOW.ANKMAAINTE.R
LMR
6:
[ (0)... (0)...,,((()))
7:
for 0= do
8:
Fooowrrdd computation*/
9:
for f = 1 + do
Vie hi,e-1, Vi [n]
9i. LMR. QUERY,, hi,e-1 hi, + ((9ne)), Vi [m]
Takes o(nm²) time
12:
h sparse b Takes ( time
14:
(gt, Di.t + diag(ø' (gt,e)), V [n]
Di, sparse
16:
end for
ft < [a T hhT,,,,, h,, .......... a [.]
b Takes (nm) time
18:
Backward computation*/
19:
for = L
do
20:
U4 a Dr, ... Dieteee vi e [n]
Takes o(nLm²) time fore)- with precision /A/n
Je FASTIENSORREGGSSIONN
22:
& Algorithm 6
L.R. (gc, (1))))))
24:
end for
25:
end for
26:
end procedure
Theorem B.1 (Formal version of Theorem 1.2), Let X R and R". and let k denote the sparsity of Di, and s denote the sparsity of AD AA E [L], i E [n]. Let m denote the width of neural network, I. denote the number of layers and a denote the dual matrix multiplication exponent (Def. C. thenn the running time of Algorithm is
O +T .Tter))
where
Tiiit = 0(m + L)),
Ther Ō((m + m(s (s k))L² m² mL-
Therefore, the cost per iteration of Algorithm is
Ō((m 1ta + m(s k)) 2 m nLL)
Proof. We analyze Tini and Titer separately.
Initialization time. We will first initialize (L 1) m matrices and one x d matrix, which takes O(m²L) L) time. Compute W (0)2, for all E [n] takes Omm2n) time. Finally, initialize the data structure takes Omm² L) time. Hence, Tiii. O 2 L)).
Cost per iteration. For each iteration, we perform one forward computation from layer to L, then backpropagate from layer 1.
Forward computation: forward computation, we first compute It. E R", which is equivalent to form 0,4+1, hence by Lemma C.5, it takes O(s + l: + m²) time. Compute hi, and Drd takes ((m) time. These computations would be performed for all L layers, hence the overall runtime of forward computation is O((s k + T.) time.
15

---

Under review as a conference paper at ICLR 2022
Backpropagation: In backpropagation, first compute u4, R" which takes (m + + m²)) time. Then, we call Algorithm to solve the Gram regres- sion problem, which due to Theorem D.14 takes Omm n²) time. Note that even want high probability version of the solver with e log mL failure probability so that we union bound over all layers and all iterations, term in running time, which absorbed by the O(-) notation. Finally, the update takes (mm²aan) time by Lemma C.2. Sum over all L layers, we get an overall running time of (( + m(s k)) + m² (LL) time.
only pay extra log nL
This concludes the proof of our Theorem.
Corollary B.2. Suppose the network width m is chosen as in F.55 and the shift parameter is chosen in F.6, then the cost per iteration of Algorithm is
Ō(m m 1.69 n.).
Remark B.3. As long as the neural network is wide enough, as in F:25 and we choose the shifted threshold properly, as in F6,, then we can make sure that both sparsity parameters i and s to be o(m), and we achieve subquadratic cost per iteration.
C Low RANK MAINTENANCE AND EFFICIENT COMPUTATION OF THE CHANGE
In this seciton, design data structure to maintain the low rank representation of change of weights, and then we show how to efficiently compute the low rank representation using this data structure.
C. Low RANK MAINTENANCE
In this short section, design simple data structure maintain the low rank representation of change of weights, and show that the matrix-vector product can implemented efficiently. Before moving, we define some notions related to rectangular matrix multiplication.
Definition C.1 (Williams (2012); Gall & Urrutia (2018)). Let w be the matrix multiplication expo- nent such that it takes +o)) time to multiply two n matrices.
Let a be the dual exponent of the matrix multiplication which is the supremum among all a 2 0uch that it takes 2+) time multiply an n x n by n. matrix.
Additionally, we define the function w()) where (b) denotes the exponent of multiplying an matrix by an n² matrix. Hence, we have w(1) w and w(a)= =
The overall idea of our low rank maintenance data structure is follows: we keep accumulating the low rank change, when the rank of the change reaches a certain threshold (m²), then we restart the data structure and update the weight matrix.
16

---

Under review as a conference paper at ICLR 2022
Algorithm 3 Low rank maintenance data structure
1: data structure LOWRANKMANTENNNCE
Lemma C.2
2: members
3:
TE, Vl E [Z]
D denotes the accumulated rank of the change
4: We, VE E [Z]
(We)e=1 E (Rmxmja
5: AW2 vl (T.
(AWe) 1 E (Rmxm)L
6:
end members
7:
8: procedures
9:
INIT( Orr([W; [Wr (O)..... (())))
Initialize the data structure
UPDATE(UE, Vi)
Update the low rank representation D Compute the matrix-vector product between A AW; and
11:
QUERY
12: end procedures
13: end data structure
Algorithm Procedures LRM data structure
procedure (W))( (Z))))
Lemma C.2
2: W; t (0)
3: AW; - O, AA E [L]
4: Te + E [L]
5: end procedure
6:
7: procedure UPDATE(UE E Rm Vi E Rmxx)
Lemma C.2
8:
AW, AW, U.V. without forming the product and sum the matrices
9:
TE t re + rf m² where w2) then
11:
We + We + AW,
b Takes ((m²) time
rt 6- 0 AW: <
13:
end
15:
end procedure
16:
17: procedure QUERY(É E [L], E RT)
Lemma C.2
18: + We +WW
> Takes O(nnz(y) m + mrr) time
19: return
20: end procedure
Lemma C.2. There exists a deterministic data structure (Algorithm 3) such that maintains
AW,
AWL such that
The procedure INIT (Algoritlum 4) takes O(m² L) time.
The procedure UPDATE (Algorithm 4) takes O(rm²~a) amortized time, where a w2)
The procedure QUERY (Algorithm 4) takes O(m (nnz(y) +rz)) time, where Tt is the rank of AW; when QUERY called.
Proof. The runtime for INIT is obvious, for QUERY, notice that we are multiplying vector y with (possibly) dense matrix We mm which takes O(nnz(y) m) time, and accumulated low rank matrix AW; with rank тz. By using the low rank decomposition AW; = UV with U, V m the time to multiply y with AW is O(mrz). Combine them together, get running time of O(m (nnz(y) + rt)).
It remains to analyze the amortized cost of UPDATE. Note that if rz < m then we just pay O(1) time to update corresponding variables in the data structure. r; = m", then we will explicitly form the x matrix AW2. To form it, notice have accumulated rz/n different sums of
17

---

Under review as a conference paper at ICLR 2022
rank-n decompositions, which can be represented
U U..
Imxr, [Ve(1), V:(2)...., .... Ve(re/n)] E E Imxr,
and AW; UVT, which takes (m²)) time to compute since rı m and (2)) Finally, note that this update of We only happens once per re/n number of calls to UPDATE, therefore we can charge each step by O m² time.
O(m²~nn) (m (mm²~nn)) "n), arrives at our final amortized running
refn
Remark C.3. Currently, the dual matrix multiplication exponent a R 0.31 (Gall Urrutia, 2018), hence the amortized time for UPDATE is O(nm* .9). If m 2 n n0 then we achieve an update time of o(m²). Similarly, the time for QUERY O(m (nnz(y) + ré)) = O unz(y) + m²+aa = O(m nnz(y) + m m,))), ), long mnz(y) o(m), then its running time is also o(m²). In our application of training neural networks, we will make sure that the inputted vector y is sparse.
C.2 EFFICIENTLY COMPUTE (t) AND (t,)
In this section, we show that how to compute the vectors Wi,e, Vi, IRm using the low rank structure of the change of weights combined. Recall definition of these vectors:
(r,e(t) a (t) ... (t e (t,e(t) = (t)) E R
Before proceeding, we list the assumptions we will be using:
For any E [4]. (t,(()) (t) is sp-sparse, where sp = k k is the sparsity of (1, and s the sparsity of Di,e(t) (0).
For any & E [L], the change of the weight matrix We AWe(t) = Wt(t) We(0). of low-rank. That is, AW&(t) Σラ Ve
For any E [n]. Wi (0)z; is pre-computed.
We first note that as direct consequence of Di,((0) is k-sparse, h,,e(0) is k-sparse well. Simi- larly, (t) - (0,)(0) has sparsity s. Hence (t) has sparsity bounded by sp.
Compute (t,)(t). Compute (t)) is equivalent to compute the following vector:
(e)()((2+(())) (W)+(()) (1)++(() (Wz(0) + AWr(1) (i)a
First, we know that (t)a E Th' an sp-sass vector, and it takes O(sp) time. The next matrix is (W) which gives two terms: ((D,,.())))) and (Wr( (De,,())) For the first term, since (t). sp-sparse, it takes (MmD)) -timme. For the second term, have
AW Σ AL.VEL = j=1
Σ 1=1
aLs D(z()a).
Each inner-product takes (ssp)-time and it takes O(mr, + sDr() = ((mr)) -time in total. Hence, in O(m(sD -tmmm, we compute the vector Wy (z()) (t). Note that we do not assume the sparsity of a.
Thus, by repeating this process for the L - intermediate matrices WJ (t),,,(t), we can obtain the vector
II c=(+1
WJ (DDD)))))))
time ((LL-)m(sp)) + r:)). Finally, by multiplying a sparse diagonal matrix (,,,(t), we get the desired vector (t).
18

---

Under review as a conference paper at ICLR 2022
Compute (t). Note that (t) is essentially ht)e- so consider how to compute ha,e(t) for general & L]. Recall that
(t,e(t) ((WV(())= (W())))))
since ha (1) is sparss,, the product W2(0h,,t (()) can be computed in ((msp) time. For the product (W)(t),,e-11( can be computed use the low rank decomposition, which takes ((mr)) time. Apply the threshold ReLU takes (m) time. Hence, the total time is O(m(re sD)) time.
The running time results are summarized in the following lemma:
Lemma C.4. For l E [L] and i [n]. suppose ( S k. Ler O. Suppose the change of D is sparse, iee, [Da,(tt) - Di.e(0)]lo 8. For l [L]. E [n], any > 0, suppose the change of We is of low-rank, i.e., AW&(t) = 1 We further assume that {ye.5 2t,j (eEe)tt()))) and (W)()));)))) Jie(m) are pre-computed.
Then, for any E [L] and i E [n], the vectors (t), (t,(t) E R can be computed O(mL(s + k: + -ime..
As a direct consequence, we combine Lemma C.2 and Lemma C.4, then we get the following corollary:
Corollary C.5. For é E [L] and E [ni], we can compute vt,,(t,, (1)e(l) E I as in Algorithm with the following time bound:
Compute (t) in time 00 + k+
Compute (t) in time (mm(s + k k +r)).
D FAST TENSOR PRODUCT REGRESSION
In this section, we design generic algorithm for solving the following type of regression task: Given two matrices U [u] ..... u IT V
[o] ..... 1 IR with > n, consider the
matrix J Raxm" formed by
vec(u1 (1( 2 (222 T
vec( (u "
We are also given a vector c E R", the goal is to solve the following regression task:
min
||J cll2.
#E
Our main theorem for this section is as follows:
Theorem D.1 (Restatement of Theorem D.14), Given two n X m matrices U and V, and a target vector e E g" Let = [vec(u]u]] ...., vec(un J] E Ruxm", There an algorithm (Algo- rithm 6) takes Ömm + n² (log(k/e) + log(m/e0)) 2) + n²) time and outputs a vector & E R=" such that
||J - c/22 S e||c|12
holds with probability at least 1 - 8, and k is the condition number of
D.1 APPROXIMATE, J VIA TensorSketch
We introduce the notion of TensorSketch for two vectors:
Definition D.2. Let hi, h2 [m]] - [s] be 3-wise independent hash functions, also let a [m] + (+1} a 4-wise independent random sign function. The degree two TensorSketch transform, S: I x R - TR is defined as follows: for any i, E [m] and r E [s].
Sr,3. = =i o2(1) hh (i) hz()) mod s).
19

---

Under review as a conference paper at ICLR 2022
Remark D.3. Apply S to vectors x; I can be implemented in time Oss nnz(z) + mnz(y)).
We introduce one key technical lemma from Avron et al. (2014):
Lemma D.4 (Theorem of Avron et al. (2014)), Let 5 E Raxm' be the TensorSketch matrix, consider fixed n-dimensional subspace V. If S(n /(+")), then with probability least 6, ||Sx||2 = (1 + > simultaneously for all V.
Now are ready to prove the main lemma of this section:
Lemma D.S. Let e, E (0, 1) denote two parameters. Lei J E R be matrix such that i-th row defined as rec(u;t) for some ui, v, R.. Then, we can compute a matrix J E R such that for any vector E IR,,, with probability at least 1 - ô, we have
|| xl12 ( 2112.
where s = ((()) The time to compute J is (nns log + nnz(U) + mnz(V)).
Proof. Notice that the row space matrix J can be viewed as an n-dimensional subspace, hence, by Lemma D.4, the TensorSketch matrix S with = O(n² (c)6)) can preserve the length of all vectors in the subspace generated by J with probability
to multiplicative factor of + 6.
The running time part is to apply the FFT algorithm to each row of J with total of rows. For each row, takes (s log s m) time, henee the overall running time is Onn(s + m)).
D
D.2 APPROXIMATE J VIA TensorSRHT
We note that the dependence on the target dimension of sketching is 0(1/8) for TensorSketch. We introduce another kind of sketching technique for tensor, called TensorSRHT. The tradeoff is lose input sparsity runtime of matrices U and V.
Definition D.6. We define the TensorSRHT S RR S P. X HD2), mm
where each of P E {0, 1) mxxd contains only at random coordinate, one view P as a sampling matrix. H is a d x d Hadamard matrix, and D1, D2 are two x d independent diag- onal matrices with diagonals that are each independently (uniform in {11
be a Rademacher random variable
Remark D.7. By using FFT algorithm, apply S to two vectors x, E. R takes time ((m m-
We again introduce a technical lemma for TensorSRHT.
Lemma D.8 (Theorem 3 of Ahle al. (2020)). Let S Rexm² be the TensorSRHT matrix, sider a fixed n-dimensional subspace V. If = Q(n (og'log(mmm (nm//o) then with probability at least 1 8, ||Sa||2
(1 e)][==] [||l0|12 12 simultaneously for all V.
Lemma D.9. Let 8 E (0,1) denote two parameters.
Given a list of vectors Let J E Ruxm² be a matrix with i-th row of J is defined as compute a matrix J gaxs such that for any vector E 12", with
21 Umm U1.-. Em R",
vec(u,v,T), Then, probability at least 8, we have
|| = T2112
where 5 Q(n (nm/ - 2). The time to compute J O(n(mlogm + s)).
Proof. The correctness follows directly from Lemma D.8. The running time follows from the FFT algorithm to each row of J, each application takes O(m log m + s) time, and we need to apply it to rows.
20

---

Under review as a conference paper at ICLR 2022
D.3 TENSOR TRICK AS PRECONDITIONER
In this section, we use TensorSketch and TensorSRHT as preconditioner to solve a regression task involving JJ This is important step implement the Gram-Newton-GGuss, iteration Cai al. (2019); Brand et al. (2021).
Before proceeding, we introduce the notion of subspace embedding:
Definition D.10 (Subspace Embedding. Sarlos (2006)). Let A E R we say matrix S E R*x a + c) l2 subspace embedding for if for any x E R², we have |S = (1 + S
Ax e Equivalently, | - U S SU| & where U is orthonormal basis for the column space of A.
We will mainly utilize efficient subspace embedding.
Definition D.11 (Lu al. (2013); Woodruff (2014)). Given a matrix A E RRNn with N poly(k), then we can compute an S E I (k/))))k)))))))))))))))))))))))) / XN such that with probability least 1 1 6, we have
||Ax||2 = (1 Ax]x
hols for all x E RR Moreover, SA can be computed in ONNi: log log N)/e)) time.
Algorithm Fast Regression algorithm of Brand et al. (2021))
procedure FASTREGRESSION(A.PM v)
Lemma D.12
2:
b A is full rank, E (0, 1/2)
3:
Compute a subspace embedding SA
>> S (pooly(Ing N
4:
Compute R such that SAR has orthonormal columns via QR decomposition D R E R
5:
0g E R
6:
f
7:
while ||AT ARa, 2/1k do
8:
2++1 f ++1
(R A AR) AAR) (RT A AR2; - R y)
9:
10: end while
return R=
12:
end procedure
Lemma D.12 (Lemma 4.2 of Brand et al. (2021)). Let N (kpolyy(oog k)). Given a matrix
A E R Nxk, let denote its condition number. Consider the following regression task:
min
||ATAr
Ar
yl12-
TER
Using the procedure FASTREGRESSION (Algorithm 5), with probability at least - 5, we can
compute an e-approximate solution satisfying
||A A vll2 S el|v1l2
in time O(Nklog(1/)) +*).
Our algorithm is similar to the ridge regression procedure in Song et al. (2021), where they first bootstrapping reduce the dimension of the original matrix,
apply their sketching algorithm then use another subspace embedding to proceed and get stronger guarantee.
We shall first prove a useful lemma.
Lemma D.13. Let A E B.N suppose SA is subspace embedding for A (Def. D.1I), then have for any x E RR*,, with probability at least 5,
I[SSA] SAr bll2 (1 + A Ar -112-
Proof. Throughout the proof, we condition on the event that S preserves the length of all vectors in the column space of A.
Note that
((AA) SAx b112 [|(SA) SAx | 2((SA) (((S))TSAass SAx,
21

---

Under review as a conference paper at ICLR 2022
We will first bound the norm of (SA) SAA, then the inner product term.
Bounding ((
Let U E I x be orthonormal basis of A, then use the equivalent definition of subspace embed- T SU - I|l I this means all the eigenvalues U S SU lie in the range c) Let V denote the matrix U s SU, then know that all eigenvalues of V lie in range [1 -d=
ding. we have | T
of of [(1 ()), (1
c)*, (1 + e) Setting c as c/4, we arrive at I| V--VV- -I| I c. This shows
that for any E IR we have II(SA) SAzl|2 (1 + Ax1/2-
Bounding ((SA) SAr,b) b) Note that
((SA) SAx, b) = A SAb)
-1/2- -11/S//a + ||SAb|| ||SA(z - -))1))
=1/2- <) + |bbb| ||A( b)ll))
(1 (A Av., b).
Combining these two terms, we conclude that, with probability at least 1 - 6,
[I(SA) SAx
bll2 (1 + A AA Ar
b/12-
Algorithm 6 Fast Regression via tensor trick
procedure FASTTTENSORREEGGESSSSN((u(
mg fwili Rmxx , C RR.)
Theorem D.14
b J= =J=[vved [vec(((()))) vec(u2u)) ...., vec(un T E IR pxm
2:
3: 81 + On
4:
$2 + O((n log m) lo
5:
Let S1 E xm be a sketching matrix
D S1 can be TensorSketch or TensorSRHT
6:
Compute J J= JST via FFT algorithm
D J Rux
7: Let S2 E R=2xs be a sketching matrix defined in Definition D.11
8:
Compute subspace embedding S².JT
9:
Compute R such that S2J R has orthonormal columns via QR decomposition
RE Ruxn
20
+ 02 E R.
11:
++ 0
12:
while
|[JJ Ra: - c/22 E do
13:
24+1
t 21 - (R JJ R) (R JJ Res - R c)
+
15:
end while
return Rz, end procedure
17:
Theorem D.14. Given two n m matrices U and V, and target vector E. Let J (vec(uro]) T)T ...., vec(u, o E Roxm², There is an algorithm (Algorithm 6) takes Ö(nm + n² (log(1 k/e) (og(m//)) + n²") time and outputs vector & E RD such that
||JJ T c12 S c||c|2
holds with probability at least 1 - 5, and k is the condition number of J.
Proof. We can decompose Algorithm into two parts:
Applying S, to efficiently form matrix J to approximate and reduce its dimension, notice here we only need for this part to be small constant, pick 0.1 suffices.
Using S2 as a preconditioner and solve the regression problem iteratively.
22

---

Under review as a conference paper at ICLR 2022
Let @ denote the solution found by the iterative regime. We will prove this statement in two-folds:
First, we will show that ||JJ - c//2 e||c|/2 with probability at least ó;
Then, we will show that ||J म - c/| (1 0.1) cll2 with probability least
1- d.
Combining these two statements, we can show that
||J d|2 (1 c|2
S 1.1e||ell2
Setting e to e/1.1 and ó 8/2, we conclude our proof. It remains to prove these two parts.
Part |JJ - clla S <||c|l2 We observe the iterative procedure is essentially the same as running FASTREGRESSION on input J M, e, hence by Lemma D.12, we have with probability at least 1 1- 6, |JJ 2 - c||2 - <||c||2-
Part 2. ||JJ c|l2 0.1 cll2 To prove this part, note that by Lemma D.8, know that JJ subspace embedding for J Hence, we can utilize Lemma D.13 and get that, with probability at least we have [JJT cll2 = (1 + 0.1)
ell2-
Combining these two parts, have proven the correctness of the theorem. It remains justify the running time. Note that running time can be decomposed into two parts: 1). The time to generate J, 2). The time to compute via iterative scheme.
Part 1. Generate To generate J, we apply S₁ E Rm" which is TensorSRHT. By Lemma D.9, it takes (m logm + 81)) time to compute J, plug in = O(n log" (nm/ the time is Onn)).
Part 2. Compute & compute ã, essentially we ru FASTREGRESSION on JT,c, c, hence by Lemma D.12, takes Ö(22 log(n/e) + n²) time, with 82 = O(nn + log m) logn) and x is the condition number of J, which has the guarantee x (1 + e)x(J). Hence, the overall running time of this part is (n² log(k//) n
Put things together, the overall running time is Ö(nm n² log(x</e)+**)
Remark D.15. Due to the probability requirement (union bounding over all layers and all data points), here only prove by using TensorSRHT. One can similar strategy obtain input sparsity time version using TensorSketch. We remark that this framework similar to the ap- proach Song et al. (2021) takes to solve kernel ridge regression, where first uses shallow but fast sketch to bootstrap, then use another sketching to proceed with the main task.
E SPECTRAL PROPERTIES OF OVER-PARAMETRIZER DEEP NEURAL NETWORK
In this section, we study the spectral properties of our Gram matrix and connects to multi-layer NTK.
E. BOUNDS ON THE LEAST EIGENVALUE KERNEL INITIALIZATION
We first define the Gram matrices for multiple layer neural network.
Definition E.1 (Multiple layer Gram matrix). The Gram matrices K2 E Raxn for l (O...., L} of an L-layer neural network are defined as follows:
(Ko)
xx.
23

---

Under review as a conference paper at ICLR 2022
1
(K--1) b,j
For l. O, let Et.igj =
E 1222 for any (i,j) [n] x [n]. Then,
(K 3,2
K 8- 1anl
(Ke) = (2..2))
EE
[0(x1)((22] ve [L- 1],
E (K (KL) == (6).2))-(((221.1....))
(0((1))((xz))
Lef År. := Amin i.) to be the minimum eigenvalue of the NTK kernel K..-
In the following lemma, we generalize Lemma C.3 in Brand et al. (2021) (also Lemma 3 in Cai et al. (2019)) into multiple layer neural networks.
Lemma E.2 (Bounds on the least cigenvalue at initialization, multiple layer version of Lemma C.3 in Brand al. (2021)). Let de denote the minimum eigenvalue of NTK defined for E-th layer of neural networks. Suppose m = 2
2 n² log(n/8)), then with probability 5, we have
3 de,
Amin (G,(G)
VLE [T]
4
Proof. For any E [L], have G2 = JrJ,
E ERnx Hence, for any i, [n].
(Gr) = vec( af(W, x4))
vvc( (f(W.x))
oW,
OW.
T
vec Diee II
W Di.cahi
Dia II W Draahhse- s k=(+1
vec
k=(+1
Diee II W Dssa) 1) (hi,e-1 @ Ami) Die II 1m8+1
L
(hii-1 & Ame)
W Dj
k=8+1
Di. II k=(+1
WT Desaa)
(hhee-1 @ Ime)(ejee-11 @ Imee II k=t+1
Dgi
W Diaaa
L
L II k=e+1
Di,t
II k=l+1
Wil Dikk)
Dj.e
Wt Dz.e a 1
where
t-1 Wa.
h4.--
In particular,
(Gr)ss a D2, Di,aa hi.1-1hhL-11 1 Φ' (wuf (' (b,,
(4)
m
We will prove that ||Gr K L blo small, which implies that Amin (GL) is close to de. The proof idea is similar Du al. (2019a) via induction on l.
For = 1, recall (9),;) = Eoe(m (Wi),,, for k E [m]] Hence, for any k E [m].
E((g1,)) =
b.
Σ
B( (2:)6(13))
(Wi)&, ~N(0, )))
6e[m]
Σ (2))0(z;))
m
be ml]
2 4 Zz..
24

---

Under review as a conference paper at ICLR 2022
Then, have
Bhh
M1 =
Σ
TH (((1,,)) ((115)))))
km(m]
Σ
ke[m]
Σ
E
[o(u)((v)]
(u, 20 F1
ke(m]
E
(o(())(v)
(u,t)~A(00
2 )
E
(o(1')((1')
(wr,,"))~N0 (0.25...)
(Ki)..
Next, we will show that 10h1,)) (h1 & (0r,)) s ((Wi),,,, m.)l 1(Wa)),,,;))) Since (W, ~ NN0, ) by the concentration of Gaussian distribution, m
hi Ai. concentrates around its expectation. First, for any & E [m],
(((W)),,,, +,)1 S Vc Vk E [m]. e [n]
holds with probability least --mme--mm cm/4, Conditioned on this event, have ( S e for all j E [n] and & E [m]. Then, by Hoeffding's 's inequality, we have for any (i,j) E [nx] x [m],
Pr [ih - (Ki)..) 2 t] s exp (-2m (2c)² = (xp(-O((/((mc2))
Hence, by union bound, we get that
|1, - (Ki)isl S1
max (L)))(n)((()
with probability at least
1 mn exp(-{(mm)) ni² exp(-S((t)/)))) (mm)))
1/2 polylog(nI./H)) 8), we have with probability least we choose e log(mnL/8) and = m-1/2
m
1 L
1h IT (Ki)is) s 2mm -1/2).
max (b)))(())((]]
Suppose for L,... (h- L), have
hh IhT (Ke)apl S (mm-m11)
max
Consider / /= + By a similar computation, we have
2 2 f 1, m
E [(Gee:)(
W.
Define a new covariance matrix
h h h h
hT h
Ees ==
v(i.j) E [n] [n].
h he-151
We have
JE [hi W.
E (m.))~ Seni))
[o(u)((o)]
ke(m)
E (1..0))-N(0,25..0
(o(u')((())
== (Ke)is
25

---

Under review as a conference paper at ICLR 2022
Hence, we have with probability at least
i.
max ((J) x(M)
he (Ke)as) (1m)).
(5)
It remains upper bound the difference || Kelloo.
K Ki
max [x)))(
E E [d(u)((o)] (1.))~W(0.28...... 28.)))
(0.0)) [(())(v))
Recall that
(Kg-1)i, (Ke-1)js
K 1 (Kg-1)3s1
Een =
v(i,)) E [m] x [m].
and hence, by the induction hypothesis, we have
|Sess - Leaaloo S
I Me-1hee-15 (Ke- 0mm-1/2)
max [x)))
Notice that Eegu can be written as
1 hhe-
(oos(BL...)
1
Moreover, when is the ReLU function, have
E
[d(u)(()]]
(lla.F((sa)) F(0e.a))
where
cos(#)
F(0)
(u.r)~NNo (()))(v) with Σ(0) = cos(())
1
We note that F(0) has the following analytic form:
1 (sin(0) + (x cos(()) 1/2].
F(0)
(6)
2T
Similarly,
E (1,0)~
[Q(u)((v]] 22 (Ke-i)i,(Ke--Diz F(Tesi))
where Te.,j = cos
(Ki-)hi
By the induction hypothesis, we have (Kg)..)
(k
he, 1/2) for i, E [n]. By Lemma F.7, we also have |/e,,llz 1+te for all 8 E [L] and Q(me²/L), (((m)))-e0((m))2)) They implies that cos cos(TE,i.)) cos(O) (m
i [n] with probability 1-O(nL)-e
Thus, by Taylor's 's theorem, we have
[F(0rs.)) - F(TE, S 1122.
Therefore, have
E
[o(u)á(1)] -
E (u.)) [o(u)((v)]
= 2
(Ke-)) (Te,.))
5 (m2m112)
That is,
K. Kelloo 5 :0mm-111 112)..
26

---

Under review as a conference paper at ICLR 2022
Combining Eqs. (5) and (7) together, get that
ber - (Ke)ial S 1(m-mmmm))) 1/2)
max (kd))[[)x(()
holds with probability at least - for l = 1
By induction, we have proved that for the first L 1 layers, the intermediate correlation h hegi is close to the intermediate Gram matrix (Ke)a.- Now, consider the last layer. Recall G is defined by Eq. (4), which has the same form as the correlation matrix of a two-layer over parapaaraeteriezed neural network with input data 1,).e)))) Define
Kk == h h--11,
E
[d'(w" hL-1,)0'(w" he-1,))
(0.2
Then, by the analysis of the two-layer case (see for example Song & Yang (2019); eu al. (2019b)), we have
A
||G, Kr S
2(2 n2 log(n/5)), where dL = Amin (K L). It remains to bound || lently, for any (i.j)= E [n] x [n].
kL Equiva-
max (KJ) [
E ((..5
1 (u))o'(v))
E
1 (u)o
(1.
The expectation has the following analytic form:
E 1 (=1) 1 sin (p)
p²2 Ppg ppq
with Σ =
(2
2x
q²
By the analysis of the (L - 1) layer, we know that |PL,E,j - PL,cs < (mm21/2) where PL. = cos(TL.))) and PL.Y cos(OL, Also, notice that cos( F(TL-155)) 10, 1/2] by Eq. (6). Hence, the derivative of the expectation is bounded, and by Taylor's theorem, we have
K2 K2 S Ō(m
It implies that [ KLll <
dz. 8. which further implies that
A 4
|GL K&ll s
Equivalently, we get that
3 \ Amin (GL)
with probability at least 8.
The lemma is then proved.
E.2 BOUNDS ON THE LEAST EIGENVALUE DURING OPTIMIZATION
In this section, we generalize the Lemma C.5 Brand et al. (2021) into multiple layer neural net- work
Lemma E.3 (Bounds on the least eigenvalue, multiple layer neural network version of Lemma C.5 of Brand et al. (2021)). Suppose m 2(A 2n2 log(n/8)), with probability least 1 8, for any set of weights W,. ... WL satisfying
I Wc(0)] S R.
then the following holds
[[) (r((ll)) A1/2.
27

---

Under review as a conference paper at ICLR 2022
Proof. Recall (GL) = (/ For simplicity, let 2g, == (WL), hI (WL),h,LL- and (0, = (W2(0)))
(GL(W)) (W))a))
, - Σ
h , -- o' ((L))d (1)
(0) (0)
m Σ o' (0))
< | - h4, (O)TAa,..-1(0))
Σ S (z) (z().))
+ +1,),1-(00) (0)
1 Σ
お (zir)) )o' (2jrr) Ø' (Zi, (0))
(0) + AL + Õ(m -1 1/2)
< -,,,
Σ of (()))))))))))(x.x)(())((..)
m
where the last step follows from o' (x) (0, and Lemma E.2.
To upper bound the above two terms, we first need to bound the move of h For any E [L-1]] we have
|l,a, - hs,r(0)ll2 -ld(Ver-=)--((((())..-0000
S (
+ |\x(0])) (1/(0)]))|1(4,.-111
S Wl)W2 W,(0)]|
||1.1111111111 (0)
+ ||W; -
< (R + -h,,
R(L + E),
where ew == |[W]( S 3 by the well-known deviations bounds concerning the singular values of Gaussian random matrices (Rudelson & Vershynin (2010)). Also, when E = 0, we have ||4o, (0))())l22 = 0. Hence, we get that for all E IL - 11, [n)]
1 (0)]l2 S cR(ccw)
since R < by our choice of R.
Hence, it implies that
| hi, -1)
< (hi,1-1 - (0)) hh,,,-1111 + 1 (h,0000(001-11,--1000000
- (0))
2(1 + R)
Similarly,
A + Õ(m 1/2)
m
Σ 8 -(aaa
m
b + (mm-1/2) 1/2)
1 (WW.). (.
m
By our assumption, we have ( r.), (Wr S R. We also know that || - hx.L-10)) ||2 < +I+ce 1--1 Then, we can follow the proof in Du et al. (2019b); Song & Yang (2019) and define the event A., == Sw, h w|l (Wr(0))/l1z R. ||- h1,1-1 (0) S O(R) such that 1 TA>O
h>0 * [(W2 ho We have A,,- happens if and only if (((((()))) (0)),
Aw) T (0) < O(R(R m 1/2 for some fixed vector Au of length at most R. It implies that
Pr[Air]
Pr
|||| < R o(R) O(R).
s~N((011)
28

---

Under review as a conference paper at ICLR 2022
Hence, using the same proof in the previous work, get that
Ay + (m 1/2)
m
('(4,))0(22))) - o' (((a.(0))) sA2 m m-
m
Putting them together, we have
(GLL(W)),s (G(())))) s (2(1+ 222ow) AL)R. And by our choice of R, we get that
(W) Gr. (W(()) & (2(1+ (2ol 1-1 + ALDTR S A1/2.
The lemma is then proved.
F CONVERGENCE ANALYSIS OF ALGORITHM 2
In this section, we analyze the convergence behavior of Algorithm 2.
F.1 PRELIMINARY
We recall the initialization of our neural network. Definition F1 (Initialization). Let for all l E [L]. Let mo = d. We assume weights are initialized
Each entry of weight vector a E mR is i..dd sampled from {-
(-vmy-
vm, -
_ uniformly
at random.
Each entry of weight matrices We E Rmmemme sampled from (0, 2/me).
We also restate the architecture of our neural network here.
Definition F.2 (Architccture). Our neural network is a standard L-layer feed-forward neural network, with the activation functions defined as scaled version of shifted ReLU activation: 6(x) /o, [x > /2/mbbaæ, where c := (2(1 4(b) + bo(b)))11/2 Here b is threshold value we will pick later. At last layer, we use a scaled version of a vector with its entry being Rademacher random variables. We define the neural network function m + R as
a @(W WWWW&&WWWWl (
We measure the loss of the neural network via squared-los: function:
C(W) 2 1=1
(f(x.)) - -1)2
use
T mn Ra denote the prediction of our network:
fx(X) = (1)) ((W)(()))))
We state two assumptions here.
Assumption F.3 (Small Spectral Norm), Let E fo...., T) and let S be a parameter. We assume
max [[a(t) W2(0)|I S R.
RE[L.]
Later, we will invoke this assumption by specifying the choice of R. Assumption E.4 (Sparsity). Let {O....,T) ,,T and let 2 1 be an integer parameter. We assume
|lleell & s, VL [T), E [n]]
Later, we will invoke this assumption by specifying the choice of 8.
Finally, throughout this entire section, will assume me for any E [T.].
29

---

Under review as a conference paper at ICLR 2022
F.2 TECHNICAL LEMMAS
We first show that during initialization, by using our threshold ReLU activation, the vector hi, is sparse, hence the diagonal matrix D..e is sparse well.
Lemma F.5 (Sparse initialization). Let 0b(x)) max <2 b,0 0} be the threshold ReLU activation
with threshold 0. After initialization, with probability 1 nL Q-2mme- (mm//) it holds for all i [n] and E [L]. |/hu,ello S - bb/44/
Proof. We fix i [n] and [L], since we will union bound over all and at last. Let uy R be a fixed vector and Wer to denote the r-th row of Wi, then by the concentration of Gaussian, we have
Pr(os((Wen)) us)) 0]
Pr =~~[(0 )
[a b) <
Let S be the following index set S = {r E [m] : (eer,, un) > b}, the above reasoning means that
for the indicator random variable r E 5), we have
E[1[r E S]] < exp(-80m//4
Use Bernstein's 's inequality (Lemma A.3) we have that for all > 0,
f//2
Pr[[S] > + S exp(
k + t/3
where k == exx((='pmmm//l By picking k, have
Pr[[S| > 2k] < exp( 3k
Note that |S] is essentially the quantity ||/r,ello. hence union bound over all and i and with probability at least
nL
exp(b"m/4))).
we have |Hu,elo S 2m exp(-b2m/44)
Remark F.6. The above lemma shows that by using the shifted ReLU activation, we make sure that all hie are sparse after initialization. Specifically, we use k == m cxp(-b3m//) as a sparsity parameter. Later, we might re-scale b so that the probability becomes exp(- /2).. stress that such re-sculing does not affect the sparsity of our initial vectors. If we re-scale and choose it as V0.4log m, then it = 0.8 and hence with high probability. |[hs,eloo S (m...))
As direct consequence, note that all initial Di.t are k-sparse well.
We state a lemma that handles the l2 norm of hi, when one uses truncated Gaussian distribution instead. Due to the length and the delicacy of the proof, we defer it Section G.
Lemma F.7 (Restatement of Lemma G.6). Let b > be a fixed scalar. Let the activation function ((x) = /e,1 [x: 22/mbba, where es (2(1 (bb) + bo((b)) Let (0,1). then over the randomness of W W(O), with probability least (mI.)) expl- (0ppp((() /2)e² (T?)), have
|hs,el2 E [1 -, c], Vi [n], & [L].
We remark that the parameter -b2/2 captures the sparsity during the initialization.
The second lemma handles the consecutive product that appears naturally in the gradient computa- tion.
Lemma E8 (Variant of Lemma 7.3 Allen-Zhu al. (2019a)). Suppose 2 M(nL log(nL)) With probability at least 1 e 22(/113) over the randomness of initializations W (0)...., (.) e mm for all E [m) and 1 & & b L.,
30

---

Under review as a conference paper at ICLR 2022
(a) WH,Ds,o-WW6-11 .r, ((VL).
(b) WW..D.6-W6-.. ... Da, W. - 2|100|2. for all vectors with Ieelo S ( log m
(c) || D.DDD.-WW--1 ... Di,a W. W.ll S ((1)[lull2, for for all vectors u with ||ullo O Log T.
(d) For any intergers 1
ooggm with probability at least 1 e --((llggm) over the randomness of initializations W W1 (z E m Iu WoD W S |||||||||||2 O for all vectors u, v with | |Plllo S 8.
The proof is similarly the original proof of the corresponding lemma in Allen-Zhu al. (2019a), however we replace the bound on hig with our Lemma F.7. We highlight this does not change the bound, merely in expense of worse probability.
The next lemma concerns the bound on the product being used in backpropagation.
Lemma F9 (Variant of Lemma 7.4 Allen-Zhu al. (2019a)). Suppose m z (oggl))
Then with probability at least 1 e e-toloo" m), for all E [m]. & [T.].
log m VL). mm
||a D ... Daeleel2 00(
Proof. By Lemma F8 part (a), we know that [[ Well S O(VL), consequently,
[ Well S O(vL) since |D s 1. This means for any vector a E R"" with probability least 1 e (kkkL²), have
D ... DellW S {(VL))
Conditioning on this event and using the randomness of Rademacher vector a, we have
Pr(la DI, Ded S t. VT[lu]la] 2 1 2exp( mt²
pick
log
Vm we know that
"illaDDLLWL.. DirWeu log m Vm
log² m
VLllu|lal 2 2exp(
2
This means with probability at least 1 - 1-e e- -SQgo m)) we have
||a D.. ... Dielalla s 0 log VL).
m
The next several lemmas bound norms after small perturbation.
Lemma E10 (Lemma 8.2 in Allen-Zhu et al. (2019a)). Suppose Assumption F3 is satisfied with R < O 107= m. With probability at least e- (3
(a) Agit be written Agiz Agi.e1 + Agina where |Agie S O(RL3/2) and
Agg,e2loo 5 O RL3/2
Alggm
mm
(b) |ADaello S ((MR2/33)) L) and ((A O(R13/2)
(c) Agr.elle AAhiell2 S 5(RT,512 (oog.
Remark F.11. Lemma F.10 establishes the connection between parameter and s of Assump- tion F.3 and F.4. As long R is small, then have O(mR273 Such relation enables us to pick R to our advantage and ensure the sparsity of AD; is sublinear in m, and hence the update time per iteration is subquadratic in
31

---

Under review as a conference paper at ICLR 2022
Lemma E.12 (Lemma 8.6 in Allen-Zhie et al. (2019a)). Suppose Assumption F.4 is satisfied with m 1 S S O E log m and m 2 InnL (n)( with probability at least - --2 log over the randomness of (O)). every i E [n], 1 S S b < L, we have
(a) [w=(0]().o-00)) AD
(0) AD, a, W. S ((VL).
(b) [[(W(0)) + ADa) + 1 (Ds,,) OVVL) if Assumption F.3 is satisfied with R S O z.
(W. AW.)l -
Corollary F.13. Suppose Assumption F3 is satisfied with R S O (z and m z
MinL (ng(n))). with probability at least 1 - e 221log m) over the randomness of W (0) and a, for any [m] and E [L], we have
11a (0) (O) AWL) ... (We-1 (W + AD, < <((LL/))
Proof. We first note that by Lemma F,10 and the choice of R, we know that Assumption F.4 satisfied with 1 s s S O 1.=log" log" m
The proceeding proof is identical to Lemma F.9, use denote the product (D,, + (0) + AWL).. (We+1(0) + (D,,)) + AD;,). Per Lemma F.12, we know that with the conditions stated, we have
|P| SO(V))
conditioning on this even, for any vector E IR"", with probability at least 1 eQ(s m)), we know that
|uull2 S OVVL)(())) [u|l2,
Use the randomness of vector we have
ml2
Pr[[a Pul St VT.[/ul22] 11 2exp( 1-22expl
2
Pick log m vm
we are done.
Lemma E.14 (Variant of Lemma 8.7 in Allen-Zhu et al. (2019a)). Suppose Assumption F.3 sat- 1 isfied with R S O Las log m and m 2 (2nL log(n))), with probability at least e Qloog m)
over the randomness of W (0) and for all i E [m), E [L]. Define
u a T (D), (W, + (W) (We+ (D), + AD;.).
Ug a (t,,()) ... We+1 (o) satisfies that
|/u - uoll2 S LL/m)
Proof. Note that we can invoke Corollary F.13 to give an upper bound on ||u||2 with probability 1 R(log m)
||u|| - L.(LL/))
Similarly, we can use Lemma F.9 to give upper bound on ||ua||a with probability - - m).
|| 2 L(VLm/))
Finally, by triangle inequality, we get
||u - uoll2 5 LLm),
with the desired probability.
32

---

Under review as a conference paper at ICLR 2022
F.3 BOUNDS ON INITIALIZATION
In the following lemma, we generalize the Lemma C.2 in Brand et al. (2021) into multiple layer neural networks.
Lemma F.15 (Bounds on initialization, multiple layer version of Lemma C.2 in Brand et al. (2021)).
Suppose [. log(nI)), then with probability ([T)) -QQ(//Z2 we have the following
(WW,,,) = O Vi E [n].
With probability at least 1 - eQ(log""" m), we have || = Ō( ve E [L]. Vi [n].
Remark F.16 The bound of in Brand al. (2021) has linear dependence because they need to bound [| (0)ll. If we do not need to bound [| (0)ll, then m does not depend linearly on n.
Proof. We will prove the two parts of the statement separately.
Part 1: By definition, for any E [n:]. we have
f(W, mx) = a (oV
We shall make use Lemma F.7 here: with probability at least 1 ((L)) exp((S((k/L2), we have Ille,zll ll2 [0.9, 1.1]. Recall that a E R" has each of its entry being a Rademacher random variable scaled by 1/m, hence it's 1 l/m-subgaussian. Using the concentration of subgaussian (Lemma A.5), we have that
m
Pr(l 2 4) < eexp(
2 1
finally by noticing that t ((1) and | E [0.9, 1.1), conclude that ( æz) 2 O(1) - 2 exp((=((m)),
union bounding over i, we conclude that
Pr[f(W,x:) 2 O(1) s (n)) exp(-S((m)), vi [n].
Part 2: We will combine Lemma F.7 and F.9, with probability at least 1 - ((L)) .E ( (2(/233)
have |ll,, [0.9, 1.1] and with probability least 1-e e 2(logg m)) we have ||a T D ... Wellz S
L/). Combine them together, we know that with probability at least 1 e -2(oo m) ||Je,,ill Ō(VL./m).
L
Remark F.17. By utilizing the structure of vector a, can show that with a weaker probability (1 e -Qtlogg m²), f(W, has even smaller magnitude (Ö( (m)). For purpose, f(W,x.)
O(1) suffices.
F.4 BOUNDS ON SMALL PERTURBATION
In the following, we generalize the Lemma C.4 in Brand al. (2021) into multiple layer neural network. For the simplicity of notation, we set all my to be m.
Lemma F.18 (multiple layer version of Lemma C.4 in Brand et al. (2021)). Suppose < T.A. T.4. Tog and m Q(nL log(nL), With probability at least 1 1- e Aillg" m) over the random initialization of (0) (0), (2))), ... Wr(0)), the following holds for any set of weights Wi.... Wr, satisfying
O
W - W = R, E [L]
*IJwe.a; Jw.(o).e,112 Ō(vT./m)
33

---

Under review as a conference paper at ICLR 2022
|Jwe Jw.(o)]l= Ö(n' L/m).
[Jwellr O0.2//2
Proof. Part 1. To simplify the notation, we ignore the subscripts below.
Consider the following computation:
Jw,,, (w), 2l
||uh -nhh
1| (0) + 1h+-1) (0)l-100)1l=
= I[(u
< [[(u +0) I|
< ||uu uoll2 (10 ||ul22 -
11Ahe-1ll2-
where
L
u (De(0) AD,)
II
(W&(0) + AWA) (D&(0) + AD&)
I.
D2(0)
II
(a(0))
k=t+1
By Lemma F.14, we know that with probability at least 1 - e Qllog m), we have
u uoll2 S Ó(VLm
By Lemma F.7, we know that with probability at least O(n)) k/L2 have [ E [0.9, 1.1].
By Corollary F.13, we know that with probability at least e fillo"" m), we have
|ull2 S L(VL/m).
By Lemma F.10, we know that with probability at least 1 c (1m we have
Ah S )
Note that due to the choice of Rt we know that
| ||2 S
Taking a union bound over all events, with probability at least - e =2(log" m), we achieve the following bound:
|J Jw ll2 S Ō(VL/m).
Part 2. Note that the squared Frobenious norm is just the sum of all squared & norm of rows, hence
Jww, Jw s Ō(n 1/2 / /L/m).
Part 3. We will prove by triangle inequality:
[w 1 - [ + [|w,, - Jw
< 5(n +L L//)
+
02 LL/m).
34

---

Under review as a conference paper at ICLR 2022
F.5 PUTTING IT ALL TOGETHER
In this section, we will prove the following core theorem that analyzes the convergence behavior of Algorithm 2:
Theorem F.19 (Formal version of Theorem 1.1). Suppose the width of the neural network satis- fies m O 2 L2), then with probability at least 1 e- -R(log m) over the randomness of the initialization of the neural network and the randomness of the algorithm, Algorithm 2 satisfies
1 22111: yl12.
|e++ yll2 S
Before moving on, we introduce several definitions and prove some useful facts related to them. Definition F.20 (function J). We define
L II k=(+1
= De,(Zz)
Z Di
Zz-1))
R Xm
where
(z) (Z,h:, Z1-1))); (Z xx))))))
E R
1.(Z...... Zz-1):= (6) ..l.
R
Fact F21. Let Je denote the function be defined as Definition F200, For any E (0....,,y. T), we have
se+1 - fc Σ EE
(L Je((1 - s) ((11-))()))
sW(t 1))ds) vee((Wr(t+ + We(()),
t=1
Proof. For i [n], consider the i-th coordinate.
(f41 - fi) = S/" f((1 - st)(t) + sW + 2;)''
I.
Of
- ())( + sW ((1-s)W(1))
+ 1),
voc(Ws(t +t Wx(t))ds
Ow,
Σ
(((Ls) sW 1)).ds)
vecc(W((t+1) W,(O),
Thus, complete the proof.
Fact F22.. For any E (O....,T), T}. we have Je ( WL( Jr,--
Proof. In order to simplify the notation, we drop the term & below. We note that for E [n]) the i-th row of matrix Jr,: is defined as
I Dr,(( II WT Dr.a)ah k=8+1
where D., diag(d' ('((Woh,..11)) and h4.8- o(We-1 ( this is essen-
tially the same as h, (WI..... We-1) and Dw,) This completes the proof.
We state the range we require for parameter R: Definition F23. We choose so that
71
$ R
VLL
d
log
35

---

Under review as a conference paper at ICLR 2022
Remark F.24. Recall that the sparsity parameter is directly related to R: O(mR2/3L) hence to ensure the sparsity is small, we shall pick as small as possible. Specifically, if we pick R to be VmLAj then = 0 "00,/'(/mmm))..). n)mL),, as long as m > X 2n" then Re ml)
Next, we pick the value of m: Definition F.25. We choose to be
m 2 1 1,/033/3/3))
Remark F.26. The choice of m here makes sure that, long as we pick R matching its lower bound, then the sparsity s )) this matches the other sparsity parameter k, which is also in the order of Om...)
We induction prove the following two claims recursively.
Definition F.27 (Induction hypothesis 1). Let E [T] be a fixed integer. have
[[a((t) Ws(0)]| < R
holds for any é [T.) Definition F.28 (Induction Hypothesis 2). Let E [T] be fixed integer. We have
[I u|12 - fl wl12.
Suppose the above two claims hold up to t, we prove they continue to hold for time t + 1. The second claim is more delicate, going to prove it first and we define
J1, = S// Je (e((1 s)W, + (WH1)dss
where Jg is defined as Definition F.20.
Lemma F229. Let 9 = (,,)
7 tfi y). We have
ll - 1l12 s life - y- Σ Jew
1
+
(1rr Jea, 1 Je
+
Σ
[IJe, Jeatta 9/)1l2
Proof. Consider the following computation:
||441 8/12
= =Ift- - (ftt1 fell2
|lfe Σ
Jeast1 vec(We.+1--We 2
t=1
L
= =1|s- - y Σ
Je, J.geelll2
t=1
= =Iff- -
Σ
Je. +
Σ
Jesti - Σ Jes.tIE
1=1
Em1
1m1
<ilf- - Σ
Jr Je J..
[1JJ, Je,,44 1
E=1
t=1
L
< < - Σ
Je.Je.
+
Σ
[1(Je Je,,41))
Σ
[1Jr. Je...e1) gē ll2
The second step follows from the definition of Je, and simple calculus.
36

---

Under review as a conference paper at ICLR 2022
Claim F.30 (Ist term in Eq. (8)). We have
II y) Σ Jew eecell
- yll2.
1
6
Proof. We have
1
Il(f - y) Σ JeaSenge
Σ
(ft 8) -Je,ee
L
t=1
4
$ Σ
(ft - y) - Je.ee
t=1
llft د1ا2
since gr,, is an (Eo < ) approximate solution to regression problem
min Je I (fe- -l)l
Claim F.31 (2nd term in Eq. (8)). We have
Σ I1rr Je Jeaa41) De Ellse 2l12.
Proof. For the second term in Eq. (8), for any l e [L], we have
I1(Je, - Je.,,--)) < 1e Je,,,++1l - 1
1 [ 1)ll2
|| Jea,till (J
< -|Je, Jeas-ill (1 1lfe v112.
(10)
We bound these term separately. First.
Je,, Jell,+11l (.(W.)) 1/ Je((1 s)W, s )ds
- 1/ ([e(WW)) - Je((1 s)W,+ + sW
1 ([((W))) - Je(W) + (Wo(W)) Je((1 - s)W, + sW(+1)l ds
< ([)((W)) - Je(Woll++ S. - - -))(11s)), + sW ds
02 / L/m).
(11)
where by Fact F.22, know (),(W.) (W) inequality is by Lemma F.18. For the second term, we have
/(n L/mm) the last
[[1 - s) vec(Wa())) + vec(We(t 1)) - voc(We(0))
(1 s) .IIve((W<(t)) - ve) + s-/| + 1)) vee (W(()))
(1 s) -[[We(t) - We(0)] W
We(t 1
O().
37

---

Under review as a conference paper at ICLR 2022
This means the perturbation of (1 [ with respect to (c(0) is small, for any E [4], hence ()e(WW) Je((1 - ([[[-s)W[++30) sW O(n LL/m).
Furthermore, we have
I (J m (J,,) S 22/2L=
(12)
where the second inequality follows from Tmin (Jr,e) Lemma E.3).
Vemin(JeadEdd 2
Vrr/2 (see
Combining Eq. (10), (11) and (12), we have
I1 Jeset1 Ilo s <O(/nL/m) 7i I1.e v12
1=11
$ 1l
111 - v112,
(13)
where the last step follows from choice of m (Definition F.25). Claim F.32 (3rd term in Eq. (8)). We have
Σ I1 J 22)l2 -llfe v1l2
1
Proof. We can show
[| - Jeaat1) 1e, gillla S 11 Jematill [| lIge, gêll2.
(14)
Moreover, one has
AL 11,,r - gõll2 S Amin (J,n(Je l19r, oi ll2
2
1l Je.
||Je,Je [ 2)
VAr/n L
-1.fe ullz.
(15)
The first step comes from Amin (J = Amin (Gr,)) 2 A2/2 (see Lemma E.3) and the last step comes from St is an eo approximate solution. The fourth step follows from Eq. (15) and the fact that [I(Je,J -1 2/Az. The last step follows from 9r is eo (Eo VAL/n) approximate solution to the regression.
Consequently, have
[[]e,, Je...ta) ge )l2 S [|Je, Je,,,41l -1 11ge, gEll2
1/2 //////LL//// LL/m) O(n (22m/2/L//m) VL/m). Wind llfe vll2 s 1/2
2
= Ō vn
2
ll.f v112.
(16)
Vx
The second step follows from Eq. (11) and (15) and the fact that |e,ll < D(VnL/m) (see Lemma F.18).
Finally, we have
Σ [J 1e 11)122 < ~0( V(L 2
11ʃe 2|12
m VdL
1
<
- [ y112.
6
38

---

Under review as a conference paper at ICLR 2022
The last step follows from choice of m (Definition F.25).
Lemma F.33 (Putting it all together). We have
1 1| u112 511 y112.
(17)
Proof. Combining Eq. (9), (13), and (??), we have proved the second claim, i.e.,
Ilfi+1 Ul1 2 fl vl12.
F.6 WEIGHTS ARE MOVING TOO FAR
Lemma E34. Let R be chosen as in Definition F:23, then the following holds:
[[W 1) Wi(O0)ll S R, Vl E [L].
Proof. It remains to show that Wi does not move far away from Wo. First, we have 1 ll2 - 119 + llge,t 12112
L (,, (fc 2
v)ll2 + 11,6, gē ll2
< I (1,e,,Je,,)) -1 11(st - -))le + 11ge, - ge ll
2
2 Il.s 2/12
LXr -1.fc- Tv/nd
2 TXL
[1 - -yl2
(18)
where the third step follows from Eq. (15) and the last step follows from the obvious fact that 1/vnd s 1/Xt.
Then
Wi(k + 1) -)<+)(())] = |
< Je 11ge,xl12
- σ nL/m)
1| yll
LλL
< v0lVvo/Tm (m/Tm)
Il - vll
2k
7 < 5
VLm
2*\1
The third step uses the fact that || < uses the fact that both [|olll and ||v|| are in the order of O(1). Consequently, we have
|| Olv
Lm
by Lemma F.18, and the last step
[W&( + W2(0)l S
We(k + 11) Wc(k)ll
k=0
1
VLm
X 2
k=0
n
A By the choice of R (Definition F.23), we know this is upper bounded by R. This concludes proof.
Lm
39

---

Under review as a conference paper at ICLR 2022
G PROOF OF LEMMA F.7
In this section, we prove a technical lemma (Lemma F.7) involving truncated gaussian distribution, which correlates the shifted ReLU activation we use.
Definition G.1 (Truncated Gaussian distribution). Suppose X (O o²). Letb E R. Then, a random variable follows from a truncated Gaussian distribution (o(0, ifY = x|x 2 b. The probability density function for No,0,, o²) is as follows:
-e-V1020 y (0),
f(y) = o(1- (bbaa)
V2T
where (-)) is the standard Gaussian distribution's CDE
Fact G.2 (Properties of truncated Gaussian distribution). For b IR suppose X N2), Le
3 = b/a. Then, we have
ob(3) [E[]] (-)(3)) where d(x)
1 _ 1²12 2
Var[X]= = o2 (6(0))((11110)) (069)((100))1))(3))))))))
X/o ~ (0,
When X is C(b 1)-suggaussian, where C >> 0 is an absolute constant.
Lemma G.3 (Concentration inequality for b-truncated chi-square distribution). For b E R, 0. let X Xb.ni that is, X Ex= Y where Yi,... Yn (o, 1) are independent b-truncated Gaussian random variables. Then, there exist two constants C1, C2 such that for any t > 0,
by(b)
Pr Pr[
x m01 xx1+
2 nt I exp (-Cint exp eexp(-Cynt//8f
1 2bb)
In particular, we have
Pr[|X - -11 b(b 1))1 2 S exp (-C11/(mb))) + exp (-C21/b2).
Proof. Since we know that Y ~ (o 1) is C(b + 1)-subgaussian, it implies that Y? is a sub- exponential random variable with parameters (4 (c+12) (6++1))) Hence, by the standard concentration of exponennttt random variables, we have
exp
nt² (611 ml 2-40*
nf S (C2 1)2
Pr
Σ Y? n B|Y?
nt 5
2 exx
otherwise
$ 2 exp
(-Cint² /ontt/b")
exx (-C2nt/b²).
Fact G.4. Let h E Rp be fixed vectors and + 0, let > 0 be a fixed scalar. W E Rmxp be random matrix with i.i.d. entries Waj (O, 2) and vector R. defined U, (Wh))) I(Wh); 2 b),, Then
1el follows i.i.d. from the following distribution: with probability 1 e - (m/mm// |val 0, and with probubility (MIl)(( |va follows from truncuted Gaussian distribution N0, [/ 12)
mlleel M|M
is in distribution identical to X3 (V-ttrunctted chi-square distribution of order w)
where follows from binomial distribution B(m, 6 and U /mm/2, b. T/A12
Proof. We assume each vector W; is generated by first generating gaussian vector g (0 2 and then setting W, = =g where the sign is chosen with half-half probability. Now, KWW, h/ (g, bi) only depends g, and is in distribution identical to (0, 2 1b13). Next, after the sign is
40

---

Under review as a conference paper at ICLR 2022
determined, the indicator 1(W,h); 2 bj is 1 with probability c-h"mm(44)AA()) and 0 with probability e m/(AI]]] "), Therefore, |0a]] satisfies the aforementioned distribution. As for ||=||2 letting w E
{0, 10.1....,77 m) be the variable indicates how many indicators are 1, then w. ~ B e - -m/m((11)))))) b.
mlln 2 X aw where D'
vm/2
and
Fact G.5 (Gaussian tail bound). For any b > 0, we have
12
C(b+1) S $(b) Se-
where C' an absolute constant.
We prove a truncated Gaussian version of Lemma 7.1 of Allen- et al. (2019a).
Lemma G.6. Let b > 0 be fixed scalar. Let the activation function be defined as Φ(z)
v x > /2/mbla, where es (2(1 {bb) bob(b)) 11/2, Let E (0,1), then over the
randomness of (0), with probability least (nL)) exp( ( /2) /T2)),
have
||s,ell2 E [1 -, e), vi E [m], E [4].
Proof. We only prove for fixed E [m] and E at the end. Below, we drop the subseript for notational convenience, and write hi and I; as he and respectively.
} because we can apply union bound
According to Fact G.4, fixing any he-1 * and letting We the only source of randomness, we have
m 2
X6/IIA with w ~ B(m, 1 - -(bb)),
2
where V' = b/llhe-lll2
We first consider the case. Then, we have |lhe-1l2 ll| 1, and b' b. Let Po = 1- -bb). By Chernoff bound, for any 8 E (0, 1), we have
Pr[w E (1 + 8)m 2 - exppeexp(200 P.m)). In the following proof, we condition on this event. By Fact G.5,
(1 - 8) e //
E (1 + 6) = e
C(b+1) (1 + (xxxp(_b3/22))
By Lemma G.3, we have
(1+ [1
bob(b)
)
Pr
2 h 12
>t] I exp (-S2(t/(ub*))) + exp (-S(t/b²))
Pb
Note that
bo(b)
bo(b)
wi
1
E (1 + 0) + (1 + 8)mP,
(1 + OP + bo(b)) + m.
Pb
P.
Let o -1 = 2(Ps + bb)()) be the normalization constant. Then, we have
Pr[lea||ha113 (1 + 5) > 21c/mm] < exp (-S((P/(wb))) + exp (-S(t//b3)
We want 2tcs/m m 06)). i.c.. 0(8c, m). (1)c)lm). Then, we have wt m0) b². Hence, by
Lemma G.3, we actually have
Pr|[cs|[hn 1122 - (1 5)1 > 0()) I exp (-226m/(ccbb)))
By taking 8 =eLL, we get that that
| |1 c/L, 1/L,1+c/I)
41

---

Under review as a conference paper at ICLR 2022
with probability at least
- exp(-2(e Pxp(-O0(ePammL))) - exp
2 exp
where the last step follows from
Po+bolb)
(1)).
We can inductively prove the l > case. Since the blowup of the norm of h1 is 1 + e/ T,, the con-
centration bound is roughly the same for he for & 2 Thus, by carefully choosing the parameters, can achieve | 12 [(1 e/) (1
+/ with high probability.
this end, by union bound over all the layers E [4] and all the input data i [nd], we get that
||a,elle [] c, + c]
with probability at least
- (00
which completes the proof of the lemma.
42


Reviews:
[{"review_id": "rH9iV3TobUe", "reviewer": "Reviewer_aph9", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1: [motivation and claims about the literature] \n\n> In the first paragraph of introduction, the authors claim that “most theoretical works focus on training on-hidden layer networks”. This is not the case. \n>>a), quite a few of the cited papers are theoretical analyses on *deep* neural networks. For example, Jacot et al (2018), Du et al (2019a), and Allen-zhu et al (2019 a, b), which are listed just above this claim in the paper. \n\n>>b), the paper misses several theoretical works on *deep* neural networks. For example, see the references: Lee et al (2019), Liu et al (2020 a,b).\nThere should be more works on deep neural networks, and I suggest the authors have a more complete query of the literature before making the claims.\n\n> The authors claim in page 1 that “ convergence rate (of first-order methods) is typically slow in non-convex settings …”. I cannot agree with this. For example, in both Du et al (2019a) and Liu et al (2020a), exponential convergence rates on over-parameterized neural networks are obtained for gradient descent methods. Typically, the setting is non-convex, see Liu et al (2020a). For comparison, the proposed method in this paper does not have a faster rate than those two.\n\n> I am not clear how the author got this claim: “it is increasingly evident that first-order methods are becoming a real bottleneck for many practical applications”. The authors should provide evidence. \n\n> The paper claims a running time of O(m^2) for first-order methods, like gradient descent. However, with parallelism this running time can be greatly reduced.\n\n2: [Paper structure] \nOne uncomfortable thing about this paper is that all the formal statements (e.g., theorems, assumptions) are not shown in the main text, but are deferred to the appendix. Even the formal version of the proposed algorithm is also not displayed in the main pages. In the first 9 pages, there are only informal statements and intuition. I strongly suggest the authors put the main theorems and algorithm, as well as key analysis, in the main text, and just put tedious proof or unimportant details into the appendix. \n\n3: [practicality]\nIn practical cases, the network width m is never too large, compared to the data size n. In these cases, the proposed algorithm seems to have no advantage compared to gradient descent methods, in terms of the computation cost. Moreover, the complicated design of this algorithm may make it not favorable to practitioners. \n\n\nReferences:\n\nLee, et al. Wide neural networks of any depth evolve as linear models under gradient descent. NeurIPS, 2019.\n\nLiu, et al (a). Toward a theory of optimization for over-parameterized systems of non-linear equations: the lessons of deep learning, Arxiv: 2003.00307, 2020.\n\nLiu, et al (b). On the linearity of large non-linear models: when and why the tangent kernel is constant. NeurIPS, 2020.", "summary_of_the_paper": "This paper aims to propose a second-order algorithm for training of neural networks with very large widths. The main claim is that the computation cost of each iteration of the proposed methods is sub-quadratic in terms of the network width.  ", "summary_of_the_review": "The main text of the paper is not self-containing and not justifiable (suggest to reorganize the paper). \nSeveral claims about the literature are not accurate.\nNo advantage compared to gradient descent methods for practical scenarios.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "msMQmPPPWBg", "reviewer": "Reviewer_xVuq", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper is clearly written. The detailed proofs are left in the appendix and the authors explained the high-level proof techniques in the main text. The pros and cons are as follows.\n\nPros. \n1. As far as I know, this paper is the first work that studied the convergence of second-order methods in multi-layer neural networks. Some of the analysis techniques can be useful in future studies. \n2. This paper reduced the per-iteration cost of second-order methods by combining the Gram-Gauss-Newton method, tensor-based sketching, and preconditioning. It's a crucial challenge in second-order methods to reduce the per-iteration cost. The ideas in this paper can inspire new algorithms that enjoy lower per-iteration costs.\n\nCons.\n1. The convergence proof of second-order methods seems very similar to the NTK proof for first-order methods. I understand there are some differences in proof details, but I still wonder if there is any conceptual difference in the proof. \n2. The linear convergence has been established for first-order methods in similar settings. It seems to me that this paper does not show any benefits of second-order methods over first-order methods in this over-parameterized network setting.\n3. The convergence analysis is restricted to the NTK regime. However, in practice, neural networks are not trained in the NTK regime and also achieved better performance in test data than the neural tangent kernel.\n4. The reduction of the per-iteration cost relies on the assumption that $m\\gg n$, which is hardly true in practice. Therefore, I am concerned about the practicability of the proposed algorithm. It can be helpful if there are empirical experiments that verifies the performance of the algorithms in standard architecture and realistic data set. ", "summary_of_the_paper": "1. This paper proved that the second-order method can minimize the training loss in linea rate on multi-layer over-parameterized neural networks. This analysis relies on the connection between neural tangent kernel and over-parameterized neural networks.\n2. This paper also reduced the per-iteration cost of second-order methods to $\\widetilde{o}(m^2)$ where $m$ is the hidden layer width, by combining Gram-Gauss-Newton method, tensor-based sketching, and efficient data structures that maintain low-rank representations. \n3. This paper also designed an algorithm to efficiently solve a regression problem in which the matrix has its rows being tensor products of two vectors. This algorithm may be of general interest. ", "summary_of_the_review": "In general, I appreciate the technical depth of this work as a theory paper. And I think the techniques for reducing per-iteration cost can be useful in future studies. My main concern is the novelty of the convergence analysis and the practicability of the proposed algorithm. Above all, I think this paper is marginally above the acceptance threshold. \n\n-----------------\n\nThanks for the response. After reading the response and other reviews, I decided to keep my original evaluation. \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "Kjl9-GO-ZE7", "reviewer": "Reviewer_djYB", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper provides an algorithm for training neural networks in the over-parametrized regime. The authors go to great length to show that under good implementation which exploits various structural properties (low rank in particular), the runtime complexity of the algorithm should make it practical. It would be nice to provide empirical validation of the method. The results are not surprising: the gains come from careful attention to these structural properties, and extend Cai et al 2019. I consider this a strength for this paper: providing a precise practical guideline on implementation, while providing theoretical convergence guarantees, is a welcome contribution. The paper's diffusion would of course profit a lot from empirical validation + a code release. Do the authors have such an implementation handy?\nI believe this paper is a good contribution to the field -- on the other hand, it would be made much stronger with such an implementation. I'm therefore on the fence on acceptance at this time: I want to encourage authors to provide this implementation.\n\nComments:\n- Can you give the meaning of $m$, $n$, $L$, $\\epsilon$ before they are used in the introduction?\n- What definition of \"over-parametrized\" NN are you using? It seems that the rank of the jacobian is at most $n$ where $n$ is the size of the mini-batch, in practice? Is this correct?\n- The paper (and in particular Theorem 1.1) uses the L2 error -- is this result extendable for any (convex? smooth?) loss function? Can you explain in the main text why the L2 error and Relu activations are necessary, and where the analysis/algorithm break in the more general case? \n- The limits of the method are hardly discussed in the main text.\n- It might be easier for the reader to follow if the contributions were given before the informal theorem statements.\n- Other practical second order methods have been proposed for neural network optimization. In particular, how does this method compare to AdaHessian? https://arxiv.org/pdf/2006.00719.pdf AdaHessian uses a diagonal approximation to the Hessian. How is it related to K-fac ? https://arxiv.org/pdf/1503.05671.pdf                                            \n- This sentence \"A smarter way is to instead of using the n × n Gram matrix for neural tangent kernel, however this would still incur a cost of O(nm2) per iteration, since each gradient is an m×m matrix and the Jacobian consists of n such gradients.\" isn't clear. What are you proposing here?\n- There are many typos and sentences which are difficult to follow. e.g. Step 2 paragraph (page 6). \"ethnic implications\" rather than \"ethics implications\". It would benefit from an attentive proof-reading.\n", "summary_of_the_paper": "This paper proposes a second-order algorithm for training neural networks, in the L2 regression setting. It provides a theoretical analysis of its complexity in the over-parametrized regime. It does not provide empirical validation of the method, or an implementation.", "summary_of_the_review": "I believe this paper provides good theoretical results on a new second order method for optimizing neural networks in the over-parametrized setting. Content-wise, I'd strongly appreciate empirical validation / implementation. On the formal side, the paper is a bit hard to follow, and contains many typos. It would benefit from careful proof-reading, and correction. I do not recommend acceptation in the current state of the paper, but could be convinced otherwise.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "7pRxNmsleP6", "reviewer": "Reviewer_91V6", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "My first concern is regarding the intuition of this paper. The author claims in the introduction that “it is increasingly evident that first-order methods are becoming a real bottleneck for many practical applications” because the convergence rate is typically slow, i.e., requiring $\\mathrm{poly}(n, L, log(1/\\epsilon))$ to converge, while second-order method enjoys a much faster convergence rate ($\\log(1/\\epsilon)$). This is true in terms of the iteration complexity, but it is also true that the per-iteration complexity of the first-order method is much lower (could be lower by $\\mathrm{poly}(m)$, while $m$ itself is a high-degree polynomial of $n$), then it is hard to directly claim that first-order method is not as efficient as second-order method.\n\nIn order to demonstrate the efficiency of the second-order method, the authors may need to characterize the total computational complexity of both first-order and second-order algorithms, rather than the \\textit{cost per iteration}. This could be a more fair comparison to justify the better performance of second-order methods.\n\nIt is known that when the data distribution is good, the neural network does not need to be that large to achieve perfect training accuracy. For example, [1,2] show that when the data are well separated by the neural tangent kernel, polylogarithmic widths are sufficient. My suggestion is that the analysis should not be restricted in the setting where $m$ is super large, but consider a wide range of different $m$, and identify under which conditions on $m$, the proposed algorithm could outperform previous ones and vice versa.\n\n1. Ji and Telgarsky, Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow ReLU networks. ICLR 2020.\n2. Chen et. al., How Much Over-parameterization Is Sufficient to Learn Deep ReLU Networks?. ICLR 2021.\n\n\nMinor comments:\n\nCould you elaborate on why the update rule in Step 1 can be solved via minimizing the second equation after Step 1?\n\nCan the same algorithm be applied to classification problems, like [3, 4]? \n\n3. Zou et al., Gradient descent optimizes over-parameterized deep ReLU networks. Machine Learning  2019\n4. Cao and Gu, Generalization Bounds of Stochastic Gradient Descent for Wide and Deep Neural Networks, NeurIPS 2019\n\n", "summary_of_the_paper": "This paper studies the training algorithms for multi-layer over-parameterized neural networks. In particular, this paper starts from gauss-newton-methods and incorporates tensor-based sketching techniques and preconditioning to improve the per-iteration computational complexity. As a result, the proposed algorithm can find the global solution within the time that is subquadratic in the network width.", "summary_of_the_review": "This paper is theoretically sound but could be further improved based on my comments.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies the performance of second-order algorithms on training multi-layers over-parameterized neural networks. The authors propose an algorithm based on the Gram-Gauss-Newton method, tensor-based sketching techniques, and preconditioning to train such a network, whose runtime is subquadratic in the width of the neural network. While some reviewers provide some weak support, none of them are in strong support, even after the author's response. I think one of the reasons is the lack of empirical experiments. Since the main claim of this paper is an efficient second-order algorithm, some experiments are necessary to back up this claim. Unfortunately, the authors did not try to add such an experiment during the rebuttal. I would suggest the authors add such experiments in the revision.
Decision: Reject