Original: T RAINING M ULTI -L AYER O VER -P ARAMETRIZED N EURAL N ETWORK IN S UBQUADRATIC T IME
Revised: T RAINING M ULTI -L AYER O VER -P ARAMETRIZED N EURAL N ETWORK IN S UBQUADRATIC T IME

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: In this work, we initiate the study of the performance of second-order algorithm on multiple-hidden layers over-parametrized neural network.
Revised: In this work, we initiate the study of the performance of second-order algorithm on multiple-hidden layers over-parametrized neural network.

Original: We propose a novel algorithm to train such network, in time subquadratic in the width of the neural network.
Revised: We propose a novel algorithm to train such network, in time subquadratic in the width of the neural network.

Original: Our algorithm combines the GramGauss-Newton method, tensor-based sketching techniques and preconditioning.
Revised: Our algorithm combines the GramGauss-Newton method, tensor-based sketching techniques and preconditioning.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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;
Revised: 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;

Original: 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.
Revised: 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.

Original: Du et al. (2019b); Allen-Zhu et al.
Revised: Du et al. (2019b); Allen-Zhu et al.

Original: (2019a;b);
Revised: (2019a;b);

Original: Du et al. (2019a); Song & Yang (2019);
Revised: Du et al. (2019a); Song & Yang (2019);

Original: Brand et al.
Revised: Brand et al.

Original: In order to analyze the dynamics of neural networks and obtain provable guarantees, using over-parametrization has been a growing trend. However, due to the complexity of analysis, most theoretical works focus on training one-hidden layer network.
Revised: ; Zou et al.

Original: This limits their contributions, since in practice, neural networks are often deep.
Revised: (2018); Cao & Gu (2019); Lee et al. (2019a); Liu et al.

Original: Such a lack of understanding has been further exacerbated when one tries to design efﬁcient algorithms to train such networks.
Revised: Chen et al. In order to analyze the dynamics of neural networks and obtain provable guarantees, using over-parametrization has been a growing trend.

Original: In terms of understanding the convergence behavior of over-parametrized networks, most of the attentions have been directed to the study of ﬁrst-order method such as gradient descent or stochastic gradient descent.
Revised: In terms of understanding the convergence behavior of over-parametrized networks, most of the attentions have been directed to the study of ﬁrst-order method such as gradient descent or stochastic gradient descent.

Original: The widespread use of ﬁrst-order method is explained, to a large degree, by its computational efﬁciency, 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.
Revised: The widespread use of ﬁrst-order method is explained, to a large degree, by its computational efﬁciency, 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.

Original: One of the major drawbacks of ﬁrst-order methods is their convergence rate is typically slow in  non-convex setting ( poly( n, L, log(1 /(cid:15) )) iterations) as shown in Allen-Zhu et al.(2019a)), which is often the case of a deep over-parameterized neural network.
Revised: One of the major drawbacks of ﬁrst-order methods is their convergence rate is typically slow in many non-convex settings ( poly( n, L, log(1 /(cid:15) )) iterations, where n is the number of training samples, L is the number of layers and (cid:15) is the precision of training), e.g., deep neural network with ReLU activation, as shown in Allen-Zhu et al.

Original: It is increasingly evident that ﬁrstorder methods are becoming a real bottleneck for many practical applications.
Revised: (2019a)), which is often the case of a deep over-parameterized neural network.

Original: Second-order method (which employs the information of the Hessian matrix) on the other hand enjoys a much faster convergence rate (only log(1 /(cid:15) )
Revised: Second-order method (which employs the information of the Hessian matrix) on the other hand enjoys a much faster convergence rate (only log(1 /(cid:15) )

Original: iterations (Zhang et al., 2019), but not poly( n ) log(1 /(cid:15) ) iterations) and exploits the local geometry of the loss function to overcome the pathological curvature issues that are critical in ﬁrst-order method.
Revised: iterations (Zhang et al., 2019), but not poly( n ) log(1 /(cid:15) ) iterations) and exploits the local geometry of the loss function to overcome the pathological curvature issues that are critical in ﬁrst-order method.

Original: Another clear advantage of second-order method over ﬁrst-order method is it does not require the tuning of learning rate.
Revised: Another clear advantage of second-order method over ﬁrst-order method is it does not require the tuning of learning rate.

Original: 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.
Revised: 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.

Original: Given any weight matrix ofintermediate layer of size m × m , its Hessian matrix has size m 2 × m 2 , which makes any naive implementation of second-order algorithm takes at least O ( m 4 ) time since one needs to write down the Hessian.
Revised: Given any weight matrix of size m × m ( m is the width of network), its Hessian matrix has size m 2 × m 2 , which makes any naive implementation of second-order algorithm takes at least O ( m 4 ) time since one needs to write down the Hessian.

Original: 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);
Revised: 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);

Original: Lee et al. (2015); Cohen et al.

Original: ; Lee et al.
Revised: Lee et al.

Original: ; Jiang et al.

Original: (2020b;a)).

Original: Recent works (Cai et al.
Revised: Recent works (Cai et al.

Original: ; Zhang et al.
Revised: ; Zhang et al.

Original: (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.
Revised: (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.

Original: Let m denote the width of networks and n denote the number of samples, they achieve a running time of O ( mn 2 ) per iteration.
Revised: Speciﬁcally, they achieve a running time of O ( mn 2 ) per iteration.

Original: Their methods are essentially variants of Gauss-Newton method combined with using neural tangent kernels (Jacot et al.
Revised: Their methods are essentially variants of Gauss-Newton method combined with using neural tangent kernels (Jacot et al.

Original: to prove the convergence.
Revised: to prove the convergence.

Original: By using the idea of randomized linear algebra (Clarkson & Woodruff (2013); Woodruff (2014)) anduse a clever sketching matrix as a preconditioner, Brand et al.
Revised: By using the idea of randomized linear algebra (Clarkson & Woodruff (2013); Woodruff (2014)) and a clever sketching matrix as a preconditioner, Brand et al.

Original: (2021) further improves this running time to (cid:101) O
Revised: (2021) further improves the running time to (cid:101) O

Original: ( mn ) per iteration.
Revised: ( mn ) per iteration.

Original: However, all of these algorithms are for training an over-parametrized neural network with onehidden 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.
Revised: 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.

Original: In the seminal work of Allen-Zhu et al. (2019a), they showed that as long as the networks are over-parametrized, ﬁrst-order methods such as gradient descent and stochastic gradient descent won’t encounter such problem.
Revised: In the seminal work of Allen-Zhu et al. (2019a), they showed that as long as the networks are over-parametrized, ﬁrst-order methods such as gradient descent and stochastic gradient descent won’t encounter such problem.

Original: But does this still hold for second-order method?
Revised: But does this still hold for second-order method?

Original: Can we provably show that second-order method has a good performance when training deep over-parametrized networks?
Revised: Can we provably show that second-order method has a good performance when training deep over-parametrized networks?

Original: Second, even the fastest of them (Brand et al.
Revised: Second, even the fastest of them (Brand et al.

Original: (2021)) would incur a running time of (cid:101) O ( m 2 nL ) per iteration, which seems unavoidable due to the size of intermediate weight matrices is m × m .
Revised: (2021)) would incur a running time of (cid:101) O ( m 2 nL ) per iteration, which seems unavoidable due to the size of intermediate weight matrices is m × m .

Original: In this work, we take the ﬁrst 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.
Revised: In this work, we take the ﬁrst 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.

Original: 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 ≥ n 8 L 12 , (Zou & Gu, 2019)).
Revised: 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 ≥ n 8 L 12 , (Zou & Gu, 2019)).

Original: Our work can be decomposed into two parts: algorithmically and analytically.
Revised: Our work can be decomposed into two parts: algorithmically and analytically.

Original: From an algorithmic perspective, our method builds upon a variant of Gauss-Newton method (Björck (1996)) called
Revised: From an algorithmic perspective, our method builds upon a variant of Gauss-Newton method (Björck (1996)) called

Original: Gram-Gauss-Newton method (Cai et al.
Revised: Gram-Gauss-Newton method (Cai et al.

Original: ; Brand et al.
Revised: ; Brand et al.

Original: 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).
Revised: 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).

Original: The gradient is low rank (rank n ).
Revised: The gradient is low rank (rank n ).

Original: The gradient can be formulated as the outer product of two vectors.
Revised: The gradient can be formulated as the outer product of two vectors.

Original: From an analytical perspective, our work is inspired by Allen-Zhu et al.
Revised: From an analytical perspective, our work is inspired by Allen-Zhu et al.

Original: (2019a).
Revised: (2019a).

Original: 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.
Revised: 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.

Original: (2019a))
Revised: (2019a))

Original: and establish a connection between a
Revised: and establish a connection between a

Original: Gram matrix we compute at each iteration and the NTK matrix.
Revised: Gram matrix we compute at each iteration and the NTK matrix.

Revised: Our Contributions.

Revised: We summarize our technical contributions below.

Original: 1. O UR R ESULT
Revised: • 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 (cid:101) o ( m 2 ) .

Revised: Our algorithm makes use of Gram-Gauss-Newton method, tensor-based sketching techniques, and data structures that maintains a low rank representation efﬁciently.

Revised: • By combining fast tensor algebra techniques and sketching-based preconditioning, we devise an algorithm to efﬁciently solve a regression problem where the involved matrix has its rows being tensor product of vectors.

Revised: 1. O UR R ESULT

Original: 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 efﬁcient 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.
Revised: 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 efﬁcient 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.

Original: 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.
Revised: 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.

Original: We use f t ∈ R n to denote the prediction of neural network at time t .
Revised: We use f t ∈ R n to denote the prediction of neural network at time t .

Original: Our ﬁrst theorem demonstrates the fast convergence rate of our algorithm.
Revised: Our ﬁrst theorem demonstrates the fast convergence rate of our algorithm.

Original: Theorem 1.1 (Convergence, informal version of Theorem F.19) .
Revised: Theorem 1.1 (Convergence, informal version of Theorem F.19) .

Original: Suppose the width of the neural network satisﬁes 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 2 m ) , we have
Revised: Suppose the width of the neural network satisﬁes 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 2 m ) , we have

Original: [Equation]
Revised: [Equation]

Original: where f t ∈ R n is the the prediction produced by neural network at time t .
Revised: where f t ∈ R n is the the prediction produced by neural network at time t .

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: The next theorem concerns the cost per iteration of our second-order algorithm.
Revised: The next theorem concerns the cost per iteration of our second-order algorithm.

Original: Theorem 1.2 (Runtime, informal version of Theorem B.1) .
Revised: Theorem 1.2 (Runtime, informal version of Theorem B.1) .

Original: There exists a randomized algorithm (Algorithm 1) that trains a multi-layer neural network of width m with the cost per training iteration being
Revised: There exists a randomized algorithm (Algorithm 1) that trains a multi-layer neural network of width m with the cost per training iteration being

Original: O ( m 2 − Ω(1) ) .
Revised: [Equation]

Original: We improve the overall training time of multi-layer over-parametrized networks with second-order method from T init + T · O ( m 2 ) to T init + T · o ( m 2 ) , where T init is the initialization time of training, typically takes O ( m 2 ) .
Revised: We improve the overall training time of multi-layer over-parametrized networks with second-order method from T init + T · O ( m 2 ) to T init + T · o ( m 2 ) , where T init is the initialization time of training, typically takes O ( m 2 ) .

Original: As we have argued before, multi-layer over-parametrized networks require m to be in the order of n 8 , hence improving the cost per iteration from quadratic to subquadratic is an important gain in speeding up training.
Revised: As we have argued before, multi-layer over-parametrized networks require m to be in the order of n 8 , hence improving the cost per iteration from quadratic to subquadratic is an important gain in speeding up training.

Original: Its advantage is even more evident when one seeks a high precision solution , and hence the number of iterations T is large.
Revised: Its advantage is even more evident when one seeks a high precision solution , and hence the number of iterations T is large.

Original: 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 ( m 2 ) time!
Revised: 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 ( m 2 ) time!

Original: Also, the 
Revised: This means that even for ﬁrst-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.

Original:  Jacobian matrixis of size n × m 2 ,  forming it naively will cost O ( nm 2 ) time, which is prohibitively large.
Revised: In our case, we have also a Jacobian matrix of size n × m 2 , so forming it naively will cost O ( nm 2 ) time, which is prohibitively large.

Original: Finally, note thatforming the update matrix will also take O ( m 2 ) time.
Revised: Finally, note that the update matrix is also an m × m matrix.

Original: 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.
Revised: 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.

Original: 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 efﬁciently.
Revised: 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 efﬁciently.

Original: As a byproduct, we have the following technical theorem:
Revised: As a byproduct, we have the following technical theorem:

Original: Theorem 1.3 (Fast Tensor Regression, informal version of Theorem D.14) .
Revised: Theorem 1.3 (Fast Tensor Regression, informal version of Theorem D.14) .

Original: Given two n × m matrices U and V with m (cid:29) n and a target vector c ∈ R n .
Revised: Given two n × m matrices U and V with m (cid:29) n and a target vector c ∈ R n .

Original: Let J = [vec( u 1 v (cid:62) 1 ) (cid:62) , . .
Revised: Let J = [vec( u 1 v (cid:62) 1 ) (cid:62) , . .

Original: , vec( u n v (cid:62) n ) (cid:62) ] ∈
Revised: , vec( u n v (cid:62) n ) (cid:62) ] ∈

Original: R n × m 2 where u
Revised: R n × m 2 where u

Original: i is the i -th row of matrix U ∀ i ∈ [ n ] .
Revised: i is the i -th row of matrix U ∀ i ∈ [ n ] .

Original: There is a randomized algorithm that takes (cid:101)
Revised: There is a randomized algorithm that takes (cid:101)

Original: O
Revised: O

Original: ( nm + n 2 (log( κ/(cid:15) )
Revised: ( nm + n 2 (log( κ/(cid:15) )

Original: + log( m/δ ))
Revised: + log( m/δ ))

Original: + n ω ) time and outputs a vector (cid:98) x ∈ R n such that
Revised: + n ω ) time and outputs a vector (cid:98) x ∈ R n such that

Original: [Equation]
Revised: [Equation]

Original: holds with probability at least 1 − δ , and κ is the condition number of J .
Revised: holds with probability at least 1 − δ , and κ is the condition number of J .

