Year: 2021

Abstract:
In the recent years of development of theoretical machine learning, overparametrization 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 GramGauss-Newton method, tensor-based sketching techniques and preconditioning.

heading: 1 INTRODUCTION
text: Deep neural networks have been playing a 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 et 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 to analyze the dynamics of neural networks and obtain provable guarantees, using over-parametrization has been a growing trend.
In terms of understanding the convergence behavior of over-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, to a large degree, by its 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-based method such as minibatch. 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/ )) iterations, where n is the number of training samples, L is the number of layers and is the precision of training), e.g., deep neural network with ReLU activation, as shown in Allen-Zhu et al. (2019a)), which is often the case of a deep over-parameterized neural network.
Second-order method (which employs the information of the Hessian matrix) on the other hand enjoys a much faster convergence rate (only log(1/ ) iterations (Zhang et al., 2019), but not poly(n) log(1/ ) 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 second-order method over first-order method is it does not require the tuning of learning rate. The expense of using second-order method is its prohibitive cost per iteration, as it is imperative to invert a dynamically-changing Hessian matrix or equivalently, solving a regression task involving Hessian. Given any weight matrix of size m × m (m is the width of network), its Hessian matrix has size m2 ×m2, which makes any naive implementation of second-order algorithm takes at least O(m4) 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 as 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 et al. (2019b); Jiang et al. (2020b;a); Song & Yu (2021)).
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 overparametrized networks with smooth activation functions. Specifically, they achieve a running time ofO(mn2) per iteration. Their methods are essentially variants of Gauss-Newton method combined with using neural tangent kernels (Jacot et al. (2018)) to prove the convergence. By using the idea of randomized linear algebra (Clarkson & Woodruff (2013); Woodruff (2014)) and a clever sketching matrix as a preconditioner, Brand et al. (2021) further improves the running time to Õ(mn) 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 as gradient descent and stochastic gradient descent won’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 et al. (2021)) would incur a running time of Õ(m2nL) 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 in training deep 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 ≥ n8L12, (Zou & Gu, 2019)). Our work can be decomposed into two parts: algorithmically and analytically. From an algorithmic perspective, our method builds upon a variant of Gauss-Newton method (Björck (1996)) called Gram-Gauss-Newton method (Cai et al. (2019); Brand et al. (2021)). In order to achieve a feasible running time, we exploit two features of the gradient, which is the key ingredient to form the Jacobian 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 to their proof which is a straightforward analysis of the gradient, we make use of the multi-layer neural tangent kernels (Du et al. (2019a)) and establish a connection between a Gram matrix we compute at each iteration and the NTK matrix.
Our Contributions. We summarize our technical contributions below.
• We develop an analytical framework for the convergence behavior of second-order method on training multi-layer over-parametrized neural network. To facilitate the analysis, we exploit the equivalence between neural tangent kernels and our over-parametrized network.
• We design a second-order algorithm to train such networks, and achieve a cost per iteration of õ(m2). Our algorithm makes use of Gram-Gauss-Newton method, tensor-based sketching techniques, and data structures that maintains a low rank representation efficiently.
• By combining fast tensor algebra techniques and sketching-based preconditioning, we devise an algorithm to efficiently solve a regression problem where the involved matrix has its rows being tensor product of vectors.

heading: 1.1 OUR RESULT
text: Our main result can be summarized in the following three theorems, with one analyzing the convergence behavior of a general Gram-based optimization framework, one designing an efficient algorithm to realize this second-order optimization scheme, and the other is a novel algorithm to solve tensor-based regression in 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 L to denote the number of layers of the network. We use ft ∈ Rn 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 ≥ poly(n,L), then there exists an algorithm (Algorithm 1) such that, over
the randomness of initialization of the network and the algorithm, with probability at least 1 − e−Ω(log 2m), we have
‖ft+1 − y‖2 ≤ 1
2 ‖ft − y‖2,
where ft ∈ Rn is the the prediction produced by neural network at time t.
The above theorem establishes the linear convergence behavior of our second-order method, which is a standard convergence result for second-order method, as well as the same behavior as in onehidden layer over-parametrized networks (Brand et al. (2021)). However, compared to one-hidden layer case, our analysis is much more sophisticated since we 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 our second-order algorithm.
Theorem 1.2 (Runtime, informal version of Theorem B.1). There exists a randomized algorithm (Algorithm 1) that trains a multi-layer neural network of widthm with the cost per training iteration being
O(m2−Ω(1)).
We improve the overall training time of multi-layer over-parametrized networks with second-order method from Tinit +T ·O(m2) to Tinit +T · o(m2), where Tinit is the initialization time of training, typically takes O(m2). As we have argued before, multi-layer over-parametrized networks require m to be in the order of n8, hence improving the cost per iteration from quadratic to subquadratic is an important gain in speeding up training. Its advantage is even more evident when one seeks a high precision solution, and hence the number of iterations T is large.
We highlight that it is non-trivial to obtain a subquadratic running time per iteration: If not handled properly, computing the matrix-vector product with weight matrices will take O(m2) time! This means that even for first-order methods such as gradient descent, it is not clear how to achieve a subquadratic running time, since one has to multiply the weight matrix with a vector in both forward evaluation and backpropagation. In our case, we have also a Jacobian matrix of size n × m2, so forming it naively will cost O(nm2) time, which is prohibitively large. Finally, note that the update matrix is also an m×m matrix. In order to circumvent these problems, we exploit the fact that the gradient is of low rank (rank n), hence one can compute a rank-n factorization and use it to support fast matrix-vector product. We also observe that each row of the Jacobian matrix can be formulated as a 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:
Theorem 1.3 (Fast Tensor Regression, informal version of Theorem D.14). Given two n×m matrices U and V with m n and a target vector c ∈ Rn. Let J = [vec(u1v>1 )>, . . . , vec(unv>n )>] ∈ Rn×m2 where ui is the i-th row of matrix U ∀i ∈ [n]. There is a randomized algorithm that takes Õ(nm+ n2(log(κ/ ) + log(m/δ)) + nω) time and outputs a vector x̂ ∈ Rn such that
‖JJ>x̂− c‖2 ≤ ‖c‖2
holds with probability at least 1− δ, and κ is the condition number of J .
From a high level, the algorithm proceeds as follows: given matrices U and V , it forms an approximation J̃ ∈ Rn×n log(m/δ), where each row is generated by applying fast tensor sketching technique to ui and vi (Ahle et al. (2020)). Then, it uses another sketching matrix for J̃ to obtain a good preconditioner R for J̃ . Subsequently, it runs a gradient descent to solve the regression.
To understand this runtime better, we note that nm 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 m2 tensors to length O(n log(m/ δ)). All subsequent operations are performed on these much smaller vectors. Finally, computing the preconditioner takes Õ(nω) time, and running the gradient descent takes Õ(n2 log(κ/ )) time.

heading: 1.2 RELATED WORK
text: Second-order Method in Optimization. Though not as prevalent as first-order method in deep learning, second-order methods are one of the most popular in convex setting, such as linear programming (Vaidya (1989); Daitch & Spielman (2008); Lee et al. (2015); Cohen et al. (2019)), empirical risk minimization (Lee et al. (2019b)), cutting plane method (Jiang et al. (2020b)) and semidefinite programming (Jiang et al. (2020a)). Due to the prohibitive high cost of implementing one step of second-order 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 et al. (2017); Pilanci & Wainwright (2017); Agarwal et al. (2017); Bernacchia et al. (2018); Cai et al. (2019); Zhang et al. (2019); Brand et al. (2021); Yao et al. (2021)) that try to improve the practicality of second-order method and adapt them to train deep neural networks. As shown in Cai et al. (2019), it is possible to exploit the equivalence between over-parametrized networks and neural tangent kernel to optimize an n×nmatrix instead of anm2×m2 matrix, which is an important breakthrough in gaining speedup for second-order method. Sketching and sampling-based 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), in the sense that our optimization also works on an n × n Gram matrix. Our algorithm also makes use of sketching and sampling, as in Pilanci & Wainwright (2017); Brand et al. (2021).
Over-parameterized Neural Networks. In recent deep learning literature, understanding the geometry and convergence behavior of various optimization algorithms on over-parameterized neural networks has received a 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 et 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 ≥ Ω(n4)), then the optimization dynamic on a neural network is equivalent to that on a NTK.
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-Solve” (Clarkson & Woodruff (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 as a preconditioner (Woodruff (2014); Brand et al. (2021)) to obtain a high precision solution.
Roadmap. In Section 2, we give a preliminary view of the training setup we consider in this paper. In Section 2.1, we introduce the notations that will be used throughout this paper. In Section 2.2, we consider the training setup. In Section 3, we 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 to achieve subquadratic cost per iteration. In Section 4, we summarize the results in this paper and point out some future directions.

heading: 2 PRELIMINARIES

heading: 2.1 NOTATIONS
text: For any positive integer n, we use [n] to denote the set {1, 2, · · · , n}. We use E[·] to denote expectation and Pr[·] for probability. We use ‖x‖2 to denote the `2 norm of a vector x. We use ‖A‖ to denote the spectral norm of matrix A. We use ‖A‖F 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 ×m. For matrix A or vector x, we use ‖A‖0, ‖x‖0 to denote the number of nonzero entries of A and x respectively. Note that ‖ · ‖0 is a semi-norm since it satisfies triangle inequality. Given a real square matrix A, we use λmax(A) and λmin(A) to denote its largest and smallest eigenvalues respectively. Given a real matrix A, we use σmax(A) and σmin(A) to denote its largest and smallest singular
values respectively. We useN (µ, σ2) to denote the Gaussian distribution with mean µ and variance σ2. We use Õ(f(n)) to 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 to two matrices, this means 〈A,B〉 = tr[A>B] where tr[A] denote the trace of matrix A.

heading: 2.2 PROBLEM SETUP
text: Let X ∈ Rm0×n denote the data matrix with n data points and m0 features. Without loss of generality, we assume ‖xi‖2 = 1,∀i ∈ [n]. Consider an L layer neural network with one vector a ∈ RmL and L matrices WL ∈ RmL×mL−1 , · · · , W2 ∈ Rm2×m1 and W1 ∈ Rm1×m0 . We will use W`(t) to denote the weight matrix at layer ` at time t, and ∇W`(t) to denote its gradient. We also use W (t) = {W1(t), . . . ,WL(t)} to denote the collection of weight matrices at time t. Architecture. We first describe our network architecture. The network consists of L hidden layers, each represented by a weight matrix W` ∈ Rm`×m`−1 for any ` ∈ [L]. The output layer consists of a vector a ∈ RmL . We define the neural network prediction function f : Rm0 → R as follows:
f(W,x) = a>φ(WL(φ(· · ·φ(W1x)))), where φ : R → R is the shifted ReLU activation function (σb(x) = max{x − b, 0}) applied coordinate-wise to a vector.
We measure the loss via squared-loss function:
L(W ) = 1 2 n∑ i=1 (yi − f(W,xi))2.
This is also the objective function for our training.
We define the prediction function ft : Rm0×n → Rn as
ft(X) = [f(W (t), x1) f(W (t), x2) · · · f(W (t), xn)] > .
Initialization. Our neural networks are initialized as follows:
• For each ` ∈ [L], the initial weight matrix W`(0) ∈ Rm`×m`−1 is initialized such that each entry is sampled from N (0, 2m` ). • Each entry of a is an i.i.d. sample from {− 1√ mL ,+ 1√ mL } uniformly at random.
Gradient. In order to write gradient in an elegant way, we define some artificial variables: gi,1 = W1xi, hi,1 = φ(W1xi), ∀i ∈ [n] gi,` = W`hi,`−1, hi,` = φ(W`hi,`−1), ∀i ∈ [n],∀` ∈ [L]\{1} (1) Di,1 = diag ( φ′(W1xi) ) , ∀i ∈ [n]
Di,` = diag ( φ′(W`hi,`−1) ) , ∀i ∈ [n],∀` ∈ [L]\{1}
Using the definitions of f and h, we have
f(W,xi) = a >hi,L, ∈ R, ∀i ∈ [n]
We can compute the gradient of L in terms of W` ∈ Rm`×m`−1 , for all ` ≥ 2
∂L(W ) ∂W` = n∑ i=1 (f(W,xi)− yi) Di,`︸︷︷︸ m`×m`  L∏ k=`+1 W>k︸︷︷︸ mk−1×mk Di,k︸︷︷︸ mk×mk  a︸︷︷︸ mL×1 h>i,`−1︸ ︷︷ ︸ 1×m`−1
(2)
Note that the gradient for W1 ∈ Rm1×m0 (recall that m0 = d) is slightly different and can not be written by general form. By the chain rule, we can compute the gradient with respect to W1,
∂L(W ) ∂W1 = n∑ i=1 (f(W,xi)− yi) Di,1︸︷︷︸ m1×m1  L∏ k=2 W>k︸︷︷︸ mk−1×mk Di,k︸︷︷︸ mk×mk  a︸︷︷︸ mL×1 x>i︸︷︷︸ 1×m0
(3)
It is worth noting that the gradient matrix is of rank n, since it’s a sum of n rank-1 matrices.
Jacobian. For each layer ` ∈ [L] and time t ∈ [T ], we define the Jacobian matrix J`,t ∈ Rn×m`m`−1 via the following formulation:
J`,t := [ vec(∂f(W (t),x1)∂W`(t) ) vec( ∂f(W (t),x2) ∂W`(t) ) · · · vec(∂f(W (t),xn)∂W`(t) ) ]> .
The Gram matrix at layer ` and time t is then defined as G`,t = J`,tJ>`,t ∈ Rn×n whose (i, j)-th entry is 〈∂f(W (t),xi)∂W` , ∂f(W (t),xj) ∂W` 〉.

heading: 3 TECHNIQUE OVERVIEW
text: In this section, we give an overview of the techniques employed in this paper. In Section 3.1, we showcase our algorithm and explain various techniques being used to obtain a subquadratic cost per iteration. In Section 3.2, we give an overview of the proof to show the convergence of our algorithm. To give a simpler and cleaner presentation, we assume m` = m for all ` ∈ [L].

heading: 3.1 SUBQUADRATIC TIME
text: In this section, we 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 a Hessian matrix of size m2 × m2, we use an n × n Gram matrix derived from the 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.
We start by demonstrating our algorithm:
Algorithm 1 Informal version of our algorithm. 1: procedure OURALGORITHM(f, {xi, yi}i∈[n]) . Theorem 1.1,1.2 2: /*Initialization*/ 3: Initialize W`(0), ∀` ∈ [L] 4: Store W1(0)xi in memory, ∀i ∈ [n] . Takes O(nm2) time 5: Store hi,`(0)← φ(W`(0)hi,`−1(0)), ∀` ∈ [L], i ∈ [n] in memory . Takes O(nLm2) time 6: for t = 0→ T do 7: /*Forward computation*/ 8: for ` = 1→ L do 9: vi,` ← hi,`−1, ∀i ∈ [n] 10: hi,` ← φ((W`(0) + ∆W`)(hi,`−1)), ∀i ∈ [n] . Takes O(n2m) + o(nm2) time 11: . hi,` is sparse 12: Di,` ← diag(φ′((W`(0) + ∆W`)hi,`−1)), ∀i ∈ [n] . Takes O(nm) time 13: . Di,` is sparse 14: end for 15: ft ← [a>h1,L, . . . , a>hn,L]> . Takes O(nm) time 16: /*Backward computation*/ 17: for ` = L→ 1 do 18: ui,` ← a>Di,LWL(t) . . . Di,`+1W`+1(t)Di,` . Takes o(nLm2) time 19: Form J̃`,t that approximates J`,t using {ui,`}ni=1, {vi,`}ni=1 20: . Takes Õ(mn) time, J̃`,t ∈ Rn×s where s = Õ(n) 21: Compute g` that approximates (J̃`,tJ̃>`,t)
−1c . Takes Õ(nm) time 22: Form J>`,tg` via low rank factorization ∑n i=1 g`,iui,`v > i,`
23: Implicitly update ∆W` ← ∆W` + ∑n i=1 g`,iui,`v > i,` and store it in memory 24: end for 25: end for 26: end procedure
Step 1: Invert Gram by solving regression. Recall the update rule of generic algorithm is given by
W`(t+ 1)←W`(t)− J>`,t(J`,tJ>`,t)c,
where c is ft− y after proper scaling. Naively forming the Gram matrix J`,tJ>`,t will take O(n2m2) time and inverting it 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 g`,t ∈ Rn such that
‖J`,tJ>`,tg`,t − c‖22 is minimized. This enables us 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 a good preconditioner via sketching then apply iterative method to solve it (Brand et al. (2021)). Roughly speaking, we first use a random matrix S ∈ Rs×m2 that has the subspace embedding property (Sarlos (2006)) to reduce the number of rows of J>, then we run a QR decomposition on matrix SJ>. This gives us a matrixR such that SJ>R has orthonormal columns. We then use gradient descent to optimize the objective ‖JJ>Rzt − y‖22. Since S is a subspace embedding for J>, we can make sure that the condition number of the matrix J>R is small (O(1)), hence the gradient descent converges after log(κ/ ) iterations, where κ is the condition number of J . However, in order to implement the gradient descent, we still need to multiply an m2 × n matrix with a length n vector, in the worst case this will incur a time of O(nm2). 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 ∂f(W,xi)∂W` ∈ R
m×m, it can be written as (for simplicity, we use hi,0 to denote xi):
∂f(W,xi) ∂W` = hi,`−1︸ ︷︷ ︸
vi∈Rm×1 a>Di,LWL . . . Di,`+1Wi,`+1Di,`︸ ︷︷ ︸ u>i ∈R1×m .
This means the gradient is essentially an outer product of two vectors, and hence has rank one. This has several interesting consequences: for over-parametrized networks, the gradient is merely of rank n instead of m. When using first-order method such as gradient descent or stochastic gradient descent, the weight is updated via a 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 our attention to design a fast approximation algorithm to the Jacobian matrix. Recall that Jacobian matrix J`,t ∈ Rn×m 2
is an n × m2 matrix, therefore writing down the matrix will take O(nm2) time. However, it is worth noticing that each row of J`,t is vec(uiv>i )
>, ∀i ∈ [n], or equivalently, ui ◦ vi where ◦ denotes the tensor product between two vectors. Suppose we are given the collection of {u1, . . . , un} ∈ (Rm)n and {v1, . . . , vn} ∈ (Rm)n, then we can compute the tensor product ui ◦ vi via tensor-based 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 O(m2) time. Since it suffices to preserve the length of all vectors in the column space of J>`,t, the target dimension s can be chosen as O( −2n · poly(log(m/ δ))). Use J̃`,t ∈ Rn×s to denote this approximation of J`,t, 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 5: Efficient update via low rank factorization. The low rank structure of the gradient can further be utilized to represent the change on weight matrices ∆W in a way such that any matrix-vector product involving ∆W can be performed fast. Let g` ∈ Rn denote the solution to the regression problem posed in Step 1. Note that by the update rule of our method, we shall use J>`,tg` ∈ Rm×m to update the weight matrix, but writing down the matrix will already take O(m2) time. Therefore,
it is instructive to find a succinct representation for the update. The key observation is that each column of J>`,t is a tensor product of two vectors: ui ◦ vi or equivalently, uiv>i . The update matrix can be rewritten as ∑n i=1 g`,iuiv > i , and we can use this representation for the update on the weight, instead of adding it directly. Let
U` := [ | | . . . | g`,1u1 g`,2u2 . . . g`,nun | | . . . | ] ∈ Rm×n, V` := [ | | . . . | v1 v2 . . . vn | | . . . | ] ∈ Rm×n,
then the update can be represented as U`V >` . Consider multiplying a vector y ∈ Rm with this representation, we first multiply y with V >` ∈ Rn×m, which takes O(mn) time. Then we multiply V >` y ∈ Rn with U` ∈ Rm×n which takes O(mn) time. This drastically reduces the cost of multiplying the weight matrix with a vector from O(m2) to O(mn).
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 a runtime ofO(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 andO(Tmn) might be in the order ofO(m2). To circumvent this problem, we design the data structure so that it will exactly compute the m×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, α (Gall & Urrutia, 2018), which means it takes O(m2) time to multiply an m×m by an m×mα matrix. Hence, we restart the data structure after around mα/n updates. Therefore, we achieve an amortized o(m2) time, which is invariant even though the number of iterations T grows larger and larger.

heading: 3.2 CONVERGENCE ANALYSIS
text: In this section, we demonstrate the strategy to prove that our second-order method achieves a linear convergence rate on the training loss.
Step 1: Initialization. Let W (0) be the random initialization. We first show that for any data point xi, we have f(W (0), xi) = O(1). The analysis draws inspiration from Allen-Zhu et al. (2019a). The general idea is, given a fixed unit length vector x, multiplying it with a random Gaussian matrix W will make sure that ‖Wx‖22 ≈ 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)‖2 ≈ 1 as well as φ(Wx) being sparse. Inductively applying this idea to each layer and carefully controlling the error occurring at each layer, we can show that with good probability, f(W (0), xi) is a constant.
We also bound the spectral norm of ∂f(W (0),xi)∂W` by Õ( √ L/m). One of the key part of this matrix is the consecutive product DLWL . . . D`+1W`+1D`. 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( √ L). Finally, we make use the fact that each entry of a is a Rademacher random variable scaled by 1/ √ m,
hence the norm of a>DLWL . . . D`+1W`+1D` is bounded by Õ( √ L/m) with good probability.
Furthermore, we show that the Gram matrix for the multiple-layer over-parametrized neural network, which is defined as J`,0J>`,0, has a nontrivial minimum eigenvalue after the initialization. In particular, we adapt the neural tangent kernel (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 a 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 `2 norm) and the consecutive product (in terms of spectral norm) is also small. Finally, using the fact that a is a Rademacher vector with scaling 1/ √ m, we can show that the change ∂f(W (0) + ∆W,xi)
∂(W` + ∆W`) − ∂f(W (0), xi) ∂W`
has its spectral norm being bounded by Õ( √ L/m) for any layer ` ∈ [L] and input data i ∈ [n].
Consequently, the Frobenious norm of the Jacobian matrix is bounded by Õ( √ nL/m).
Step 3: Connect everything via a double induction. Put things together, we use a 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 ft−1 and y. By carefully bounding various terms and exploiting the fact the Jacobian matrix always has a relative small spectral norm (Õ( √ nL/m)), 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 ‖ft − y‖2.

heading: 4 DISCUSSION AND FUTURE DIRECTIONS
text: In this work, we propose and analyze a second-order method to train multi-layer over-parametrized neural networks. Our algorithm achieves a linear convergence rate in terms of training loss, and achieves a subquadratic (o(m2)) 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 a significant improvement fromO(m2n) time 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-based 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 a nearly linear cost per iteration on m as in the case of training one-hidden layer over-parametrized networks (Brand et al. (2021))? In particular, can this runtime be achieved under the current best width of multi-layer over-parametrized networks (m ≥ n8)? We note that the major limitation in our method is the sparsity of the change of the diagonal matrices (∆D) 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 in fact imposes on a lower bound on ‖∆W‖. It seems to us that, in order to achieve a 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 we maintain the changes of weight more efficiently? In our work, we achieve speedup in the neural network training process by observing that the change of the weights are 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 program (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 orthogonal direction to maintain the change is to design an initialization setup such that while we still have enough randomness to obtain provable guarantees, the matrix-vector product with the initial weight matrix can be performed faster than O(m2) by sparsifying the Gaussian matrix as in Dereziń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 activation 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 on such networks. Though the convergence of first-order method on over-parametrized multi-layer RNN has been established, it is unclear whether such analysis can be extended to second-order method.
Ethics Statement. This is a theory paper that proposed efficient algorithm to train multi-layer overparametrized neural networks, hence it 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, see 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.

heading: A PRELIMINARIES AND PROBABILITY TOOLS
text: 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.
Notations. For any positive integer n, we use [n] to denote the set {1, 2, · · · , n}. We use E[·] to denote expectation and Pr[·] for probability. We use ‖x‖2 to denote the `2 norm of a vector x. We use ‖A‖ to denote the spectral norm of matrix A. We use ‖A‖F 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 ×m. For matrix A or vector x, we use ‖A‖0, ‖x‖0 to denote the number of nonzero entries of A and x respectively. Note that ‖ · ‖0 is a semi-norm since it satisfies triangle inequality. Given a real square matrix A, we use λmax(A) and λmin(A) to denote its largest and smallest eigenvalues respectively. Given a real matrix A, we use σmax(A) and σmin(A) to denote its largest and smallest singular values respectively. We use N (µ, σ2) to denote the Gaussian distribution with mean µ and variance σ2. We use Õ(f(n)) to 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 to two matrices, this means 〈A,B〉 = tr[A>B] where tr[A] denote the trace of matrix A. Lemma A.1 (Chernoff bound Chernoff (1952)). Let X = ∑n i=1Xi, where Xi = 1 with probability
pi and Xi = 0 with probability 1−pi, and all Xi are independent. Let µ = E[X] = ∑n i=1 pi. Then 1. Pr[X ≥ (1 + δ)µ] ≤ exp(−δ2µ/3), ∀δ > 0 ; 2. Pr[X ≤ (1− δ)µ] ≤ exp(−δ2µ/2), ∀0 < δ < 1. Lemma A.2 (Hoeffding bound Hoeffding (1963)). Let X1, · · · , Xn denote n independent bounded variables in [ai, bi]. Let X = ∑n i=1Xi, then we have
Pr[|X − E[X]| ≥ t] ≤ 2 exp ( − 2t
2∑n i=1(bi − ai)2
) .
Lemma A.3 (Bernstein inequality Bernstein (1924)). Let X1, · · · , Xn be independent zero-mean random variables. Suppose that |Xi| ≤M almost surely, for all i. Then, for all positive t,
Pr [ n∑ i=1 Xi > t ] ≤ exp ( − t 2/2∑n j=1 E[X2j ] +Mt/3 ) .
Lemma A.4 (Anti-concentration of Gaussian distribution). Let X ∼ N (0, σ2), then
Pr[|X| ≤ t] ∈ (2 3 t σ , 4 5 t σ ).
Lemma A.5 (Concentration of subgaussian random variables). Let a ∈ Rn be a vector where each entry of a is sampled from a subgaussian distribution of parameter σ2, then for any vector x ∈ Rn,
Pr[|〈a, x〉| ≥ t · ‖x‖2] ≤ 2 exp(− t2
2σ2 ).
Lemma A.6 (Small ball probability). Let a ∈ Rn be a vector such that |ai| ≥ δ for all i ∈ [n], x1, . . . , xn are n i.i.d. Rademacher random variables. Then, there exist absolute constants C1, C2 such that for any t > 0,
Pr[|〈a, x〉| ≤ t] ≤ min { C1t
‖a‖2 , C2t δ √ n
} .

heading: B COMPLETE ALGORITHM AND ITS RUNTIME ANALYSIS
text: In this section, we first present our complete algorithm, then give a runtime analysis of it.
Algorithm 2 Complete version of our algorithm.
1: procedure COMPLETEALGORITHM(X ∈ Rd×n, y ∈ Rn) . Theorem B.1 2: /*Initialization*/ 3: Initialize W`(0), ∀` ∈ [L] 4: Store W1(0)xi in memory, ∀i ∈ [n] . Takes O(nm2) time 5: LOWRANKMAINTENANCE LMR . Algorithm 3 6: LMR.INIT({W1(0) . . . ,WL(0)}) 7: for t = 0→ T do 8: /*Forward computation*/ 9: for ` = 1→ L do 10: vi,` ← hi,`−1, ∀i ∈ [n] 11: gi,` ← LMR.QUERY(`, hi,`−1) . Takes o(nm2) time 12: hi,` ← φ(gi,`), ∀i ∈ [n] 13: . hi,` is sparse 14: Di,` ← diag(φ′(gi,`)), ∀i ∈ [n] . Takes O(nm) time 15: . Di,` is sparse 16: end for 17: ft ← [a>h1,L, . . . , a>hn,L]> . Takes O(nm) time 18: /*Backward computation*/ 19: for ` = L→ 1 do 20: ui,` ← a>Di,LWL(t) . . . Di,`+1W`+1(t)Di,`,∀i ∈ [n] . Takes o(nLm2) time 21: g` ← FASTTENSORREGRESSION({ui,`}ni=1, {vi,`}ni=1, c) with precision √ λ/n 22: . Algorithm 6 23: LMR.UPDATE({g`,iui,`}ni=1, {vi,`}ni=1) 24: end for 25: end for 26: end procedure
Theorem B.1 (Formal version of Theorem 1.2). Let X ∈ Rd×n and y ∈ Rn, and let k denote the sparsity of Di,` and s denote the sparsity of ∆Di,`, ∀` ∈ [L], i ∈ [n]. Let m denote the width of neural network, L denote the number of layers and α denote the dual matrix multiplication exponent (Def. C.1),then the running time of Algorithm 2 is
O(Tinit + T · Titer), where
Tinit = O(m2(n+ L)), Titer = Õ((m1+α +m(s+ k))L2 +m2−αnL).
Therefore, the cost per iteration of Algorithm 2 is
Õ((m1+α +m(s+ k))L2 +m2−αnL).
Proof. We analyze Tinit and Titer separately. Initialization time. We will first initialize (L − 1) m ×m matrices and one m × d matrix, which takes O(m2L) time. Compute W1(0)xi for all i ∈ [n] takes O(m2n) time. Finally, initialize the data structure takes O(m2L) time. Hence, Tinit = O(m2(n+ L)). Cost per iteration. For each iteration, we perform one forward computation from layer 1 to L, then backpropagate from layer L to 1.
• Forward computation: In forward computation, we first compute gi,` ∈ Rm, which is equivalent to form vi,`+1, hence by Lemma C.5, it takes O(s + k + mα) time. Compute hi,` and Di,` takes O(m) time. These computations would be performed for all L layers, hence the overall runtime of forward computation is O((s+ k +mα)L) time.
• Backpropagation: In backpropagation, we first compute ui,` ∈ Rm, which takes O(mL(s + k + mα)) time. Then, we call Algorithm 6 to solve the Gram regression problem, which due to Theorem D.14 takes Õ(mn + nω) time. Note that even we want a high probability version of the solver with e− log
2 nL failure probability so that we can union bound over all layers and all iterations, we only pay extra log2 nL term in running time, which is absorbed by the Õ(·) notation. Finally, the update takes O(m2−αn) time by Lemma C.2. Sum over all L layers, we get an overall running time of Õ((m1+α +m(s+ k))L2 +m2−αnL) time.
This concludes the proof of our Theorem.
Corollary B.2. Suppose the network width m is chosen as in F.25 and the shift parameter b is chosen as in F.6, then the cost per iteration of Algorithm 2 is
Õ(m1.8L2 +m1.69nL).
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 F.6, then we can make sure that both sparsity parameters k and s to be o(m), and we achieve subquadratic cost per iteration.

heading: C LOW RANK MAINTENANCE AND EFFICIENT COMPUTATION OF THE CHANGE
text: In this seciton, we design a 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.

heading: C.1 LOW RANK MAINTENANCE
text: In this short section, we design a simple data structure to maintain the low rank representation of change of weights, and show that the matrix-vector product can be implemented efficiently.
Before moving, we define some notions related to rectangular matrix multiplication.
Definition C.1 (Williams (2012); Gall & Urrutia (2018)). Let ω be the matrix multiplication exponent such that it takes nω+o(1) time to multiply two n× n matrices. Let α be the dual exponent of the matrix multiplication which is the supremum among all a ≥ 0 such that it takes n2+o(1) time to multiply an n× n by n× na matrix. Additionally, we define the function ω(·) where ω(b) denotes the exponent of multiplying an n × n matrix by an n× nb matrix. Hence, we have ω(1) = ω and ω(α) = 2.
The overall idea of our low rank maintenance data structure is as 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.
Algorithm 3 Low rank maintenance data structure 1: data structure LOWRANKMAINTENANCE . Lemma C.2 2: members 3: r`,∀` ∈ [L] . r` denotes the accumulated rank of the change 4: W`,∀` ∈ [L] . {W`}L`=1 ∈ (Rm×m)L 5: ∆W`,∀` ∈ [L] . {∆W`}L`=1 ∈ (Rm×m)L 6: end members 7: 8: procedures 9: INIT({W1(0), . . .WL(0)}) . Initialize the data structure 10: UPDATE(U`, V`) . Update the low rank representation 11: QUERY(`, y) . Compute the matrix-vector product between ∆W` and y 12: end procedures 13: end data structure
Algorithm 4 Procedures of LRM data structure 1: procedure INIT({W1(0), . . . ,WL(0)}) . Lemma C.2 2: W` ←W`(0) 3: ∆W` ← 0,∀` ∈ [L] 4: r` ← 0,∀` ∈ [L] 5: end procedure 6: 7: procedure UPDATE(U` ∈ Rm×n, V` ∈ Rm×n) . Lemma C.2 8: ∆W` ← ∆W` + U`V >` without forming the product and sum the two matrices 9: r` ← r` + n 10: if r` = ma where a = ω(2) then 11: W` ←W` + ∆W` . Takes O(m2) time 12: r` ← 0 13: ∆W` ← 0 14: end if 15: end procedure 16: 17: procedure QUERY(` ∈ [L], y ∈ Rm) . Lemma C.2 18: z ←W` · y + ∆W` · y . Takes O(nnz(y) ·m+mr`) time 19: return z 20: end procedure
Lemma C.2. There exists a deterministic data structure (Algorithm 3) such that maintains ∆W1, . . . ,∆WL such that
• The procedure INIT (Algorithm 4) takes O(m2L) time.
• The procedure UPDATE (Algorithm 4) takes O(nm2−α) amortized time, where α = ω(2)
• The procedure QUERY (Algorithm 4) takes O(m · (nnz(y) + r`)) time, where r` is the rank of ∆W` when QUERY is called.
Proof. The runtime for INIT is obvious, for QUERY, notice that we are multiplying vector y with a (possibly) dense matrix W` ∈ Rm×m, which takes O(nnz(y) · m) time, and an accumulated low rank matrix ∆W` with rank r`. By using the low rank decomposition ∆W` = UV > with U, V ∈ Rm×r` , the time to multiply y with ∆W is O(mr`). Combine them together, we get a running time of O(m · (nnz(y) + r`)). It remains to analyze the amortized cost of UPDATE. Note that if r` < ma, then we just pay O(1) time to update corresponding variables in the data structure. If r` = ma, then we will explicitly form the m × m matrix ∆W`. To form it, notice we have accumulated r`/n different sums of
rank-n decompositions, which can be represented as
U = [U`(1), U`(2), . . . , U`(r`/n)] ∈ Rm×r` , V = [V`(1), V`(2), . . . , V`(r`/n)] ∈ Rm×r` ,
and ∆W` = UV >, which takes O(m2) time to compute since r` = ma and a = ω(2). Finally, note that this update of W` only happens once per r`/n number of calls to UPDATE, therefore we can charge each step by O( m 2
r`/n ) = O(m2−an) = O(m2−αn), arrives at our final amortized running
time.
Remark C.3. Currently, the dual matrix multiplication exponent α ≈ 0.31 (Gall & Urrutia, 2018), hence the amortized time for UPDATE is O(nm1.69). If m ≥ n10/3, then we achieve an update time of o(m2). Similarly, the time for QUERY is O(m · (nnz(y) + r`)) = O(m · nnz(y) + m1+α) = O(m · nnz(y) + m1.31), as long as nnz(y) = o(m), then its running time is also o(m2). In our application of training neural networks, we will make sure that the inputted vector y is sparse.
C.2 EFFICIENTLY COMPUTE ui,`(t) AND vi,`(t)
In this section, we show that how to compute the vectors ui,`, vi,` ∈ Rm using the low rank structure of the change of weights combined. Recall the definition of these vectors:
ui,`(t) > = a>Di,L(t)WL(t) . . . Di,`+1(t)W`+1(t)Di,`(t) ∈ R1×m,
vi,`(t) = hi,`−1(t) ∈ Rm.
Before proceeding, we list the assumptions we will be using:
• For any ` ∈ [L], Di,`(t) is sD-sparse, where sD := k + s, k is the sparsity of Di,`(0) and s is the sparsity of Di,`(t)−Di,`(0).
• For any ` ∈ [L], the change of the weight matrix W`, ∆W`(t) := W`(t) −W`(0), is of low-rank. That is, ∆W`(t) = ∑rt j=1 y`,jz > `,j .
• For any i ∈ [n], W1(0)xi is pre-computed.
We first note that as a direct consequence of Di,`(0) is k-sparse, hi,`(0) is k-sparse as well. Similarly, hi,`(t)− hi,`(0) has sparsity s. Hence hi,`(t) has sparsity bounded by sD.
Compute ui,`(t). Compute ui,`(t) is equivalent to compute the following vector:
Di,`(t)(W`+1(0) + ∆W`+1(t)) >Di,`+1(t) · · · (WL(0) + ∆WL(t))>Di,L(t)a.
First, we know thatDi,L(t)a ∈ Rm is an sD-sparse vector, and it takesO(sD) time. The next matrix is (WL(0)+∆WL(t))>, which gives two terms: WL(0)>(Di,L(t)a) and ∆WL(t)>(Di,L(t)a). For the first term, since Di,L(t)a is sD-sparse, it takes O(msD)-time. For the second term, we have
∆WL(t) >(Di,L(t)a) = rt∑ j=1 zL,jy > L,j(Di,L(t)a) = rt∑ j=1 zL,j · 〈yL,j , Di,L(t)a〉.
Each inner-product takes O(sD)-time and it takes O(mrt + sDrt) = O(mrt)-time in total. Hence, in O(m(sD + rt))-time, we compute the vector WL(t)>Di,L(t)a. Note that we do not assume the sparsity of a.
Thus, by repeating this process for the L − ` intermediate matrices W>j (t)Di,j(t), we can obtain the vector  L∏
j=`+1
W>j (t)Di,j(t)  a in time O((L− `)m(sD + rt)). Finally, by multiplying a sparse diagonal matrix Di,`(t), we get the desired vector ui,`(t).
Compute vi,`(t). Note that vi,`(t) is essentially hi,`−1(t), so we consider how to compute hi,`(t) for general ` ∈ [L]. Recall that
hi,`(t) = φ((W`(0) + ∆W`(t))hi,`−1(t)),
since hi,`−1(t) is sD-sparse, the productW`(0hi,`−1(t))) can be computed inO(msD) time. For the product ∆W`(t)hi,`−1(t) can be computed use the low rank decomposition, which takes O(mrt) time. Apply the threshold ReLU takes O(m) time. Hence, the total time is O(m(rt + sD)) time.
The running time results are summarized in the following lemma: Lemma C.4. For ` ∈ [L] and i ∈ [n], suppose ‖Di,`(0)‖0 ≤ k. Let t > 0. Suppose the change of Di,` is sparse, i.e., ‖Di,`(t)−Di,`(0)‖0 ≤ s. For ` ∈ [L], i ∈ [n], for any t > 0, suppose the change of W` is of low-rank, i.e., ∆W`(t) = ∑rt j=1 y`,jz > `,j . We further assume that {y`,j , z`,j}`∈[L],j∈[rt] and {W1(0)xi}i∈[n] are pre-computed.
Then, for any ` ∈ [L] and i ∈ [n], the vectors u`,i(t), v`,i(t) ∈ Rm can be computed in O(mL(s+ k + rt))-time.
As a direct consequence, if we combine Lemma C.2 and Lemma C.4, then we get the following corollary: Corollary C.5. For ` ∈ [L] and i ∈ [n], we can compute vi,`(t), ui,`(t) ∈ Rm as in Algorithm 2 with the following time bound:
• Compute ui,`(t) in time O(mL(s+ k + r`)).
• Compute vi,`(t) in time O(m(s+ k + r`)).

heading: D FAST TENSOR PRODUCT REGRESSION
text: In this section, we design a generic algorithm for solving the following type of regression task:
Given two matrices U = [u>1 , . . . , u > n ] >, V = [v>1 , . . . , v > n ] ∈ Rm×n with m n, consider the matrix J ∈ Rn×m2 formed by
J =  vec(u1v > 1 ) > vec(u2v > 2 ) >
... vec(unv > n ) >  . We are also given a vector c ∈ Rn, the goal is to solve the following regression task:
min x∈Rn
‖JJ>x− c‖22.
Our main theorem for this section is as follows: Theorem D.1 (Restatement of Theorem D.14). Given two n ×m matrices U and V , and a target vector c ∈ Rn. Let J = [vec(u1v>1 )>, . . . , vec(unv>n )>] ∈ Rn×m 2
. There is an algorithm (Algorithm 6) takes Õ(nm + n2(log(κ/ ) + log(m/ δ) −2) + nω) time and outputs a vector x̂ ∈ Rn such that
‖JJ>x̂− c‖2 ≤ ‖c‖2 holds with probability at least 1− δ, and κ is the condition number of J .
D.1 APPROXIMATE J VIA TensorSketch
We introduce the notion of TensorSketch for two vectors: Definition D.2. Let h1, h2 : [m] → [s] be 3-wise independent hash functions, also let σ : [m] → {±1} be a 4-wise independent random sign function. The degree two TensorSketch transform, S : Rm × Rm → Rs is defined as follows: for any i, j ∈ [m] and r ∈ [s],
Sr,(i,j) = σ1(i) · σ2(j) · 1[h1(i) + h2(j) = r mod s].
Remark D.3. Apply S to two vectors x, y ∈ Rm can be implemented in time O(s log s+ nnz(x) + nnz(y)).
We introduce one key technical lemma from Avron et al. (2014):
Lemma D.4 (Theorem 1 of Avron et al. (2014)). Let S ∈ Rs×m2 be the TensorSketch matrix, consider a fixed n-dimensional subspace V . If s = Ω(n2/( 2δ)), then with probability at least 1− δ, ‖Sx‖2 = (1± )‖x‖2 simultaneously for all x ∈ V .
Now we are ready to prove the main lemma of this section:
Lemma D.5. Let , δ ∈ (0, 1) denote two parameters. Let J ∈ Rn×m2 be a matrix such that i-th row of J is defined as vec(uiv>i ) for some ui, vi ∈ Rm. Then, we can compute a matrix J̃ ∈ Rn×s such that for any vector x ∈ Rn, with probability at least 1− δ, we have
‖J̃>x‖2 = (1± )‖J>x‖2,
where s = Ω(n2/( 2δ)). The time to compute J̃ is O(ns log s+ nnz(U) + nnz(V )).
Proof. Notice that the row space of matrix J can be viewed as an n-dimensional subspace, hence, by Lemma D.4, the TensorSketch matrix S with s = Ω(n2/( 2δ)) can preserve the length of all vectors in the subspace generated by J> with probability 1− δ, to a multiplicative factor of 1± . The running time part is to apply the FFT algorithm to each row of J with a total of n rows. For each row, it takes O(s log s+m) time, hence the overall running time is O(n(s log s+m)).
D.2 APPROXIMATE J VIA TensorSRHT
We note that the dependence on the target dimension of sketching is O(1/δ) for TensorSketch. We introduce another kind of sketching technique for tensor, called TensorSRHT. The tradeoff is we lose input sparsity runtime of matrices U and V .
Definition D.6. We define the TensorSRHT S : Rd × Rd → Rm as S = 1√ m P · (HD1 ×HD2), where each row of P ∈ {0, 1}m×d2 contains only one 1 at a random coordinate, one can view P as a sampling matrix. H is a d × d Hadamard matrix, and D1, D2 are two d × d independent diagonal matrices with diagonals that are each independently set to be a Rademacher random variable (uniform in {−1, 1}). Remark D.7. By using FFT algorithm, apply S to two vectors x, y ∈ Rm takes timeO(m logm+s).
We again introduce a technical lemma for TensorSRHT.
Lemma D.8 (Theorem 3 of Ahle et al. (2020)). Let S ∈ Rs×m2 be the TensorSRHT matrix, consider a fixed n-dimensional subspace V . If s = Ω(n log3(nm/ δ) −2), then with probability at least 1− δ, ‖Sx‖2 = (1± )‖x‖2 simultaneously for all x ∈ V . Lemma D.9. Let , δ ∈ (0, 1) denote two parameters. Given a list of vectors u1, · · · , um, v1, · · · , vm ∈ Rm. Let J ∈ Rn×m 2
be a matrix with i-th row of J is defined as vec(uiv > i ). Then, we can compute a matrix J̃ ∈ Rn×s such that for any vector x ∈ Rn, with probability at least 1− δ, we have
‖J̃>x‖2 = (1± )‖J>x‖2,
where s = Ω(n log3(nm/( δ)) −2). The time to compute J̃ is O(n(m logm+ 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 logm+ s) time, and we need to apply it to n rows.

heading: D.3 TENSOR TRICK AS A PRECONDITIONER
text: In this section, we use TensorSketch and TensorSRHT as a preconditioner to solve a regression task involving JJ>. This is an important step to implement the Gram-Newton-Gauss iteration Cai et al. (2019); Brand et al. (2021).
Before proceeding, we introduce the notion of subspace embedding: Definition D.10 (Subspace Embedding, Sarlos (2006)). LetA ∈ RN×k, we say a matrix S ∈ Rs×N is a (1± )− `2 subspace embedding for A if for any x ∈ Rk, we have ‖SAx‖22 = (1± )‖Ax‖22. Equivalently, ‖I − U>S>SU‖ ≤ where U is an orthonormal basis for the column space of A.
We will mainly utilize efficient subspace embedding. Definition D.11 (Lu et al. (2013); Woodruff (2014)). Given a matrixA ∈ RN×k withN = poly(k), then we can compute an S ∈ Rkpoly(log(k/δ))/ 2×N such that with probability at least 1−δ, we have
‖SAx‖2 = (1± )‖Ax‖2 hols for all x ∈ Rk. Moreover, SA can be computed in O(Nk log((k logN)/ )) time.
Algorithm 5 Fast Regression algorithm of Brand et al. (2021)] 1: procedure FASTREGRESSION(A, y, ) . Lemma D.12 2: . A ∈ RN×k is full rank, ∈ (0, 1/2) 3: Compute a subspace embedding SA . S ∈ Rkpoly(log k) ×N 4: Compute R such that SAR has orthonormal columns via QR decomposition . R ∈ Rk×k 5: z0 = 0k ∈ Rk 6: t← 0 7: while ‖A>ARzt − y‖2 ≥ do 8: zt+1 ← zt − (R>A>AR)>(R>A>ARzt −R>y) 9: t← t+ 1 10: end while 11: return Rzt 12: end procedure
Lemma D.12 (Lemma 4.2 of Brand et al. (2021)). Let N = Ω(kpoly(log k)). Given a matrix A ∈ RN×k, let κ denote its condition number. Consider the following regression task:
min x∈Rk
‖A>Ax− y‖2.
Using the procedure FASTREGRESSION (Algorithm 5), with probability at least 1 − δ, we can compute an -approximate solution x̂ satisfying
‖A>Ax̂− y‖2 ≤ ‖y‖2
in time Õ(Nk log(κ/ ) + kω).
Our algorithm is similar to the ridge regression procedure in Song et al. (2021), where they first apply their sketching algorithm as a bootstrapping to reduce the dimension of the original matrix, then use another subspace embedding to proceed and get stronger guarantee.
We shall first prove a useful lemma. Lemma D.13. Let A ∈ RN×k, suppose SA is a subspace embedding for A (Def. D.11), then we have for any x ∈ Rk, with probability at least 1− δ,
‖(SA)>SAx− b‖2 = (1± )‖A>Ax− b‖2.
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
‖(SA)>SAx− b‖22 = ‖(SA)>SAx‖22 + ‖b‖22 − 2〈(SA)>SAx, b〉.
We will first bound the norm of (SA)>SAx, then the inner product term.

heading: Bounding ‖(SA)>SAx‖22
text: Let U ∈ RN×k be an orthonormal basis of A, then use the equivalent definition of subspace embedding, we have ‖U>S>SU − I‖ ≤ , this means all the eigenvalues of U>S>SU lie in the range of of [(1− )2, (1 + )2]. Let V denote the matrix U>S>SU , then we know that all eigenvalues of V >V lie in range [(1 − )4, (1 + )4]. Setting as /4, we arrive at ‖V >V − I‖ ≤ . This shows that for any x ∈ Rk, we have ‖(SA)>SAx‖2 = (1± )‖A>Ax‖2. Bounding 〈(SA)>SAx, b〉 Note that
〈(SA)>SAx, b〉 = 〈SAx, SAb〉 = 1/2 · (‖SAx‖22 + ‖SAb‖22 − ‖SA(x− b)‖22) = 1/2 · (1± )(‖Ax‖22 + ‖Ab‖22 − ‖A(x− b)‖22) = (1± )〈A>Ax, b〉.
Combining these two terms, we conclude that, with probability at least 1− δ,
‖(SA)>SAx− b‖2 = (1± )‖A>Ax− b‖2.
Algorithm 6 Fast Regression via tensor trick
1: procedure FASTTENSORREGRESSION({ui}ni=1 ∈ Rm×n, {vi}ni=1 ∈ Rm×n, c ∈ Rn) . Theorem D.14 2: . J = [vec(u1v>1 ) >, vec(u2v > 2 ) >, . . . , vec(unv > n ) >]> ∈ Rn×m2 3: s1 ← Θ(n log3(nm/( δ)) −2) 4: s2 ← Θ((n+ logm) log n) 5: Let S1 ∈ Rs1×m 2
be a sketching matrix . S1 can be TensorSketch or TensorSRHT 6: Compute J̃ = JS>1 via FFT algorithm . J̃ ∈ Rn×s1 7: Let S2 ∈ Rs2×s1 be a sketching matrix defined in Definition D.11 8: Compute a subspace embedding S2J̃> 9: Compute R such that S2J̃>R has orthonormal columns via QR decomposition . R ∈ Rn×n
10: z0 ← 0k ∈ Rk 11: t← 0 12: while ‖J̃ J̃>Rzt − c‖2 ≥ do 13: zt+1 ← zt − (R>J̃ J̃>R)>(R>J̃ J̃>Rzt −R>c) 14: t← t+ 1 15: end while 16: return Rzt 17: end procedure
Theorem D.14. Given two n × m matrices U and V , and a target vector c ∈ Rn. Let J = [vec(u1v > 1 ) >, . . . , vec(unv > n ) >] ∈ Rn×m2 . There is an algorithm (Algorithm 6) takes Õ(nm + n2(log(κ/ ) + log(m/δ)) + nω) time and outputs a vector x̂ ∈ Rn such that
‖JJ>x̂− c‖2 ≤ ‖c‖2 holds with probability at least 1− δ, and κ is the condition number of J .
Proof. We can decompose Algorithm 6 into two parts:
• Applying S1 to efficiently form matrix J̃ to approximate J and reduce its dimension, notice here we only need for this part to be a small constant, pick = 0.1 suffices.
• Using S2 as a preconditioner and solve the regression problem iteratively.
Let x̂ denote the solution found by the iterative regime. We will prove this statement in two-folds:
• First, we will show that ‖J̃ J̃>x̂− c‖2 ≤ ‖c‖2 with probability at least 1− δ;
• Then, we will show that ‖JJ>x̂− c‖2 = (1± 0.1)‖J̃ J̃>x̂− c‖2 with probability at least 1− δ.
Combining these two statements, we can show that
‖JJ>x̂− c‖2 = (1± 0.1)‖J̃ J̃>x̂− c‖2 ≤ 1.1 ‖c‖2
Setting to /1.1 and δ to δ/2, we conclude our proof. It remains to prove these two parts.
Part 1. ‖J̃ J̃>x̂− c‖2 ≤ ‖c‖2 We observe the iterative procedure is essentially the same as running FASTREGRESSION on input J̃>, y, , hence by Lemma D.12, we have with probability at least 1−δ, ‖J̃ J̃>x̂− c‖2 ≤ ‖c‖2.
Part 2. ‖JJ>x̂ − c‖2 = (1 ± 0.1)‖J̃ J̃>x̂ − c‖2 To prove this part, note that by Lemma D.8, we know that J̃> is a subspace embedding for J>. Hence, we can utilize Lemma D.13 and get that, with probability at least 1− δ, we have ‖JJ>x̂− c‖2 = (1± 0.1)‖J̃ J̃>x̂− c‖2. Combining these two parts, we have proven the correctness of the theorem. It remains to 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 x̂ via iterative scheme.
Part 1. Generate J̃ To generate J̃ , we apply S1 ∈ Rs1×m 2
which is a TensorSRHT. By Lemma D.9, it takes O(n(m logm + s1)) time to compute J̃ , plug in s1 = Θ(n log3(nm/δ)), the time is Õ(nm).
Part 2. Compute x̂ To compute x̂, essentially we run FASTREGRESSION on J̃>, c, , hence by Lemma D.12, it takes Õ(s2n log(κ/ ) + nω) time, with s2 = Θ((n + logm) log n) and κ is the condition number of J̃ , which has the guarantee κ = (1± )κ(J). Hence, the overall running time of this part is Õ(n2 log(κ/ ) + nω).
Put things together, the overall running time is Õ(nm+ n2 log(κ/ ) + nω).
Remark D.15. Due to the probability requirement (union bounding over all layers and all data points), here we only prove by using TensorSRHT. One can use similar strategy to obtain an input sparsity time version using TensorSketch. We remark that this framework is similar to the approach Song et al. (2021) takes to solve kernel ridge regression, where one first uses a shallow but fast sketch to bootstrap, then use another sketching to proceed with the main task.

heading: E SPECTRAL PROPERTIES OF OVER-PARAMETRIZED DEEP NEURAL NETWORK
text: In this section, we study the spectral properties of our Gram matrix and connects it to multi-layer NTK.

heading: E.1 BOUNDS ON THE LEAST EIGENVALUE OF KERNEL AT INITIALIZATION
text: We first define the Gram matrices for multiple layer neural network.
Definition E.1 (Multiple layer Gram matrix). The Gram matrices K` ∈ Rn×n for ` ∈ {0, . . . , L} of an L-layer neural network are defined as follows:
• (K0)i,j := x>i xj
• For ` > 0, let Σ`,i,j := [ (K`−1)i,i (K`−1)i,j (K`−1)j,i (K`−1)j,j ] ∈ R2×2 for any (i, j) ∈ [n]× [n]. Then,
(K`)i,j := E (x1,x2)∼N (0,2Σ`−1,i,j) [φ(x1)φ(x2)] ∀` ∈ [L− 1],
(KL)i,j := E (x1,x2)∼N (0,2ΣL−1,i,j)
[φ′(x1)φ ′(x2)]
Let λL := λmin(KL) to be the minimum eigenvalue of the NTK kernel KL.
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 eigenvalue at initialization, multiple layer version of Lemma C.3 in Brand et al. (2021)). Let λ` denote the minimum eigenvalue of NTK defined for `-th layer of neural networks. Suppose m` = Ω(λ−2` n 2 log(n/δ)), then with probability 1− δ, we have
λmin(G`(0)) ≥ 3
4 λ`, ∀` ∈ [L]
Proof. For any ` ∈ [L], we have G` = J`J>` ∈ Rn×n. Hence, for any i, j ∈ [n],
(G`)i,j = vec( ∂f(W,xi)
∂W` )>vec(
∂f(W,xj)
∂W` )
= vec ( Di,`
L∏ k=`+1 W>k Di,kah > i,`−1
)> vec ( Dj,`
L∏ k=`+1 W>k Dj,kah > j,`−1
)
= ( (hi,`−1 ⊗ Im`) ( Di,` L∏ k=`+1 W>k Di,ka ))> (hj,`−1 ⊗ Im`) ( Dj,` L∏ k=`+1 W>k Dj,ka )
= ( Di,`
L∏ k=`+1 W>k Di,ka )> (h>i,`−1 ⊗ Im`)(hj,`−1 ⊗ Im`) ( Dj,` L∏ k=`+1 W>k Dj,ka )
= a> ( Di,`
L∏ k=`+1 W>k Di,k
)>( Dj,`
L∏ k=`+1 W>k Dj,k ) a · h>i,`−1hj,`−1,
where
hi,`−1 = `−1∏ k=1 Di,kWkxi.
In particular, (GL)i,j = a >Di,LDj,La · h>i,L−1hj,L−1
= 1
m m∑ r=1 φ′(〈w(r)L , hi,L−1〉)φ ′(〈w(r)L , hj,L−1〉)h > i,L−1hj,L−1 (4)
We will prove that ‖GL −KL‖∞ is small, which implies that λmin(GL) is close to λ`. The proof idea is similar to Du et al. (2019a) via induction on `.
For ` = 1, recall (g1,i)k = ∑ b∈[m](W1)k,b(xi)b for k ∈ [m]. Hence, for any k ∈ [m],
E[(g1,i)k(g1,j)k] = ∑
b,b′∈[m]
E[(W1)k,b(W1)k,b′(xi)b(xj)b′ ]
= ∑ b∈[m] E[(W1)2k,b] · (xi)b(xj)b ((W1)k,b ∼ N (0, 2m ).)
= 2
m ∑ b∈[m] (xi)b(xj)b
= 2
m x>i xj .
Then, we have
E[h>1,ih1,j ] = ∑ k∈[m] E[(h1,i)k(h1,j)k]
= ∑ k∈[m] E[φ((g1,i)k)φ((g1,j)k)]
= ∑ k∈[m] E (u,v)∼N (0, 2m Σ1,i,j) [φ(u)φ(v)]
= E (u,v)∼N (0, 2m Σ1,i,j) [mφ(u)φ(v)]
= E (u′,v′)∼N (0,2Σ1,i,j)
[φ(u′)φ(v′)]
= (K1)i,j .
Next, we will show that h>1,ih1,j concentrates around its expectation. First, for any k ∈ [m],
|(h1,i)k(h1,j)k| ≤ |(g1,i)k(g1,j)k| ≤ |〈(W1)k,∗, xi〉| · |〈(W1)k,∗, xj〉|.
Since 〈(W1)k,∗, xi〉 ∼ N (0, 2‖xi‖ 2 2
m ), by the concentration of Gaussian distribution,
|〈(W1)k,∗, xi〉| ≤ √ c ∀k ∈ [m], i ∈ [n]
holds with probability at least 1−mne−cm/4. Conditioned on this event, we have |(h1,i)k(h1,j)k| ≤ c for all i, j ∈ [n] and k ∈ [m]. Then, by Hoeffding’s inequality, we have for any (i, j) ∈ [n]× [n],
Pr [ |h>1,ih1,j − (K1)i,j | ≥ t ] ≤ exp ( − t 2
2m · (2c)2
) = exp(−Ω(t2/(mc2))).
Hence, by union bound, we get that
max (i,j)∈[n]×[n]
|h>1,ih1,j − (K1)i,j | ≤ t
with probability at least
1−mn exp(−Ω(mc))− n2 exp(−Ω(t2/(mc2))).
If we choose c := log(mnL/δ)m and t := m −1/2 · polylog(nL/δ), we have with probability at least 1− δL ,
max (i,j)∈[n]×[n]
|h>1,ih1,j − (K1)i,j | ≤ Õ(m−1/2).
Suppose for ` = 1, . . . h (h < L), we have
max (i,j)∈[n]×[n]
|h>`,ih`,j − (K`)i,j | ≤ Õ(m−1/2).
Consider ` = h+ 1. By a similar computation, we have
E W`
[(g`,i)k(g`,j)k] = 2
m h>`−1,ih`−1,j .
Define a new covariance matrix
Σ̂`,i,j :=
[ h>`−1,ih`−1,i h > `−1,ih`−1,j
h>`−1,jh`−1,i h > `−1,jh`−1,j
] ∀(i, j) ∈ [n]× [n].
We have
E W` [h>i,`hj,`] = ∑ k∈[m] E (u,v)∼N (0, 2m Σ̂`,i,j) [φ(u)φ(v)]
= E (u′,v′)∼N (0,2Σ̂`,i,j)
[φ(u′)φ(v′)]
:= (K̂`)i,j .
Hence, we have with probability at least 1− δL ,
max (i,j)∈[n]×[n] ∣∣∣h>`,ih`,j − (K̂`)i,j∣∣∣ ≤ Õ(m−1/2). (5) It remains to upper bound the difference ‖K̂` −K`‖∞.∥∥∥K̂` −K`∥∥∥
∞ = max (i,j)∈[n]×[n] ∣∣∣∣∣ E(u,v)∼N (0,2Σ̂`,i,j)[φ(u)φ(v)]− E(u,v)∼N (0,2Σ`,i,j)[φ(u)φ(v)] ∣∣∣∣∣ .
Recall that
Σ`,i,j := [ (K`−1)i,i (K`−1)i,j (K`−1)j,i (K`−1)j,j ] ∀(i, j) ∈ [n]× [n],
and hence, by the induction hypothesis, we have
‖Σ̂`,i,j − Σ`,i,j‖∞ ≤ max (i,j)∈[n]×[n] ∣∣h>`−1,ih`−1,j − (K`−1)i,j∣∣ = Õ(m−1/2). Notice that Σ̂`,i,j can be written as[
‖h`−1,i‖22 cos(θ`,i,j)‖h`−1,i‖2‖h`−1,j‖2 cos(θ`,i,j)‖h`−1,i‖2‖h`−1,j‖2 ‖h`−1,j‖22
] .
Moreover, when φ is the ReLU function, we have
E (u,v)∼N (0,2Σ̂`,i,j) [φ(u)φ(v)] = 2‖h`−1,i‖2‖h`−1,j‖2 · F (θ`,i,j),
where
F (θ) := E (u,v)∼N (0,Σ(θ))
[φ(u)φ(v)] with Σ(θ) := [
1 cos(θ) cos(θ) 1
] .
We note that F (θ) has the following analytic form:
F (θ) = 1
2π (sin(θ) + (π − θ) cos(θ)) ∈ [0, 1/2]. (6)
Similarly,
E (u,v)∼N (0,2Σ`,i,j)
[φ(u)φ(v)] = 2 √ (K`−1)i,i(K`−1)j,j · F (τ`,i,j),
where τ`,i,j := cos−1 (
(K`−1)i,j√ (K`−1)i,i(K`−1)j,j
) . By the induction hypothesis, we have (K`)i,j ∈
h>`,ih`,j±Õ(m−1/2) for all i, j ∈ [n]. By Lemma F.7, we also have ‖h`,i‖2 ∈ 1± for all ` ∈ [L] and i ∈ [n] with probability 1−O(nL)·e−Ω(m 2/L). They implies that cos(τ`,i,j) ∈ cos(θ)±Õ(m−1/2). Thus, by Taylor’s theorem, we have
|F (θ`,i,j)− F (τ`,i,j)| ≤ Õ(m−1/2).
Therefore, we have∣∣∣∣∣ E(u,v)∼N (0,2Σ̂`,i,j)[φ(u)φ(v)]− E(u,v)∼N (0,2Σ`,i,j)[φ(u)φ(v)] ∣∣∣∣∣
= 2 ∣∣∣∣‖h`−1,i‖2‖h`−1,j‖2F (θ`,i,j)−√(K`−1)i,i(K`−1)j,jF (τ`,i,j)∣∣∣∣ ≤ Õ(m−1/2).
That is,
‖K̂` −K`‖∞ ≤ Õ(m−1/2). (7)
Combining Eqs. (5) and (7) together, we get that
max (i,j)∈[n]×[n]
|h>`,ih`,j − (K`)i,j | ≤ Õ(m−1/2)
holds with probability at least 1− δL for ` = h+ 1.
By induction, we have proved that for the first L − 1 layers, the intermediate correlation h>`,ih`,j is close to the intermediate Gram matrix (K`)i,j . Now, we consider the last layer. RecallGL is defined by Eq. (4), which has the same form as the correlation matrix of a two-layer over-parameterized neural network with input data {hL−1,i}i∈[n]. Define
(K̂L)i,j := h > L−1,ihL−1,j · E
w∼N (0,2Im)
[ φ′(w>hL−1,i)φ ′(w>hL−1,j) ] .
Then, by the analysis of the two-layer case (see for example Song & Yang (2019); Du et al. (2019b)), we have
‖GL − K̂L‖ ≤ λL 8 ,
if m = Ω(λ−2L n 2 log(n/δ)), where λL := λmin(KL). It remains to bound ‖K̂L −KL‖∞. Equivalently, for any (i, j) ∈ [n]× [n],
max (i,j)∈[n]×[n] ∣∣∣∣∣ E(u,v)∼N (0,2Σ̂L,i,j)[φ′(u)φ′(v)]− E(u,v)∼N (0,2ΣL,i,j)[φ′(u)φ′(v)] ∣∣∣∣∣ .
The expectation has the following analytic form:
E (z1,z2)∼N (0,Σ)
[φ′(z1)φ ′(z2)] =
1 4 +
sin−1(ρ)
2π with Σ = [ p2 ρpq ρpq q2 ] .
By the analysis of the (L − 1)-layer, we know that |ρL,i,j − ρ̂L,i,j | ≤ Õ(m−1/2), where ρL,i,j := cos(τL,i,j) and ρ̂L,i,j := cos(θL,i,j). Also, notice that cos(τL,i,j) = F (τL−1,i,j) ∈ [0, 1/2] by Eq. (6). Hence, the derivative of the expectation is bounded, and by Taylor’s theorem, we have
‖K̂L −KL‖∞ ≤ Õ(m−1/2).
It implies that ‖K̂L −KL‖ ≤ λL8 , which further implies that
‖GL −KL‖ ≤ λL 4 .
Equivalently, we get that
λmin(GL) ≥ 3
4 λL
with probability at least 1− δ. The lemma is then proved.

heading: E.2 BOUNDS ON THE LEAST EIGENVALUE DURING OPTIMIZATION
text: In this section, we generalize the Lemma C.5 in Brand et al. (2021) into multiple layer neural network 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` n
2 log(n/δ)), with probability least 1− δ, for any set of weights W1, · · ·WL satisfying
‖W` −W`(0)‖ ≤ R.
then the following holds
‖GL(W )−GL(W (0))‖F ≤ λL/2.
Proof. Recall (GL)(W )i,j = 1m ∑m r=1 φ
′(〈(WL)r, hi,L−1〉)φ′(〈(WL)r, hj,L−1〉)h>i,L−1hj,L−1. For simplicity, let zi,r := (WL)>r hi,L−1 and zi,r(0) := (WL(0)) > r hi,L−1(0).
|(GL(W ))i,j − (GL(W (0)))i,j |
= ∣∣∣∣∣h>i,L−1hj,L−1 1m m∑ r=1 φ′(zi,r)φ ′(zj,r)− hi,L−1(0)>hj,L−1(0) 1 m m∑ r=1 φ′(zi,r(0))φ ′(zj,r(0)) ∣∣∣∣∣ ≤ ∣∣h>i,L−1hj,L−1 − hi,L−1(0)>hj,L−1(0)∣∣ · 1m m∑ r=1 φ′(zi,r)φ ′(zj,r)
+ ∣∣hi,L−1(0)>hj,L−1(0)∣∣ · 1
m ∣∣∣∣∣ m∑ r=1 φ′(zi,r)φ ′(zj,r)− φ′(zi,r(0))φ′(zj,r(0)) ∣∣∣∣∣ ≤ ∣∣h>i,L−1hj,L−1 − hi,L−1(0)>hj,L−1(0)∣∣+ λL + Õ(m−1/2)m ∣∣∣∣∣ m∑ r=1 φ′(zi,r)φ ′(zj,r)− φ′(zi,r(0))φ′(zj,r(0))
∣∣∣∣∣ , where the last step follows from φ′(x) ∈ {0, 1} and Lemma E.2. To upper bound the above two terms, we first need to bound the move of hi,L−1. For any ` ∈ [L−1], we have
‖hi,` − hi,`(0)‖2 = ‖φ(W`hi,`−1)− φ(W`(0)hi,`−1(0))‖2 ≤ ‖φ(W`hi,`−1)− φ(W`hi,`−1(0))‖2
+ ‖φ(W`hi,`−1(0))− φ(W`(0)hi,`−1(0))‖2 ≤ (‖W` −W`(0)‖+ ‖W`(0)‖) · ‖hi,`−1 − hi,`−1(0)‖2
+ ‖W` −W`(0)‖ · ‖hi,`−1(0)‖2 ≤ (R+ cW )‖hi,`−1 − hi,`−1(0)‖2 +R(1 + ),
where cW := ‖W`(0)‖ ≤ 3 by the well-known deviations bounds concerning the singular values of Gaussian random matrices (Rudelson & Vershynin (2010)). Also, when ` = 0, we have ‖hi,0 − hi,0(0)‖2 = 0. Hence, we get that for all ` ∈ [L− 1], i ∈ [n],
‖hi,` − hi,`(0)‖2 ≤ √ 1 + R(2cW ) `,
since R 1 by our choice of R. Hence, it implies that∣∣h>i,L−1hj,L−1 − hi,L−1(0)>hj,L−1(0)∣∣
≤ ∣∣(hi,L−1 − hi,L−1(0))>hj,L−1∣∣+ ∣∣hi,L−1(0)>(hj,L−1 − hj,L−1(0))∣∣
≤ 2(1 + )R(2cW )L−1.
Similarly,
λL + Õ(m −1/2)
m
∣∣∣∣∣ m∑ r=1 φ′(zi,r)φ ′(zj,r)− φ′(zi,r(0))φ′(zj,r(0)) ∣∣∣∣∣ = λL + Õ(m −1/2)
m ∣∣1(WL)>r hi,L−1>0,(WL)>r hj,L−1>0 − 1(WL(0))>r hi,L−1(0)>0,(WL(0))>r hj,L−1(0)>0∣∣ . By our assumption, we have ‖(WL)r − (WL(0))r‖2 ≤ R. We also know that ‖hi,L−1 − hi,L−1(0)‖2 ≤ √ 1 + R(2cW )
L−1. Then, we can follow the proof in Du et al. (2019b); Song & Yang (2019) and define the event Ai,r := ∃w, h : ‖w − (WL(0))r‖2 ≤ R, ‖h− hi,L−1(0)‖2 ≤ O(R) such that 1w>h>0 6= 1(WL(0))>r hi,L−1(0)>0. We haveAi,r happens if and only if |((WL(0))r+ ∆w)>hi,L−1(0)| < O(R(R + m−1/2)) for some fixed vector ∆w of length at most R. It implies that
Pr[Ai,r] = Pr z∼N (0,1)
[|z| < R+ o(R)] ≤ O(R).
Hence, using the same proof in the previous work, we get that
λL + Õ(m −1/2)
m
∣∣∣∣∣ m∑ r=1 φ′(zi,r)φ ′(zj,r)− φ′(zi,r(0))φ′(zj,r(0)) ∣∣∣∣∣ ≤ (λL +m−1/2)R. Putting them together, we have
|(GL(W ))i,j − (GL(W (0)))i,j | ≤ (2(1 + )(2cW )L−1 + λL)R. And by our choice of R, we get that
‖GL(W )−GL(W (0))‖F ≤ (2(1 + )(2cW )L−1 + λL)nR ≤ λL/2.
The lemma is then proved.

heading: F CONVERGENCE ANALYSIS OF ALGORITHM 2
text: In this section, we analyze the convergence behavior of Algorithm 2.

heading: F.1 PRELIMINARY
text: We recall the initialization of our neural network. Definition F.1 (Initialization). Let m = m` for all ` ∈ [L]. Let m0 = d. We assume weights are initialized as
• Each entry of weight vector a ∈ RmL is i.i.d. sampled from {− 1√ mL ,+ 1√ mL } uniformly
at random.
• Each entry of weight matrices W` ∈ Rm`×m`−1 sampled from N (0, 2/m`).
We also restate the architecture of our neural network here. Definition F.2 (Architecture). Our neural network is a standard L-layer feed-forward neural network, with the activation functions defined as a scaled version of shifted ReLU activation: φ(x) = √ cb1[x > √ 2/mb]x, where cb := (2(1 − Φ(b) + bφ(b)))−1/2. Here b is a 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 f : Rm0 → R as
f(W,xi) = a >φ(WLφ(WL−1φ(. . . φ(W1xi)))).
We measure the loss of the neural network via squared-loss function:
L(W ) = 1 2 n∑ i=1 (f(xi)− yi)2.
We use ft : Rm0×n → Rn denote the prediction of our network:
ft(X) = [f(W (t), x1), . . . , f(W (t), xn)] >.
We state two assumptions here. Assumption F.3 (Small Spectral Norm). Let t ∈ {0, . . . , T} and let R ≤ 1 be a parameter. We assume
max `∈[L]
‖W`(t)−W`(0)‖ ≤ R.
Later, we will invoke this assumption by specifying the choice of R. Assumption F.4 (Sparsity). Let t ∈ {0, . . . , T} and let s ≥ 1 be an integer parameter. We assume
‖∆Di,`‖0 ≤ s,∀` ∈ [L], i ∈ [n]. Later, we will invoke this assumption by specifying the choice of s.
Finally, throughout this entire section, we will assume m` = m for any ` ∈ [L].

heading: F.2 TECHNICAL LEMMAS
text: We first show that during initialization, by using our threshold ReLU activation, the vector hi,` is sparse, hence the diagonal matrix Di,` is sparse as well.
Lemma F.5 (Sparse initialization). Let σb(x) = max{x − b, 0} be the threshold ReLU activation with threshold b > 0. After initialization, with probability 1 − nL · e−Ω(me−b
2m/4), it holds for all i ∈ [n] and ` ∈ [L], ‖hi,`‖0 ≤ O(m · e−b 2m/4).
Proof. We fix i ∈ [n] and ` ∈ [L], since we will union bound over all i and ` at last. Let ui ∈ Rm be a fixed vector and W`,r to denote the r-th row of W`, then by the concentration of Gaussian, we have
Pr[σb(〈W`,r, ui〉) > 0] = Pr z∼N (0, 2m ) [z > b] ≤ exp(−b2m/4).
Let S be the following index set S := {r ∈ [m] : 〈W`,r, ui〉 > b}, the above reasoning means that for the indicator random variable 1[r ∈ S], we have
E[1[r ∈ S]] ≤ exp(−b2m/4).
Use Bernstein’s inequality (Lemma A.3) we have that for all t > 0,
Pr[|S| > k + t] ≤ exp(− t 2/2
k + t/3 ),
where k := m · exp(−b2m/4). By picking t = k, we have
Pr[|S| > 2k] ≤ exp(−3k 8 ).
Note that |S| is essentially the quantity ‖hi,`‖0, hence we can union bound over all ` and i and with probability at least
1− nL · exp(−Ω(m · exp(b2m/4))),
we have ‖hi,`‖0 ≤ 2m · exp(−b2m/4).
Remark F.6. The above lemma shows that by using the shifted ReLU activation, we make sure that all hi,` are sparse after initialization. Specifically, we use k := m · exp(−b2m/4) as a sparsity parameter. Later, we might re-scale b so that the probability becomes exp(−b2/2). We stress that such re-scaling does not affect the sparsity of our initial vectors. If we re-scale b and choose it as√
0.4 logm, then k = m0.8 and hence with high probability, ‖hi,`‖0 ≤ O(m0.8). As a direct consequence, we note that all initial Di,` are k-sparse as well.
We state a lemma that handles the `2 norm of hi,` when one uses truncated Gaussian distribution instead. Due to the length and the delicacy of the proof, we defer it to Section G.
Lemma F.7 (Restatement of Lemma G.6). Let b > 0 be a fixed scalar. Let the activation function φ(x) := √ cb1[x > √ 2/mb]x, where cb := (2(1 − Φ(b) + bφ(b)))−1/2. Let ∈ (0, 1), then over the randomness of W (0), with probability at least 1−O(nL) · exp(−Ω(m exp(−b2/2) 2/L2)), we have
‖hi,`‖2 ∈ [1− , 1 + ], ∀i ∈ [n], ` ∈ [L].
We remark that the parameter e−b 2/2m captures the sparsity during the initialization.
The second lemma handles the consecutive product that appears naturally in the gradient computation.
Lemma F.8 (Variant of Lemma 7.3 in Allen-Zhu et al. (2019a)). Suppose m ≥ Ω(nL log(nL)). With probability at least 1− e−Ω(k/L2) over the randomness of initializations W1(0), . . . ,WL(0) ∈ Rm×m, for all i ∈ [n] and 1 ≤ a ≤ b ≤ L,
(a) ‖WbDi,b−1Wb−1 . . . Di,aWa‖ ≤ O( √ L).
(b) ‖WbDi,b−1Wb−1 . . . Di,aWav‖2 ≤ 2‖v‖2, for all vectors with ‖v‖0 ≤ O( mL logm ).
(c) ‖u>WbDi,b−1Wb−1 . . . Di,aWa‖2 ≤ O(1)‖u‖2, for all vectors uwith ‖u‖0 ≤ O( mL logm ).
(d) For any intergers 1 ≤ s ≤ O( mL logm ), with probability at least 1 − e−Ω(s logm) over the randomness of initializations W1(0), . . . ,WL(0) ∈ Rm×m, |u>WbDi,b−1Wb−1 . . . Di,aWav| ≤ ‖u‖2‖v‖2 · O( mL logm ) for all vectors u, v with ‖u‖0, ‖v‖0 ≤ s.
The proof is similarly to the original proof of the corresponding lemma in Allen-Zhu et al. (2019a), however we replace the bound on hi,` with our Lemma F.7. We highlight this does not change the bound, merely in expense of a worse probability.
The next lemma concerns the bound on the product being used in backpropagation. Lemma F.9 (Variant of Lemma 7.4 in Allen-Zhu et al. (2019a)). Suppose m ≥ Ω(nL log(nL)). Then with probability at least 1− e−Ω(log2m), for all i ∈ [n], ` ∈ [L],
‖a>Di,LWL . . . Di,`W`‖2 ≤ O( logm√ m
√ L).
Proof. By Lemma F.8 part (a), we know that ‖WL . . . Di,`W`‖ ≤ O( √ L), consequently, ‖Di,LWL . . . Di,`W`‖ ≤ O( √ L) since ‖Di,L‖ ≤ 1. This means for any vector u ∈ Rm, with probability at least 1− e−Ω(k/L2), we have
‖Di,LWL . . . Di,`W`u‖2 ≤ O( √ L)‖u‖2.
Conditioning on this event and using the randomness of Rademacher vector a, we have
Pr[|a>Di,LWL . . . Di,`W`u| ≤ t · √ L‖u‖2] ≥ 1− 2 exp(− mt2
2 ),
pick t = logm√ m , we know that
Pr[|a>Di,LWL . . . Di,`W`u| ≤ logm√ m · √ L‖u‖2] ≥ 1− 2 exp(− log2m 2 ).
This means with probability at least 1− e−Ω(log2m), we have
‖a>Di,LWL . . . Di,`W`‖2 ≤ O( logm√ m
√ L).
The next several lemmas bound norms after small perturbation. Lemma F.10 (Lemma 8.2 in Allen-Zhu et al. (2019a)). Suppose Assumption F.3 is satisfied with R ≤ O( 1
L9/2 log3m ). With probability at least 1− e−Ω(mR2/3L),
(a) ∆gi,` can be written as ∆gi,` = ∆gi,`,1 + ∆gi,`,2 where ‖∆gi,`,1‖2 ≤ O(RL3/2) and ‖∆gi,`,2‖∞ ≤ O(RL 5/2√logm√ m ).
(b) ‖∆Di,`‖0 ≤ O(mR2/3L) and ‖(∆Di,`)gi,`‖2 ≤ O(RL3/2). (c) ‖∆gi,`‖2, ‖∆hi,`‖2 ≤ O(RL5/2 √
logm). Remark F.11. Lemma F.10 establishes the connection between parameter R and s of Assumption F.3 and F.4. As long as R is small, then we have s = O(mR2/3L). Such a relation enables us to pickR to our advantage and ensure the sparsity of ∆Di,` is sublinear inm, and hence the update time per iteration is subquadratic in m.
Lemma F.12 (Lemma 8.6 in Allen-Zhu et al. (2019a)). Suppose Assumption F.4 is satisfied with 1 ≤ s ≤ O( mL3 logm ) and m ≥ Ω(nL log(nL)), with probability at least 1 − e
−Ω(s logm) over the randomness of W (0), for every i ∈ [n], 1 ≤ a ≤ b ≤ L, we have
(a) ‖Wb(0)(Di,b−1(0) + ∆Di,b−1) . . . (Di,a(0) + ∆Di,a)Wa(0)‖ ≤ O( √ L).
(b) ‖(Wb(0) + ∆Wb)(Di,b−1(0) + ∆Di,b−1) . . . (Di,a(0) + ∆Di,a)(Wa(0) + ∆Wa)‖ ≤ O( √ L) if Assumption F.3 is satisfied with R ≤ O( 1L1.5 ).
Corollary F.13. Suppose Assumption F.3 is satisfied with R ≤ O( 1 L4.5 log3m ) and m ≥ Ω(nL log(nL)), with probability at least 1 − e−Ω(log2m) over the randomness of W (0) and a, for any i ∈ [n] and ` ∈ [L], we have
‖a>(Di,L(0) + ∆Di,L(0))(WL(0) + ∆WL) . . . (W`+1(0) + ∆W`+1)(Di,`(0) + ∆Di,`)‖2 ≤ Õ( √ L/m).
Proof. We first note that by Lemma F.10 and the choice of R, we know that Assumption F.4 is satisfied with 1 ≤ s ≤ O( m
L2 log2m ).
The proceeding proof is identical to Lemma F.9, use P to denote the product (Di,L(0) + ∆Di,L(0))(WL(0) + ∆WL) . . . (W`+1(0) + ∆W`+1)(Di,`(0) + ∆Di,`). Per Lemma F.12, we know that with the conditions stated, we have
‖P‖ ≤ O( √ L),
conditioning on this even, for any vector u ∈ Rm, with probability at least 1−e−Ω(s logm), we know that
‖Pu‖2 ≤ O( √ L)‖u‖2.
Use the randomness of vector a, we have
Pr[|a>Pu| ≤ t · √ L‖u‖2] ≥ 1− 2 exp(− mt2
2 ).
Pick t = logm√ m , we are done.
Lemma F.14 (Variant of Lemma 8.7 in Allen-Zhu et al. (2019a)). Suppose Assumption F.3 is satisfied with R ≤ O( 1
L4.5 log3m ) and m ≥ Ω(nL log(nL)), with probability at least 1 − e−Ω(log2m)
over the randomness of W (0) and a, for all i ∈ [n], ` ∈ [L]. Define
u := a>(Di,L(0) + ∆Di,L)(WL(0) + ∆WL(0)) . . . (W`+1(0) + ∆W`+1)(Di,`(0) + ∆Di,`),
u0 := a >Di,L(0)WL(0) . . .W`+1(0)Di,`(0).
It satisfies that
‖u− u0‖2 ≤ Õ( √ L/m).
Proof. Note that we can invoke Corollary F.13 to give an upper bound on ‖u‖2 with probability 1− e−Ω(log2m),
‖u‖2 ≤ Õ( √ L/m).
Similarly, we can use Lemma F.9 to give an upper bound on ‖u0‖2 with probability 1− e−Ω(log 2m), ‖u0‖2 ≤ Õ( √ L/m).
Finally, by triangle inequality, we get ‖u− u0‖2 ≤ Õ( √ L/m),
with the desired probability.

heading: F.3 BOUNDS ON INITIALIZATION
text: 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 m = Ω(nL log(nL)), then with probability 1−O(nL) · e−Ω(k/L2), we have the following
• f(W,xi) = O(1), ∀i ∈ [n]. • With probability at least 1 − e−Ω(log2m), we have ‖J`,0,i‖ = Õ( √ L/m), ∀` ∈ [L], ∀i ∈
[n].
Remark F.16. The bound of m in Brand et al. (2021) has a linear dependence on n because they need to bound ‖W (0)‖2. If we do not need to bound ‖W (0)‖2, 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 i ∈ [n], we have
f(W,xi) = a >φ(WL(φ(· · ·φ(W1xi)))).
We shall make use Lemma F.7 here: with probability at least 1−O(nL) · exp(−Ω(k/L2)), we have ‖hi,L‖2 ∈ [0.9, 1.1]. Recall that a ∈ Rm has each of its entry being a Rademacher random variable scaled by 1/m, hence it’s 1/m-subgaussian. Using the concentration of subgaussian (Lemma A.5), we have that
Pr[|f(W,xi)| ≥ t] ≤ 2 exp(− t2m
2‖hi,L‖2 ),
finally by noticing that t = Θ(1) and ‖hi,L‖2 ∈ [0.9, 1.1], we conclude that
Pr[|f(W,xi)| ≥ O(1)] ≤ 2 exp(−Ω(m)),
union bounding over all i, we conclude that
Pr[f(W,xi) ≥ O(1)] ≤ O(n) · exp(−Ω(m)),∀i ∈ [n].
Part 2: We will combine Lemma F.7 and F.9, with probability at least 1−O(nL) · e−Ω(k/L2), we have ‖hi,`‖2 ∈ [0.9, 1.1] and with probability at least 1−e−Ω(log
2m), we have ‖a>Di,L . . .W`‖2 ≤ Õ( √ L/m). Combine them together, we know that with probability at least 1 − e−Ω(log2m),
‖J`,0,i‖ ≤ Õ( √ L/m).
Remark F.17. By utilizing the structure of vector a, we can show that with a weaker probability (1 − e−Ω(log2m)), f(W,xi) has even smaller magnitude (Õ( 1√m )). For our purpose, f(W,xi) = O(1) suffices.

heading: F.4 BOUNDS ON SMALL PERTURBATION
text: In the following, we generalize the Lemma C.4 in Brand et al. (2021) into multiple layer neural network. For the simplicity of notation, we set all m` to be m. Lemma F.18 (multiple layer version of Lemma C.4 in Brand et al. (2021)). Suppose R ≤ O( 1
L4.5 log3m ) and m = Ω(nL log(nL)). With probability at least 1− e−Ω(log2m) over the random initialization of W (0) = {W1(0),W2(0), · · ·WL(0)}, the following holds for any set of weights W1, · · · ,WL satisfying
‖W` −W`(0)‖ ≤ R,∀` ∈ [L]
• ‖JW`,xi − JW`(0),xi‖2 = Õ( √ L/m).
• ‖JW` − JW`(0)‖F = Õ(n1/2 √ L/m).
• ‖JW`‖F = Õ(n1/2 √ L/m).
Proof. Part 1. To simplify the notation, we ignore the subscripts i below.
Consider the following computation:
‖JW`,xi − JW`(0),xi‖2 = ‖uh > `−1 − u0h>`−1(0)‖F
= ‖u(h`−1(0) + ∆h`−1)> − u0h>`−1(0)‖F = ‖(u− u0)h>`−1(0) + u(∆h>`−1)‖F ≤ ‖(u− u0)h>`−1(0)‖F + ‖u(∆h>`−1)‖F ≤ ‖u− u0‖2 · ‖h`−1(0)‖2 + ‖u‖2 · ‖∆h`−1‖2.
where
u := (D`(0) + ∆D`)
( L∏
k=`+1
(Wk(0) + ∆Wk) >(Dk(0) + ∆Dk) ) a
u0 := D`(0)
( L∏
k=`+1
Wk(0) >Dk(0) ) a
By Lemma F.14, we know that with probability at least 1− e−Ω(log2m), we have ‖u− u0‖2 ≤ Õ( √ L/m).
By Lemma F.7, we know that with probability at least 1−O(nL) · e−k/L2 , we have ‖h`−1(0)‖2 ∈ [0.9, 1.1].
By Corollary F.13, we know that with probability at least 1− e−Ω(log2m), we have ‖u‖2 ≤ Õ( √ L/m).
By Lemma F.10, we know that with probability at least 1− e−Ω(mR2/3L), we have
‖∆h`−1‖2 ≤ Õ(RL2.5).
Note that due to the choice of R, we know that
‖∆h`−1‖2 ≤ 1.
Taking a union bound over all events, with probability at least 1 − e−Ω(log2m), we achieve the following bound:
‖JW`,xi − JW`(0),xi‖2 ≤ Õ( √ L/m).
Part 2. Note that the squared Frobenious norm is just the sum of all squared `2 norm of rows, hence ‖JW` − JW`(0)‖F ≤ Õ(n 1/2 √ L/m).
Part 3. We will prove by triangle inequality:
‖JW`‖F ≤ ‖JW`(0)‖F + ‖JW` − JW`(0)‖F ≤ Õ(n1/2 √ L/m) + Õ(n1/2 √ L/m)
= Õ(n1/2 √ L/m).

heading: F.5 PUTTING IT ALL TOGETHER
text: 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 satisfies m = Ω(λ−2L n
2L2), then with probability at least 1 − e−Ω(log2m) over the randomness of the initialization of the neural network and the randomness of the algorithm, Algorithm 2 satisfies
‖ft+1 − y‖2 ≤ 1
2 ‖ft − y‖2.
Before moving on, we introduce several definitions and prove some useful facts related to them.
Definition F.20 (function J). We define
J`(Z1, . . . , ZL)i := Di,`(Z`) L∏
k=`+1
Z>k Di,k(Zk)a(hi(Z1, . . . , Z`−1)) > ∈ Rm`×m`−1
where
Di,`(Z`) := diag(φ ′(Z`hi(Z1, . . . , Z`−1))), ∈ Rm`×m`
hi(Z1, . . . , Z`−1) := φ(Z`−1(φ(Z`−2 · · · (φ(Z1xi))))) ∈ Rm`−1
Fact F.21. Let J` denote the function be defined as Definition F.20. For any t ∈ {0, . . . , T}, we have
ft+1 − ft = L∑ `=1 (∫ 1 0 J`((1− s)W (t) + sW (t+ 1))ds )> · vec(W`(t+ 1)−W`(t)),
Proof. For i ∈ [n], consider the i-th coordinate. (ft+1 − ft)i = ∫ 1
0
f((1− s)W (t) + sW (t+ 1), xi)′ds
= ∫ 1 0 L∑ `=1 ( ∂f ∂W` ((1− s)W (t) + sW (t+ 1), xi) )> · vec(W`(t+ 1)−W`(t))ds
= L∑ `=1 (∫ 1 0 J`((1− s)W (t) + sW (t+ 1))ids )> · vec(W`(t+ 1)−W`(t)),
Thus, we complete the proof.
Fact F.22. For any t ∈ {0, . . . , T}, we have J`(W1(t), . . . ,WL(t)) = J`,t.
Proof. In order to simplify the notation, we drop the term t below.
We note that for i ∈ [n], the i-th row of matrix J`,t is defined as
Di,`( L∏ k=`+1 W>k Di,k)ah > i,`−1,
where Di,` = diag(φ′(W`hi,`−1)) and hi,`−1 = φ(W`−1(φ(W`−2 . . . (φ(W1xi))))), this is essentially the same as hi(W1, . . . ,W`−1) and Di,`(W`). This completes the proof.
We state the range we require for parameter R:
Definition F.23. We choose R so that n√ mL · 1 λL ≤ R ≤ 1 L4.5 log3m .
Remark F.24. Recall that the sparsity parameter s is directly related toR: s = O(mR2/3L), hence to ensure the sparsity is small, we shall pick R as small as possible. Specifically, if we pick R to be n√ mLλL , then s = O(λ−1L n √ mL), as long as m λ−2L n2L then s ≈ O( √ m).
Next, we pick the value of m: Definition F.25. We choose m to be
m ≥ Ω(λ−10/3L n 10/3L5/3).
Remark F.26. The choice of m here makes sure that, as long as we pick R matching its lower bound, then the sparsity s = O(m0.8), this matches the other sparsity parameter k, which is also in the order of O(m0.8).
We use induction to prove the following two claims recursively. Definition F.27 (Induction hypothesis 1). Let t ∈ [T ] be a fixed integer. We have
‖W`(t)−W`(0)‖ ≤ R holds for any ` ∈ [L]. Definition F.28 (Induction Hypothesis 2). Let t ∈ [T ] be a fixed integer. We have
‖ft − y‖2 ≤ 1
2 ‖ft−1 − y‖2.
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, we are going to prove it first and we define
J`,t,t+1 := ∫ 1 0 J` ( (1− s)Wt + sWt+1 ) ds,
where J` is defined as Definition F.20. Lemma F.29. Let g?` := (J`,tJ>`,t)−1 · 1L (ft − y). We have
‖ft+1 − y‖2 ≤ ‖ft − y − L∑ `=1 J`,tJ > `,tg`,t‖2
+ L∑ `=1 ‖(J`,t − J`,t,t+1)J>`,tg?` ‖2
+ L∑ `=1 ‖(J`,t − J`,t,t+1)J>`,t(g`,t − g?` )‖2. (8)
Proof. Consider the following computation: ‖ft+1 − y‖2
= ‖ft − y + (ft+1 − ft)‖2
= ‖ft − y + L∑ `=1 J`,t,t+1 · vec(W`,t+1 −W`,t)‖2
= ‖ft − y − L∑ `=1 J`,t,t+1 · J>`,tg`,t‖2
= ‖ft − y − L∑ `=1 J`,tJ > `,tg`,t + L∑ `=1 J`,tJ > `,tg`,t − L∑ `=1 J`,t,t+1J > `,tg`,t‖2
≤ ‖ft − y − L∑ `=1 J`,tJ > `,tg`,t‖2 + L∑ `=1 ‖(J`,t − J`,t,t+1)J>`,tg`,t‖2
≤ ‖ft − y − L∑ `=1 J`,tJ > `,tg`,t‖2 + L∑ `=1 ‖(J`,t − J`,t,t+1)J>`,tg?` ‖2 + L∑ `=1 ‖(J`,t − J`,t,t+1)J>`,t(g`,t − g?` )‖2,
The second step follows from the definition of J`,t,t+1 and simple calculus.
Claim F.30 (1st term in Eq. (8)). We have
‖(ft − y)− L∑ `=1 J`,tJ > `,tg`,t‖2 ≤ 1 6 ‖ft − y‖2.
Proof. We have
‖(ft − y)− L∑ `=1 J`,tJ > `,tg`,t‖2 = ‖ L∑ `=1 ( 1 L (ft − y)− J`,tJ>`,tg`,t)‖2
≤ L∑ `=1 ‖ 1 L · (ft − y)− J`,tJ>`,tg`,t‖2
≤ 1 6 ‖ft − y‖2, (9)
since g`,t is an 0 ( 0 ≤ 16 ) approximate solution to regression problem
min g ‖J`,tJ>`,tg −
1 L (ft − y)‖2.
Claim F.31 (2nd term in Eq. (8)). We have L∑ `=1 ‖(J`,t − J`,t,t+1)J>`,tg?` ‖2 ≤ 1 6 ‖ft − y‖2.
Proof. For the second term in Eq. (8), for any ` ∈ [L], we have
‖(J`,t − J`,t,t+1)J>`,tg?` ‖2 ≤ ‖J`,t − J`,t,t+1‖ · ‖J>`,tg?` ‖2
= ‖J`,t − J`,t,t+1‖ · ‖J>`,t(J`,tJ>`,t)−1 · 1
L (ft − y)‖2
≤ 1 L · ‖J`,t − J`,t,t+1‖ · ‖J>`,t(J`,tJ>`,t)−1‖ · ‖ft − y‖2. (10)
We bound these term separately. First,
‖J`,t − J`,t,t+1‖ = ∥∥∥∥J`(Wt)− ∫ 1
0
J`((1− s)Wt + sWt+1)ds ∥∥∥∥
≤ ∫ 1
0
‖J`(Wt)− J`((1− s)Wt + sWt+1)‖ ds
≤ ∫ 1
0
‖J`(Wt)− J`(W0)‖+ ‖J`(W0)− J`((1− s)Wt + sWt+1)‖ ds
≤ ‖J`(Wt)− J`(W0)‖+ ∫ 1
0
‖J`(W0)− J`((1− s)Wt + sWt+1)‖ ds
≤ Õ(n1/2 √ L/m), (11)
where by Fact F.22, we know ‖J`(Wt)−J`(W0)‖ = ‖JW`(t)−JW`(0)‖ ≤ Õ(n1/2 √ L/m), the last inequality is by Lemma F.18. For the second term, we have
‖(1− s) · vec(W`(t)) + s · vec(W`(t+ 1))− vec(W`(0))‖2 ≤ (1− s) · ‖vec(W`(t))− vec(W`(0))‖2 + s · ‖vec(W`(t+ 1))− vec(W`(0))‖2 = (1− s) · ‖W`(t)−W`(0)‖F + s · ‖W`(t+ 1)−W`(0)‖F ≤ O(R).
This means the perturbation of (1− s)W`(t) + sW`(t+ 1) with respect to W`(0) is small, for any ` ∈ [L], hence ‖J`(W0)− J`((1− s)Wt + sWt+1)‖ = Õ(n1/2 √ L/m).
Furthermore, we have
‖J>`,t(J`,tJ>`,t)−1‖ = 1 σmin(J>`,t) ≤ √ 2/λL, (12)
where the second inequality follows from σmin(J`,t) = √ λmin(J`,tJ>`,t) ≥ √ λL/2 (see Lemma E.3).
Combining Eq. (10), (11) and (12), we have L∑ `=1 ‖(J`,t − J`,t,t+1)J>`,tg?‖2 ≤ Õ( √ nL/m) · λ−1/2L · ‖ft − y‖2
≤ 1 6 ‖ft − y‖2, (13)
where the last step follows from choice of m (Definition F.25).
Claim F.32 (3rd term in Eq. (8)). We have L∑ `=1 ‖(J`,t − J`,t,t+1)J>`,t(g`,t − g?` )‖2 ≤ 1 6 ‖ft − y‖2
Proof. We can show
‖(J`,t − J`,t,t+1)J>`,t(g`,t − g?` )‖2 ≤ ‖J`,t − J`,t,t+1‖ · ‖J>`,t‖ · ‖g`,t − g?` ‖2. (14)
Moreover, one has λL 2 ‖g`,t − g?` ‖2 ≤ λmin(J`,tJ>`,t) · ‖g`,t − g?` ‖2
≤ ‖J`,tJ>`,tg`,t − J`,tJ>`,tg?` ‖2
= ‖J`,tJ>`,tg`,t − 1
L (ft − y)‖2 ≤ √ λL/n
L · ‖ft − y‖2. (15)
The first step comes from λmin(J`,tJ>`,t) = λmin(G`,t) ≥ λL/2 (see Lemma E.3) and the last step comes from gt is an 0 approximate solution. The fourth step follows from Eq. (15) and the fact that ‖(J`,tJ>`,t)−1‖ ≤ 2/λL. The last step follows from gt is an 0 ( 0 ≤ √ λL/n) approximate solution to the regression.
Consequently, we have
‖(J`,t − J`,t,t+1)J>`,t(g`,t − g?` )‖2 ≤ ‖J`,t − J`,t,t+1‖ · ‖J>`,t‖ · ‖g`,t − g?` ‖2 ≤ Õ(n1/2 √ L/m) · Õ(n1/2 √ L/m) · 2
L √ nλL · ‖ft − y‖2
= Õ(
√ n
m ) 2√ λL · ‖ft − y‖2. (16)
The second step follows from Eq. (11) and (15) and the fact that ‖J`,t‖ ≤ Õ( √ nL/m) (see Lemma F.18).
Finally, we have L∑ `=1 ‖(J`,t − J`,t,t+1)J>`,t(g`,t − g?` )‖2 ≤ Õ( √ nL m ) 2√ λL · ‖ft − y‖2
≤ 1 6 ‖ft − y‖2.
The last step follows from choice of m (Definition F.25).
Lemma F.33 (Putting it all together). We have
‖ft+1 − y‖2 ≤ 1
2 ‖ft − y‖2. (17)
Proof. Combining Eq. (8), (9), (13), and (??), we have proved the second claim, i.e.,
‖ft+1 − y‖2 ≤ 1
2 ‖ft − y‖2.

heading: F.6 WEIGHTS ARE NOT MOVING TOO FAR
text: Lemma F.34. Let R be chosen as in Definition F.23, then the following holds: ‖W`(t+ 1)−W`(0)‖ ≤ R, ∀` ∈ [L].
Proof. It remains to show that Wt does not move far away from W0. First, we have ‖g`,t‖2 ≤ ‖g?` ‖2 + ‖g`,t − g?` ‖2
= 1
L ‖(J`,tJ>`,t)−1(ft − y)‖2 + ‖g`,t − g?` ‖2
≤ 1 L ‖(J`,tJ>`,t)−1‖ · ‖(ft − y)‖2 + ‖g`,t − g?` ‖2 ≤ 2 LλL · ‖ft − y‖2 + 2 L √ nλL · ‖ft − y‖2
. 1
LλL · ‖ft − y‖2 (18)
where the third step follows from Eq. (15) and the last step follows from the obvious fact that 1/ √ nλL ≤ 1/λL.
Then ‖W`(k + 1)−W`(k)‖ = ‖J>`,kg`,k‖
≤ ‖J`,k‖ · ‖g`,k‖2 ≤ Õ( √ nL/m) · 1
LλL · ‖fk − y‖2
≤ Õ( √ n/Lm) · 1
2kλL · ‖f0 − y‖2
≤ Õ( n√ Lm ) · 1 2kλL .
The third step uses the fact that ‖J`,k‖ ≤ ‖J`,k‖F ≤ O( √ n Lm ) by Lemma F.18, and the last step uses the fact that both ‖f0‖2 and ‖y‖2 are in the order of O(1). Consequently, we have
‖W`(t+ 1)−W`(0)‖ ≤ t∑
k=0
‖W`(k + 1)−W`(k)‖
≤ t∑
k=0
Õ( n√ Lm ) · 1 λL 1 2k
≤ Õ( n√ Lm ) · 1 λL .
By the choice of R (Definition F.23), we know this is upper bounded by R. This concludes our proof.

heading: G PROOF OF LEMMA F.7
text: In this section, we prove a technical lemma (Lemma F.7) involving truncated gaussian distribution, which correlates to the shifted ReLU activation we use. Definition G.1 (Truncated Gaussian distribution). SupposeX ∼ N (0, σ2). Let b ∈ R. Then, we say a random variable Y follows from a truncated Gaussian distribution Nb(0, σ2) if Y = X|X ≥ b. The probability density function for Nb(0, σ2) is as follows:
f(y) = 1 σ(1− Φ(b/σ)) · 1√ 2π e−y 2/(2σ2) y ∈ [b,∞),
where Φ(·) is the standard Gaussian distribution’s CDF. Fact G.2 (Properties of truncated Gaussian distribution). For b ∈ R, suppose X ∼ Nb(0, σ2). Let β := b/σ. Then, we have
• E[X] = σφ(β)1−Φ(β) , where φ(x) := 1√ 2π e−x 2/2.
• Var[X] = σ2(1 + βφ(β)/(1− Φ(β))− (φ(β)/(1− Φ(β)))2).
• X/σ ∼ Nb/σ(0, 1).
• When σ = 1, X is C(b+ 1)-subgaussian, where C > 0 is an absolute constant. Lemma G.3 (Concentration inequality for b-truncated chi-square distribution). For b ∈ R, n > 0, let X ∼ χ2b,n; that is, X = ∑n i=1 Y 2 i where Y1, . . . , Yn ∼ Nb(0, 1) are independent b-truncated Gaussian random variables. Then, there exist two constants C1, C2 such that for any t > 0,
Pr [∣∣∣∣X − n(1 + bφ(b)1− Φ(b) ) ∣∣∣∣ ≥ nt] ≤ exp (−C1nt2/b4)+ exp (−C2nt/b2) .
In particular, we have Pr [|X − n(1 + b(b+ 1))| ≥ t] ≤ exp ( −C1t2/(nb4) ) + exp ( −C2t/b2 ) .
Proof. Since we know that Yi ∼ Nb(0, 1) is C(b + 1)-subgaussian, it implies that Y 2i is a subexponential random variable with parameters (4 √ 2C2(b+1)2, 4C2(b+1)2). Hence, by the standard concentration of sub-exponential random variables, we have
Pr [∣∣∣∣∣ n∑ i=1 Y 2i − nE[Y 2i ] ∣∣∣∣∣ ≥ nt ] ≤ 2 exp ( − nt 2 2·32C4(b+1)4 ) if nt ≤ 8C2(b+ 1)2 2 exp ( − nt2·4C2(b+1)2 ) otherwise
≤ 2 exp ( −C1nt2/b4 ) + 2 exp ( −C2nt/b2 ) .
Fact G.4. Let h ∈ Rp be fixed vectors and h 6= 0, let b > 0 be a fixed scalar, W ∈ Rm×p be random matrix with i.i.d. entries Wi,j ∼ N (0, 2m ) and vector v ∈ R
m defined as vi = φ((Wh)i) = 1[(Wh)i ≥ b](Wh)i. Then
• |vi| follows i.i.d. from the following distribution: with probability 1−e−b 2m/(4‖h‖2), |vi| =
0, and with probability e−b 2m/(4‖h‖2), |vi| follows from truncated Gaussian distribution Nb(0, 2m‖h‖ 2 2).
• m‖v‖ 2 2
2‖h‖22 is in distribution identical to χ2b′,ω (b ′-truncated chi-square distribution of order ω)
where ω follows from binomial distribution B(m, e−b2m/(4‖h‖2)) and b′ = √ m/2
‖h‖2 b.
Proof. We assume each vector Wi is generated by first generating a gaussian vector g ∼ N (0, 2mI) and then setting Wi = ±g where the sign is chosen with half-half probability. Now, |〈Wi, h〉| = |〈g, h〉| only depends on g, and is in distribution identical to Nb(0, 2m‖h‖ 2 2). Next, after the sign is
determined, the indicator 1[(Wih)i ≥ b] is 1 with probability e−b 2m/(4‖h‖2) and 0 with probability 1−e−b2m/(4‖h‖2). Therefore, |vi| satisfies the aforementioned distribution. As for ‖v‖22, letting ω ∈ {0, 1, . . . ,m} be the variable indicates how many indicators are 1, then ω ∼ B(m, e−b2m/(4‖h‖2)) and m‖v‖ 2 2
2‖h‖22 ∼ χ2b′,ω , where b′ =
√ m/2
‖h‖2 b.
Fact G.5 (Gaussian tail bound). For any b > 0, we have
e−b 2/2
C(b+ 1) ≤ 1− Φ(b) ≤ e−b 2/2,
where C is an absolute constant.
We prove a truncated Gaussian version of Lemma 7.1 of Allen-Zhu et al. (2019a). Lemma G.6. Let b > 0 be a fixed scalar. Let the activation function be defined as φ(x) :=√ cb1[x > √ 2/mb]x, where cb := (2(1 − Φ(b) + bφ(b)))−1/2. Let ∈ (0, 1), then over the randomness of W (0), with probability at least 1 − O(nL) · exp(−Ω(m exp(−b2/2) 2/L2)), we have
‖hi,`‖2 ∈ [1− , 1 + ], ∀i ∈ [n], ` ∈ [L].
Proof. We only prove for a fixed i ∈ [n] and ` ∈ {0, 1, 2, . . . , L} because we can apply union bound at the end. Below, we drop the subscript i for notational convenience, and write hi,` and xi as h` and x respectively.
According to Fact G.4, fixing any h`−1 6= 0 and letting W` be the only source of randomness, we have
m 2 ‖h`‖22 ∼ χ2b/‖h‖2,ω, with ω ∼ B(m, 1− Φ(b ′)),
where b′ := b/‖h`−1‖2. We first consider the ` = 1 case. Then, we have ‖h`−1‖2 = 1, and b′ = b. Let Pb := 1− Φ(b). By Chernoff bound, for any δ ∈ (0, 1), we have
Pr[ω ∈ (1± δ)mPb] ≥ 1− exp(−Ω(δ2Pbm)).
In the following proof, we condition on this event. By Fact G.5,
ω ∈ (1± δ)Pbm ⇐⇒ ω ∈ [ (1− δ) e −b2/2
C(b+ 1) m, (1 + δ) exp(−b2/2)m
] .
By Lemma G.3, we have
Pr [∣∣∣∣m2 ‖h1‖22 − ω ( 1 + bφ(b)
Pb )∣∣∣∣ > t] ≤ exp (−Ω(t2/(ωb4)))+ exp (−Ω(t/b2)) Note that
ω ( 1 + bφ(b)
Pb
) ∈ (1± δ)mPb + (1± δ)mPb · bφ(b)
Pb = (1± δ)(Pb + bφ(b)) ·m.
Let c−1b := 2(Pb + bφ(b)) be the normalization constant. Then, we have Pr[|cb‖h1‖22 − (1± δ)| > 2tcb/m] ≤ exp ( −Ω(t2/(ωb4)) ) + exp ( −Ω(t/b2) ) .
We want 2tcb/m = O(δ), i.e., t = O(δc−1b m). Then, we have ωt = m Ω(1) > b2. Hence, by Lemma G.3, we actually have
Pr[|cb‖h1‖22 − (1± δ)| > O(δ)] ≤ exp ( −Ω(δm/(cbb2)) ) .
By taking δ = /L, we get that
‖h1‖22 ∈ [1− /L, 1 + /L]
with probability at least 1− exp(−Ω( 2Pbm/L2))− exp ( −Ω( m/(cbb2L)) ) ≥ 1− exp(−Ω( 2Pbm/L2)),
where the last step follows from 1cbb2 = Pb+bφ(b) b2 = Θ(Pb).
We can inductively prove the ` > 1 case. Since the blowup of the norm of h1 is 1 ± /L, the concentration bound is roughly the same for h` for ` ≥ 2. Thus, by carefully choosing the parameters, we can achieve ‖h`‖22 ∈ [(1− /L)`, (1 + /L)`] with high probability. In this end, by a union bound over all the layers ` ∈ [L] and all the input data i ∈ [n], we get that
‖hi,`‖2 ∈ [1− , 1 + ]
with probability at least
1−O(nL) exp(−Ω( 2Pbm/L2)),
which completes the proof of the lemma.

Review summaries:
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. 2. 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.  3. 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. 

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.

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.