Original: From a high level, the algorithm proceeds as follows: given matrices U and V , it forms an approximation (cid:101) J ∈ R n × n log( m/δ ) , where each row is generated by applying fast tensor sketching technique to u i and v i (Ahle et al.
Revised: From a high level, the algorithm proceeds as follows: given matrices U and V , it forms an approximation (cid:101) J ∈ R n × n log( m/δ ) , where each row is generated by applying fast tensor sketching technique to u i and v i (Ahle et al.

Original: Then, it uses another sketching matrix for (cid:101) J to obtain a good preconditioner R for (cid:101) J .
Revised: Then, it uses another sketching matrix for (cid:101) J to obtain a good preconditioner R for (cid:101) J .

Original: Subsequently, it runs a gradient descent to solve the regression.
Revised: Subsequently, it runs a gradient descent to solve the regression.

Original: 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.
Revised: 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.

Original: The algorithm then uses tensor-based sketching techniques (Ahle et al.
Revised: The algorithm then uses tensor-based sketching techniques (Ahle et al.

Original: to squash length m 2 tensors to length
Revised: to squash length m 2 tensors to length

Original: O ( n log( m/(cid:15)δ )) .
Revised: O ( n log( m/(cid:15)δ )) .

Original: All subsequent operations are performed on these much smaller vectors.
Revised: All subsequent operations are performed on these much smaller vectors.

Original: Finally, computing the preconditioner takes (cid:101) O ( n ω ) time, and running the gradient descent takes (cid:101) O ( n 2 log( κ/(cid:15) )) time.
Revised: Finally, computing the preconditioner takes (cid:101) O ( n ω ) time, and running the gradient descent takes (cid:101) O ( n 2 log( κ/(cid:15) )) time.

Original: Our Contributions.

Original: We summarize our technical contributions below.

Original: • We develop an analytical framework for the convergence behavior of second-order method on training multi-layer over-parametrized neural network.

Original: To facilitate the analysis, we exploit the equivalence between neural tangent kernels and our over-parametrized network.

Original: • We design an algorithm to implement the second-order method on training such networks, and achieve a cost per iteration of (cid:101) o

Original: ( m 2 ) .

Original: Our algorithm makes use of Gram-Gauss-Newton method, tensor-based sketching techniques, and data structures that maintains a low rank representation efﬁciently.

Original: • By combining fast tensor algebra techniques and sketching-based preconditioning, we devise an algorithm to efﬁciently solve a regression problem where the involved matrix has its rows being tensor product of vectors.

Original: 1. R ELATED W ORK
Revised: 1. R ELATED W ORK

Original: Second-order Method in Optimization.
Revised: Second-order Method in Optimization.

Original: Though not as prevalent as ﬁrst-order method in deep learning, second-order methods are one of the most popular in convex setting, such assolving linear programming (Vaidya (1989); Daitch & Spielman (2008); Lee et al. (2015)
Revised: Though not as prevalent as ﬁrst-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)

Original: ; Cohen et al.
Revised: ; Cohen et al.

Original: empirical risk minimization (Lee et al.
Revised: empirical risk minimization (Lee et al.

Original:  , cutting plane method (Jiang et al. (2020b)) and semideﬁnite programming (Jiang et al.
Revised: (2019b)) , cutting plane method (Jiang et al. (2020b)) and semideﬁnite programming (Jiang et al.

Original: (2020a)).
Revised: (2020a)).

Original: 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 .
Revised: 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 .

Original: In non-convex setting, there’s a vast body of ongoing works (Botev et al.
Revised: In non-convex setting, there’s a vast body of ongoing works (Martens & Grosse (2015); Botev et al.

Original: ; Pilanci & Wainwright (2017); Agarwal et al.
Revised: (2017); Pilanci & Wainwright (2017); Agarwal et al.

Original: ; Bernacchia et al.
Revised: ; Bernacchia et al.

Original: ; Cai et al.
Revised: ; Cai et al.

Original: ; Zhang et al.
Revised: ; Zhang et al.

Original:  Brand et al.
Revised: ; Brand et al.

Revised: ; Yao et al.

Original: that try to improve the practicality of second-order method and adapt them to train deep neural networks.
Revised: that try to improve the practicality of second-order method and adapt them to train deep neural networks.

Original: As shown in Cai et al.
Revised: As shown in Cai et al.

Original: (2019), it is possible to exploit the equivalence between over-parametrized  and neural tangent kernel to optimize an n × n matrix instead of an m 2 × m 2 matrix, which is an important breakthrough in gaining speedup for second-order method.
Revised: (2019), it is possible to exploit the equivalence between over-parametrized networks and neural tangent kernel to optimize an n × n matrix instead of an m 2 × m 2 matrix, which is an important breakthrough in gaining speedup for second-order method.

Original: Sketching and sampling-based methods can also be used to accelerate the computation of inverses of the Hessian matrix (Pilanci & Wainwright (2017)).
Revised: Sketching and sampling-based methods can also be used to accelerate the computation of inverses of the Hessian matrix (Pilanci & Wainwright (2017)).

Original: In spirit, our work resembles most with Cai et al.
Revised: In spirit, our work resembles most with Cai et al.

Original: and Brand et al.
Revised: and Brand et al.

Original: (2021), in the sense that our optimization also works on an n × n Gram matrix.
Revised: (2021), in the sense that our optimization also works on an n × n Gram matrix.

Original: Our algorithm also makes use of sketching and sampling, as in Pilanci & Wainwright (2017); Brand et al.
Revised: Our algorithm also makes use of sketching and sampling, as in Pilanci & Wainwright (2017); Brand et al.

Original: Over-parameterized Neural Networks.
Revised: Over-parameterized Neural Networks.

Original: 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.
Revised: 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.

Original: (2019a;b); Du et al. (2019a); Song & Yang (2019); Ji & Telgarsky (2020)).
Revised: (2019a;b); Du et al. (2019a); Song & Yang (2019); Ji & Telgarsky (2020); Zou et al. ; Cao & Gu (2019); Liu et al.

Original: The seminal work of Jacot et al.
Revised: The seminal work of Jacot et al.

Original: (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 ≥ Ω( n 4 )) , then the optimization dynamic on a neural network is equivalent to that on a NTK.
Revised: (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 ≥ Ω( n 4 )) , then the optimization dynamic on a neural network is equivalent to that on a NTK.

Original: Sketching.
Revised: Sketching.

Original: Using randomized linear algebra to reduce the dimension of the problem and speedup the algorithms for various problems has received a lot of attention (Sarlos (2006); Clarkson & Woodruff (2013); Woodruff (2014)) due to its wide range of applications to various machine learning problems, especially the efﬁcient approximation of kernel matrices.
Revised: 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 efﬁcient approximation of kernel matrices (Avron et al.

Original: (Avron et al.

Original: Ahle et al.
Revised: Ahle et al.

Original: ; Woodruff & Zandieh (2020))
Revised: ; Woodruff & Zandieh (2020)).

Original: 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 efﬁciency.
Revised: 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 efﬁciency.

Original: Another line of work is to use sketching as a preconditioner (Woodruff (2014); Brand et al.
Revised: Another line of work is to use sketching as a preconditioner (Woodruff (2014); Brand et al.

Original: to obtain a high precision solution.
Revised: to obtain a high precision solution.

Original: Roadmap.
Revised: Roadmap.

Original: In Section 2, we give a preliminary view of the training setup we consider in this paper.
Revised: In Section 2, we give a preliminary view of the training setup we consider in this paper.

Original: In Section 2.1, we introduce the notations that will be used throughout this paper.
Revised: In Section 2.1, we introduce the notations that will be used throughout this paper.

Original: In Section 2.2, we consider the training setup.
Revised: In Section 2.2, we consider the training setup.

Original: In Section 3, we overview the techniques employed in this paper.
Revised: In Section 3, we overview the techniques employed in this paper.

Original: In Section 3.2, we demonstrate various techniques to prove the convergence of our second-order method.
Revised: In Section 3.2, we demonstrate various techniques to prove the convergence of our second-order method.

Original: In Section 3.1, we examine the algorithmic tools utilized in this paper to achieve subquadratic cost per iteration.
Revised: In Section 3.1, we examine the algorithmic tools utilized in this paper to achieve subquadratic cost per iteration.

Original: In Section 4, we summarize the results in this paper and point out some future directions.
Revised: In Section 4, we summarize the results in this paper and point out some future directions.

Original:  P RELIMINARIES 2. N OTATIONS
Revised:  P RELIMINARIES 2. N OTATIONS

Original: For any positive integer n , we use [ n ] to denote the set { 1 , 2 , · · · , n } .
Revised: For any positive integer n , we use [ n ] to denote the set { 1 , 2 , · · · , n } .

Original: We use E [ · ] to denote expectation and Pr[ · ] for probability.
Revised: We use E [ · ] to denote expectation and Pr[ · ] for probability.

Original: We use (cid:107) x
Revised: We use (cid:107) x

Original: (cid:107) 2 to denote the (cid:96) 2 norm of a vector x .
Revised: (cid:107) 2 to denote the (cid:96) 2 norm of a vector x .

Original: We use (cid:107) A (cid:107) to denote the spectral norm of matrix A .
Revised: We use (cid:107) A (cid:107) to denote the spectral norm of matrix A .

Original: We use (cid:107) A (cid:107) F to denote the Frobenius norm of A .
Revised: We use (cid:107) A (cid:107) F to denote the Frobenius norm of A .

Original: We use
Revised: We use

Original: A (cid:62) to denote the transpose of matrix A .
Revised: A (cid:62) to denote the transpose of matrix A .

Original: We use I m to denote the identity matrix of size m × m .
Revised: We use I m to denote the identity matrix of size m × m .

Original: For matrix A or vector x , we use (cid:107) A (cid:107) 0 , (cid:107) x (cid:107) 0 to denote the number of nonzero entries of A and x respectively.
Revised: For matrix A or vector x , we use (cid:107) A (cid:107) 0 , (cid:107) x (cid:107) 0 to denote the number of nonzero entries of A and x respectively.

Original: Note that (cid:107) · (cid:107) 0 is a semi-norm since it satisﬁes triangle inequality.
Revised: Note that (cid:107) · (cid:107) 0 is a semi-norm since it satisﬁes triangle inequality.

Original: 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.
Revised: 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.

Original: We use N ( µ, σ 2 ) to denote the Gaussian distribution with mean µ and variance σ 2 .
Revised: We use N ( µ, σ 2 ) to denote the Gaussian distribution with mean µ and variance σ 2 .

Original: We use (cid:101) O
Revised: We use (cid:101) O

Original: ( f ( n )) to denote O ( f ( n ) · poly log( f ( n )) .
Revised: ( f ( n )) to denote O ( f ( n ) · poly log( f ( n )) .

Original: We use (cid:104)· , ·(cid:105) 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 (cid:104) A, B (cid:105) =
Revised: We use (cid:104)· , ·(cid:105) 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 (cid:104) A, B (cid:105) =

Original: tr[ A (cid:62) B ] where tr[ A ] denote the trace of matrix A .
Revised: tr[ A (cid:62) B ] where tr[ A ] denote the trace of matrix A .

Original: 2. P ROBLEM S ETUP
Revised: [Equation]

Original: Let X ∈ R m 0 × n denote the data matrix with n data points and m 0 features.
Revised: Let X ∈ R m 0 × n denote the data matrix with n data points and m 0 features.

Original: Without loss of generality, we assume (cid:107) x
Revised: Without loss of generality, we assume (cid:107) x

Original: i (cid:107) 2 = 1 , ∀ i ∈ [ n ] .
Revised: i (cid:107) 2 = 1 , ∀ i ∈ [ n ] .

Original: Consider an L layer neural network with one vector a ∈ R m L and L matrices W L ∈ R m L × m L − 1 , · · · , W 2 ∈ R m 2 × m 1 and W 1 ∈ R m 1 × m 0 .
Revised: Consider an L layer neural network with one vector a ∈ R m L and L matrices W L ∈ R m L × m L − 1 , · · · , W 2 ∈ R m 2 × m 1 and W 1 ∈ R m 1 × m 0 .

Original: We will use
Revised: We will use

Original: W (cid:96) ( t ) to denote the weight matrix at layer (cid:96) at time t , and ∇ W (cid:96) ( t ) to denote its gradient.
Revised: W (cid:96) ( t ) to denote the weight matrix at layer (cid:96) at time t , and ∇ W (cid:96) ( t ) to denote its gradient.

Original: We also use W ( t ) =
Revised: We also use W ( t ) =

Original: { W 1 ( t ) , . . .
Revised: { W 1 ( t ) , . . .

Original: , W L ( t ) } to denote the collection of weight matrices at time t .
Revised: , W L ( t ) } to denote the collection of weight matrices at time t .

Original: Architecture.
Revised: Architecture.

Original: We ﬁrst describe our network architecture.
Revised: We ﬁrst describe our network architecture.

Original: The network consists of L hidden layers, each represented by a weight matrix W (cid:96) ∈
Revised: The network consists of L hidden layers, each represented by a weight matrix W (cid:96) ∈

Original: R m (cid:96) ×
Revised: R m (cid:96) ×

Original: m
Revised: m

Original: (cid:96) − 1 for any (cid:96) ∈
Revised: (cid:96) − 1 for any (cid:96) ∈

Original: [ L ] .
Revised: [ L ] .

Original: The output layer consists of a vector a ∈ R m L .
Revised: The output layer consists of a vector a ∈ R m L .

Original: We deﬁne the neural network prediction function f : R m 

Revised: We deﬁne the neural network prediction function f : R m

Original: → R as follows:
Revised: → R as follows:

Original: [Equation]
Revised: [Equation]

Original:  where φ : R → R is the shifted ReLU activation function ( σ b ( x ) = max { x − b, 0 } , it is common and widely used in practical training setting) applied coordinate-wise to a vector.
Revised: , where φ : R → R is the shifted ReLU activation function ( σ b ( x ) = max { x − b, 0 } ) applied coordinate-wise to a vector.

Original: We measure the loss via squared-loss function:
Revised: We measure the loss via squared-loss function:

Original: [Equation]
Revised: [Equation]

Original: This is also the objective function for our training.
Revised: This is also the objective function for our training.

Original: We deﬁne the prediction function f t : R m 0 × n → R n asfollows:
Revised: We deﬁne the prediction function f t : R m 0 × n → R n as

Original: [Equation]
Revised: [Equation]

Original: Initialization.
Revised: Initialization.

Original: Our neural networks are initialized as follows:
Revised: Our neural networks are initialized as follows:

Original: • For each (cid:96) ∈
Revised: • For each (cid:96) ∈

Original: [ L ] , the initial weight matrix W (cid:96) (0)
Revised: [ L ] , the initial weight matrix W (cid:96) (0)

Original: ∈
Revised: ∈

Original: R m (cid:96) ×
Revised: R m (cid:96) ×

Original: m (cid:96) − 1 is initialized such that each entry is sampled from N (0 , 2 m (cid:96) ) .
Revised: m (cid:96) − 1 is initialized such that each entry is sampled from N (0 , 2 m (cid:96) ) .

Original: • Each entry of a is an i.i
Revised: • Each entry of a is an i.i.d.

Original: d. sample from {− 1  m L , + 1  m L } uniformly at random.
Revised: sample from {− 1 √ m L , + 1 √ m L } uniformly at random.

Original: Gradient.
Revised: Gradient.

Original: In order to write gradient in an elegant way, we deﬁne some artiﬁcial variables as follows g
Revised: In order to write gradient in an elegant way, we deﬁne some artiﬁcial variables:

Original: [Equation]
Revised: [Equation]

Original: Thus, using the deﬁnitions of f and h , we have
Revised: Using the deﬁnitions of f and h , we have

Original: [Equation]
Revised: [Equation]

Original: = a (cid:62) h

Original: i,L , ∈ R , ∀ i ∈ [ n ]

Original: We can compute the gradient of L in terms of W (cid:96) ∈
Revised: We can compute the gradient of L in terms of W (cid:96) ∈

Original: R m (cid:96) ×
Revised: R m (cid:96) ×

Original: m
Revised: m

Original: (cid:96) − 1 , for all (cid:96) ≥ [Equation] W [Equation]
Revised: (cid:96) − 1 , for all (cid:96) ≥∂ L ( W ) ∂W

Original: (cid:124)(cid:123)(cid:122)(cid:125)

Original: [Equation]
Revised: [Equation]

Original:  

Original: [Equation]

Original: (cid:96)

Original: [Equation]

Original: [Equation]

Original: (cid:124)(cid:123)(cid:122)(cid:125)

Original: m k − 1 × m k

Original: D i,k (cid:124)(cid:123)(cid:122)(cid:125)

Original: m k

Original: [Equation]

Original: k

Original: [Equation]

Original:  a (cid:124)(cid:123)(cid:122)(cid:125) m

Original: L ×

Original: [Equation]

Original: (cid:62) i,(cid:96) −(cid:124)

Original: (cid:123)(cid:122) (cid:125) 1 × [Equation] (cid:96) −(2) Note that the gradient for W 1 ∈ R m 1 × m 0 (recall that m 0 = d ) is slightly different and can not be written by general form.
Revised: Note that the gradient for W 1 ∈ R m 1 × m 0 (recall that m 0 = d ) is slightly different and can not be written by general form.

Original: By the chain rule, we can compute the gradient with respect to W 1 , ∂ L ( W ) ∂W= n (cid:88) i = [Equation] i [Equation]
Revised: By the chain rule, we can compute the gradient with respect to W 1 , ∂ L ( W ) ∂W

Original: (cid:124)(cid:123)(cid:122)(cid:125) m

Original: [Equation]

Original:  

Original: [Equation]

Original: k

Original: [Equation]

Original: [Equation]

Original: k (cid:124)(cid:123)(cid:122)(cid:125)

Original: m k − 1 × m k

Original: D i,k (cid:124)(cid:123)(cid:122)(cid:125)

Original: [Equation]
Revised: [Equation]

Original: It is worth noting that the gradient matrix is of rank n , since it’s a sum of n matrices, where each matrix is of rank 1.
Revised: It is worth noting that the gradient matrix is of rank n , since it’s a sum of n rank-1 matrices.

Original: Jacobian.
Revised: Jacobian.

Original: For each layer (cid:96) ∈
Revised: For each layer (cid:96) ∈

Original: [ L ] and time t ∈ [ T ] , we deﬁne the Jacobian matrix J (cid:96),t ∈ R n × m (cid:96) m (cid:96) [Equation]
Revised: [ L ] and time t ∈ [ T ] , we deﬁne the Jacobian matrix J (cid:96),t ∈ R n × m (cid:96) m (cid:96) −via

Original: via the following formulation:
Revised: the following formulation:

Original: [Equation]
Revised: [Equation]

Original: The Gram matrix at layer (cid:96) and time t is then deﬁned as G (cid:96),t =
Revised: The Gram matrix at layer (cid:96) and time t is then deﬁned as G (cid:96),t =

Original: J (cid:96),t J (cid:62) (cid:96),t ∈ R n × n whose ( i, j ) -th entry is (cid:104)
Revised: J (cid:96),t J (cid:62) (cid:96),t ∈ R n × n whose ( i, j ) -th entry is (cid:104)

Original: ∂f ( W ( t ) ,x i ) ∂W
Revised: ∂f ( W ( t ) ,x i ) ∂W

Original: (cid:96) , ∂f ( W ( t ) ,x j ) ∂W (cid:96) (cid:105) .
Revised: (cid:96) [Equation]

Original:  T ECHNIQUE O VERVIEW
Revised:  T ECHNIQUE O VERVIEW

Original: In this section, we give an overview of the techniques employed in this paper.
Revised: In this section, we give an overview of the techniques employed in this paper.

Original: In Section 3.1, we showcase our algorithm and explain various techniques being used to obtain a subquadratic cost per iteration.
Revised: In Section 3.1, we showcase our algorithm and explain various techniques being used to obtain a subquadratic cost per iteration.

Original: In Section 3.2, we give an overview of the proof to show the convergence of our algorithm.
Revised: In Section 3.2, we give an overview of the proof to show the convergence of our algorithm.

Original: To give a simpler and cleaner presentation, we assume m (cid:96) = m for all (cid:96)
Revised: To give a simpler and cleaner presentation, we assume m (cid:96) = m for all (cid:96)

Original: ∈
Revised: ∈

Original: [ L ] .
Revised: [ L ] .

Original: 3. S UBQUADRATIC TIME
Revised: 3. S UBQUADRATIC TIME

Original: In this section, we overview the techniques deployed in our implementation of the second-order method.
Revised: In this section, we overview the techniques deployed in our implementation of the second-order method.

Original: Our main focuspoint is to achieve subquadratic cost per iteration.
Revised: Our main focus is to achieve subquadratic cost per iteration.

Original: Note that the most naive implementation of Newton’s method will take at least O ( m 4 ) time, since the size of the Hessian is m 2 × m 2 . A smarter way is to instead of using the n × n Gram matrix for neural tangent kernel, however this would still incur a cost of O ( nm 2 ) per iteration, since each gradient is an m × m matrix
Revised: Instead of using a Hessian matrix of size m 2 × m 2 , we use an n × n Gram matrix derived from the neural tangent kernel. However this would still incur a cost of O ( nm 2 ) per iteration, since each gradient is an m × m matrix

Original: and the Jacobian consists of n such gradients.
Revised: and the Jacobian consists of n such gradients.

Original: We start by demonstrating our algorithm:
Revised: We start by demonstrating our algorithm:

Original: Step 1: Invert Gram by solving regression.

Original: Recall the update rule of generic algorithm is given by

Original: [Equation]

Original: where c is f t − y after proper scaling.

Original: Naively forming the Gram matrix J (cid:96),t J (cid:62) (cid:96),t will take O ( n 2 m 2 ) time and inverting it will take O ( n ω ) time.

Original: To avoid the quadratic cost at this step, we instead solve the regression task, or a linear system since the Gram matrix has full rank:

Original: ﬁnd the vector g (cid:96),t ∈ R n such that

Original: [Equation]

Original: is minimized.

Original: This enables us to utilize the power of sketching to solve the regression efﬁciently.

Original: Step 2: Solve Gram regression via preconditioning.

Original: In order to solve this regression task, we adapt the idea of obtaining a good preconditioner via sketching then apply iterative method to solve it (Brand et al.

Original: Roughly speaking, we ﬁrst use a random matrix S ∈ R s × m 2 called subspace embedding (Sarlos (2006)) to squash down the number of rows of J (cid:62) and we use QR decomposition on matrix SJ (cid:62) to obtain a matrix R such that SJ (cid:62) R has orthonormal columns.

Original: We then use gradient descent to iterative optimize over the objective (cid:107) JJ (cid:62)

Original: Rz t − y (cid:107) 22 , since S is a subspace embedding for J (cid:62) , we can make sure that the condition number of the matrix J (cid:62) R is small ( O (1) ), hence the gradient descent converges after log( κ/(cid:15) ) iterations, where κ is the condition number of J .

Original: However, in order to implement the gradient descent, we will still have to multiply an m 2 × n matrix with a length n vector, in the worst case this will incur a time of O ( nm 2 ) .

Original: In order to bypass this barrier, we need to exploit extra structure properties of the matrices, which will be demonstrated in the following steps.

Original: Step 3: Low rank structure of the gradient.

Original: Instead of studying Jacobian directly, we ﬁrst try to understand the low rank structure of the gradient.

Original: Consider ∂f ( W,x i

Original: [Equation]

Original: ∂W

Original: [Equation]

Original: ∈ R m × m , it can be written as (for simplicity of notation, we use h i, 0 to denote x i ):

Original: [Equation]

Original: 1: procedure O UR A LGORITHM ( f, { x i , y i } i ∈ [ n ] ) (cid:46) Theorem 1.1,1.
Revised: 1: procedure O UR A LGORITHM ( f, { x i , y i } i ∈ [ n ] ) (cid:46) Theorem 1.1,1.

Original: 2: /*Initialization*/ 3:
Revised: 2: /*Initialization*/ 3:

Original: Initialize W (cid:96) (0) , ∀ (cid:96) ∈
Revised: Initialize W (cid:96) (0) , ∀ (cid:96) ∈

Original: [ L ] 4:
Revised: [ L ] 4:

Original: Store W 1 (0) x i in memory, ∀ i ∈ [ n ] (cid:46) Takes O ( nm 2 ) time 5:
Revised: Store W 1 (0) x i in memory, ∀ i ∈ [ n ] (cid:46) Takes O ( nm 2 ) time 5:

Original: Store h i,(cid:96) (0) ← φ ( W (cid:96) (0) h i,(cid:96) − 1 (0)) , ∀ (cid:96) ∈
Revised: Store h i,(cid:96) (0) ← φ ( W (cid:96) (0) h i,(cid:96) − 1 (0)) , ∀ (cid:96) ∈

Original: [ L ] , i ∈ [ n ] in memory (cid:46) Takes O ( nLm 2 ) time 6: for t = 

Revised: [ L ] , i ∈ [ n ] in memory (cid:46) Takes O ( nLm 2 ) time 6: for t =

Original: → T do 7: /*Forward computation*/ 8: for (cid:96) = 

Revised: → T do 7: /*Forward computation*/ 8: for (cid:96) =

Original: → L do 9: v i,(cid:96) ← h i,(cid:96) − 1 , ∀ i ∈ [ n ] 10: h i,(cid:96) ← φ (( W (cid:96) (0)
Revised: → L do 9: v i,(cid:96) ← h i,(cid:96) − 1 , ∀ i ∈ [ n ] 10: h i,(cid:96) ← φ (( W (cid:96) (0)

Original: + ∆ W (cid:96) )
Revised: + ∆ W (cid:96) )

Original: ( h i,(cid:96) − 1 )) , ∀ i ∈ [ n ] (cid:46) Takes O ( n 2 m )
Revised: ( h i,(cid:96) − 1 )) , ∀ i ∈ [ n ] (cid:46) Takes O ( n 2 m )

Original: + o ( nm 2 ) time 11: (cid:46) h i,(cid:96) is sparse 12:
Revised: + o ( nm 2 ) time 11: (cid:46) h i,(cid:96) is sparse 12:

Original: D i,(cid:96) ← diag( φ (cid:48) (( W (cid:96) (0) + ∆ W (cid:96) )
Revised: D i,(cid:96) ← diag( φ (cid:48) (( W (cid:96) (0) + ∆ W (cid:96) )

Original: h i,(cid:96) − 1 )) , ∀ i ∈ [ n ] (cid:46) Takes O ( nm ) time 13: (cid:46) D i,(cid:96) is sparse 14: end for 15: f t ← [ a (cid:62) h 1 ,L , . . .
Revised: h i,(cid:96) − 1 )) , ∀ i ∈ [ n ] (cid:46) Takes O ( nm ) time 13: (cid:46) D i,(cid:96) is sparse 14: end for 15: f t ← [ a (cid:62) h 1 ,L , . . .

Original: , a (cid:62) h n,L ] (cid:62) (cid:46) Takes O ( nm ) time 16: /*Backward
Revised: , a (cid:62) h n,L ] (cid:62) (cid:46) Takes O ( nm ) time 16: /*Backward

Original: computation*/ 17: for (cid:96) =
Revised: computation*/ 17: for (cid:96) =

Original: L → 1 do 18: u i,(cid:96) ← a (cid:62) D i,L W L ( t ) . . .
Revised: L → 1 do 18: u i,(cid:96) ← a (cid:62) D i,L W L ( t ) . . .

Original: D i,(cid:96)
Revised: D i,(cid:96)

Original: +1 W (cid:96) +1 ( t ) D i,(cid:96) (cid:46) Takes o ( nLm 2 ) time 19:
Revised: +1 W (cid:96) +1 ( t ) D i,(cid:96) (cid:46) Takes o ( nLm 2 ) time 19:

Original: Form (cid:101) J (cid:96),t that approximates J (cid:96),t using { u i,(cid:96) } ni =1 , { v i,(cid:96) } ni = 20: (cid:46) Takes (cid:101) O ( mn ) time, (cid:101) J (cid:96),t ∈ R n × s where s = (cid:101) O ( n ) 21:
Revised: Form (cid:101) J (cid:96),t that approximates J (cid:96),t using { u i,(cid:96) } ni =1 , { v i,(cid:96) } ni = 20: (cid:46) Takes (cid:101) O ( mn ) time, (cid:101) J (cid:96),t ∈ R n × s where s = (cid:101) O ( n ) 21:

Original: Compute g (cid:96) that approximates ( (cid:101) J (cid:96),t (cid:101) J (cid:62) (cid:96),t ) − 1 c (cid:46) Takes (cid:101) O ( nm ) time 22:
Revised: Compute g (cid:96) that approximates ( (cid:101) J (cid:96),t (cid:101) J (cid:62) (cid:96),t ) − 1 c (cid:46) Takes (cid:101) O ( nm ) time 22:

Original: Form J (cid:62) (cid:96),t g (cid:96) via low rank factorization (cid:80) ni
Revised: Form J (cid:62) (cid:96),t g (cid:96) via low rank factorization (cid:80) ni

Original: =1 g (cid:96),i u i,(cid:96) v (cid:62) i,(cid:96) 23:
Revised: =1 g (cid:96),i u i,(cid:96) v (cid:62) i,(cid:96) 23:

Original: Implicitly update ∆ W (cid:96) ← ∆ W (cid:96) + (cid:80) ni =1 g (cid:96),i u i,(cid:96) v (cid:62) i,(cid:96) and store it in memory 24: end for 25: end for 26: end procedure 
Revised: Implicitly update ∆ W (cid:96) ← ∆ W (cid:96) + (cid:80) ni =1 g (cid:96),i u i,(cid:96) v (cid:62) i,(cid:96) 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 [Equation]

Revised: where c is f t − y after proper scaling.

Revised: Naively forming the Gram matrix J (cid:96),t J (cid:62) (cid:96),t will take O ( n 2 m 2 ) time and inverting it will take O ( n ω ) time.

Original: This means the gradient is essentially an outer product of two vectors, and hence has rank one.
Revised: To avoid the quadratic cost at this step, we instead solve a regression, or a linear system since the Gram matrix has full rank: ﬁnd the vector g [Equation] ,t ∈ R n such that

Revised: [Equation]

Revised: is minimized.

Revised: This enables us to utilize the power of sketching to solve the regression efﬁciently.

Revised: Step 2: Solve Gram regression via preconditioning.

Revised: 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.

Revised: Roughly speaking, we ﬁrst use a random matrix S ∈ R s × m 2 that has the subspace embedding property (Sarlos (2006)) to reduce the number of rows of J (cid:62) , then we run a QR decomposition on matrix SJ (cid:62) .

Revised: This gives us a matrix R such that SJ (cid:62) R has orthonormal columns.

Revised: We then use gradient descent to optimize the objective (cid:107) JJ (cid:62)

Revised: Rz t − y (cid:107) 2 2 .

Revised: Since S is a subspace embedding for J (cid:62) , we can make sure that the condition number of the matrix J (cid:62) R is small ( O (1) ), hence the gradient descent converges after log( κ/(cid:15) ) iterations, where κ is the condition number of

Revised: J .

Revised: However, in order to implement the gradient descent, we still need to multiply an m 2 × n matrix with a length n vector, in the worst case this will incur a time of O ( nm 2 ) .

Revised: In order to bypass this barrier, we need to exploit extra structural properties of the Jocobian, which will be demonstrated in the following steps.

Revised: Step 3: Low rank structure of the gradient.

Revised: Instead of studying Jacobian directly, we ﬁrst try to understand the low rank structure of the gradient.

Revised: Consider ∂f ( W,x i ) ∂W (cid:96) ∈

Revised: R m × m , it can be written as (for simplicity, we use h i, 0 to denote x i ):

Revised: [Equation]

Revised: This means the gradient is essentially an outer product of two vectors, and hence has rank one.

Original: This has several interesting consequences: for over-parametrized networks, the gradient is merely of rank n instead of m .
Revised: This has several interesting consequences: for over-parametrized networks, the gradient is merely of rank n instead of m .

Original: This means when using ﬁrst-order method such as gradient descent or stochastic gradient descent, the weight is updated via a low rank matrix.
Revised: When using ﬁrst-order method such as gradient descent or stochastic gradient descent, the weight is updated via a low rank matrix.

Original: To some extent, this explains why the weight does not move too far from initialization in over-parametrized networks when using ﬁrstorder method to train.
Revised: To some extent, this explains why the weight does not move too far from initialization in over-parametrized networks when using ﬁrst-order method to train.

Original: Also, as we will illustrate below, this enables the efﬁcient approximation of Jacobian matrices and maintenance of the change.
Revised: Also, as we will illustrate below, this enables the efﬁcient approximation of Jacobian matrices and maintenance of the change.

Original: Step 4: Fast approximation of the Jacobian matrix.
Revised: Step 4: Fast approximation of the Jacobian matrix.

Original: We now turn our attention to design a fast approximation  to the Jacobian matrix.
Revised: We now turn our attention to design a fast approximation algorithm to the Jacobian matrix.

Original: Recall that Jacobian matrix J (cid:96),t ∈ R n × m 2 is an n × mmatrix, therefore writing down the matrix will take O ( nm 2 ) matrix.
Revised: Recall that Jacobian matrix J (cid:96),t ∈ R n × m 2 is an n × m 2 matrix, therefore writing down the matrix will take O ( nm 2 ) time.

Original: However, it is worth noticing that each row of J (cid:96),t is vec( u i v (cid:62) i ) (cid:62) , ∀ i ∈ [ n ] , or equivalently, u i ◦ v i where ◦ denotes the tensor product between two vectors.
Revised: However, it is worth noticing that each row of J (cid:96),t is vec( u i v (cid:62) i ) (cid:62) , ∀ i ∈ [ n ] , or equivalently, u i ◦ v i where ◦ denotes the tensor product between two vectors.

Original: Suppose we are given the collection of { u 1 , . .
Revised: Suppose we are given the collection of { u 1 , . .

Original: . , u n } ∈ ( R m ) n and { v 1 , . .
Revised: . , u n } ∈ ( R m ) n and { v 1 , . .

Original: . , v n } ∈ ( R m ) n , then we can compute the tensor product u i ◦ v i via tensor-based sketching techniques, such as TensorSketch (Avron et al.
Revised: . , v n } ∈ ( R m ) n , then we can compute the tensor product u i ◦ v i via tensor-based sketching techniques, such as TensorSketch (Avron et al.

Original: Diao et al.
Revised: Diao et al.

Original: (2017; 2019)) or
Revised: (2017; 2019)) or

Original: TensorSRHT (Ahle et al.
Revised: TensorSRHT (Ahle et al.

Original: ; Woodruff & Zandieh (2020)) in time nearly linear in m and the targeted sketching dimension s , in contrast to the naive O ( m 2 ) time.
Revised: ; Woodruff & Zandieh (2020)) in time nearly linear in m and the targeted sketching dimension s , in contrast to the naive O ( m 2 ) time.

Original: Since it sufﬁces to preserve the length of all vectors in the column space of J (cid:62) (cid:96),t , the target dimension s can be chosen as
Revised: Since it sufﬁces to preserve the length of all vectors in the column space of J (cid:62) (cid:96),t , the target dimension s can be chosen as

Original: O ( (cid:15) − 2 n · poly(log( m/(cid:15)δ ))) .
Revised: O ( (cid:15) − 2 n · poly(log( m/(cid:15)δ ))) .

Original: Use (cid:101) J (cid:96),t ∈ R n × s to denote this approximation of J (cid:96),t , we perform the preconditioned gradient descent we described above on this smaller matrix.
Revised: Use (cid:101) J (cid:96),t ∈ R n × s to denote this approximation of J (cid:96),t , we perform the preconditioned gradient descent we described above on this smaller matrix.

Original: This enables to lower the overall cost of the regression step to be subquadratic in m .
Revised: This enables to lower the overall cost of the regression step to be subquadratic in m .

Original: Step 5: Efﬁcient update via low rank factorization.
Revised: Step 5: Efﬁcient update via low rank factorization.

Original: 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.
Revised: 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.

Original: Let g (cid:96) ∈
Revised: Let g (cid:96) ∈

Original: R
Revised: R

Original: n denote the solution to the regression problem posed in Step 1.
Revised: n denote the solution to the regression problem posed in Step 1.

Original: Note that by the update rule of our method, we shall use J (cid:62) (cid:96),t g (cid:96)
Revised: Note that by the update rule of our method, we shall use J (cid:62) (cid:96),t g (cid:96)

Original: ∈
Revised: ∈

Original: R m × m to update the weight matrix, but writing down the matrix will already take O ( m 2 ) time.
Revised: R m × m to update the weight matrix, but writing down the matrix will already take O ( m 2 ) time.

Original: Therefore, it is instructive to ﬁnd a succinct representation for the update.
Revised: Therefore, it is instructive to ﬁnd a succinct representation for the update.

Original: The key observation is that each column of J (cid:62) (cid:96),t is a tensor product of two vectors: u i ◦ v i or equivalently, u i v (cid:62) i .
Revised: The key observation is that each column of J (cid:62) (cid:96),t is a tensor product of two vectors: u i ◦ v i or equivalently, u i v (cid:62) i .

Original: The update matrix can be rewritten as (cid:80) ni
Revised: The update matrix can be rewritten as (cid:80) ni

Original: =1 g (cid:96),i u i v (cid:62) i , andthis implies that we can use this representation for the update on the weight, instead of adding it directly to the original weight.
Revised: =1 g (cid:96),i u i v (cid:62) i , and we can use this representation for the update on the weight, instead of adding it directly.

Original: Let
Revised: Let

Original: U (cid:96) := (cid:34) | | . . . | g (cid:96), 1 ug (cid:96), 2 u. . . g (cid:96),n u n | | . . . | (cid:35) ∈
Revised: [Equation]

Original: R m × n , V (cid:96) :

Original: = (cid:34) | | . . .

Original: | vv. . . v n | | . . . | (cid:35) ∈ R m × n , then the update can be represented as U (cid:96) V (cid:62) (cid:96) .
Revised: then the update can be represented as U (cid:96) V (cid:62) (cid:96) .

Original: Consider multiplying a vector y ∈ R m with this representation, we ﬁrst multiply y with V (cid:62) (cid:96) ∈ R n × m , which takes O ( mn ) time.
Revised: Consider multiplying a vector y ∈ R m with this representation, we ﬁrst multiply y with V (cid:62) (cid:96) ∈ R n × m , which takes O ( mn ) time.

Original: Then we multiply V (cid:62) (cid:96) y ∈ R n with U (cid:96) ∈
Revised: Then we multiply V (cid:62) (cid:96) y ∈ R n with U (cid:96) ∈

Original: R m × n which takes O ( mn ) time.
Revised: R m × n which takes O ( mn ) time.

Original: This drastically reduces the cost of multiplying the weight matrix with a vector from O ( m 2 ) to O ( mn ) .
Revised: This drastically reduces the cost of multiplying the weight matrix with a vector from O ( m 2 ) to O ( mn ) .

Original: 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 of O ( Tmn ) .
Revised: 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 of O ( Tmn ) .

Original: This is ﬁne when T is relatively small, however, if one looks for a high precision solution which requires a large number of iterations, then T might be too large and O ( Tmn ) might be in the order of O ( m 2 ) .
Revised: This is ﬁne when T is relatively small, however, if one looks for a high precision solution which requires a large number of iterations, then T might be too large and O ( Tmn ) might be in the order of O ( m 2 ) .

Original: 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.
Revised: 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.

Original: This can be viewed as a “restart” of the data structure.
Revised: This can be viewed as a “restart” of the data structure.

Original: 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 ( m 2 ) time to multiply an m × m by an m × m α matrix.
Revised: 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 ( m 2 ) time to multiply an m × m by an m × m α matrix.

Original: Hence, we restart the data structure after around m α /n updates.
Revised: Hence, we restart the data structure after around m α /n updates.

Original: Therefore, we achieve an amortized o ( m 2 ) time, which is invariant even though the number of iterations T grows larger and larger.
Revised: Therefore, we achieve an amortized o ( m 2 ) time, which is invariant even though the number of iterations T grows larger and larger.

Original: 3. C ONVERGENCE ANALYSIS
Revised: 3. C ONVERGENCE ANALYSIS

Original: In this section, we demonstrate the strategy to prove that by using a generic algorithm, one can achieve a linear convergence rate on  training loss.
Revised: In this section, we demonstrate the strategy to prove that our second-order method achieves a linear convergence rate on the training loss.

Original: Step 1: Initialization.
Revised: Step 1: Initialization.

Original: Let W (0) be the random initialization.
Revised: Let W (0) be the random initialization.

Original: We ﬁrst show that for any data point x i , we have f ( W (0) , x i ) =
Revised: We ﬁrst show that for any data point x i , we have f ( W (0) , x i ) =

Original: O (1) .
Revised: O (1) .

Original: The analysis draws inspiration from Allen-Zhu et al. (2019a).
Revised: The analysis draws inspiration from Allen-Zhu et al. (2019a).

Original: The general idea is, given a ﬁxed unit length vector x , multiplying it with a random Gaussian matrix
Revised: The general idea is, given a ﬁxed unit length vector x , multiplying it with a random Gaussian matrix

Original: W will make sure that (cid:107) Wx (cid:107) 

Revised: W will make sure that (cid:107) Wx (cid:107)

Original: ≈ 2 .
Revised: ≈ 2 .

Original: Again, since W is a random Gaussian matrix, applying  ReLU activation gives a random vector with a truncated Gaussian distribution conditioned on a binomial random variable indicating which neurons are activated.
Revised: 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.

Original: We will end up with (cid:107) φ ( Wx ) (cid:107) 2 ≈ 1 as well as φ ( Wx ) being sparse.
Revised: We will end up with (cid:107) φ ( Wx ) (cid:107) 2 ≈ 1 as well as φ ( Wx ) being sparse.

Original: 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) , x i ) is a constant.
Revised: 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) , x i ) is a constant.

Original: We also bound the spectral norm of [Equation]
Revised: We also bound the spectral norm of ∂f ( W (0) ,x i ) ∂W (cid:96) by (cid:101) O

Original: by (cid:101) O ( (cid:112) L/m ) .
Revised: ( (cid:112) L/m ) .

Original: One of the key part of this matrix is the consecutive product D 
Revised: One of the key part of this matrix is the consecutive product D L W L . . .

Original: [Equation]
Revised: D (cid:96) +1 W (cid:96) +1 D (cid:96) .

Original: By studying the distribution of its product with a ﬁxed vector, one can show that the spectral norm of this consecutive product is bounded by O ( √ L ) .
Revised: By studying the distribution of its product with a ﬁxed vector, one can show that the spectral norm of this consecutive product is bounded by O ( √ L ) .

Original: 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 (cid:62) D L W L . . .
Revised: 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 (cid:62) D L W L . . .

Original: D (cid:96) +1 W (cid:96) +1 D (cid:96) is bounded by (cid:101) O
Revised: D (cid:96) +1 W (cid:96) +1 D (cid:96) is bounded by (cid:101) O

Original: ( (cid:112) L/m ) with good probability.
Revised: ( (cid:112) L/m ) with good probability.

Original: Furthermore, we show that the Gram matrix for the multiple-layer over-parametrized neural network, which is deﬁned as J (cid:96), 0 J (cid:62) (cid:96), 0 , has a nontrivial minimum eigenvalue after the initialization.
Revised: Furthermore, we show that the Gram matrix for the multiple-layer over-parametrized neural network, which is deﬁned as J (cid:96), 0 J (cid:62) (cid:96), 0 , has a nontrivial minimum eigenvalue after the initialization.

Original: In particular, we adapt the neural tangent kernel (NTK) for multiple-layer neural networks deﬁned by
Revised: In particular, we adapt the neural tangent kernel (NTK) for multiple-layer neural networks deﬁned by

Original: Du et al. (2019a) into our setting by analyzing the corresponding Gaussian process with  ReLU activation function.
Revised: Du et al. (2019a) into our setting by analyzing the corresponding Gaussian process with shifted ReLU activation function.

Original: 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.
Revised: 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.

Original: Step 2: Small perturbation.
Revised: Step 2: Small perturbation.

Original: 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.
Revised: 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.

Original: As long as the perturbation is small enough, it is possible to show that the change of the h vector (in terms of (cid:96) 2 norm) and the consecutive product (in terms of spectral norm) is also small.
Revised: As long as the perturbation is small enough, it is possible to show that the change of the h vector (in terms of (cid:96) 2 norm) and the consecutive product (in terms of spectral norm) is also small.

Original: Finally, using the fact that a is a Rademacher vector with scaling 1 / √ m , we can show that the change
Revised: Finally, using the fact that a is a Rademacher vector with scaling 1 / √ m , we can show that the change

Original: [Equation]
Revised: [Equation]

Original: has its spectral norm being bounded by (cid:101)
Revised: has its spectral norm being bounded by (cid:101)

Original: O ( (cid:112) L/m ) for any layer (
Revised: O ( (cid:112) L/m ) for any layer (

Original: cid:96) ∈
Revised: cid:96) ∈

Original: [ L ] and input data i ∈ [ n ] .
Revised: [ L ] and input data i ∈ [ n ] .

Original: Consequently, the Frobenious norm of the Jacobian matrix is bounded by (cid:101) O ( (cid:112) nL/m ) .
Revised: Consequently, the Frobenious norm of the Jacobian matrix is bounded by (cid:101) O ( (cid:112) nL/m ) .

Original: Step 3: Connect everything via a double induction.
Revised: Step 3: Connect everything via a double induction.

Original: Put things together, we use a double induction argument, where we assume the perturbation of weight matrix is small and the gap between f t and y is at most 1 / 2 of the gap between f t − 1 and y .
Revised: Put things together, we use a double induction argument, where we assume the perturbation of weight matrix is small and the gap between f t and y is at most 1 / 2 of the gap between f t − 1 and y .

Original: By carefully bounding various terms and exploiting the fact the Jacobian matrix always has a relative small spectral norm ( (cid:101) O
Revised: By carefully bounding various terms and exploiting the fact the Jacobian matrix always has a relative small spectral norm ( (cid:101) O

Original: ( (cid:112) nL/m ) )
Revised: ( (cid:112) nL/m ) )

Original: , we ﬁrst show that the weights are not moving too far from the initialization, then use this fact to derive a ﬁnal convergence bound for (cid:107)
Revised: , we ﬁrst show that the weights are not moving too far from the initialization, then use this fact to derive a ﬁnal convergence bound for (cid:107)

Original: f t − y (cid:107) 2 .
Revised: f t − y (cid:107) 2 .

Original:  D ISCUSSION AND F UTURE D IRECTIONS
Revised:  D ISCUSSION AND F UTURE D IRECTIONS

Original: In this work, we propose and analyze a second-order method to train multi-layer over-parametrized neural networks.
Revised: In this work, we propose and analyze a second-order method to train multi-layer over-parametrized neural networks.

Original: Our algorithm achieves a linear convergence rate in terms of training loss, and achieves a subquadratic ( o ( m 2 ) ) cost per training iteration.
Revised: Our algorithm achieves a linear convergence rate in terms of training loss, and achieves a subquadratic ( o ( m 2 ) ) cost per training iteration.

Original: From an analytical perspective, we greatly extend the analysis of (Allen-Zhu et al.
Revised: From an analytical perspective, we greatly extend the analysis of (Allen-Zhu et al.

Original: (2019a))
Revised: (2019a))

Original: 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)).
Revised: 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)).

Original: From an algorithmic perspective, we achieve a subquadratic cost per iteration, which is a signiﬁcant improvement from O ( m 2 n ) time per iteration due to the prohibitively large network width m .
Revised: From an algorithmic perspective, we achieve a subquadratic cost per iteration, which is a signiﬁcant improvement from O ( m 2 n ) time per iteration due to the prohibitively large network width m .

Original: 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.
Revised: 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.

Original: Our algorithm is especially valuable when one requires a high precision solution on training loss, and hence the number of iterations is large.
Revised: Our algorithm is especially valuable when one requires a high precision solution on training loss, and hence the number of iterations is large.

Original: 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.
Revised: 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.

Original: In particular, can this runtime be achieved under the current best width of multi-layer over-parametrized networks ( m ≥ n 8 )?
Revised: In particular, can this runtime be achieved under the current best width of multi-layer over-parametrized networks ( m ≥ n 8 )?

Original: 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 ( (cid:107) ∆ W (cid:107) ).
Revised: 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 ( (cid:107) ∆ W (cid:107) ).

Original: In our analysis of convergence, we go through a careful double induction argument, which in fact imposes on a lower bound on (cid:107) ∆ W (cid:107) .
Revised: In our analysis of convergence, we go through a careful double induction argument, which in fact imposes on a lower bound on (cid:107) ∆ W (cid:107) .

Original: 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.
Revised: 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.

Original: A related question is, how can we maintain the changes of weight more efﬁciently?
Revised: A related question is, how can we maintain the changes of weight more efﬁciently?

Original: In our work,intuitively, we achieve speedup in the neural network training process by observing that the change of the weights are small in each iteration.
Revised: 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.

Original: Similar phenomenon also appears in some classical optimization problem (e.g., solving linear programming (Cohen et al., 2019; Jiang et al., 2021) and solving semideﬁnite programming (Jiang et al., 2020a)) and they achieved 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.
Revised: Similar phenomenon also appears in some classical optimization problem (e.g., solving linear program (Cohen et al., 2019; Jiang et al., 2021) and solving semideﬁnite 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.

Original: Can we adapt their techniques to neural network training?
Revised: Can we adapt their techniques to neural network training?

Original: 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 matrixvector product with the initial weight matrix can be performed faster than O ( m 2 ) by sparsifying the Gaussian matrix as in Derezi´nski et al.
Revised: 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 ( m 2 ) by sparsifying the Gaussian matrix as in Derezi´nski et al.

Original: or imposing extra structural assumption such as using circulant Gaussian (Rauhut et al., 2012; Nelson & NguyÅn, 2013; Krahmer et al., 2014).
Revised: or imposing extra structural assumption such as using circulant Gaussian (Rauhut et al., 2012; Nelson & NguyÅn, 2013; Krahmer et al., 2014).

Original: Another question concerns activation functions.
Revised: Another question concerns activation functions.

Original: In this paper, we consider the threshold ReLU activation and design our algorithm and analysis based on its properties.
Revised: In this paper, we consider the shifted ReLU activation and design our algorithm and analysis around its properties.

Original: Is it possible to generalize our algorithm and analysis to various other activation functions, such as sigmoid, tanh or leaky ReLU?
Revised: Is it possible to generalize our algorithm and analysis to various other activation functions, such as sigmoid, tanh or leaky ReLU?

Original: If one chooses a smooth activation, can we get a better result in terms of convergence rate?
Revised: If one chooses a smooth activation, can we get a better result in terms of convergence rate?

Original: Can we leverage this structure to design faster algorithms?
Revised: Can we leverage this structure to design faster algorithms?

Original: Finally, the network architecture considered in this paper is the standard feedforward network.
Revised: Finally, the network architecture considered in this paper is the standard feedforward network.

Original: Is it possible to extend our analysis and algorithm to other architectures, such as recurrent neural networks (RNN)?
Revised: Is it possible to extend our analysis and algorithm to other architectures, such as recurrent neural networks (RNN)?

Original: For RNN, the weight matrices for each layer are the same, hence it is trickier to analyze the training dynamics on such networks.
Revised: For RNN, the weight matrices for each layer are the same, hence it is trickier to analyze the training dynamics on such networks.

Original: Though the convergence of ﬁrst-order method on over-parametrized multi-layer RNN has been established, it is unclear whether such analysis can be extended to second-order method.
Revised: Though the convergence of ﬁrst-order method on over-parametrized multi-layer RNN has been established, it is unclear whether such analysis can be extended to second-order method.

Original: Ethics Statement.
Revised: Ethics Statement.

Original: This is a theory paper that proposed efﬁcient algorithm to train multi-layer overparametrized neural networks, hence it does not have direct ethnic implications.
Revised: This is a theory paper that proposed efﬁcient algorithm to train multi-layer overparametrized neural networks, hence it does not have direct ethics implications.

Original: Reproducibility Statement.
Revised: Reproducibility Statement.

Original: All results in this paper can be directly veriﬁed and reproduced via reading into the proofs.
Revised: All results in this paper can be directly veriﬁed and reproduced via reading into the proofs.

Original: For complete algorithm description and its runtime analysis, see Section B.
Revised: For complete algorithm description and its runtime analysis, see Section B.

Original: For speciﬁc 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.
Revised: For speciﬁc 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.

Original: A P RELIMINARIES AND P ROBABILITY TOOLS
Revised: A P RELIMINARIES AND P ROBABILITY TOOLS

Original: 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.
Revised: 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.

Original: Notations.
Revised: Notations.

Original: For any positive integer n , we use [ n ] to denote the set { 1 , 2 , · · · , n } .
Revised: For any positive integer n , we use [ n ] to denote the set { 1 , 2 , · · · , n } .

Original: We use E [ · ] to denote expectation and Pr[ · ] for probability.
Revised: We use E [ · ] to denote expectation and Pr[ · ] for probability.

Original: We use (cid:107) x
Revised: We use (cid:107) x

Original: (cid:107) 2 to denote the (cid:96) 2 norm of a vector x .
Revised: (cid:107) 2 to denote the (cid:96) 2 norm of a vector x .

Original: We use (cid:107) A (cid:107) to denote the spectral norm of matrix A .
Revised: We use (cid:107) A (cid:107) to denote the spectral norm of matrix A .

Original: We use (cid:107) A (cid:107) F to denote the Frobenius norm of
Revised: We use (cid:107) A (cid:107) F to denote the Frobenius norm of

Original: A .
Revised: A .

Original: We use A (cid:62) to denote the transpose of matrix A .
Revised: We use A (cid:62) to denote the transpose of matrix A .

Original: We use I m to denote the identity matrix of size m × m .
Revised: We use I m to denote the identity matrix of size m × m .

Original: For matrix A or vector x , we use (cid:107) A (cid:107) 0 , (cid:107) x (cid:107) 0 to denote the number of nonzero entries of
Revised: For matrix A or vector x , we use (cid:107) A (cid:107) 0 , (cid:107) x (cid:107) 0 to denote the number of nonzero entries of

Original: A and x respectively.
Revised: A and x respectively.

Original: Note that (cid:107) · (cid:107) 0 is a semi-norm since it satisﬁes triangle inequality.
Revised: Note that (cid:107) · (cid:107) 0 is a semi-norm since it satisﬁes triangle inequality.

Original: 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.
Revised: 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.

Original: We use N ( µ, σ 2 ) to denote the Gaussian distribution with mean µ and variance σ 2 .
Revised: We use N ( µ, σ 2 ) to denote the Gaussian distribution with mean µ and variance σ 2 .

Original: We use (cid:101) O
Revised: We use (cid:101) O

Original: ( f ( n )) to denote O ( f ( n ) · poly log( f ( n )) .
Revised: ( f ( n )) to denote O ( f ( n ) · poly log( f ( n )) .

Original: We use (cid:104)· , ·(cid:105) 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 (cid:104) A, B (cid:105) =
Revised: We use (cid:104)· , ·(cid:105) 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 (cid:104) A, B (cid:105) =

Original: tr[ A (cid:62) B ] where tr[ A ] denote the trace of matrix A .
Revised: tr[ A (cid:62) B ] where tr[ A ] denote the trace of matrix A .

Original: Lemma A.1 (Chernoff bound Chernoff (1952)) .
Revised: Lemma A.1 (Chernoff bound Chernoff (1952)) .

Original: Let X = (cid:80) ni
Revised: Let X = (cid:80) ni

Original: =1 X i , where X i = 1 with probability p i and X i = 0 with probability 1 − p
Revised: =1 X i , where X i = 1 with probability p i and X i = 0 with probability 1 − p

Original: i , and all X i are independent.
Revised: i , and all X i are independent.

Original: Let µ = E
Revised: Let µ = E

Original: [ X ] = (cid:80) ni
Revised: [ X ] = (cid:80) ni

Original: =1 p i .
Revised: =1 p i .

Original: Then
Revised: Then

Original: Pr[ X ≥ (1 + δ ) µ ] ≤ exp( − δ 2 µ/ 3) , ∀ δ > 0 ; 2.
Revised: Pr[ X ≥ (1 + δ ) µ ] ≤ exp( − δ 2 µ/ 3) , ∀ δ > 0 ; 2.

Original: Pr[ X ≤ (1 − δ ) µ ] ≤ exp( − δ 2 µ/ 2) , ∀ 0 < δ < 1 .
Revised: Pr[ X ≤ (1 − δ ) µ ] ≤ exp( − δ 2 µ/ 2) , ∀ 0 < δ < 1 .

Original: Lemma A.2 (Hoeffding bound Hoeffding (1963)) .
Revised: Lemma A.2 (Hoeffding bound Hoeffding (1963)) .

Original: Let X 1 , · · · , X n denote n independent bounded variables in [ a i , b i ] .
Revised: Let X 1 , · · · , X n denote n independent bounded variables in [ a i , b i ] .

Original: Let X = (cid:80) ni
Revised: Let X = (cid:80) ni

Original: =1 X i , then we have
Revised: =1 X i , then we have

Original: [Equation]
Revised: [Equation]

Original: Lemma A.3 (Bernstein inequality Bernstein (1924)) .
Revised: Lemma A.3 (Bernstein inequality Bernstein (1924)) .

Original: Let X 1 , · · · , X n be independent zero-mean random variables.
Revised: Let X 1 , · · · , X n be independent zero-mean random variables.

Original: Suppose that | X i | ≤ M almost surely, for all i .
Revised: Suppose that | X i | ≤ M almost surely, for all i .

Original: Then, for all positive t ,
Revised: Then, for all positive t ,

Original: Pr (cid:34) n (cid:88) i =
Revised: Pr [Equation]

Original: X i>
Revised: i

Original: [Equation]
Revised: [Equation]

Original: Lemma A.4 (Anti-concentration of Gaussian distribution) .
Revised: Lemma A.4 (Anti-concentration of Gaussian distribution) .

Original: Let X ∼ N (0 , σ 2 ) , then
Revised: Let X ∼ N (0 , σ 2 ) , then

Original: Pr[ 
Revised: Pr[ | X | ≤ [Equation]

Original: [Equation]
Revised: [Equation]

Revised: t σ ) .

Original: Lemma A.5 (Concentration of subgaussian random variables) .
Revised: Lemma A.5 (Concentration of subgaussian random variables) .

Original: Let a ∈ R n be a vector where each entry of a is sampled from a subgaussian distribution of parameter σ 2 , then for any vector x ∈ R n , 
Revised: Let a ∈ R n be a vector where each entry of a is sampled from a subgaussian distribution of parameter σ 2 , then for any vector x ∈ R n , [Equation]

Revised: Lemma A.6 (Small ball probability) .

Original: Pr[
Revised: Let a ∈ R n be a vector such that | a i | ≥ δ for all i ∈ [ n ] , x 1 , . . . , x n are n i.i.d. Rademacher random variables. Then, there exist absolute constants C 1 , Csuch that for any t > 0 ,

Original: [Equation]
Revised: [Equation]

Original: B C OMPLETE A LGORITHM AND ITS R UNTIME A NALYSIS
Revised: B C OMPLETE A LGORITHM AND ITS R UNTIME A NALYSIS

Original: In this section, we ﬁrst present our complete algorithm, then give a runtime analysis of it.
Revised: In this section, we ﬁrst present our complete algorithm, then give a runtime analysis of it.

Revised: Algorithm 2 Complete version of our algorithm.

Original: 1: procedure C OMPLETE A LGORITHM ( X ∈ R d × n , y ∈ R n ) (cid:46) Theorem B.
Revised: 1: procedure C OMPLETE A LGORITHM ( X ∈ R d × n , y ∈ R n ) (cid:46) Theorem B.

Original: 2: /*Initialization*/ 3:
Revised: 2: /*Initialization*/ 3:

Original: Initialize W (cid:96) (0) , ∀ (cid:96) ∈
Revised: Initialize W (cid:96) (0) , ∀ (cid:96) ∈

Original: [ L ] 4:
Revised: [ L ] 4:

Original: Store W 1 (0) x i in memory, ∀ i ∈ [ n ] (cid:46) Takes O ( nm 2 ) time 5:
Revised: Store W 1 (0) x i in memory, ∀ i ∈ [ n ] (cid:46) Takes O ( nm 2 ) time 5:

Original: L OW R ANK M AINTENANCE LMR (cid:46)
Revised: L OW R ANK M AINTENANCE LMR (cid:46)

Original: Algorithm6:
Revised: Algorithm6:

Original: LMR.I NIT ( { W 1 (0) . . .
Revised: LMR.I NIT ( { W 1 (0) . . .

Original: , W L (0) } ) 7: for t = 

Revised: , W L (0) } ) 7: for t =

Original: → T do 8: /*Forward computation*/ 9: for (cid:96) = 

Revised: → T do 8: /*Forward computation*/ 9: for (cid:96) =

Original: → L do 10: v i,(cid:96) ← h i,(cid:96) − 1 , ∀ i ∈ [ n ] 11: g i,(cid:96)
Revised: → L do 10: v i,(cid:96) ← h i,(cid:96) − 1 , ∀ i ∈ [ n ] 11: g i,(cid:96)

Original: ← LMR.Q UERY
Revised: ← LMR.Q UERY

Original: (cid:96), h i,(cid:96) − 1 ) (cid:46) TakesO ( n 2 m ) + o ( nm 2 ) time 12: h i,(cid:96) ← φ ( g i,(cid:96) ) , ∀ i ∈ [ n ] 13: (cid:46) h i,(cid:96) is sparse 14:
Revised: (cid:96), h i,(cid:96) − 1 ) (cid:46) Takes o ( nm 2 ) time 12: h i,(cid:96) ← φ ( g i,(cid:96) ) , ∀ i ∈ [ n ] 13: (cid:46) h i,(cid:96) is sparse 14:

Original: D i,(cid:96) ← diag( φ (cid:48) ( g i,(cid:96) )) , ∀ i ∈ [ n ] (cid:46) Takes O ( nm ) time 15: (cid:46) D i,(cid:96) is sparse 16: end for 17: f t ← [ a (cid:62) h 1 ,L , . . .
Revised: D i,(cid:96) ← diag( φ (cid:48) ( g i,(cid:96) )) , ∀ i ∈ [ n ] (cid:46) Takes O ( nm ) time 15: (cid:46) D i,(cid:96) is sparse 16: end for 17: f t ← [ a (cid:62) h 1 ,L , . . .

Original: , a (cid:62) h n,L ] (cid:62) (cid:46) Takes O ( nm ) time 18: /*Backward computation*/ 19: for (cid:96) =
Revised: , a (cid:62) h n,L ] (cid:62) (cid:46) Takes O ( nm ) time 18: /*Backward computation*/ 19: for (cid:96) =

Original: L → 1 do 20: u i,(cid:96) ← a (cid:62) D i,L W L ( t ) . . .
Revised: L → 1 do 20: u i,(cid:96) ← a (cid:62) D i,L W L ( t ) . . .

Original: D i,(cid:96)
Revised: D i,(cid:96)

Original: +1 W (cid:96) +1 ( t ) D i,(cid:96)  (cid:46) Takes o ( nLm 2 ) time 21: g (cid:96)
Revised: +1 W (cid:96) +1 ( t ) D i,(cid:96) , ∀ i ∈ [ n ] (cid:46) Takes o ( nLm 2 ) time 21: g (cid:96)

Original: ← F AST T ENSOR R EGRESSION ( { u i,(cid:96) } ni =1 , { v i,(cid:96) } ni =1 , c ) with precision (cid:112) λ/n 22: (cid:46)
Revised: ← F AST T ENSOR R EGRESSION ( { u i,(cid:96) } ni =1 , { v i,(cid:96) } ni =1 , c ) with precision (cid:112) λ/n 22: (cid:46)

Original: Algorithm23:
Revised: Algorithm23:

Original: LMR.U PDATE ( { g (cid:96),i u i,(cid:96) } ni =1 , { v i,(cid:96) } ni =1 ) 24: end for 25: end for 26: end procedure
Revised: LMR.U PDATE ( { g (cid:96),i u i,(cid:96) } ni =1 , { v i,(cid:96) } ni =1 ) 24: end for 25: end for 26: end procedure

Original: Theorem B.1 (Formal version of Theorem 1.2) .
Revised: Theorem B.1 (Formal version of Theorem 1.2) .

Original: Let X ∈ R d × n and y ∈ R n , and let k denote the sparsity of D i,(cid:96) and s denote the sparsity of ∆ D i,(cid:96) , ∀ (cid:96) ∈
Revised: Let X ∈ R d × n and y ∈ R n , and let k denote the sparsity of D i,(cid:96) and s denote the sparsity of ∆ D i,(cid:96) , ∀ (cid:96) ∈

Original: [ L ] , i ∈ [ n ] , then the running time of
Revised: [ L ] , i ∈ [ n ] .

Revised: 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

Original:  Algorithm 2 is 
Revised: the running time of Algorithm 2 is O

Original: [Equation]
Revised: [Equation]

Original:  where 
Revised: init [Equation] T · T iter ) , where T

Original: [Equation]
Revised: [Equation]

Original: Therefore, the cost per iteration of Algorithm 2 is
Revised: Therefore, the cost per iteration of Algorithm 2 is

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We analyze T init and T iter separately.
Revised: We analyze T init and T iter separately.

Original: Initialization time.
Revised: Initialization time.

Original: We will ﬁrst initialize ( L − 1) m × m matrices and one m × d matrix, which takes O ( m 2 L ) time.
Revised: We will ﬁrst initialize ( L − 1) m × m matrices and one m × d matrix, which takes O ( m 2 L ) time.

Original: Compute W 1 (0) x i for all i ∈ [ n ] takes O ( m 2 n ) time.
Revised: Compute W 1 (0) x i for all i ∈ [ n ] takes O ( m 2 n ) time.

Original: Finally, initialize the data structure takes O ( m 2 L ) time.
Revised: Finally, initialize the data structure takes O ( m 2 L ) time.

Original: Hence, T init =
Revised: Hence, T init =

Original: O ( m  ( n + L )) .
Revised: O ( m 2 ( n + L )) .

Original: Cost per iteration.
Revised: Cost per iteration.

Original: For each iteration, we perform one forward computation from layer 1 to L , then backpropagate from layer L to 1.
Revised: For each iteration, we perform one forward computation from layer 1 to L , then backpropagate from layer L to 1.

Original: • Forward computation:
Revised: • Forward computation:

Original: In forward computation, we ﬁrst compute g i,(cid:96) ∈
Revised: In forward computation, we ﬁrst compute g i,(cid:96) ∈

Original: R m , which is equivalent to form v i,(cid:96) +1 , hence by Lemma C.5, it takes O ( s + k + m α ) time.
Revised: R m , which is equivalent to form v i,(cid:96) +1 , hence by Lemma C.5, it takes O ( s + k + m α ) time.

Original: Compute h i,(cid:96) and D i,(cid:96) takes O ( m ) time.
Revised: Compute h i,(cid:96) and D i,(cid:96) takes O ( m ) time.

Original: These computations would be performed for all L layers, hence the overall runtime of forward computation is O
Revised: These computations would be performed for all L layers, hence the overall runtime of forward computation is O

Original: (( s + k + m α ) L ) time.
Revised: (( s + k + m α ) L ) time.

Original: • Backpropagation:
Revised: • Backpropagation:

Original: In backpropagation, we ﬁrst compute u i,(cid:96) ∈ R m , which takes O ( mL ( s + k
Revised: In backpropagation, we ﬁrst compute u i,(cid:96) ∈ R m , which takes O ( mL ( s + k

Original: + m α )) time.
Revised: + m α )) time.

Original: Then, we call Algorithm 6 to solve the Gram regression problem, which due to Theorem D.14 takes (cid:101) O
Revised: Then, we call Algorithm 6 to solve the Gram regression problem, which due to Theorem D.14 takes (cid:101) O

Original: ( mn + n ω ) time.
Revised: ( mn + n ω ) time.

Original: 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 log 2 nL term in running time, which is absorbed by the (cid:101)
Revised: 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 log 2 nL term in running time, which is absorbed by the (cid:101)

Original: O ( · ) notation.
Revised: O ( · ) notation.

Original: Finally, the update takes O ( m 2 − α n ) time.
Revised: Finally, the update takes O ( m 2 − α n ) time by Lemma C.2.

Original: Sum over all L layers, we get an overall running time of (cid:101) O
Revised: Sum over all L layers, we get an overall running time of (cid:101) O

Original: (( m 1+ α + m ( s + k ))
Revised: (( m 1+ α + m ( s + k ))

Original: L 2 + m 2 − α nL ) time.
Revised: L 2 + m 2 − α nL ) time.

Original: This concludes the proof of our Theorem.
Revised: This concludes the proof of our Theorem.

Original: Corollary B.2.
Revised: Corollary B.2.

Original: 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
Revised: 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

Original: [Equation]
Revised: [Equation]

Original: Remark B.3.
Revised: Remark B.3.

Original: As long as the neural network is wide enough, as in F.

Revised: As long as the neural network is wide enough, as in F.

Original: 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.
Revised: 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.

Original: C L OW RANK MAINTENANCE AND EFFICIENT COMPUTATION OF THE CHANGE
Revised: C L OW RANK MAINTENANCE AND EFFICIENT COMPUTATION OF THE CHANGE

Original: In this seciton, we design a data structure to maintain the low rank representation of change of weights, and then we show how to efﬁciently compute the low rank representation using this data structure.
Revised: In this seciton, we design a data structure to maintain the low rank representation of change of weights, and then we show how to efﬁciently compute the low rank representation using this data structure.

Original: C. L OW RANK MAINTENANCE
Revised: C. L OW RANK MAINTENANCE

Original: 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 efﬁciently.
Revised: 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 efﬁciently.

Original: Before moving, we deﬁne some notions related to rectangular matrix multiplication.
Revised: Before moving, we deﬁne some notions related to rectangular matrix multiplication.

Original: Deﬁnition C.1.
Revised: Deﬁnition C.1 (Williams (2012); Gall & Urrutia (2018)) .

Original: Let ω be the matrix multiplication exponent such that it takes n ω + o (1) time to multiply two n × n matrices.
Revised: Let ω be the matrix multiplication exponent such that it takes n ω + o (1) time to multiply two n × n matrices.

Original: Let α be the dual exponent of the matrix multiplication which is the supremum among all a ≥ 0 such that it takes n 2+ o (1) time to multiply an n × n by n × n a matrix.
Revised: Let α be the dual exponent of the matrix multiplication which is the supremum among all a ≥ 0 such that it takes n 2+ o (1) time to multiply an n × n by n × n a matrix.

Original: Additionally, we deﬁne the function ω ( · ) where ω ( b ) denotes the exponent of multiplying an n × n matrix by an n × n b matrix.
Revised: Additionally, we deﬁne the function ω ( · ) where ω ( b ) denotes the exponent of multiplying an n × n matrix by an n × n b matrix.

Original: Hence, we have ω (1) =
Revised: Hence, we have ω (1) =

Original: ω and ω ( α )
Revised: ω and ω ( α )

Original: = 2 .
Revised: = 2 .

Original: 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
Revised: 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.

Original: 1: data structure L OW R ANK M AINTENANCE (cid:46) Lemma C.
Revised: 1: data structure L OW R ANK M AINTENANCE (cid:46) Lemma C.

Original: 2: members 3: r (cid:96) , ∀ (cid:96) ∈
Revised: 2: members 3: r (cid:96) , ∀ (cid:96) ∈

Original: [ L ] (cid:46) r (cid:96) denotes the accumulated rank of the change 4:
Revised: [ L ] (cid:46) r (cid:96) denotes the accumulated rank of the change 4:

Original: W (cid:96) , ∀ (cid:96) ∈
Revised: W (cid:96) , ∀ (cid:96) ∈

Original: [ L ] (cid:46) { W (cid:96) } L(cid:96)
Revised: [ L ] (cid:46) { W (cid:96) } L(cid:96)

Original: =1 ∈ ( R m × m )
Revised: =1 ∈ ( R m × m )

Original: L 5: ∆ W (cid:96) , ∀ (cid:96) ∈
Revised: L 5: ∆ W (cid:96) , ∀ (cid:96) ∈

Original: [ L ] (cid:46) { ∆ W (cid:96) } L(cid:96) =
Revised: [ L ] (cid:46) { ∆ W (cid:96) } L(cid:96) =

Original: 1 ∈ ( R m × m )
Revised: 1 ∈ ( R m × m )

Original: L 6: end members 7: 8: procedures 9:
Revised: L 6: end members 7: 8: procedures 9:

Original: I NIT ( { W 1 (0) , . . .
Revised: I NIT ( { W 1 (0) , . . .

Original: W L (0) } )
Revised: W L (0) } )

Original: (cid:46) Initialize the data structure 10:
Revised: (cid:46) Initialize the data structure 10:

Original: U PDATE ( U (cid:96) , V (cid:96) )
Revised: U PDATE ( U (cid:96) , V (cid:96) )

Original: (cid:46) Update the low rank representation 11:
Revised: (cid:46) Update the low rank representation 11:

Original: Q UERY ( (cid:96), y ) (cid:46) Compute the matrix-vector product between ∆ W (cid:96) and y 12: end procedures 13: end data structure
Revised: Q UERY ( (cid:96), y ) (cid:46) Compute the matrix-vector product between ∆ W (cid:96) and y 12: end procedures 13: end data structure

Original: 1: procedure I NIT ( { W 1 (0) , . .
Revised: 1: procedure I NIT ( { W 1 (0) , . .

Original: , W L (0) } )
Revised: , W L (0) } )

Original: (cid:46) Lemma C.
Revised: (cid:46) Lemma C.

Original: W (cid:96) ← W (cid:96) (0) 3: ∆ W (cid:96) ← 0 , ∀ (cid:96) ∈
Revised: W (cid:96) ← W (cid:96) (0) 3: ∆ W (cid:96) ← 0 , ∀ (cid:96) ∈

Original: [ L ] 4: r (cid:96) ← 0 , ∀ (cid:96) ∈
Revised: [ L ] 4: r (cid:96) ← 0 , ∀ (cid:96) ∈

Original: [ L ] 5: end procedure 6: 7: procedure U PDATE ( U (cid:96)
Revised: [ L ] 5: end procedure 6: 7: procedure U PDATE ( U (cid:96)

Original: ∈
Revised: ∈

Original: R m × n , V (cid:96) ∈
Revised: R m × n , V (cid:96) ∈

Original: R m × n ) (cid:46) Lemma C.
Revised: R m × n ) (cid:46) Lemma C.

Original: 8: ∆ W (cid:96) ← ∆ W (cid:96) + U (cid:96) V (cid:62) (cid:96) without forming the product and sum the two matrices 9: r (cid:96) ← r (cid:96) + n 10: if r (cid:96) =
Revised: 8: ∆ W (cid:96) ← ∆ W (cid:96) + U (cid:96) V (cid:62) (cid:96) without forming the product and sum the two matrices 9: r (cid:96) ← r (cid:96) + n 10: if r (cid:96) =

Original: m a where a = ω (2) then 11:
Revised: m a where a = ω (2) then 11:

Original: W (cid:96) ← W (cid:96) + ∆ W (cid:96) (cid:46) Takes O ( m 2 ) time 12: r (cid:96) ←13: ∆ W (cid:96) ←14: end if 15: end procedure 16: 17: procedure Q UERY (
Revised: W (cid:96) ← W (cid:96) + ∆ W (cid:96) (cid:46) Takes O ( m 2 ) time 12: r (cid:96) ←13: ∆ W (cid:96) ←14: end if 15: end procedure 16: 17: procedure Q UERY (

Original: (cid:96) ∈
Revised: (cid:96) ∈

Original: [ L ] , y ∈ R m ) (cid:46) Lemma C.
Revised: [ L ] , y ∈ R m ) (cid:46) Lemma C.

Original:  z ←( W (cid:96)  + ∆ W (cid:96) ) y (cid:46) Takes O (nnz( y ) · m + mr (cid:96) ) time 19: return z 20: end procedure
Revised: 18: z ← W (cid:96) · y + ∆ W (cid:96) · y (cid:46) Takes O (nnz( y ) · m + mr (cid:96) ) time 19: return z 20: end procedure

Original: Lemma C.2.
Revised: Lemma C.2.

Original: There exists a deterministic data structure (Algorithm 3) such that maintains ∆ W 1 , . . .
Revised: There exists a deterministic data structure (Algorithm 3) such that maintains ∆ W 1 , . . .

Original: , ∆ W L such that
Revised: , ∆ W L such that

Original: • The procedure I NIT (Algorithm 4) takes O ( m 2 L ) time.
Revised: • The procedure I NIT (Algorithm 4) takes O ( m 2 L ) time.

Original: • The procedure U PDATE (Algorithm 4) takes O ( nm 2 − α ) amortized time, where α = ω (2) • The procedure Q UERY (Algorithm 4) takes O ( m · (nnz( y )
Revised: • The procedure U PDATE (Algorithm 4) takes O ( nm 2 − α ) amortized time, where α = ω (2) • The procedure Q UERY (Algorithm 4) takes O ( m · (nnz( y )

Original: + r (cid:96) )) time, where r (cid:96) is the rank of ∆ W (cid:96) when Q UERY is called.
Revised: + r (cid:96) )) time, where r (cid:96) is the rank of ∆ W (cid:96) when Q UERY is called.

Original: Proof.
Revised: Proof.

Original: The runtime for I NIT is obvious, for Q UERY , notice that we are multiplying vector y with a (possibly) dense matrix W (cid:96) ∈
Revised: The runtime for I NIT is obvious, for Q UERY , notice that we are multiplying vector y with a (possibly) dense matrix W (cid:96) ∈

Original: R m × m , which takes O (nnz( y ) · m ) time, and an accumulated low rank matrix ∆ W (cid:96) with rank r (cid:96) .
Revised: R m × m , which takes O (nnz( y ) · m ) time, and an accumulated low rank matrix ∆ W (cid:96) with rank r (cid:96) .

Original: By using the low rank decomposition ∆ W (cid:96) =
Revised: By using the low rank decomposition ∆ W (cid:96) =

Original: UV (cid:62) with
Revised: UV (cid:62) with

Original: U, V ∈ R m × r (cid:96) , the time to multiply y with ∆ W is O ( mr (cid:96) ) .
Revised: U, V ∈ R m × r (cid:96) , the time to multiply y with ∆ W is O ( mr (cid:96) ) .

Original: Combine them together, we get a running time of O ( m · (nnz( y )
Revised: Combine them together, we get a running time of O ( m · (nnz( y )

Original: + r (cid:96) )) .
Revised: + r (cid:96) )) .

Original: It remains to analyze the amortized cost of U PDATE .
Revised: It remains to analyze the amortized cost of U PDATE .

Original: Note that if r (cid:96) <
Revised: Note that if r (cid:96) <

Original: m a , then we just pay O (1) time to update corresponding variables in the data structure.
Revised: m a , then we just pay O (1) time to update corresponding variables in the data structure.

Original: If r (cid:96) =
Revised: If r (cid:96) =

Original: m a , then we will explicitly form the m × m matrix ∆ W (cid:96) .
Revised: m a , then we will explicitly form the m × m matrix ∆ W (cid:96) .

Original: To form it, notice we have accumulated r (cid:96) /n different sums of rank- n decompositions, which can be represented as
Revised: To form it, notice we have accumulated r (cid:96) /n different sums of rank- n decompositions, which can be represented as

Original: U =
Revised: U =

Original: [Equation] U [Equation] 1) , U (cid:96) (2) , . . 
Revised: [ U (cid:96) (1) , U (cid:96) (2) , . . .

Original: , U (cid:96) ( r (cid:96) /n )]
Revised: , U (cid:96) ( r (cid:96) /n )]

Original: ∈ R m × r (cid:96) , V =
Revised: ∈ R m × r (cid:96) , V =

Original: [ V (cid:96) (1) , V (cid:96) (2) , . .
Revised: [ V (cid:96) (1) , V (cid:96) (2) , . .

Original: , V (cid:96) ( r (cid:96) /n )]
Revised: , V (cid:96) ( r (cid:96) /n )]

Original: ∈ R m × r (cid:96) , and ∆ W (cid:96) = UV (cid:62) , which takes O ( m 2 ) time to compute since r (cid:96) =
Revised: ∈ R m × r (cid:96) , and ∆ W (cid:96) = UV (cid:62) , which takes O ( m 2 ) time to compute since r (cid:96) =

Original: m a and a = ω (2) .
Revised: m a and a = ω (2) .

Original: Finally, note that this update of W (cid:96) only happens once per r (cid:96) /n number of calls to U PDATE , therefore we can charge each step by O ( m 2 r (cid:96) /n )
Revised: Finally, note that this update of W (cid:96) only happens once per r (cid:96) /n number of calls to U PDATE , therefore we can charge each step by O ( m 2 r (cid:96) /n )

Original: =
Revised: =

Original: O ( m 2 − a n ) =
Revised: O ( m 2 − a n ) =

Original: O ( m 2 − α n ) , arrives at our ﬁnal amortized running time.
Revised: O ( m 2 − α n ) , arrives at our ﬁnal amortized running time.

Original: Remark C.3.
Revised: Remark C.3.

Original: Currently, the dual matrix multiplication exponent α ≈ 0 .
Revised: Currently, the dual matrix multiplication exponent α ≈ 0 .

Original: 31 (Gall & Urrutia, 2018), hence the amortized time for U PDATE is O ( nm 1 . 69 ) 
Revised: 31 (Gall & Urrutia, 2018), hence the amortized time for U PDATE is O ( nm 1 . 69 ) .

Original: = o ( m 2 ) .
Revised: If m ≥ n 10 / 3 , then we achieve an update time of o ( m 2 ) .

Original: Similarly, the time for Q UERY is O ( m · (nnz( y )
Revised: Similarly, the time for Q UERY is O ( m · (nnz( y )

Original: + r (cid:96) ))
Revised: + r (cid:96) ))

Original: =
Revised: =

Original: O ( m · nnz( y )+ m 1+ α ) =
Revised: O ( m · nnz( y ) + m 1+ α ) =

Original: O ( m · nnz( y )+ m 1 . 31 ) , as long as nnz( y ) =
Revised: O ( m · nnz( y ) + m 1 . 31 ) , as long as nnz( y ) =

Original: o ( m ) , then its running time is also o ( m 2 ) .
Revised: o ( m ) , then its running time is also o ( m 2 ) .

Original: In our application of training neural networks, we will make sure that the inputted vector y is sparse.
Revised: In our application of training neural networks, we will make sure that the inputted vector y is sparse.

Original: C. E FFICIENTLY COMPUTE u i,(cid:96) ( t ) AND v i,(cid:96) ( t )
Revised: C. E FFICIENTLY COMPUTE u i,(cid:96) ( t ) AND v i,(cid:96) ( t )

Original: In this section, we show that how to compute the vectors u i,(cid:96) , v i,(cid:96) ∈
Revised: In this section, we show that how to compute the vectors u i,(cid:96) , v i,(cid:96) ∈

Original: R m using the low rank structure of the change of weights combined.
Revised: R m using the low rank structure of the change of weights combined.

Original: Recall the deﬁnition of these vectors:
Revised: Recall the deﬁnition of these vectors:

Original: [Equation]
Revised: [Equation]

Original: Before proceeding, we list the assumptions we will be using:
Revised: Before proceeding, we list the assumptions we will be using:

Original: • For any (cid:96) ∈
Revised: • For any (cid:96) ∈

Original: [ L ] , D i,(cid:96) ( t ) is s D -sparse, where s D :=
Revised: [ L ] , D i,(cid:96) ( t ) is s D -sparse, where s D :=

Original: k + s , k is the sparsity of D i,(cid:96) (0) and s is the sparsity of D i,(cid:96)
Revised: k + s , k is the sparsity of D i,(cid:96) (0) and s is the sparsity of D i,(cid:96)

Original: ( t ) − D i,(cid:96) (0) .
Revised: ( t ) − D i,(cid:96) (0) .

Original: • For any (cid:96) ∈
Revised: • For any (cid:96) ∈

Original: [ L ] , the change of the weight matrix W (cid:96) , ∆ W (cid:96) ( t ) :
Revised: [ L ] , the change of the weight matrix W (cid:96) , ∆ W (cid:96) ( t ) :

Original: = W (cid:96) ( t ) − W (cid:96) (0) , is of low-rank.
Revised: = W (cid:96) ( t ) − W (cid:96) (0) , is of low-rank.

Original: That is, ∆ W (cid:96) ( t ) =
Revised: That is, ∆ W (cid:96) ( t ) =

Original: (cid:80) [Equation] j 
Revised: (cid:80) r t j =1 y

Original: [Equation] 1 y [Equation] ,j z (cid:62) (cid:96),j .
Revised: (cid:96),j z (cid:62) (cid:96),j .

Original: • For any i ∈ [ n ] , W 1 (0) x
Revised: • For any i ∈ [ n ] , W 1 (0) x

Original: i is pre-computed.
Revised: i is pre-computed.

Original: We ﬁrst note that as a direct consequence of D i,(cid:96) (0) is k -sparse, h i,(cid:96) (0) is k -sparse as well.
Revised: We ﬁrst note that as a direct consequence of D i,(cid:96) (0) is k -sparse, h i,(cid:96) (0) is k -sparse as well.

Original: Similarly, h i,(cid:96) ( t ) − h i, [Equation] (0) has sparsity s .
Revised: Similarly, h i,(cid:96) ( t ) − h i,(cid:96) (0) has sparsity s .

Original: Hence h i,(cid:96) ( t ) has sparsity bounded by s D .
Revised: Hence h i,(cid:96) ( t ) has sparsity bounded by s D .

Original: Compute u i,(cid:96) ( t ) .
Revised: Compute u i,(cid:96) ( t ) .

Original: Compute u i,(cid:96) ( t ) is equivalent to compute the following vector:
Revised: Compute u i,(cid:96) ( t ) is equivalent to compute the following vector:

Original: [Equation]
Revised: [Equation]

Original: {+(cid:96) ( t )( W (cid:96) +
+}
Revised: {+(cid:96) ( t )( W (cid:96) +
+}

Revised: + ∆ W (cid:96) +1 ( t ))

Revised: (cid:62) D i,(cid:96)

Revised: +1 ( t ) · · · ( W L (0) + ∆ W L ( t )) (cid:62) D

Original:  First, we know that D i,L ( t ) a ∈ R m is an s D -sparse vector, and it takes O ( s D ) time.
Revised: i,L ( t ) a. First, we know that D i,L ( t ) a ∈ R m is an s D -sparse vector, and it takes O ( s D ) time.

Original: The next matrix is ( W L (0)+∆ W L ( t )) (cid:62) , which gives two terms: W L (0) (cid:62) ( D i,L ( t ) a ) and ∆ W L ( t ) (cid:62) ( D i,L ( t ) a ) .
Revised: The next matrix is ( W L (0)+∆ W L ( t )) (cid:62) , which gives two terms: W L (0) (cid:62) ( D i,L ( t ) a ) and ∆ W L ( t ) (cid:62) ( D i,L ( t ) a ) .

Original: For the ﬁrst term, since D i,L ( t ) a is s D -sparse, it takes O ( ms D ) -time.
Revised: For the ﬁrst term, since D i,L ( t ) a is s D -sparse, it takes O ( ms D ) -time.

Original: For the second term, we have∆
Revised: For the second term, we have

Original: [Equation]
Revised: [Equation]

Revised: z

Revised: [Equation]

Revised: y

Revised: [Equation]

Revised: L

Revised: [Equation]

Revised: ( D i

Revised: [Equation]

Revised: ( t

Revised: [Equation]

Revised: a

Revised: [Equation]

Revised: r

Revised: [Equation]

Revised: (cid:88)

Revised: j = z L

Revised: [Equation]

Revised: (cid:104)

Revised: y L

Revised: [Equation]

Revised: D i

Revised: [Equation]

Revised: L ( t ) a (cid:105) .

Original: Each inner-product takes O ( s D ) -time
Revised: Each inner-product takes O ( s D ) -time

Original: and it takes O ( mr t + s D r t ) =
Revised: and it takes O ( mr t + s D r t ) =

Original: O ( mr t ) -time in total.
Revised: O ( mr t ) -time in total.

Original: Hence, in O ( m ( s D + r t )) -time, we compute the vector W L ( t ) (cid:62) D
Revised: Hence, in O ( m ( s D + r t )) -time, we compute the vector W L ( t ) (cid:62) D

Original: i,L ( t ) a .
Revised: i,L ( t ) a .

Original: Note that we do not assume the sparsity of a .
Revised: Note that we do not assume the sparsity of a .

Original: Thus, by repeating this process for the L − (cid:96) intermediate matrices W (cid:62) j ( t ) D i,j ( t ) , we can obtain the vector [Equation] 
Revised: Thus, by repeating this process for the L − (cid:96) intermediate matrices W (cid:62) j ( t ) D i,j ( t ) , we can obtain the vector 

Original: [Equation]
Revised: [Equation]

Original: in time O
Revised: in time O

Original: (( L − (cid:96) ) m
Revised: (( L − (cid:96) ) m

Original: ( s D + r t )) .
Revised: ( s D + r t )) .

Original: Finally, by multiplying a sparse diagonal matrix D i,(cid:96) ( t ) , we get the desired vector u i,(cid:96) ( t ) .
Revised: Finally, by multiplying a sparse diagonal matrix D i,(cid:96) ( t ) , we get the desired vector u i,(cid:96) ( t ) .

Original: Compute v i,(cid:96) ( t ) .
Revised: Compute v i,(cid:96) ( t ) .

Original: Note that v i,(cid:96) ( t ) is essentially h i,(cid:96) − 1 ( t ) , so we consider how to compute h i,(cid:96) ( t ) for general (cid:96)
Revised: Note that v i,(cid:96) ( t ) is essentially h i,(cid:96) − 1 ( t ) , so we consider how to compute h i,(cid:96) ( t ) for general (cid:96)

Original: ∈
Revised: ∈

Original: [ L ] .
Revised: [ L ] .

Original: Recall that
Revised: Recall that

Original: [Equation]
Revised: [Equation]

Original: since h i,(cid:96) − 1 ( t ) is s D -sparse, the product W (cid:96) (0 h i,(cid:96) − 1 ( t ))) can be computed in O ( ms D ) time.
Revised: since h i,(cid:96) − 1 ( t ) is s D -sparse, the product W (cid:96) (0 h i,(cid:96) − 1 ( t ))) can be computed in O ( ms D ) time.

Original: For the product DeltaW (cid:96) ( t ) h i,(cid:96) − 1 ( t ) can be computed use the low rank decomposition, which takes O ( mr t ) time.
Revised: For the product ∆ W (cid:96) ( t ) h i,(cid:96) − 1 ( t ) can be computed use the low rank decomposition, which takes O ( mr t ) time.

Original: Apply the threshold ReLU takes O ( m ) time.
Revised: Apply the threshold ReLU takes O ( m ) time.

Original: Hence, the total time is O ( m ( r t + s D )) time.
Revised: Hence, the total time is O ( m ( r t + s D )) time.

Original: The running time results are summarized in the following lemma:
Revised: The running time results are summarized in the following lemma:

Original: Lemma C.4.
Revised: Lemma C.4.

Original: For (cid:96) ∈
Revised: For (cid:96) ∈

Original: [ L ] and i ∈ [ n ] , suppose (cid:107) D i,(cid:96) (0) (cid:107) 0 ≤ k .
Revised: [ L ] and i ∈ [ n ] , suppose (cid:107) D i,(cid:96) (0) (cid:107) 0 ≤ k .

Original: Let t > 0 .
Revised: Let t > 0 .

Original: Suppose the change of
Revised: Suppose the change of

Original: D i,(cid:96) is sparse, i.e., (cid:107) D i,(cid:96) ( t ) − D i,(cid:96) (0) (cid:107) 0 ≤ s .
Revised: D i,(cid:96) is sparse, i.e., (cid:107) D i,(cid:96) ( t ) − D i,(cid:96) (0) (cid:107) 0 ≤ s .

Original: For (cid:96) ∈
Revised: For (cid:96) ∈

Original: [ L ] , i ∈ [ n ] , for any t > 0 , suppose the change of W (cid:96) is of low-rank, i.e., ∆ W (cid:96) ( t )
Revised: [ L ] , i ∈ [ n ] , for any t > 0 , suppose the change of W (cid:96) is of low-rank, i.e., ∆ W (cid:96) ( t )

Original: = (cid:80) r t j =1 y
Revised: = (cid:80) r t j =1 y

Original: (cid:96),j z (cid:62) (cid:96),j .
Revised: (cid:96),j z (cid:62) (cid:96),j .

Original: We further assume that { y (cid:96),j , z (cid:96),j } (cid:96) ∈
Revised: We further assume that { y (cid:96),j , z (cid:96),j } (cid:96) ∈

Original: [ L ] ,j ∈ [ r t ] and { W 1 (0) x
Revised: [ L ] ,j ∈ [ r t ] and { W 1 (0) x

Original: i } i ∈ [ n ] are pre-computed.
Revised: i } i ∈ [ n ] are pre-computed.

Original: Then, for any (cid:96) ∈
Revised: Then, for any (cid:96) ∈

Original: [ L ] and i ∈ [ n ] , the vectors u (cid:96),i ( t ) , v (
Revised: [ L ] and i ∈ [ n ] , the vectors u (cid:96),i ( t ) , v (

Original: cid:96),i ( t ) ∈
Revised: cid:96),i ( t ) ∈

Original: R m can be computed in O ( mL ( s + k + r t ))
Revised: R m can be computed in O ( mL ( s + k + r t ))

Original: -time.
Revised: -time.

Original: As a direct consequence, if we combine Lemma C.2 and Lemma C.4, then we get the following corollary:
Revised: As a direct consequence, if we combine Lemma C.2 and Lemma C.4, then we get the following corollary:

Original: Corollary C.5.
Revised: Corollary C.5.

Original: For (cid:96) ∈
Revised: For (cid:96) ∈

Original: [ L ] and i ∈ [ n ] , we can compute v i,(cid:96) ( t ) , u i,(cid:96) ( t )
Revised: [ L ] and i ∈ [ n ] , we can compute v i,(cid:96) ( t ) , u i,(cid:96) ( t )

Original: ∈
Revised: ∈

Original: R m as in Algorithmwith the following time bound:
Revised: R m as in Algorithmwith the following time bound:

Original: • Compute u i,(cid:96) ( t ) in time O ( mL ( s + k
Revised: • Compute u i,(cid:96) ( t ) in time O ( mL ( s + k

Original: + r (cid:96) )) .
Revised: + r (cid:96) )) .

Original: • Compute v i,(cid:96) ( t ) in time O
Revised: • Compute v i,(cid:96) ( t ) in time O

Original: ( m ( s + k + r (cid:96) )) .
Revised: ( m ( s + k + r (cid:96) )) .

Original: D F AST T ENSOR P RODUCT R EGRESSION
Revised: D F AST T ENSOR P RODUCT R EGRESSION

Original: In this section, we design a generic algorithm for solving the following type of regression task:
Revised: In this section, we design a generic algorithm for solving the following type of regression task:

Original: Given two matrices U =
Revised: Given two matrices U =

Original: [ u (cid:62) 1 , . . .
Revised: [ u (cid:62) 1 , . . .

Original: , u (cid:62) n ] (cid:62) , V =
Revised: , u (cid:62) n ] (cid:62) , V =

Original: [ v (cid:62) 1 , . .
Revised: [ v (cid:62) 1 , . .

Original: . , v (cid:62) n ] ∈ R m × n with m (cid:29) n , consider the matrix J ∈ R n × m 2 formed by
Revised: . , v (cid:62) n ] ∈ R m × n with m (cid:29) n , consider the matrix J ∈ R n × m 2 formed by

Original: [Equation]
Revised: [Equation]

Original: We are also given a vector c ∈ R n , the goal is to solve the following regression task:
Revised: We are also given a vector c ∈ R n , the goal is to solve the following regression task:

Original: [Equation]
Revised: [Equation]

Original: Our main theorem for this section is as follows:
Revised: Our main theorem for this section is as follows:

Original: Theorem D.1 (Restatement of Theorem D.14) .
Revised: Theorem D.1 (Restatement of Theorem D.14) .

Original: Given two n × m matrices U and V , and a target vector c ∈ R n .
Revised: Given two n × m matrices U and V , and a target vector c ∈ R n .

Original: Let J = [vec( u 1 v (cid:62) 1 ) (cid:62) , . .
Revised: Let J = [vec( u 1 v (cid:62) 1 ) (cid:62) , . .

Original: , vec( u n v (cid:62) n ) (cid:62) ]
Revised: , vec( u n v (cid:62) n ) (cid:62) ]

Original: ∈ R n × m 2 .
Revised: ∈ R n × m 2 .

Original: There is an algorithm (Algorithm 6) takes (cid:101)
Revised: There is an algorithm (Algorithm 6) takes (cid:101)

Original: O
Revised: O

Original: ( nm + n 2 (log( κ/(cid:15) )
Revised: ( nm + n 2 (log( κ/(cid:15) )

Original: + log( m/(cid:15)δ )
Revised: + log( m/(cid:15)δ )

Original: (cid:15) − 2 ) + n ω ) time and outputs a vector (cid:98) x ∈ R n such that
Revised: (cid:15) − 2 ) + n ω ) time and outputs a vector (cid:98) x ∈ R n such that

Original: [Equation]
Revised: [Equation]

Original: holds with probability at least 1 − δ , and κ is the condition number of J .
Revised: holds with probability at least 1 − δ , and κ is the condition number of J .

Original: D. A PPROXIMATE J VIA TensorSketch
Revised: D. A PPROXIMATE J VIA TensorSketch

Original: We introduce the notion of TensorSketch for two vectors:
Revised: We introduce the notion of TensorSketch for two vectors:

Original: Deﬁnition D.2.
Revised: Deﬁnition D.2.

Original: Let h 1 , h 2 : [ m ] → [ s ] be 3-wise independent hash functions, also let σ : [ m ] → {± 1 } be a 4-wise independent random sign function.
Revised: Let h 1 , h 2 : [ m ] → [ s ] be 3-wise independent hash functions, also let σ : [ m ] → {± 1 } be a 4-wise independent random sign function.

Original: The degree two TensorSketch transform,
Revised: The degree two TensorSketch transform,

Original: S : R m × R m → R s is deﬁned as follows: for any i, j ∈ [ m ] and r ∈ [ s ] , 
Revised: S : R m × R m → R s is deﬁned as follows: for any i, j ∈ [ m ] and r ∈ [ s ] , [Equation] j

Original: [Equation]
Revised: [Equation]

Revised: mod s ] .

Original: Remark D.3.
Revised: Remark D.3.

Original: Apply S to two vectors x, y ∈
Revised: Apply S to two vectors x, y ∈

Original: R m can be implemented in time O
Revised: R m can be implemented in time O

Original: ( s log s + nnz( x ) + nnz( y )) .
Revised: ( s log s + nnz( x ) + nnz( y )) .

Original: We introduce one key technical lemma from Avron et al.
Revised: We introduce one key technical lemma from Avron et al.

Original: Lemma D.4 (Theorem 1 of Avron et al.
Revised: Lemma D.4 (Theorem 1 of Avron et al.

Original: Let S ∈ R s × m 2 be the TensorSketch matrix, consider a ﬁxed n -dimensional subspace V .
Revised: Let S ∈ R s × m 2 be the TensorSketch matrix, consider a ﬁxed n -dimensional subspace V .

Original: If s = Ω( n 2 / ( (cid:15) 2 δ )) , then with probability at least 1 − δ , (cid:107) Sx (cid:107) 2 = (1 ± (cid:15) ) (cid:107) x (cid:107) 2 simultaneously for all x ∈ V .
Revised: If s = Ω( n 2 / ( (cid:15) 2 δ )) , then with probability at least 1 − δ , (cid:107) Sx (cid:107) 2 = (1 ± (cid:15) ) (cid:107) x (cid:107) 2 simultaneously for all x ∈ V .

Original: Now we are ready to prove the main lemma of this section:
Revised: Now we are ready to prove the main lemma of this section:

Original: Lemma D.5.
Revised: Lemma D.5.

Original: Let (cid:15), δ ∈ (0 , 1) denote two parameters.
Revised: Let (cid:15), δ ∈ (0 , 1) denote two parameters.

Original: Let J ∈ R n × m 2 be a matrix such that i -th row of J is deﬁned as vec( u i v (cid:62) i ) for some u i , v i ∈
Revised: Let J ∈ R n × m 2 be a matrix such that i -th row of J is deﬁned as vec( u i v (cid:62) i ) for some u i , v i ∈

Original: R m .
Revised: R m .

Original: Then, we can compute a matrix (cid:101) J ∈ R n × s such that for any vector x ∈ R n , with probability at least 1 − δ , we have
Revised: Then, we can compute a matrix (cid:101) J ∈ R n × s such that for any vector x ∈ R n , with probability at least 1 − δ , we have

Original: [Equation]
Revised: [Equation]

Original: where s = Ω( n 2 / ( (cid:15) 2 δ )) .
Revised: where s = Ω( n 2 / ( (cid:15) 2 δ )) .

Original: The time to compute (cid:101) J is O ( ns log s + nnz( U ) + nnz( V )) .
Revised: The time to compute (cid:101) J is O ( ns log s + nnz( U ) + nnz( V )) .

Original: Proof.
Revised: Proof.

Original: 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 = Ω( n 2 / ( (cid:15) 2 δ )) can preserve the length of all vectors in the subspace generated by J (cid:62) with probability 1 − δ , to a multiplicative factor of 1 ± (cid:15) .
Revised: 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 = Ω( n 2 / ( (cid:15) 2 δ )) can preserve the length of all vectors in the subspace generated by J (cid:62) with probability 1 − δ , to a multiplicative factor of 1 ± (cid:15) .

Original: The running time part is to apply the FFT algorithm to each row of J with a total of n rows.
Revised: The running time part is to apply the FFT algorithm to each row of J with a total of n rows.

Original: For each row, it takes O ( s log s + m ) time, hence the overall running time is O ( n ( s log s + m )) .
Revised: For each row, it takes O ( s log s + m ) time, hence the overall running time is O ( n ( s log s + m )) .

Original: D. A PPROXIMATE J VIA TensorSRHT
Revised: D. A PPROXIMATE J VIA TensorSRHT

Original: We note that the dependence on the target dimension of sketching is O (1 /δ ) for TensorSketch .
Revised: We note that the dependence on the target dimension of sketching is O (1 /δ ) for TensorSketch .

Original: We introduce another kind of sketching technique for tensor, called TensorSRHT .
Revised: We introduce another kind of sketching technique for tensor, called TensorSRHT .

Original: The tradeoff is we lose input sparsity runtime of matrices U and V .
Revised: The tradeoff is we lose input sparsity runtime of matrices U and V .

Original: Deﬁnition D.6.
Revised: Deﬁnition D.6.

Original: We deﬁne the TensorSRHT S : R d × R d → R m as S = 1 √ m P · ( HD 1 × HD 2 ) , where each row of P ∈ { 0 , 1 } m × d  contains only one 1 at a random coordinate, one can view P as a sampling matrix.
Revised: We deﬁne the TensorSRHT S : R d × R d → R m as S = 1 √ m P · ( HD 1 × HD 2 ) , where each row of P ∈ { 0 , 1 } m × d 2 contains only one 1 at a random coordinate, one can view P as a sampling matrix.

Original: H is a d × d Hadamard matrix, and D 1 , D 2 are two d × d independent diagonal matrices with diagonals that are each independently set to be a Rademacher random variable (uniform in {− 1 , 1 } ).
Revised: H is a d × d Hadamard matrix, and D 1 , D 2 are two d × d independent diagonal matrices with diagonals that are each independently set to be a Rademacher random variable (uniform in {− 1 , 1 } ).

Original: Remark D.7.
Revised: Remark D.7.

Original: By using FFT algorithm, apply S to two vectors x, y ∈
Revised: By using FFT algorithm, apply S to two vectors x, y ∈

Original: R m takes time O
Revised: R m takes time O

Original: ( m log m + s ) .
Revised: ( m log m + s ) .

Original: We again introduce a technical lemma for TensorSRHT .
Revised: We again introduce a technical lemma for TensorSRHT .

Original: Lemma D.8 (Theorem 3 of Ahle et al.
Revised: Lemma D.8 (Theorem 3 of Ahle et al.

Original: Let S ∈ R s × m 2 be the TensorSRHT matrix, consider a ﬁxed n -dimensional subspace V .
Revised: Let S ∈ R s × m 2 be the TensorSRHT matrix, consider a ﬁxed n -dimensional subspace V .

Original: If s = Ω( n log 3 ( nm/(cid:15)δ ) (cid:15) − 2 ) , then with probability at least 1 − δ , (cid:107) Sx (cid:107) [Equation] = ( [Equation]
Revised: If s = Ω( n log 3 ( nm/(cid:15)δ ) (cid:15) − 2 ) , then with probability at least 1 − δ , (cid:107) Sx (cid:107) 2 = (1 ± (cid:15) )

Original: (cid:15) ) (cid:107) [Equation] (cid:107) 2 simultaneously for all x ∈ V .
Revised: (cid:107) x (cid:107) 2 simultaneously for all x ∈ V .

Original: Lemma D.9.
Revised: Lemma D.9.

Original: Let (cid:15), δ ∈ (0 , 1) denote two parameters.
Revised: Let (cid:15), δ ∈ (0 , 1) denote two parameters.

Original: Given a list of vectors u 1 , · · · , u m , v 1 , · · · , v m ∈ R m .
Revised: Given a list of vectors u 1 , · · · , u m , v 1 , · · · , v m ∈ R m .

Original: Let J ∈ R n × m 2 be a matrix with i -th row of J is deﬁned as vec( u i v (cid:62) i ) .
Revised: Let J ∈ R n × m 2 be a matrix with i -th row of J is deﬁned as vec( u i v (cid:62) i ) .

Original: Then, we can compute a matrix (cid:101) J ∈ R n × s such that for any vector x ∈ R n , with probability at least 1 − δ , we have
Revised: Then, we can compute a matrix (cid:101) J ∈ R n × s such that for any vector x ∈ R n , with probability at least 1 − δ , we have

Original: [Equation]
Revised: [Equation]

Original: where s = Ω( n log 3 ( nm/ ( (cid:15)δ ))
Revised: where s = Ω( n log 3 ( nm/ ( (cid:15)δ ))

Original: (cid:15) − 2 ) .
Revised: (cid:15) − 2 ) .

Original: The time to compute (cid:101) J is O ( n ( m log m + s )) .
Revised: The time to compute (cid:101) J is O ( n ( m log m + s )) .

Original: Proof.
Revised: Proof.

Original: The correctness follows directly from Lemma D.8.
Revised: The correctness follows directly from Lemma D.8.

Original: The running time follows from the FFT algorithm to each row of J , each application takes O ( m log m + s ) time, and we need to apply it to n rows.
Revised: The running time follows from the FFT algorithm to each row of J , each application takes O ( m log m + s ) time, and we need to apply it to n rows.

Original: D. T ENSOR TRICK AS A PRECONDITIONER
Revised: D. T ENSOR TRICK AS A PRECONDITIONER

Original: In this section, we use TensorSketch and TensorSRHT as a preconditioner to solve a regression task involving JJ (cid:62) .
Revised: In this section, we use TensorSketch and TensorSRHT as a preconditioner to solve a regression task involving JJ (cid:62) .

Original: This is an important step to implement the Gram-Newton-Gauss iteration Cai et al.
Revised: This is an important step to implement the Gram-Newton-Gauss iteration Cai et al.

Original: (2019); Brand et al.
Revised: (2019); Brand et al.

Original: Before proceeding, we introduce the notion of subspace embedding :
Revised: Before proceeding, we introduce the notion of subspace embedding :

Original: Deﬁnition D.10 (Subspace Embedding, Sarlos (2006)) .
Revised: Deﬁnition D.10 (Subspace Embedding, Sarlos (2006)) .

Original: Let A ∈ R N × k , we say a matrix S ∈ R s × N is a (1 ± (cid:15) ) − (cid:96) 2 subspace embedding for A if for any x ∈ R k , we have (cid:107) SAx (cid:107) 

Revised: Let A ∈ R N × k , we say a matrix S ∈ R s × N is a (1 ± (cid:15) ) − (cid:96) 2 subspace embedding for A if for any x ∈ R k , we have (cid:107) SAx (cid:107)

Original: = (1 ± (cid:15) )
Revised: = (1 ± (cid:15) )

Original: (cid:107) Ax (cid:107) 22 .
Revised: (cid:107) Ax (cid:107) 22 .

Original: Equivalently, (cid:107) I − U (cid:62) S (cid:62) SU (cid:107) ≤ (cid:15) where U is an orthonormal basis for the column space of A .
Revised: Equivalently, (cid:107) I − U (cid:62) S (cid:62) SU (cid:107) ≤ (cid:15) where U is an orthonormal basis for the column space of A .

Original: We will mainly utilize efﬁcient subspace embedding.
Revised: We will mainly utilize efﬁcient subspace embedding.

Original: Deﬁnition D.11 (Lu et al.
Revised: Deﬁnition D.11 (Lu et al.

Original: ; Woodruff (2014)) .
Revised: ; Woodruff (2014)) .

Original: Given a matrix A ∈ R N × k with N = poly( k ) , then we can compute an S ∈ R k poly(log( k/δ )) /(cid:15)
Revised: Given a matrix A ∈ R N × k with N = poly( k ) , then we can compute an S ∈ R k poly(log( k/δ )) /(cid:15)

Original: 2 × N such that with probability at least 1 − δ , we have (cid:107) SAx (cid:107) 2 =
Revised: 2 × N such that with probability at least 1 − δ , we have (cid:107) SAx (cid:107) 2 =

Original: (1 ± (cid:15) )
Revised: (1 ± (cid:15) )

Original: (cid:107) Ax (cid:107)hols for all x ∈ R k .
Revised: (cid:107) Ax (cid:107)hols for all x ∈ R k .

Original: Moreover, SA can be computed in O ( Nk log(( k log N ) /(cid:15) )) time.
Revised: Moreover, SA can be computed in O ( Nk log(( k log N ) /(cid:15) )) time.

Original: Lemma D.12 (Lemma 4.2 of Brand et al. (2021)) .

Original: Let N = Ω

Original: ( k poly(log k )) .

Original: Given a matrix

Original: A ∈ R N × k , let κ denote its condition number.

Original: Consider the following regression task:

Original: [Equation]

Original: 1: procedure F AST R EGRESSION ( A, y, (cid:15) )
Revised: 1: procedure F AST R EGRESSION ( A, y, (cid:15) )

Original: (cid:46) Lemma D.
Revised: (cid:46) Lemma D.

Original: 2: (cid:46)
Revised: 2: (cid:46)

Original: A ∈ R N × k is full rank, (cid:15) ∈ (0 , 1 / 2) 3:
Revised: A ∈ R N × k is full rank, (cid:15) ∈ (0 , 1 / 2) 3:

Original: Compute a subspace embedding SA (cid:46) S ∈ R k poly(log k ) × N 4:
Revised: Compute a subspace embedding SA (cid:46) S ∈ R k poly(log k ) × N 4:

Original: Compute R such that SAR has orthonormal columns via QR decomposition (cid:46)
Revised: Compute R such that SAR has orthonormal columns via QR decomposition (cid:46)

Original: R ∈ R k ×
Revised: R ∈ R k ×

Original: k 5:
Revised: k 5:

Original: z 0 = 

Revised: z 0 =

Original: k ∈ R k 6: t ←7: while (cid:107) A (cid:62) ARz t − y (cid:107) 2 ≥ (cid:15) do 8: z t +1 ← z t − ( R (cid:62) A (cid:62) AR ) (cid:62) ( R (cid:62) A (cid:62) ARz t − R (cid:62) y ) 9: t ← t +10: end while 11: return Rz t 12: end procedure 
Revised: k ∈ R k 6: t ←7: while (cid:107) A (cid:62) ARz t − y (cid:107) 2 ≥ (cid:15) do 8: z t +1 ← z t − ( R (cid:62) A (cid:62) AR ) (cid:62) ( R (cid:62) A (cid:62) ARz t − R (cid:62) y ) 9: t ← t +10: end while 11: return Rz t 12: end procedure Lemma D.12 (Lemma 4.2 of Brand et al. (2021)) . Let N = Ω ( k poly(log k )) .

Revised: Given a matrix

Original:  Using the procedure F AST R EGRESSION (Algorithm 5), with probability at least 1 − δ , we can compute an (cid:15) -approximate solution (cid:98) x satisfying
Revised: A ∈ R N × k , let κ denote its condition number. Consider the following regression task: [Equation] Using the procedure F AST R EGRESSION (Algorithm 5), with probability at least 1 − δ , we can compute an (cid:15) -approximate solution (cid:98) x satisfying

Original: [Equation]
Revised: [Equation]

Original: in time (cid:101) O
Revised: in time (cid:101) O

Original: ( Nk log( κ/(cid:15) )
Revised: ( Nk log( κ/(cid:15) )

Original: + k ω ) .
Revised: + k ω ) .

Original: Our algorithm is similar to the ridge regression procedure in Song et al.
Revised: Our algorithm is similar to the ridge regression procedure in Song et al.

Original: (2021), where they ﬁrst 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.
Revised: (2021), where they ﬁrst 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.

Original: We shall ﬁrst prove a useful lemma.
Revised: We shall ﬁrst prove a useful lemma.

Original: Lemma D.13.
Revised: Lemma D.13.

Original: Let A ∈ R N × k , suppose SA is a subspace embedding for A (Def. D.11), then we have for any x ∈ R k , with probability at least 1 − δ ,
Revised: Let A ∈ R N × k , suppose SA is a subspace embedding for A (Def. D.11), then we have for any x ∈ R k , with probability at least 1 − δ ,

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Throughout the proof, we condition on the event that S preserves the length of all vectors in the column space of A .
Revised: Throughout the proof, we condition on the event that S preserves the length of all vectors in the column space of A .

Original: Note that
Revised: Note that

Original: [Equation]
Revised: [Equation]

Original: We will ﬁrst bound the norm of ( SA ) (cid:62) SAx , then the inner product term.
Revised: We will ﬁrst bound the norm of ( SA ) (cid:62) SAx , then the inner product term.

Original: Bounding (cid:107) ( SA ) (cid:62) SAx (cid:107)
Revised: Bounding (cid:107) ( SA ) (cid:62) SAx (cid:107)

Original: Let U ∈ R N × k be an orthonormal basis of A , then use the equivalent deﬁnition of subspace embedding, we have (cid:107) U (cid:62) S (cid:62) SU − I (cid:107) ≤ (cid:15) , this means all the eigenvalues of U (cid:62) S (cid:62) SU lie in the range of of [(1 − (cid:15) ) 2 , (1 + (cid:15) ) 2 ] .
Revised: Let U ∈ R N × k be an orthonormal basis of A , then use the equivalent deﬁnition of subspace embedding, we have (cid:107) U (cid:62) S (cid:62) SU − I (cid:107) ≤ (cid:15) , this means all the eigenvalues of U (cid:62) S (cid:62) SU lie in the range of of [(1 − (cid:15) ) 2 , (1 + (cid:15) ) 2 ] .

Original: Let V denote the matrix U (cid:62) S (cid:62) SU , then we know that all eigenvalues of
Revised: Let V denote the matrix U (cid:62) S (cid:62) SU , then we know that all eigenvalues of

Original: V (cid:62) V lie in range [(1 − (cid:15) ) 4 , (1 + (cid:15) ) 4 ] .
Revised: V (cid:62) V lie in range [(1 − (cid:15) ) 4 , (1 + (cid:15) ) 4 ] .

Original: Setting (cid:15) as (cid:15)/ 4 , we arrive at (cid:107) V (cid:62) V − I (cid:107) ≤ (cid:15) .
Revised: Setting (cid:15) as (cid:15)/ 4 , we arrive at (cid:107) V (cid:62) V − I (cid:107) ≤ (cid:15) .

Original: This shows that for any x ∈ R k , we have (cid:107) ( SA ) (cid:62) SAx (cid:107) 2 ={+
+}
Revised: This shows that for any x ∈ R k , we have (cid:107) ( SA ) (cid:62) SAx (cid:107) {+
+}

Original: [Equation]
Revised: = (1 ± (cid:15) )

Original: [Equation] A (cid:62) Ax (cid:107) 2 .
Revised: (cid:107) A (cid:62) Ax (cid:107) 2 .

Original: Bounding (cid:104) ( SA ) (cid:62) SAx, b (cid:105) Note that [Equation] Combining these two terms, we conclude that, with probability at least 1 − δ ,
Revised: Bounding (cid:104) ( SA ) (cid:62) SAx, b (cid:105) Note that (cid:104) ( SA ) (cid:62) SAx, b (cid:105) = (cid:104) SAx, SAb (cid:105) =

Original: [Equation]
Revised: [Equation]

Original: ( SA

Original: [Equation]
Revised: [Equation]

Original: SAx
Revised: SAx

Original: [Equation]
Revised: [Equation]

Original: 1 ±

Original: [Equation]
Revised: [Equation]

Original: A
Revised: Ax

Original: [Equation]
Revised: [Equation]

Original: b

Original: [Equation]
Revised: [Equation]

Original: 1: procedure F AST T ENSOR R EGRESSION ( { u
Revised: (cid:107)

Original: [Equation]
Revised: [Equation]

Original: ni =
Revised: Combining these two terms, we conclude that, with probability at least 1 − δ ,

Original: [Equation]
Revised: [Equation]

Original:  ) (cid:46)
Revised: 1: procedure F AST T ENSOR R EGRESSION ( { u i } ni =1 ∈ R m × n , { v i } ni =1 ∈ R m × n , c ∈ R n ) (cid:46)

Original: Theorem D.
Revised: Theorem D.

Revised: 2: (cid:46)

Revised: [Equation]

Revised: m

Original: [Equation]
Revised: [Equation]

Original: 3: s 1 ← Θ( n log 3 ( nm/ ( (cid:15)δ ))
Revised: 3: s 1 ← Θ( n log 3 ( nm/ ( (cid:15)δ ))

Original: (cid:15) − 2 ) 4: s 2 ← Θ(( n + log m ) log n ) 5:
Revised: (cid:15) − 2 ) 4: s 2 ← Θ(( n + log m ) log n ) 5:

Original: Let S 1 ∈ R s 1 ×
Revised: Let S 1 ∈ R s 1 ×

Original: m 2 be a sketching matrix (cid:46) S 1 can be TensorSketch or TensorSRHT 6:
Revised: m 2 be a sketching matrix (cid:46) S 1 can be TensorSketch or TensorSRHT 6:

Original: Compute (cid:101) J = JS (cid:62) 1 via FFT algorithm (cid:46) (cid:101) J ∈ R n × s7:
Revised: Compute (cid:101) J = JS (cid:62) 1 via FFT algorithm (cid:46) (cid:101) J ∈ R n × s7:

Original: Let S 2 ∈ R s 2 × s 1 be a sketching matrix deﬁned in Deﬁnition D.
Revised: Let S 2 ∈ R s 2 × s 1 be a sketching matrix deﬁned in Deﬁnition D.

Original: Compute a subspace embedding S 2 (cid:101) J (cid:62) 9:
Revised: Compute a subspace embedding S 2 (cid:101) J (cid:62) 9:

Original: Compute R such that S 2 (cid:101) J (cid:62) R has orthonormal columns via QR decomposition (cid:46)
Revised: Compute R such that S 2 (cid:101) J (cid:62) R has orthonormal columns via QR decomposition (cid:46)

Original: R ∈ R n × n 10: z 0 ← 0 k ∈ R k 11: t ←12: while (cid:107) (cid:101) J (cid:101) J (cid:62)
Revised: R ∈ R n × n 10: z 0 ← 0 k ∈ R k 11: t ←12: while (cid:107) (cid:101) J (cid:101) J (cid:62)

Original: Rz t − c (cid:107) 2 ≥ (cid:15) do 13:
Revised: Rz t − c (cid:107) 2 ≥ (cid:15) do 13:

Original: z t +1 ← z t − ( R (cid:62) (cid:101) J (cid:101) J (cid:62) R ) (cid:62) ( R (cid:62) (cid:101) J (cid:101) J (cid:62)
Revised: z t +1 ← z t − ( R (cid:62) (cid:101) J (cid:101) J (cid:62) R ) (cid:62) ( R (cid:62) (cid:101) J (cid:101) J (cid:62)

Original: Rz t − R (cid:62) c ) 14: t ← t +15: end while 16: return Rz t 17: end procedure
Revised: Rz t − R (cid:62) c ) 14: t ← t +15: end while 16: return Rz t 17: end procedure

Original: Theorem D.14.
Revised: Theorem D.14.

Original: Given two n × m matrices U and V , and a target vector c ∈ R n .
Revised: Given two n × m matrices U and V , and a target vector c ∈ R n .

Original: Let J = [vec( u 1 v (cid:62) 1 ) (cid:62) , . .
Revised: Let J = [vec( u 1 v (cid:62) 1 ) (cid:62) , . .

Original: , vec( u n v (cid:62) n ) (cid:62) ]
Revised: , vec( u n v (cid:62) n ) (cid:62) ]

Original: ∈ R n × m 2 .
Revised: ∈ R n × m 2 .

Original: There is an algorithm (Algorithm 6) takes (cid:101)
Revised: There is an algorithm (Algorithm 6) takes (cid:101)

Original: O
Revised: O

Original: ( nm + n 2 (log( κ/(cid:15) )
Revised: ( nm + n 2 (log( κ/(cid:15) )

Original: + log( m/δ ))
Revised: + log( m/δ ))

Original: + n ω ) time and outputs a vector (cid:98) x ∈ R n such that
Revised: + n ω ) time and outputs a vector (cid:98) x ∈ R n such that

Original: [Equation]
Revised: [Equation]

Original: holds with probability at least 1 − δ , and κ is the condition number of J .
Revised: holds with probability at least 1 − δ , and κ is the condition number of J .

Original: Proof.
Revised: Proof.

Original: We can decompose Algorithm 6 into two parts:
Revised: We can decompose Algorithm 6 into two parts:

Original: • Applying S 1 to efﬁciently form matrix (cid:101)
Revised: • Applying S 1 to efﬁciently form matrix (cid:101)

Original: J to approximate J and reduce its dimension, notice here we only need (cid:15) for this part to be a small constant, pick (cid:15)
Revised: J to approximate J and reduce its dimension, notice here we only need (cid:15) for this part to be a small constant, pick (cid:15)

Original: = 0 .
Revised: = 0 .

Original: 1 sufﬁces.
Revised: 1 sufﬁces.

Original: • Using S 2 as a preconditioner and solve the regression problem iteratively.
Revised: • Using S 2 as a preconditioner and solve the regression problem iteratively.

Original: Let (cid:98) x denote the solution found by the iterative regime.
Revised: Let (cid:98) x denote the solution found by the iterative regime.

Original: We will prove this statement in two-folds:
Revised: We will prove this statement in two-folds:

Original: • First, we will show that (cid:107) (cid:101) J (cid:101) J (cid:62) (cid:98) x − c (cid:107) 2 ≤ (cid:15) (cid:107) c (cid:107) 2 with probability at least 1 − δ ; •
Revised: • First, we will show that (cid:107) (cid:101) J (cid:101) J (cid:62) (cid:98) x − c (cid:107) 2 ≤ (cid:15) (cid:107) c (cid:107) 2 with probability at least 1 − δ ; •

Original: Then, we will show that (cid:107) JJ (cid:62) (cid:98) x − c (cid:107) 2 = (1 ± 0 . 1) (cid:107) (cid:101) J (cid:101) J (cid:62) (cid:98) x − c (cid:107)  with probability at least 1 − δ .
Revised: Then, we will show that (cid:107) JJ (cid:62) (cid:98) x − c (cid:107) 2 = (1 ± 0 . 1) (cid:107) (cid:101) J (cid:101) J (cid:62) (cid:98) x − c (cid:107) 2 with probability at least 1 − δ .

Original: Combining these two statements, we can show that
Revised: Combining these two statements, we can show that

Original: [Equation]
Revised: [Equation]

Original: Setting (cid:15) to (cid:15)/ 1 . 1 and δ to δ/ 2 , we conclude our proof.
Revised: Setting (cid:15) to (cid:15)/ 1 . 1 and δ to δ/ 2 , we conclude our proof.

Original: It remains to prove these two parts.
Revised: It remains to prove these two parts.

Original: Part 1.
Revised: Part 1.

Original: (cid:107) (cid:101) J (cid:101) J
Revised: (cid:107) (cid:101) J (cid:101) J

Original: (cid:62) (cid:98) x − c (cid:107) 2 ≤ (cid:15) (cid:107) c
Revised: (cid:62) (cid:98) x − c (cid:107) 2 ≤ (cid:15) (cid:107) c

Original: (cid:107) 

Revised: (cid:107)

Original: We observe the iterative procedure is essentially the same as running
Revised: We observe the iterative procedure is essentially the same as running

Original: F AST R EGRESSION on input (cid:101) J (cid:62) , y, (cid:15) , hence by Lemma D.12, we have with probability at least 1 − δ , (cid:107) (cid:101) J (cid:101)
Revised: F AST R EGRESSION on input (cid:101) J (cid:62) , y, (cid:15) , hence by Lemma D.12, we have with probability at least 1 − δ , (cid:107) (cid:101) J (cid:101)

Original: J (cid:62) (cid:98) x − c (cid:107) 2 ≤ (cid:15) (cid:107) c (cid:107) 2 .
Revised: J (cid:62) (cid:98) x − c (cid:107) 2 ≤ (cid:15) (cid:107) c (cid:107) 2 .

Original: Part 2.
Revised: Part 2.

Original: (cid:107) JJ (cid:62) (cid:98) x − c (cid:107) 2 = (1 ± 0 .
Revised: (cid:107) JJ (cid:62) (cid:98) x − c (cid:107) 2 = (1 ± 0 .

Original: 1) (cid:107) (cid:101) J (cid:101) J
Revised: 1) (cid:107) (cid:101) J (cid:101) J

Original: (cid:62) (cid:98) x − c (cid:107) 

Revised: (cid:62) (cid:98) x − c (cid:107)

Original: To prove this part, note that by Lemma D.8, we know that (cid:101) J (cid:62) is a subspace embedding for J (cid:62) .
Revised: To prove this part, note that by Lemma D.8, we know that (cid:101) J (cid:62) is a subspace embedding for J (cid:62) .

Original: Hence, we can utilize Lemma D.13 and get that, with probability at least 1 − δ , we have (cid:107) JJ (cid:62) (cid:98) x − c (cid:107) 2 = (1 ± 0 . 1) (cid:107) [Table] c [Table]
Revised: Hence, we can utilize Lemma D.13 and get that, with probability at least 1 − δ , we have (cid:107) JJ (cid:62) (cid:98) x − c (cid:107) 2 = (1 ± 0 . 1) (cid:107) (cid:101) J (cid:101) J (cid:62) (cid:98) x − c (cid:107) 2 .

Original: Combining these two parts, we have proven the correctness of the theorem.
Revised: Combining these two parts, we have proven the correctness of the theorem.

Original: It remains to justify the running time.
Revised: It remains to justify the running time.

Original: Note that running time can be decomposed into two parts: 1).
Revised: Note that running time can be decomposed into two parts: 1).

Original: The time to generate (cid:101) J , 2).
Revised: The time to generate (cid:101) J , 2).

Original: The time to compute (cid:98) x via iterative scheme.
Revised: The time to compute (cid:98) x via iterative scheme.

Original: Part 1.
Revised: Part 1.

Original: Generate (cid:101) J
Revised: Generate (cid:101) J

Original: To generate (cid:101) J , we apply S 1 ∈ R s 1 × m 2 which is a TensorSRHT .
Revised: To generate (cid:101) J , we apply S 1 ∈ R s 1 × m 2 which is a TensorSRHT .

Original: By
Revised: By

Original: Lemma D.9, it takes O ( n ( m log m + s 1 )) time to compute (cid:101) J , plug in s 1 = Θ( n log 3 ( nm/δ )) , the time is (cid:101) O ( nm ) .
Revised: Lemma D.9, it takes O ( n ( m log m + s 1 )) time to compute (cid:101) J , plug in s 1 = Θ( n log 3 ( nm/δ )) , the time is (cid:101) O ( nm ) .

Original: Part 2.
Revised: Part 2.

Original: Compute (cid:98) x
Revised: Compute (cid:98) x

Original: To compute (cid:98) x , essentially we run F AST R EGRESSION on (cid:101) J (cid:62) , c, (cid:15) , hence by
Revised: To compute (cid:98) x , essentially we run F AST R EGRESSION on (cid:101) J (cid:62) , c, (cid:15) , hence by

Original: Lemma D.12, it takes (cid:101) O
Revised: Lemma D.12, it takes (cid:101) O

Original: ( s 2 n log( κ/(cid:15) )
Revised: ( s 2 n log( κ/(cid:15) )

Original: + n ω ) time, with s 

Revised: + n ω ) time, with s

Original: = Θ(( n + log m ) log n ) and κ is the condition number of (cid:101) J , which has the guarantee κ = (1 ± (cid:15) )
Revised: = Θ(( n + log m ) log n ) and κ is the condition number of (cid:101) J , which has the guarantee κ = (1 ± (cid:15) )

Original: κ ( J ) .
Revised: κ ( J ) .

Original: Hence, the overall running time of this part is (cid:101) O ( n 2 log( κ/(cid:15) ) + n ω ) .
Revised: Hence, the overall running time of this part is (cid:101) O ( n 2 log( κ/(cid:15) ) + n ω ) .

Original: Put things together, the overall running time is (cid:101)
Revised: Put things together, the overall running time is (cid:101)

Original: O
Revised: O

Original: ( nm + n 2 log( κ/(cid:15) )
Revised: ( nm + n 2 log( κ/(cid:15) )

Original: + n ω ) .
Revised: + n ω ) .

Original: Remark D.15.
Revised: Remark D.15.

Original: Due to the probability requirement (union bounding over all layers and all data points), here we only prove by using TensorSRHT .
Revised: Due to the probability requirement (union bounding over all layers and all data points), here we only prove by using TensorSRHT .

Original: One can use similar strategy to obtain an input sparsity time version using TensorSketch .
Revised: One can use similar strategy to obtain an input sparsity time version using TensorSketch .

Original: We remark that this framework is similar to the approach Song et al.
Revised: We remark that this framework is similar to the approach Song et al.

Original: (2021) takes to solve kernel ridge regression, where one ﬁrst uses a shallow but fast sketch to bootstrap, then use another sketching to proceed with the main task.
Revised: (2021) takes to solve kernel ridge regression, where one ﬁrst uses a shallow but fast sketch to bootstrap, then use another sketching to proceed with the main task.

Original: E S PECTRAL P ROPERTIES OF O VER - PARAMETRIZED D EEP N EURAL N ETWORK
Revised: E S PECTRAL P ROPERTIES OF O VER - PARAMETRIZED D EEP N EURAL N ETWORK

Original: In this section, we study the spectral properties of our Gram matrix and connects it to multi-layer
Revised: In this section, we study the spectral properties of our Gram matrix and connects it to multi-layer

Original: NTK.
Revised: NTK.

Original: E. B OUNDS ON THE LEAST EIGENVALUE OF KERNEL AT INITIALIZATION
Revised: E. B OUNDS ON THE LEAST EIGENVALUE OF KERNEL AT INITIALIZATION

Original: We ﬁrst deﬁne the Gram matrices for multiple layer neural network.
Revised: We ﬁrst deﬁne the Gram matrices for multiple layer neural network.

Original: Deﬁnition E.1 (Multiple layer Gram matrix) .
Revised: Deﬁnition E.1 (Multiple layer Gram matrix) .

Original: The Gram matrices K (cid:96) ∈ R n × n for (cid:96) ∈ { 0 , . . .
Revised: The Gram matrices K (cid:96) ∈ R n × n for (cid:96) ∈ { 0 , . . .

Original: , L } of an L -layer neural network are deﬁned as follows:
Revised: , L } of an L -layer neural network are deﬁned as follows:

Original: • ( K 0 )
Revised: • ( K 0 )

Original: i,j := x (cid:62) i x j • For (cid:96) > 0 , let Σ (cid:96),i,j :=[Equation] ( K
Revised: i,j := x (cid:62) i x j • For (cid:96) > 0 , let Σ (cid:96),i,j :=

Original: [Equation]
Revised: [Equation]

Original: ∈ R 2 × 2 for any ( i, j ) ∈ [ n ] × [ n ] .
Revised: 2 for any ( i, j ) ∈ [ n ] × [ n ] .

Original: Then,
Revised: Then,

Original: [Equation]
Revised: [Equation]

Original: Let λ L := λ min ( K L ) to be the minimum eigenvalue of the NTK kernel K L .
Revised: Let λ L := λ min ( K L ) to be the minimum eigenvalue of the NTK kernel K L .

Original: In the following lemma [Equation] we generalize Lemma C.3 in Brand et al.
Revised: In the following lemma, we generalize Lemma C.3 in Brand et al.

Original: (2021)(also Lemma 3 in Cai et al.
Revised: (also Lemma 3 in Cai et al.

Original: into multiple layer neural networks.
Revised: into multiple layer neural networks.

Original: Lemma E.2 (Bounds on the least eigenvalue at initialization, multiple layer version of Lemma C. in Brand et al.
Revised: Lemma E.2 (Bounds on the least eigenvalue at initialization, multiple layer version of Lemma C. in Brand et al.

Original: Let λ (cid:96) denote the minimum eigenvalue of NTK deﬁned for (cid:96) -th layer of neural networks.
Revised: Let λ (cid:96) denote the minimum eigenvalue of NTK deﬁned for (cid:96) -th layer of neural networks.

Original: Suppose m (cid:96) =
Revised: Suppose m (cid:96) =

Original: Ω( λ − 2 (cid:96) n 2 log( n/δ )) , then with probability 1 − δ , we have
Revised: Ω

Original: [Equation]
Revised: [Equation]

Original: Remark E.3. Note that deﬁnition of G ∈ R n × n ,
Revised: λ

Original: [Equation]
Revised: [Equation]

Original: However, we do not need any bounds related to G . We only need to care about the bound related
Revised: [Equation]

Original: G (cid:96) (per layer based Gram matrix). For more details, we can see Du et al.
Revised: n 2 log( n/δ )) , then with probability 1 − δ , we have

Original: (2019a), page 5 right column, the Gram matrix is deﬁned for each layer, where each entry is the inner product between two derivatives. It seems to make sense if we prove for G (cid:96) .
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: For any (cid:96) ∈
Revised: For any (cid:96) ∈

Original: [ L ] , we have G (cid:96) =
Revised: [ L ] , we have G [Equation] =

Original: J
Revised: J

Original: [Equation]
Revised: [Equation]

Original: J (cid:62) (cid:96) ∈ R n × n .

Original: Hence, for any i, j ∈ [ n ] ,
Revised: Hence, for any i, j ∈ [ n ] ,

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: In particular
Revised: In particular,

Original: [Equation]
Revised: [Equation]

Original: We will prove that (cid:107) G L − K L (cid:107) ∞ is small, which implies that λ min ( G L ) is close to λ (cid:96) .
Revised: We will prove that (cid:107) G L − K L (cid:107) ∞ is small, which implies that λ min ( G L ) is close to λ (cid:96) .

Original: The proof idea is similar to Du et al. (2019a) via induction on (cid:96) .
Revised: The proof idea is similar to Du et al. (2019a) via induction on (cid:96) .

Original: For (cid:96) = 1 , recall ( g 1 ,i ) k
Revised: For (cid:96) = 1 , recall ( g 1 ,i ) k

Original: = (cid:80) b ∈ [ m ] ( W 1 ) k,b ( x i ) b for k ∈ [ m ] .
Revised: = (cid:80) b ∈ [ m ] ( W 1 ) k,b ( x i ) b for k ∈ [ m ] .

Original: Hence, for any k ∈ [ m ] ,
Revised: Hence, for any k ∈ [ m ] ,

Original: [Equation]
Revised: [Equation]

Revised: E [( W 1 ) k,b ( W 1 ) k,b (cid:48)

Revised: ( x i ) b ( x j ) b (cid:48) ]

Revised: = (cid:88)

Revised: b ∈ [ m ]

Revised: E [( W 1 ) 2 k,b ] · ( x i ) b ( x j ) b ( ( W 1 ) k,b ∼ N (0 , 2 m ) .)

Revised: =m (cid:88)

Revised: b ∈ [ m ] ( x i ) b ( x j ) b =

Revised: [Equation]

Revised: m

Revised: [Equation]

Original: Then, we have
Revised: Then, we have

Original: [Equation]
Revised: [Equation]

Original: Next, we will show that h (cid:62) 1 ,i h 1 ,j concentrates around its expectation.
Revised: Next, we will show that h (cid:62) 1 ,i h 1 ,j concentrates around its expectation.

Original: First, for any k ∈ [ m ] ,
Revised: First, for any k ∈ [ m ] ,

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: Since (cid:104) ( W 1 ) k, ∗ , x i (cid:105) ∼ N (0 , 2 (cid:107) x
Revised: Since (cid:104) ( W 1 ) k, ∗ , x i (cid:105) ∼ N (0 , 2 (cid:107) x

Original: i (cid:107) 22 m ) , by the concentration of Gaussian distribution, |(cid:104) ( W 1 )
Revised: i (cid:107) 22 m ) , by the concentration of Gaussian distribution, |(cid:104) ( W 1 )

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: k ∈ [
Revised: k ∈ [

Original: [Equation]
Revised: [Equation]

Original: ] , i ∈ [ n ] holds with probability at least 1 − mne − cm/ 4 .
Revised: ] , i ∈ [ n ] holds with probability at least 1 − mne − cm/ 4 .

Original: Conditioned on this event, we have | ( h 1 ,i ) k
Revised: Conditioned on this event, we have | ( h 1 ,i ) k

Original: ( h 1 ,j ) k | ≤ c for all i, j ∈ [ n ] and k ∈ [ m ] .
Revised: ( h 1 ,j ) k | ≤ c for all i, j ∈ [ n ] and k ∈ [ m ] .

Original: Then, by Hoeffding’s inequality, we have for any ( i, j ) ∈ [ n ] × [ n ] ,
Revised: Then, by Hoeffding’s inequality, we have for any ( i, j ) ∈ [ n ] × [ n ] ,

Original: [Equation]
Revised: [Equation]

Original: Hence, by union bound, we get that
Revised: Hence, by union bound, we get that

Original: [Equation]
Revised: [Equation]

Original: with probability at least
Revised: with probability at least

Original: [Equation]
Revised: [Equation]

Original: If we choose c := log( mnL/δ ) m and t := m − 1 / 2 · polylog( nL/δ ) , we have with probability at least 1 − δL ,
Revised: If we choose c := log( mnL/δ ) m and t := m − 1 / 2 · polylog( nL/δ ) , we have with probability at least 1 − δL ,

Original: [Equation]
Revised: [Equation]

Original: Suppose for (cid:96) = 1 , . . .
Revised: Suppose for (cid:96) = 1 , . . .

Original: h ( h < L ), we have
Revised: h ( h < L ), we have

Original: [Equation]
Revised: [Equation]

Original: Consider (cid:96) =
Revised: Consider (cid:96) =

Original: h
Revised: h

Original: + 1 .
Revised: + 1 .

Original: By a similar computation, we have
Revised: By a similar computation, we have

Original: [Equation]
Revised: [Equation]

Original: Deﬁne a new covariance matrix
Revised: Deﬁne a new covariance matrix

Original: [Equation]
Revised: [Equation]

Original: h (cid:96) − 1 ,i h (cid:62) (cid:96) − 1 ,i h (cid:96) − 1 ,j h (cid:62) (cid:96) − 1 ,j h (cid:96) − 1 ,i h (cid:62) (cid:96) − 1 ,j h (cid:96) − 1 ,j (cid:21) ∀ ( i, j ) ∈ [ n ] × [ n ] .
Revised: h (cid:96) − 1 ,i h (cid:62) (cid:96) − 1 ,i h (cid:96) − 1 ,j h (cid:62) (cid:96) − 1 ,j h (cid:96) − 1 ,i h (cid:62) (cid:96) − 1 ,j h (cid:96) − 1 ,j (cid:21) ∀ ( i, j ) ∈ [ n ] × [ n ] .

Original: We have
Revised: We have

Original: [Equation]
Revised: [Equation]

Original: Hence, we have with probability at least 1 − δL ,
Revised: Hence, we have with probability at least 1 − δL ,

Original: [Equation]
Revised: [Equation]

Original: It remains to upper bound the difference (cid:107)
Revised: It remains to upper bound the difference (cid:107)

Original: [Equation]
Revised: [Equation]

Original: K
Revised: K

Original: [Equation]
Revised: [Equation]

Original: Recall that
Revised: Recall that

Original: [Equation]
Revised: [Equation]

Original: and hence, by the induction hypothesis, we have
Revised: and hence, by the induction hypothesis, we have

Original: [Equation]
Revised: [Equation]

Original: Notice that (cid:98) Σ (cid:96),i,j can be written as
Revised: Notice that (cid:98) Σ (cid:96),i,j can be written as

Original: [Equation]
Revised: [Equation]

Original: Moreover
Revised: Moreover

Original: [Equation]
Revised: [Equation]

Original: when φ is the ReLU function
Revised: when φ is the ReLU function

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: F
Revised: F

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: We note that F ( θ ) has the following analytic form:
Revised: We note that F ( θ ) has the following analytic form:

Original: [Equation]
Revised: [Equation]

Original: Similarly,
Revised: Similarly,

Original: E
Revised: E

Original: [Equation]
Revised: [Equation]

Original: where τ (cid:96),i,j := cos −(cid:18)
Revised: where τ (cid:96),i,j := cos −(cid:18)

Original: [Equation]
Revised: [Equation]

Original: √
Revised: √

Original: [Equation]
Revised: [Equation]

Original: (cid:19) .
Revised: (cid:19) .

Original: By the induction hypothesis, we have ( K (cid:96) )
Revised: By the induction hypothesis, we have ( K (cid:96) )

Original: i,j ∈ h (cid:62) (cid:96),i h (cid:96),j ± (cid:101) O ( m − 1 / 2 ) for all i, j ∈ [ n ] .
Revised: i,j ∈ h (cid:62) (cid:96),i h (cid:96),j ± (cid:101) O ( m − 1 / 2 ) for all i, j ∈ [ n ] .

Original: By Lemma F.7, we also have (cid:107) h (cid:96),i (cid:107) 2 ∈ 1 ± (cid:15) for all (cid:96) ∈
Revised: By Lemma F.7, we also have (cid:107) h (cid:96),i (cid:107) 2 ∈ 1 ± (cid:15) for all (cid:96) ∈

Original: [ L ] and i ∈ [ n ] with probability 1 − O ( nL ) · e − Ω( m(cid:15) 2 /L ) .
Revised: [ L ] and i ∈ [ n ] with probability 1 − O ( nL ) · e − Ω( m(cid:15) 2 /L ) .

Original: They implies that cos( τ (cid:96),i,j ) ∈ cos( θ ) ± (cid:101) O ( m − 1 / 2 ) .
Revised: They implies that cos( τ (cid:96),i,j ) ∈ cos( θ ) ± (cid:101) O ( m − 1 / 2 ) .

Original: Thus, by Taylor’s theorem, we have
Revised: Thus, by Taylor’s theorem, we have

Original: [Equation]
Revised: [Equation]

Original: Therefore, we have
Revised: Therefore, we have

Original: [Equation]
Revised: [Equation]

Original: That is,
Revised: That is,

Original: [Equation]
Revised: [Equation]

Original: Combining Eqs.
Revised: Combining Eqs.

Original: (5) and (7) together, we get that
Revised: (5) and (7) together, we get that

Original: [Equation]
Revised: [Equation]

Original: holds with probability at least 1 − δL for (cid:96) =
Revised: holds with probability at least 1 − δL for (cid:96) =

Original: h + 1 .
Revised: h + 1 .

Original: By induction, we have proved that for the ﬁrst L − 1 layers, the intermediate correlation h (cid:62) (cid:96),i h (cid:96),j is close to the intermediate Gram matrix ( K (cid:96) )
Revised: By induction, we have proved that for the ﬁrst L − 1 layers, the intermediate correlation h (cid:62) (cid:96),i h (cid:96),j is close to the intermediate Gram matrix ( K (cid:96) )

Original: i,j .
Revised: i,j .

Original: Now, we consider the last layer.
Revised: Now, we consider the last layer.

Original: Recall G L is deﬁned by Eq.
Revised: Recall G L is deﬁned by Eq.

Original: (4), which has the same form as the correlation matrix of a two-layer over-parameterized neural network with input data { h L − 1 ,i } i ∈ [ n ] .
Revised: (4), which has the same form as the correlation matrix of a two-layer over-parameterized neural network with input data { h L − 1 ,i } i ∈ [ n ] .

Original: Deﬁne
Revised: Deﬁne

Original: [Equation]
Revised: [Equation]

Original: Then, by the analysis of the two-layer case (see for example Song & Yang (2019); Du et al. (2019b)), we have
Revised: Then, by the analysis of the two-layer case (see for example Song & Yang (2019); Du et al. (2019b)), we have

Original: [Equation]
Revised: [Equation]

Original: if m = Ω( λ − 2 L n 2 log( n/δ )) , where λ L := λ min ( K L ) .
Revised: if m = Ω( λ − 2 L n 2 log( n/δ )) , where λ L := λ min ( K L ) .

Original: It remains to bound (cid:107) (cid:98) K L − K L (cid:107) ∞ .
Revised: It remains to bound (cid:107) (cid:98) K L − K L (cid:107) ∞ .

Original: Equivalently, for any ( i, j ) ∈ [ n ] × [ n ] ,
Revised: Equivalently, for any ( i, j ) ∈ [ n ] × [ n ] ,

Original: [Equation]
Revised: [Equation]

Original: The expectation has the following analytic form:
Revised: The expectation has the following analytic form:

Original: [Equation]
Revised: [Equation]

Original: By the analysis of the ( L − 1) -layer, we know that | ρ L,i,j − (cid:98) ρ L,i,j | ≤ (cid:101) O ( m − 1 / 2 ) , where ρ L,i,j := cos( τ L,i,j ) and (cid:98)
Revised: By the analysis of the ( L − 1) -layer, we know that | ρ L,i,j − (cid:98) ρ L,i,j | ≤ (cid:101) O ( m − 1 / 2 ) , where ρ L,i,j := cos( τ L,i,j ) and (cid:98)

Original: ρ L,i,j := cos( θ L,i,j ) .
Revised: ρ L,i,j := cos( θ L,i,j ) .

Original: Also, notice that cos( τ L,i,j ) =
Revised: Also, notice that cos( τ L,i,j ) =

Original: F ( τ L − 1 ,i,j ) ∈ [0 , 1 / 2] by
Revised: F ( τ L − 1 ,i,j ) ∈ [0 , 1 / 2] by

Original: Eq.
Revised: Eq.

Original: Hence, the derivative of the expectation is bounded, and by Taylor’s theorem, we have
Revised: Hence, the derivative of the expectation is bounded, and by Taylor’s theorem, we have

Original: [Equation]
Revised: [Equation]

Original: It implies that (cid:107) (cid:98) K L − K L (cid:107) ≤ λ L 8 , which further implies that
Revised: It implies that (cid:107) (cid:98) K L − K L (cid:107) ≤ λ L 8 , which further implies that

Original: [Equation]
Revised: [Equation]

Original: Equivalently, we get that
Revised: Equivalently, we get that

Original: [Equation]
Revised: [Equation]

Original: with probability at least 1 − δ .
Revised: with probability at least 1 − δ .

Original: The lemma is then proved.
Revised: The lemma is then proved.

Original: [Equation]
Revised: [Equation]

Original: B OUNDS
Revised: B OUNDS

Original: In this section, we generalize the Lemma C.5 in Brand et al.
Revised: In this section, we generalize the Lemma C.5 in Brand et al.

Original: (2021) into multiple layer neural network
Revised: (2021) into multiple layer neural network

Original: Lemma E.4 (Bounds on the least eigenvalue, multiple layer neural network version of Lemma C. of Brand et al.
Revised: Lemma E.3 (Bounds on the least eigenvalue, multiple layer neural network version of Lemma C. of Brand et al.

Original: Suppose m = Ω( λ − 2 (cid:96) n 2 log( n/δ )) , with probability least 1 − δ , for any set of weights W 1 , · · · W L satisfying
Revised: Suppose m = Ω( λ − 2 (cid:96) n 2 log( n/δ )) , with probability least 1 − δ , for any set of weights W 1 , · · · W L satisfying

Original: [Equation]
Revised: [Equation]

Original: then the following holds
Revised: then the following holds

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Recall ( G L )( W ) i,
Revised: Recall ( G L )( W ) i,

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: For simplicity
Revised: For simplicity

Original: [Equation]
Revised: [Equation]

Original: let z i,r :=
Revised: let z i,r :=

Original: ( W L ) (cid:62) r h
Revised: ( W L ) (cid:62) r h

Original: i,L − 1 and z i,r (0) :
Revised: i,L − 1 and z i,r (0) :

Original: = ( W L (0))
Revised: = ( W L (0))

Original: (cid:62) r h i,L − 1 (0) .
Revised: (cid:62) r h i,L − 1 (0) .

Original: | ( G L ( W
Revised: | ( G L ( W

Original: [Equation]
Revised: [Equation]

Original: i,j −
Revised: i,j −

Original: [Equation]
Revised: [Equation]

Original: G L ( W
Revised: G L ( W

Original: [Equation]
Revised: [Equation]

Original: ) i,
Revised: ) i,

Original: [Equation]
Revised: [Equation]

Original: |
Revised: |

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: where the last step follows from φ (cid:48)
Revised: where the last step follows from φ (cid:48)

Original: ( x ) ∈ { 0 , 1 } and Lemma E.2.
Revised: ( x ) ∈ { 0 , 1 } and Lemma E.2.

Original: To upper bound the above two terms, we ﬁrst need to bound the move of h i,L − 1 .
Revised: To upper bound the above two terms, we ﬁrst need to bound the move of h i,L − 1 .

Original: For any (cid:96) ∈
Revised: For any (cid:96) ∈

Original: [ L − 1] , we have
Revised: [ L − 1] , we have

Original: [Equation]
Revised: [Equation]

Original: R
Revised: R

Original: [Equation]
Revised: [Equation]

Original: where c W := (cid:107) W (cid:96) (0) (cid:107) ≤ 3 by the well-known deviations bounds concerning the singular values of Gaussian random matrices (Rudelson & Vershynin (2010)).
Revised: where c W := (cid:107) W (cid:96) (0) (cid:107) ≤ 3 by the well-known deviations bounds concerning the singular values of Gaussian random matrices (Rudelson & Vershynin (2010)).

Original: Also, when (cid:96) = 0 , we have (cid:107) h
Revised: Also, when (cid:96) = 0 , we have (cid:107) h

Original: i, 0 − h i, 0 (0) (cid:107) 2 = 0 .
Revised: i, 0 − h i, 0 (0) (cid:107) 2 = 0 .

Original: Hence, we get that for all (cid:96) ∈
Revised: Hence, we get that for all (cid:96) ∈

Original: [ L − 1] , i ∈ [ n ] ,
Revised: [ L − 1] , i ∈ [ n ] ,

Original: [Equation]
Revised: [Equation]

Original: h
Revised: h

Original: i
Revised: i

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: ≤
Revised: ≤

Original: [Equation]
Revised: [Equation]

Original: +
Revised: +

Original: [Equation]
Revised: [Equation]

Original: since R (cid:28) 1 by our choice of R .
Revised: since R (cid:28) 1 by our choice of R .

Original: Hence, it implies that
Revised: Hence, it implies that

Original: [Equation]
Revised: [Equation]

Original: ,L − 1 h j,L − 1 − h
Revised: ,L − 1 h j,L − 1 − h

Original: i,L − 1 (0) (cid:62) h j,L − 1 (0) (cid:12)(cid:12) ≤
Revised: i,L − 1 (0) (cid:62) h j,L − 1 (0) (cid:12)(cid:12) ≤

Original: [Equation]
Revised: [Equation]

Original: ( h i,L − 1 − h i,L − 1 (0)) (cid:62) h j,L − 1 (cid:12)(cid:12) + (cid:12)(cid:12) h
Revised: ( h i,L − 1 − h i,L − 1 (0)) (cid:62) h j,L − 1 (cid:12)(cid:12) + (cid:12)(cid:12) h

Original: i,L − 1 (0) (cid:62) ( h j,L − 1 − h j,L − 1 (0))
Revised: i,L − 1 (0) (cid:62) ( h j,L − 1 − h j,L − 1 (0))

Original: (cid:12)(cid:12) ≤ 2(1 + (cid:15) )
Revised: (cid:12)(cid:12) ≤ 2(1 + (cid:15) )

Original: R (2 c W ) L − 1 .
Revised: R (2 c W ) L − 1 .

Original: Similarly,
Revised: Similarly,

Original: [Equation]
Revised: [Equation]

Original: +
Revised: +

Original: [Equation]
Revised: [Equation]

Original: O
Revised: O

Original: [Equation]
Revised: [Equation]

Original: ) m
Revised: ) m

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: m
Revised: m

Original: [Equation]
Revised: [Equation]

Original: By our assumption, we have (cid:107) ( W L ) r − ( W L (0))
Revised: By our assumption, we have (cid:107) ( W L ) r − ( W L (0))

Original: r (cid:107) 2 ≤ R .
Revised: r (cid:107) 2 ≤ R .

Original: We also know that (cid:107) h
Revised: We also know that (cid:107) h

Original: i,L − 1 − h
Revised: i,L − 1 − h

Original: i,L − 1 (0) (cid:107) 2 ≤ √ 1 + (cid:15)R (2 c W ) L − 1 .
Revised: i,L − 1 (0) (cid:107) 2 ≤ √ 1 + (cid:15)R (2 c W ) L − 1 .

Original: Then, we can follow the proof in Du et al.
Revised: Then, we can follow the proof in Du et al.

Original: (2019b)
Revised: (2019b)

Original: ; Song & Yang (2019) and deﬁne the event
Revised: ; Song & Yang (2019) and deﬁne the event

Original: A i,r := ∃ w, h : (cid:107) w − ( W L (0))
Revised: A i,r := ∃ w, h : (cid:107) w − ( W L (0))

Original: r (cid:107) 2 ≤ R, (cid:107) h − h
Revised: r (cid:107) 2 ≤ R, (cid:107) h − h

Original: i,L − 1 (0) (cid:107) 2 ≤
Revised: i,L − 1 (0) (cid:107) 2 ≤

Original: O ( R ) such that 1 w (cid:62) h> 0 (cid:54) = 

Revised: O ( R ) such that 1 w (cid:62) h> 0 (cid:54) =

Original: ( W L (0))
Revised: ( W L (0))

Original: (cid:62) r h i,L − 1 (0) > 0 .
Revised: (cid:62) r h i,L − 1 (0) > 0 .

Original: We have A i,r happens if and only if | (( W L (0))
Revised: We have A i,r happens if and only if | (( W L (0))

Original: r + ∆ w ) (cid:62) h i,L − 1 (0)
Revised: r + ∆ w ) (cid:62) h i,L − 1 (0)

Original: | <
Revised: | <

Original: O
Revised: O

Original: ( R ( R + m − 1 / 2 )) for some ﬁxed vector ∆ w of length at most R .
Revised: ( R ( R + m − 1 / 2 )) for some ﬁxed vector ∆ w of length at most R .

Original: It implies that
Revised: It implies that

Original: [Equation]
Revised: [Equation]

Original: Hence, using the same proof in the previous work, we get that
Revised: Hence, using the same proof in the previous work [Equation] we get that

Original: [Equation]
Revised: [Equation]

Original: Putting them together, we have
Revised: Putting them together, we have

Original: [Equation]
Revised: [Equation]

Original: And by our choice of R , we get that
Revised: And by our choice of R [Equation] we get that

Original: [Equation]
Revised: [Equation]

Original: The lemma is then proved.
Revised: The lemma is then proved.

Original: F C ONVERGENCE ANALYSIS OF A LGORITHM
Revised: F C ONVERGENCE ANALYSIS OF A LGORITHM

Original: In this section
Revised: In this section

Original: [Equation]
Revised: [Equation]

Original: we analyze the convergence behavior of Algorithm 2.
Revised: we analyze the convergence behavior of Algorithm [Equation]

Original: F. P RELIMINARY
Revised: F. P RELIMINARY

Original: We recall the initialization of our neural network.
Revised: We recall the initialization of our neural network.

Original: Deﬁnition F.1 (Initialization) .
Revised: Deﬁnition F.1 (Initialization) .

Original: Let m = m (cid:96) for all (cid:96) ∈
Revised: Let m = m (cid:96) for all (cid:96) ∈

Original: [ L ] .
Revised: [ L ] .

Original: Let m 0 =
Revised: Let m 0 =

Original: d .
Revised: d .

Original: We assume weights are initialized as
Revised: We assume weights are initialized as

Original: • Each entry of weight vector a ∈ R
Revised: • Each entry of weight vector a ∈ R

Original: m L is i.i
Revised: m L is i.i.d.

Original: [Equation] d. sampled from {− 1  m L , + 1  m L } uniformly at random.
Revised: sampled from {− 1 √ m L , + 1 √ m L } uniformly at random.

Original: • Each entry of weight matrices W (cid:96) ∈
Revised: • Each entry of weight matrices W (cid:96) ∈

Original: R m (cid:96) ×
Revised: R m (cid:96) ×

Original: m
Revised: m

Original: (cid:96) − 1 sampled from N (0 , 2 /m (cid:96) ) .
Revised: (cid:96) − 1 sampled from N (0 , 2 /m (cid:96) ) .

Original: We also restate the architecture of our neural network here.
Revised: We also restate the architecture of our neural network here.

Original: Deﬁnition F.2 (Architecture) .
Revised: Deﬁnition F.2 (Architecture) .

Original: Our neural network is a standard L -layer feed-forward neural network, with the activation functions deﬁned as a scaled version of threshold ReLU activation: φ ( x ) = √ c b 1 [ x > (cid:112) 2 /mb ] x , where c b :=
Revised: Our neural network is a standard L -layer feed-forward neural network, with the activation functions deﬁned as a scaled version of shifted ReLU activation: φ ( x ) = √ c b 1 [ x > (cid:112) 2 /mb ] x , where c b :=

Original: (2(1 − Φ( b ) + bφ ( b )))
Revised: (2(1 − Φ( b ) + bφ ( b )))

Original: − 1 / 2 .
Revised: − 1 / 2 .

Original: Here b is a threshold value we will pick later.
Revised: Here b is a threshold value we will pick later.

Original: At last layer, we use a scaled version of a vector with its entry being
Revised: At last layer, we use a scaled version of a vector with its entry being

Original: Rademacher random variables.
Revised: Rademacher random variables.

Original: We deﬁne the neural network function f : R m 

Revised: We deﬁne the neural network function f : R m

Original: → R as
Revised: → R as

Original: [Equation]
Revised: [Equation]

Original: We measure the loss of the neural network via squared-loss function:
Revised: We measure the loss of the neural network via squared-loss function:

Original: [Equation]
Revised: [Equation]

Original: We use f t : R m 0 × n → R n denote the prediction of our network:
Revised: We use f t : R m 0 × n → R n denote the prediction of our network:

Original: [Equation]
Revised: [Equation]

Original: We state two assumptions here.
Revised: We state two assumptions here.

Original: Assumption F.3 (Small Spectral Norm) .
Revised: Assumption F.3 (Small Spectral Norm) .

Original: Let t ∈ { 0 , . .
Revised: Let t ∈ { 0 , . .

Original: , T } and let R ≤ 1 be a parameter.
Revised: , T } and let R ≤ 1 be a parameter.

Original: We assume
Revised: We assume

Original: [Equation]
Revised: [Equation]

Original: Later, we will invoke this assumption by specifying the choice of R .
Revised: Later, we will invoke this assumption by specifying the choice of R .

Original: Assumption F.4 (Sparsity) .
Revised: Assumption F.4 (Sparsity) .

Original: Let t ∈ { 0 , . .
Revised: Let t ∈ { 0 , . .

Original: , T } and let s ≥ 1 be an integer parameter.
Revised: , T } and let s ≥ 1 be an integer parameter.

Original: We assume
Revised: We assume

Original: [Equation]
Revised: [Equation]

Original: (cid:107)

Original: [Equation]
Revised: [Equation]

Original: s, ∀ (cid:96) ∈
Revised: (cid:96) ∈

Original: [ L ] , i ∈ [ n ] .
Revised: [ L ] , i ∈ [ n ] .

Original: Later, we will invoke this assumption by specifying the choice of s.
Revised: Later, we will invoke this assumption by specifying the choice of s.

Original: Finally, throughout this entire section, we will assume m (cid:96) =
Revised: Finally, throughout this entire section, we will assume m (cid:96) =

Original: F. T ECHNICAL L EMMAS
Revised: F. T ECHNICAL L EMMAS

Original: We ﬁrst show that during initialization, by using our threshold ReLU activation, the vector h i,(cid:96) is sparse, hence the diagonal matrix D i,(cid:96) is sparse as well.
Revised: We ﬁrst show that during initialization, by using our threshold ReLU activation, the vector h i,(cid:96) is sparse, hence the diagonal matrix D i,(cid:96) is sparse as well.

Original: Lemma F.5 (Sparse initialization) .
Revised: Lemma F.5 (Sparse initialization) .

Original: Let σ b ( x ) = max { x − b, 0 } be the threshold ReLU activation with threshold b > 0 .
Revised: Let σ b ( x ) = max { x − b, 0 } be the threshold ReLU activation with threshold b > 0 .

Original: After initialization, with probability 1 − nL · e − Ω( me − b 2 m/ 4 ) , it holds for all i ∈ [ n ] and (cid:96) ∈
Revised: After initialization, with probability 1 − nL · e − Ω( me − b 2 m/ 4 ) , it holds for all i ∈ [ n ] and (cid:96) ∈

Original: [ L ] , (cid:107) h i,(cid:96) (cid:107) 0 ≤ O ( m · e − b 2 m/ 4 ) .
Revised: [ L ] , (cid:107) h i,(cid:96) (cid:107) 0 ≤ O ( m · e − b 2 m/ 4 ) .

Original: Proof.
Revised: Proof.

Original: We ﬁx i ∈ [ n ] and (cid:96) ∈
Revised: We ﬁx i ∈ [ n ] and (cid:96) ∈

Original: [ L ] , since we will union bound over all i and (cid:96) at last.
Revised: [ L ] , since we will union bound over all i and (cid:96) at last.

Original: Let u i ∈
Revised: Let u i ∈

Original: R m be a ﬁxed vector and W (cid:96),r to denote the r -th row of W (cid:96) , then by the concentration of Gaussian, we have
Revised: R m be a ﬁxed vector and W (cid:96),r to denote the r -th row of W (cid:96) , then by the concentration of Gaussian, we have

Original: [Equation]
Revised: [Equation]

Original: Let S be the following index set S :
Revised: Let S be the following index set S :

Original: = { r ∈ [ m ] : (cid:104) W (cid:96),r , u i (cid:105) >
Revised: = { r ∈ [ m ] : (cid:104) W (cid:96),r , u i (cid:105) >

Original: b } , the above reasoning means that for the indicator random variable 1 [ r ∈ S ] , we have
Revised: b } , the above reasoning means that for the indicator random variable 1 [ r ∈ S ] , we have

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: b 2 m/
Revised: b 2 m/

Original: [Equation]
Revised: [Equation]

Original: Use Bernstein’s inequality (Lemma A.3) we have that for all t > 0 ,
Revised: Use Bernstein’s inequality (Lemma A.3) we have that for all t > 0 ,

Original: [Equation]
Revised: [Equation]

Original: where k := m · exp( − b 2 m/ 4) .
Revised: where k := m · exp( − b 2 m/ 4) .

Original: By picking t =
Revised: By picking t =

Original: k , we have
Revised: k , we have

Original: [Equation]
Revised: [Equation]

Original: Note that | S | is essentially the quantity (cid:107) h i,(cid:96) (cid:107) 0 , hence we can union bound over all (cid:96) and i and with probability at least
Revised: Note that | S | is essentially the quantity (cid:107) h i,(cid:96) (cid:107) 0 , hence we can union bound over all (cid:96) and i and with probability at least

Original: [Equation]
Revised: [Equation]

Original: we have (cid:107) h i,(cid:96) (cid:107) 0 ≤ 2 m · exp( − b 2 m/ 4) .
Revised: we have (cid:107) h i,(cid:96) (cid:107) 0 ≤ 2 m · exp( − b 2 m/ 4) .

Original: Remark F.6.
Revised: Remark F.6.

Original: The above lemma shows that by using the shifted ReLU activation, we make sure that all h i,(cid:96) are sparse after initialization.
Revised: The above lemma shows that by using the shifted ReLU activation, we make sure that all h i,(cid:96) are sparse after initialization.

Original: Speciﬁcally, we use k := m · exp( − b 2 m/ 4) as a sparsity parameter.
Revised: Speciﬁcally, we use k := m · exp( − b 2 m/ 4) as a sparsity parameter.

Original: Later, we might re-scale b so that the probability becomes exp( − b 2 / 2) .
Revised: Later, we might re-scale b so that the probability becomes exp( − b 2 / 2) .

Original: We stress that such re-scaling does not affect the sparsity of our initial vectors.
Revised: We stress that such re-scaling does not affect the sparsity of our initial vectors.

Original: If we re-scale b and choose it as √ 0 .
Revised: If we re-scale b and choose it as √ 0 .

Original: 4 log m , then k
Revised: 4 log m , then k

Original: = m 0 .
Revised: = m 0 .

Original: and hence with high probability, (cid:107) h i,(cid:96) (cid:107) 

Revised: and hence with high probability, (cid:107) h i,(cid:96) (cid:107)

Original: ≤
Revised: ≤

Original: O ( m 0 . 8 ) .
Revised: O ( m 0 . 8 ) .

Original: As a direct consequence, we note that all initial D i,(cid:96) are k -sparse as well.
Revised: As a direct consequence, we note that all initial D i,(cid:96) are k -sparse as well.

Original: We state a lemma that handles the (cid:96) 2 norm of h i,(cid:96) when one uses truncated Gaussian distribution instead.
Revised: We state a lemma that handles the (cid:96) 2 norm of h i,(cid:96) when one uses truncated Gaussian distribution instead.

Original: Due to the length and the delicacy of the proof, we defer it to Section G.
Revised: Due to the length and the delicacy of the proof, we defer it to Section G.

Original: Lemma F.7 (Restatement of Lemma G.6) .
Revised: Lemma F.7 (Restatement of Lemma G.6) .

Original: Let b > 0 be a ﬁxed scalar.
Revised: Let b > 0 be a ﬁxed scalar.

Original: Let the activation function φ ( x ) := √ c b 1 [ x > (cid:112) 2 /mb ] x , where c b :=
Revised: Let the activation function φ ( x ) := √ c b 1 [ x > (cid:112) 2 /mb ] x , where c b :=

Original: (2(1 − Φ( b ) + bφ ( b )))
Revised: (2(1 − Φ( b ) + bφ ( b )))

Original: − 1 / 2 .
Revised: − 1 / 2 .

Original: Let (cid:15) ∈ (0 , 1) , then over the randomness of W (0) , with probability at least 1 − O ( nL ) · exp( − Ω( m exp( − b 2 / 2)
Revised: Let (cid:15) ∈ (0 , 1) , then over the randomness of W (0) , with probability at least 1 − O ( nL ) · exp( − Ω( m exp( − b 2 / 2)

Original: [Equation]
Revised: [Equation]

Original: 2 /L 2 )) , we have
Revised: 2 /L 2 )) , we have

Original: [Equation]
Revised: [Equation]

Original: We remark that the parameter e − b 2 / 2 m captures the sparsity during the initialization.
Revised: We remark that the parameter e − b 2 / 2 m captures the sparsity during the initialization.

Original: The second lemma handles the consecutive product that appears naturally in the gradient computation.
Revised: The second lemma handles the consecutive product that appears naturally in the gradient computation.

Original: Lemma F.8 (Variant of Lemma 7.3 in Allen-Zhu et al. (2019a)) .
Revised: Lemma F.8 (Variant of Lemma 7.3 in Allen-Zhu et al. (2019a)) .

Original: Suppose m ≥ Ω( nL log( nL )) .
Revised: Suppose m ≥ Ω( nL log( nL )) .

Original: With probability at least 1 − e − Ω( k/L 2 ) over the randomness of initializations W 1 (0) , . . .
Revised: With probability at least 1 − e − Ω( k/L 2 ) over the randomness of initializations W 1 (0) , . . .

Original: , W L (0) ∈
Revised: , W L (0) ∈

Original: R m × m , for all i ∈ [ n ] and 1 ≤ a ≤ b ≤ L ,
Revised: R m × m , for all i ∈ [ n ] and 1 ≤ a ≤ b ≤ L ,

Original: (a) (cid:107) W b D
Revised: (a) (cid:107) W b D

Original: i,b − 1 W b − 1 . . .
Revised: i,b − 1 W b − 1 . . .

Original: D i,a W a (cid:107) ≤ O ( √ L ) .
Revised: D i,a W a (cid:107) ≤ O ( √ L ) .

Original: (b) (cid:107) W b D i,b − 1 W b − 1 . . .
Revised: (b) (cid:107) W b D i,b − 1 W b − 1 . . .

Original: D i,a W a v (cid:107) 2 ≤ 2 (cid:107) v (cid:107) 2 , for all vectors with (cid:107) v (cid:107) 0 ≤ O ( m L log m ) .
Revised: D i,a W a v (cid:107) 2 ≤ 2 (cid:107) v (cid:107) 2 , for all vectors with (cid:107) v (cid:107) 0 ≤ O ( m L log m ) .

Original: (c) (cid:107) u (cid:62) W b D
Revised: (c) (cid:107) u (cid:62) W b D

Original: i,b − 1 W b − 1 . . .
Revised: i,b − 1 W b − 1 . . .

Original: D i,a W a (cid:107) 2 ≤ O (1) (cid:107) u (cid:107) 2 , for all vectors u with (cid:107) u (cid:107) 0 ≤ O ( m L log m ) .
Revised: D i,a W a (cid:107) 2 ≤ O (1) (cid:107) u (cid:107) 2 , for all vectors u with (cid:107) u (cid:107) 0 ≤ O ( m L log m ) .

Original: (d) For any intergers
Revised: (d) For any intergers

Original: ) , with probability at least 1 − e − Ω( s log m ) over the randomness of initializations W 1 (0) , . . .
Revised: ) , with probability at least 1 − e − Ω( s log m ) over the randomness of initializations W 1 (0) , . . .

Original: , W L (0)
Revised: , W L (0)

Original: ×
Revised: ×

Original: | u (cid:62) W b D i,b − 1 W b − 1 . . .
Revised: | u (cid:62) W b D i,b − 1 W b − 1 . . .

Original: D i,a W a v | ≤
Revised: D i,a W a v | ≤

Original: u (cid:107) 2 (cid:107) v (cid:107) 2 · O (
Revised: u (cid:107) 2 (cid:107) v (cid:107) 2 · O (

Original: L log m ) for all vectors u, v with (cid:107) u (cid:107) 0 , (cid:107) v (cid:107) 0 ≤ s .
Revised: L log m ) for all vectors u, v with (cid:107) u (cid:107) 0 , (cid:107) v (cid:107) 0 ≤ s .

Original: The proof is similarly to the original proof of the corresponding lemma in Allen-Zhu et al.
Revised: The proof is similarly to the original proof of the corresponding lemma in Allen-Zhu et al.

Original: (2019a)
Revised: (2019a)

Original: , however we replace the bound on h i,(cid:96) with our Lemma F.7.
Revised: , however we replace the bound on h i,(cid:96) with our Lemma F.7.

Original: We highlight this does not change the bound, merely in expense of a worse probability.
Revised: We highlight this does not change the bound, merely in expense of a worse probability.

Original: The next lemma concerns the bound on the product being used in backpropagation.
Revised: The next lemma concerns the bound on the product being used in backpropagation.

Original: Lemma F.9 (Variant of Lemma 7.4 in Allen-Zhu et al. (2019a)) .
Revised: Lemma F.9 (Variant of Lemma 7.4 in Allen-Zhu et al. (2019a)) .

Original: Suppose m ≥ Ω( nL log( nL )) .
Revised: Suppose m ≥ Ω( nL log( nL )) .

Original: Then with probability at least 1 − e − Ω(log 2 m ) , for all i ∈ [ n ] , (cid:96) ∈
Revised: Then with probability at least 1 − e − Ω(log 2 m ) , for all i ∈ [ n ] , (cid:96) ∈

Original: [ L ] ,
Revised: [ L ] ,

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: By Lemma F.8 part (a), we know that (cid:107) W L . . .
Revised: By Lemma F.8 part (a), we know that (cid:107) W L . . .

Original: D i,(cid:96) W (cid:96) (cid:107) ≤ O ( √ L ) , consequently, (cid:107) D
Revised: D i,(cid:96) W (cid:96) (cid:107) ≤ O ( √ L ) , consequently, (cid:107) D

Original: i,L W L . . .
Revised: i,L W L . . .

Original: D i,(cid:96) W (cid:96) (cid:107) ≤ O ( √ L ) since (cid:107) D i,L (cid:107) ≤ 1 .
Revised: D i,(cid:96) W (cid:96) (cid:107) ≤ O ( √ L ) since (cid:107) D i,L (cid:107) ≤ 1 .

Original: This means for any vector u ∈ R m , with probability at least 1 − e − Ω
Revised: This means for any vector u ∈ R m , with probability at least 1 − e − Ω

Original: ( k/L 2 ) , we have
Revised: ( k/L 2 ) , we have

Original: [Equation]
Revised: [Equation]

Original: Conditioning on this event and using the randomness of Rademacher vector a , we have
Revised: Conditioning on this event and using the randomness of Rademacher vector a , we have

Original: [Equation]
Revised: [Equation]

Original: pick t = log m √ m , we know that
Revised: pick t = log m √ m , we know that

Original: [Equation]
Revised: [Equation]

Original: This means with probability at least 1 − e − Ω(log 2 m ) , we have
Revised: This means with probability at least 1 − e − Ω(log 2 m ) , we have

Original: [Equation]
Revised: [Equation]

Original: The next several lemmas bound norms after small perturbation.
Revised: The next several lemmas bound norms after small perturbation.

Original: Lemma F.10 (Lemma 8.2 in Allen-Zhu et al. (2019a)) .
Revised: Lemma F.10 (Lemma 8.2 in Allen-Zhu et al. (2019a)) .

Original: Suppose Assumption F.3 is satisﬁed with
Revised: Suppose Assumption F.3 is satisﬁed with

Original: R ≤ O ( 1 L 9 / 2 log 3 m ) .
Revised: R ≤ O ( 1 L 9 / 2 log 3 m ) .

Original: With probability at least 1 − e − Ω( mR 2 / 3 L ) ,
Revised: With probability at least 1 − e − Ω( mR 2 / 3 L ) ,

Original: (a) ∆ g i,(cid:96) can be written as ∆ g i,(cid:96) =
Revised: (a) ∆ g i,(cid:96) can be written as ∆ g i,(cid:96) =

Original: ∆ g i,(cid:96), 1 + ∆ g i,(cid:96), 2 where (cid:107) ∆ g i,(cid:96), 1 (cid:107) 2 ≤ O ( RL 3 / 2 ) and (cid:107) ∆ g i,(cid:96), 2 (cid:107) ∞ ≤ O ( RL 5 / 2 √ log m √ m ) .
Revised: ∆ g i,(cid:96), 1 + ∆ g i,(cid:96), 2 where (cid:107) ∆ g i,(cid:96), 1 (cid:107) 2 ≤ O ( RL 3 / 2 ) and (cid:107) ∆ g i,(cid:96), 2 (cid:107) ∞ ≤ O ( RL 5 / 2 √ log m √ m ) .

Original: (b) (cid:107) ∆ D i,(cid:96) (cid:107) 0 ≤ O ( mR 2 / 3 L ) and (cid:107) (∆ D i,(cid:96) )
Revised: (b) (cid:107) ∆ D i,(cid:96) (cid:107) 0 ≤ O ( mR 2 / 3 L ) and (cid:107) (∆ D i,(cid:96) )

Original: g i,(cid:96) (cid:107) 2 ≤ O ( RL 3 / 2 ) .
Revised: g i,(cid:96) (cid:107) 2 ≤ O ( RL 3 / 2 ) .

Original: (c) (cid:107) ∆ g i,(cid:96) (cid:107) 2 , (cid:107) ∆ h i,(cid:96) (cid:107) 2 ≤ O ( RL 5 / 2 √ log m ) .
Revised: (c) (cid:107) ∆ g i,(cid:96) (cid:107) 2 , (cid:107) ∆ h i,(cid:96) (cid:107) 2 ≤ O ( RL 5 / 2 √ log m ) .

Original: Remark F.11.
Revised: Remark F.11.

Original: Lemma F.10 establishes the connection between parameter R and s of Assumption F.3 and F.4.
Revised: Lemma F.10 establishes the connection between parameter R and s of Assumption F.3 and F.4.

Original: As long as R is small, then we have s =
Revised: As long as R is small, then we have s =

Original: O ( mR 2 / 3 L ) .
Revised: O ( mR 2 / 3 L ) .

Original: Such a relation enables us to pick R to our advantage and ensure the sparsity of ∆ D i,(cid:96) is sublinear in m , and hence the update time per iteration is subquadratic in m .
Revised: Such a relation enables us to pick R to our advantage and ensure the sparsity of ∆ D i,(cid:96) is sublinear in m , and hence the update time per iteration is subquadratic in m .

Original: Lemma F.12 (Lemma 8.6 in Allen-Zhu et al. (2019a)) .
Revised: Lemma F.12 (Lemma 8.6 in Allen-Zhu et al. (2019a)) .

Original: Suppose Assumption F.4 is satisﬁed with 1 ≤ s ≤ O ( m L 3 log m ) and m ≥ Ω( nL log( nL )) , with probability at least 1 − e − Ω( s log m ) over the randomness of W (0) , for every i ∈ [ n ] , 1 ≤ a ≤ b ≤ L , we have
Revised: Suppose Assumption F.4 is satisﬁed with 1 ≤ s ≤ O ( m L 3 log m ) and m ≥ Ω( nL log( nL )) , with probability at least 1 − e − Ω( s log m ) over the randomness of W (0) , for every i ∈ [ n ] , 1 ≤ a ≤ b ≤ L , we have

Original: [Equation]
Revised: [Equation]

Original:  if Assumption F.3 is satisﬁed with R ≤ O ( 1 L 1 . 5 ) .
Revised: O [Equation] √ L ) if Assumption F.3 is satisﬁed with R [Equation]

Revised: [Equation]

Original: Corollary F.13.
Revised: Corollary F.13.

Original: Suppose Assumption F.3 is satisﬁed with R ≤ O ( 1 L 4 .
Revised: Suppose Assumption F.3 is satisﬁed with R ≤ O ( 1 L 4 .

Original: 5 log 3 m ) and m ≥ Ω( nL log( nL )) , with probability at least 1 − e − Ω(log 2 m ) over the randomness of W (0) and a , for any i ∈ [ n ] and (cid:96) ∈
Revised: 5 log 3 m ) and m ≥ Ω( nL log( nL )) , with probability at least 1 − e − Ω(log 2 m ) over the randomness of W (0) and a , for any i ∈ [ n ] and (cid:96) ∈

Original: [ L ] , we have
Revised: [ L ] , we have

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We ﬁrst note that by Lemma F.10 and the choice of R , we know that Assumption F.4 is satisﬁed with 1 ≤ s ≤ O ( m L 2 log 2 m ) .
Revised: We ﬁrst note that by Lemma F.10 and the choice of R , we know that Assumption F.4 is satisﬁed with 1 ≤ s ≤ O ( m L 2 log 2 m ) .

Original: The proceeding proof is identical to Lemma F.9, use P to denote the product ( D i,L (0) + ∆ D i,L (0))
Revised: The proceeding proof is identical to Lemma F.9, use P to denote the product ( D i,L (0) + ∆ D i,L (0))

Original: ( W L (0) + ∆ W L ) . . .
Revised: ( W L (0) + ∆ W L ) . . .

Original: ( W (cid:96) +1 (0) + ∆ W (cid:96) +1 )
Revised: ( W (cid:96) +1 (0) + ∆ W (cid:96) +1 )

Original: ( D i,(cid:96) (0) + ∆ D i,(cid:96) ) .
Revised: ( D i,(cid:96) (0) + ∆ D i,(cid:96) ) .

Original: Per Lemma F.12, we know that with the conditions stated, we have
Revised: Per Lemma F.12, we know that with the conditions stated, we have

Original: [Equation]
Revised: [Equation]

Original: , conditioning on this even, for any vector u ∈ R m , with probability at least 1 − e − Ω( s log m ) , we know that
Revised: conditioning on this even, for any vector u ∈ R m , with probability at least 1 − e − Ω( s log m ) , we know that

Original: [Equation]
Revised: [Equation]

Original: Use the randomness of vector a , we have
Revised: Use the randomness of vector a , we have

Original: [Equation]
Revised: [Equation]

Original: Pick t = log m √ m , we are done.
Revised: Pick t = log m √ m , we are done.

Original: Lemma F.14 (Variant of Lemma 8.7 in Allen-Zhu et al. (2019a)) .
Revised: Lemma F.14 (Variant of Lemma 8.7 in Allen-Zhu et al. (2019a)) .