Training (Overparametrized) Neural Networks in Near-Linear Time

Introduction

Understanding the dynamics of gradient-based optimization of deep neural networks has been a central focal point of theoretical machine learning in recent years [LY17, ZSJ+17, ZSD17, LL18, DZPS19, AZLS19a, AZLS19b, AZLL19, BJW19, OS19, ADH+19b, SY19, Dan20, JT20, BELM20]. This line of work led to a remarkable rigorous understanding of the generalization, robustness and convergence rate of first-order (SGD-based) algorithms, which are the standard choice for training DNNs. By contrast, the computational complexity of implementing gradient-based training algorithms (e.g., backpropagation) in such non-convex landscape is less understood, and gained traction only recently due to the overwhelming size of training data and complexity of network design [MG15, DHS11, LJH+19, CGH+19, ZMG19].

Second-order gradient algorithms (which employ information about the Hessian of the loss function), pose an intriguing computational tradeoff in this context: On one hand, they are known to converge extremely fast, at a rate independent of the input size (i.e., only $O(\log 1/\epsilon)$ iterations [ZMG19]), and offer a qualitative advantage in overcoming pathological curvature issues that arise in first-order methods, by exploiting the local geometry of the loss function. This feature implies another practical advantage of second order methods, namely, that they do not require tuning the learning rate [CGH+19, ZMG19]. On the other hand, second-order methods have a prohibitive cost per iteration, as they involve inverting a dynamically-changing dense Hessian matrix. This drawback explains the scarcity of second order methods in large scale non-convex optimization, in contrast to its popularity in the convex setting.

Following [CGH+19, ZMG19], we focus on two-layer (i.e., single hidden-layer) neural networks. While our algorithm extends to the multilayer case (with a slight comprise on the width dependence), we argue that, as far as training time, the two-layer case is not only the common case, but in fact the only interesting case for constant training error: Indeed, in the multilayer case ($L\geq 2$), we claim that the mere cost of feed-forward computation of the network’s output is already $\Omega_{\epsilon}(m^{2}nL)$. Indeed, the total number of parameters of $L$-layer networks is $M=(L-1)m^{2}+md$, and as such, feed-forward computation requires, at the very least, computing a single product of $m\times m$ (dense) matrices $W$ with a $m\times 1$ vector for each training data, which already costs $m^{2}n$ time:

1 Our Result

Our main result is a quadratic speedup to the algorithm of [CGH+19], yielding an essentially optimal training algorithm for overparametrized two-layer neural networks. Moreover, in contrast to [CGH+19], our algorithm applies to the more complex and realistic case of ReLU activation functions. Our main result is shown below (For a more comprehensive comparison, see Table 1 below and references therein).

Suppose the width of a two layer ReLU neural network satisfies

where $\lambda>0$ denotes the minimum eigenvalue of the Gram matrix (see Eq. (5) below), $n$ is the number of training data, $d$ is the input dimension. Then with probability $1-\delta$ over the random initialization of neural network and the randomness of the training algorithm, our algorithm achieves

The computational cost of each iteration is $\widetilde{O}(mnd+n^{3})$, and the running time for reducing the training loss to $\epsilon$ is $\widetilde{O}((mnd+n^{3})\log(1/\epsilon))$. Using fast matrix-multiplication, the total running time can be further reduced to $\widetilde{O}((mnd+n^{\omega})\log(1/\epsilon))$. Here, $\omega<2.373$ denotes the fast matrix-multiplication (FMM) constant for multiplying two $n\times n$ matrices [Wil12, LG14].

We stress that that our algorithm runs in (near) linear time even for networks with width $m\gtrsim n^{2}$ and in fact, under the common belief that $\omega=2$, this is true so long as $m\gtrsim n$ (!). This means that the bottleneck for linear-time training of small-width DNNs is not computational, but rather analytic: The overparametrization requirements ($m\gtrsim n^{4}$) in Theorem 1.1 stems from current-best analysis of the convergence guarantees of (S)GD-based training of ReLU networks, and any improvement on these bounds would directly yield linear-time training for thinner networks using our algorithm.

The majority of ML optimization literature on overparametrized network training is dedicated to understanding and minimizing the number of iterations of the training process [ZMG19, CGH+19] as opposed to the cost per iteration, which is the focus of our paper. Our work shows that it is possible to harness the toolbox of randomized linear algebra— which was heavily used in the past decade to reduce the cost of convex optimization tasks— in the nonconvex setting of deep learning as well. A key ingredient in our algorithm is linear sketching, where the main idea is to carefully compress a linear system underlying an optimization problem, in a way that preserves a good enough solution to the problem yet can be solved much faster in lower dimension. This is the essence of the celebrated Sketch-and-Solve (S&S) paradigm [CW13]. As we explain below, our main departure from the classic S&S framework (e.g., [PW17]) is that we cannot afford to directly solve the underlying compressed regression problem (as this approach turns out to be prohibitively slow for our application). Instead, we use sketching (or sampling) to facilitate fast preconditioning of linear systems (in the spirit of [ST04, KOSZ13, RT08, Woo14]), which in turn enables to solve the compressed regression problem to very high accuracy via first-order conjugate gradient descent. This approach essentially decouples the sketching error from the final precision error of the Gauss-Newton step, enabling a much smaller sketch size. We believe this (somewhat unconventional) approach to non-convex optimization is the most enduring message of our work.

2 Related Work

Despite the prevalence of first order methods in deep learning applications, there is a vast body of ongoing work [BRB17, BLH18, MG15, GM16, GKS18, CGH+19, ZMG19] aiming to design more scalable second-order algorithms that overcome the limitations of (S)GD for optimizing deep models. Grosse and Martens [MG15, GM16] designed the K-FAC method, where the idea is to use Kronecker-factors to approximate the Fisher information matrix, combined with natural gradient descent. This approach has been further explored and extended by [WMG+17, GLB+18, MBJ18]. Gupta et al. [GKS18] designed the “Shampoo method”, based on the idea of structure-aware preconditioning. Anil et al. [AGK+20] further validate the practical perfromance of Shampoo and incorporated it into hardware. However, despite sporadic empirical evidence of such second-order methods (e.g., K-FAC and Shampoo), these methods generally lack a provable theoretical guarantee on the performance when applied to deep neural networks. Furthermore, in the overparametrized setting, their cost per-iteration in general is at least $\Omega(mn^{2})$.

We remark that in the convex setting, theoretical guarantees for large-scale second-order algorithms have been established (e.g.,[ABH17, PW17, MNJ16, Bub15]), but such rigorous analysis in non-convex setting was only recently proposed ([CGH+19, ZMG19]). Our algorithm bears some similarities to the NewtonSketch algorithm of [PW17], which also incorporates sketching into second order Newton methods. A key difference, however, is that the algorithm of [PW17] works only for convex problems, and requires access to $(\nabla^{2}f(x))^{1/2}$ (i.e., the square-root of the Hessian). Most importantly, though, [PW17] use the standard (black-box) Sketch-and-Solve paradigm to reduce the computational cost, while this approach incurs large computation overhead in our non-convex setting. By contrast, we use sketching as a subroutine for fast preconditioning. As a by-product, in Section D we show how to apply our techniques to give a substantial improvement over [PW17] in the convex setting.

The aforementioned works of [ZMG19] and [CGH+19] are most similar in spirit to ours. Zhang et al. [ZMG19] analyzed the convergence rate of Natural gradient descent algorithms for two-layer (overparametrized) neural networks, and showed that the number of iterations is independent of the training data size $n$ (essentially $\log(1/\epsilon)$). They also demonstrate similar results for the convergence rate of K-FAC in the overparametrized regime, albeit with larger requirement on the width $m$. Another downside of K-FAC is the high cost per iteration ($\sim mn^{2}$). Cai et al. [CGH+19] analyzed the convergence rate of the so-called Gram-Gauss-Newton algorithm for training two-layer (overparametrized) neural network with smooth activation gates. They proved a quardratic (i.e., doubly-logarithnmic) convergence rate in this setting ($\log(\log(1/\epsilon))$) albeit with $O(mn^{2})$ cost per iteration. It is noteworthy that this quadratic convergence rate analysis does not readily extend to the more complex and realistic setting of ReLU activation gates, which is the focus of our work. [CGH+19] also prove bounds on the convergence of ‘batch GGN’, showing that it is possible to reduce the cost-per-iteration to $m$, at the price of $O(n^{2}\log(1/\epsilon))$ iterations, for very heavily overparametrized DNNs (currently $m=\Omega(n^{18})$).

Sketching

In the classic S&S paradigm, the underlying regression solver is treated as a black box, and the computational savings comes from applying it on a smaller compressed matrix. Since then, sketching (or sampling) has also been used in a non-black-box fashion for speeding-up optimization tasks, e.g., as a subroutine for preconditioning [Woo14, RT08, ST04, KOSZ13] or fast inverse-maintenance in Linear Programming solvers, semi-definite programming, cutting plane methods, and empirical-risk minimization [CLS19, JSWZ20, JKL+20, JLSW20, LSZ19].

Overparametrization in neural networks

A long and active line of work in recent deep learning literature has focused on obtaining rigorous bounds on the convergence rate of various local-search algorithms for optimizing DNNs [LL18, DZPS19, AZLS19a, AZLS19b, ADH+19a, ADH+19b, SY19, JT20]. The breakthrough work of Jacob et al. [JGH18] and subsequent developments For a complete list of references, we refer the readers to [ADH+19a, ADH+19b]. introduced the notion of neural tangent kernels (NTK), implying that for wide enough networks ($m\gtrsim n^{4}$), (stochastic) gradient descent provably converges to an optimal solution, with generalization error independent of the number of network parameters.

Technical Overview

We now provide a streamlined overview of our main result, Theorem 1.1. As discussed in the introduction, our algorithm extends to multi-layer ReLU networks , though we focus on the two-layer case (one-hidden layer), which is the most interesting case where one can indeed hope for linear training time.

where $(f_{t}-y)$ is the training error with respect to the network’s output and the training labels $y$. Since the Gauss-Newton method is robust to small perturbation errors (essentially [Vai89b, Vai89a]), our analysis shows that it is sufficient to find an approximate solution $g^{\prime}_{t}$ such that $J_{t}^{\top}g^{\prime}_{t}$ satisfies

Our key idea is to use dimension reduction—not to directly invert the compressed matrix—but rather to precondition it quickly. More precisely, our approach is to use a (conjugate) gradient-descent solver for the regression problem itself, with a fast preconditioning step, ensuring exponentially faster convergence to very high (polynomially small) accuracy. Indeed, conjugate gradient descent is guaranteed to find a $\gamma$-approximate solution to a regression problem $\min_{x}\|Ax-b\|_{2}$ in $O(\sqrt{\kappa}\log(1/\gamma))$ iterations, where $\kappa(A)$ is the condition number of $A$ (i.e., the ratio of maximum to minimum eigenvalue). Therefore, if we can ensure that $\kappa(G_{t})$ is small, then we can $\gamma$-solve the regression problem in $\sim mn\log(1/\gamma)=\widetilde{O}(mn)$ time, since the per-iteration cost of first-order SGD is linear ($\sim mn$).

The crucial advantage of our approach is that it decouples the sketching error from the final precision of the regression problem: Unlike the usual ‘sketch-and-solve’ method, where the sketching error $\delta$ directly affects the overall precision of the solution to (2), here $\delta$ only affects the quality of the preconditioner (i.e., the ratio of max/min singular values of the sketch $\widetilde{G}_{t}$), hence it suffices to take a constant sketching error $\delta=0.1$ (say), while letting the SGD deal with the final precision (at it has logarithmic dependence on $\gamma$). See Lemma B.1 for the formal details.

We remark that, by definition, the preconditioning step (on the JL sketch) does not preserve the eigen-spectrum of $G_{t}$, which is in fact necessary to guarantee the fast convergence of the Gauss-Newton iteration (see Lemma C.3) . The point is that this preconditioning step is only preformed as a local subroutine so as to solve the regression problem, and does not affect the convergence rate of the outer loop.

Preliminaries

We can compute the gradient of $\mathcal{L}$ in terms of $w_{r}$

We assume the least eigenvalue $\lambda$ of the kernel matrix $K$ defined in Eq. (5) satisfies $\lambda>0$.

2 Subspace embedding

Subspace embedding was first introduced by Sarlós [Sar06], it has been extensively used in numerical linear algebra field over the last decade [CW13, NN13, BW14, SWZ19]. For a more detailed survey, we refer the readers to [Woo14]. The formal definition is:

Combining Fast-JL sketching matrix [AC06, DMM06, Tro11, DMIMW12, LDFU13, PSW17] with a classical $\epsilon$-net argument [Woo14] gives subspace embedding,

Our Algorithm

Our main algorithm is shown in Algorithm 1. We have the following convergence result of our algorithm.

Suppose the width of a ReLU neural network satisfies

then with probability $1-\delta$ over the random initialization of neural network and the randomness of the training algorithm, our algorithm (procedure FasterTwoLayer in Algorithm 1) achieves

The computation cost in each iteration is $\widetilde{O}(mnd+n^{3})$, and the running time for reducing the training loss to $\epsilon$ is $\widetilde{O}((mnd+n^{3})\log(1/\epsilon))$. Using fast matrix-multiplication, the total running time can be further reduced to $\widetilde{O}((mnd+n^{\omega})\log(1/\epsilon))$.

The main difference between [CGH+19, ZMG19] and our algorithm is that we perform an approximate Newton update (see line 6). The crucial observation here is that the Newton method is robust to small loss, thus it suffices to present a fine approximation. This observation is well-known in the convex optimization but unclear to the non-convex (but overparameterized) neural network setting. Another crucial observation is that instead of directly approximating the Gram matrix, it is suffices to approximate $(J_{t}J_{t}^{\top})^{-1}g_{t}=G_{t}^{-1}g_{t}$. Intuitively, this follows from

where $(J_{t}^{\top}J_{t})^{\dagger}$ denotes the pseudo-inverse of $J_{t}^{\top}J_{t}$ and the last term is exactly the Newton update. This observation allows us to formulate the problem a regression problem (see Eq. (6)), on which we can introduce techniques from randomize linear algebra and develop fast algorithm that solves it in near linear time.

Using procedure FastRegression (in Algorithm 2), with probability $1-\delta$, we can compute an $\epsilon$-approximate solution $x^{\prime}$ satisfying

in $\widetilde{O}\left(Nk\log(\kappa/\epsilon)+k^{3}\right)$ time.

It should come as no surprise that our techniques can help accelerating a broad class of solvers in convex optimization problems as well. In the full version of this paper, we elaborate on this application, and in particular show how our technique improves the runtime of the “Newton-Sketch” algorithm of [PW17].

Conclusion and Open Problems

Our work provides a computationally-efficient (near-linear time) second-order algorithm for training sufficiently overparametrized two-layer neural network, overcoming the drawbacks of traditional first-order gradient algorithms. Our main technical contribution is developing a faster regression solver which uses linear sketching for fast preconditioning (in time independent of the network width). As such, our work demonstrates that the toolbox of randomized linear algebra can substantially reduce the computational cost of second-order methods in non-convex optimization, and not just in the convex setting for which it was originally developed (e.g., [PW17, Woo14, CLS19, JSWZ20, JKL+20, JLSW20, LSZ19]).

Finally, we remark that, while the running time of our algorithm is $\widetilde{O}(Mn+n^{3})$ (or $O(Mn+n^{\omega})$ using FMM), it is no longer (near) linear for networks with parameters $M\leq n^{2}$ (resp. $M\lesssim n^{\omega-1}$). While it is widely believed that $\omega=2$ [CKSU05], FMM algorithms are impractical at present, and it would therefore be very interesting to improve the extra additive term from $n^{3}$ to $n^{2+o(1)}$ (which seems best possible for dense $n\times n$ matrices), or even to $n^{3-\epsilon}$ using a practically viable algorithm. Faster preconditioners seem key to this avenue.

Acknowledgments

The authors would like to thank David Woodruff for telling us the tensor trick for computing kernel matrices and helping us improve the presentation of the paper. The authors would like to thank Sanjeev Arora, Simon S. Du, and Jason Lee for the suggestion of this topic. The authors would like to thank Yangsibo Huang, Shunhua Jiang, Yaonan Jin, Kai Li, Xiaoxiao Li, Zhenyu Song, Yushan Su, Fan Yi, and Hengjie Zhang for very useful discussions.

Appendix A Appendix

Organization The Appendix is organized as follows. Section A contains notations and some basic facts. In Section B we present the fast regression solver. In Section C we prove our main result for two-layer ReLU networks. Finally, in Section D we show that our optimization framework can obtain acceleration in classic convex optimization setting, improve over [PW17].

A.2 Probability Tools

Let $X=\sum_{i=1}^{n}X_{i}$, where $X_{i}=1$ with probability $p_{i}$ and $X_{i}=0$ with probability $1-p_{i}$, and all $X_{i}$ are independent. Let $\mu=\E[X]=\sum_{i=1}^{n}p_{i}$. Then 1. $\Pr[X\geq(1+\delta)\mu]\leq\exp(-\delta^{2}\mu/3)$, $\forall\delta>0$ ; 2. $\Pr[X\leq(1-\delta)\mu]\leq\exp(-\delta^{2}\mu/2)$, $\forall 0<\delta<1$.

Let $X_{1},\cdots,X_{n}$ denote $n$ independent bounded variables in $[a_{i},b_{i}]$. Let $X=\sum_{i=1}^{n}X_{i}$, then we have

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

A.3 Basic Facts

For any two matrices $A,B$, $\kappa(B)\leq\kappa(AB)\kappa(A)$.

Hence we have $\sigma_{\max}(B)\leq\sigma_{\max}(AB)/\sigma_{\min}(A)$. Similarly, we have

i.e., $\sigma_{\min}(B)\geq\sigma_{\min}(AB)/\sigma_{\max}(A)$. Thus we conclude

Appendix B Fast regression solver

We can compute an $\epsilon$-approximate solution $x^{\prime}$ satisfying

in $\widetilde{O}\left(Nk\log(\kappa/\epsilon)+k^{3}\right)$ time. Using fast matrix-multiplication, the total running time can be further reduced to $\widetilde{O}((mnd+n^{\omega})\log(1/\epsilon))$.

We choose $\epsilon_{0}=0.1$, and consider the regression problem

By lemma B.2, using gradient descent, after $t=\log(1/\epsilon)$ iterations, we can find $z_{t}$ satisfying

where $z^{\star}=(R^{\top}A^{\top}AR)^{-1}R^{\top}y$ is the optimal solution to Eq. (10). We are going to show that $x_{t}=Rz_{t}$ is an $2\kappa\epsilon$-approximate solution to the original regression problem (8), i.e.,

where the first step follows from Eq. (12) (13), the second step follows from $R$ is a square matrix and thus $\kappa(R)=\kappa(R^{\top})$, the third step follows from Fact A.4 and the last step follows from Eq. (9).

For the running time, the preconditioning time is $\widetilde{O}(Nk+k^{3})$, the number of iteration for gradient desent is $\log(\kappa/\epsilon)$, the running time per iteration is $\widetilde{O}(Nk)$, thus the total running time is

The preconditioning can be reduced to $\widetilde{O}(Nk+k^{\omega})$ when using fast matrix multiplication to compute the QR decomposition of $SA$ [DDH07]. ∎

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

The gradient at time $t$ is $B^{\top}(Bx_{t}-y)$ and $x_{t+1}=x_{t}-B^{\top}(Bx_{t}-y)$, thus we have

The second step follows from $B^{\top}Bx^{\star}=B^{\top}y$. The last step follows from the eigenvalue of $BB^{\top}$ belongs to $[\frac{9}{16},\frac{25}{16}]$ by our assumption. Thus we complete the proof.∎

Appendix C Our Algorithm

We delicate to prove the following result in this section, which is essentially Theorem 4.1.

Suppose the width of the neural network satisfies $m=\Omega(\max\{\lambda^{-4}n^{4},\lambda^{-2}n^{2}d\log(16n/\delta)\})$, then with probability $1-\delta$ over the random initialization of neural network and the randomness of the algorithm, our algorithm achieves

The computation cost in each iteration is $\widetilde{O}(mnd+n^{3})$, and the running time for reducing the training loss to $\epsilon$ is $\widetilde{O}((mnd+n^{3})\log(1/\epsilon))$. Using fast matrix multiplication, the running time is $\widetilde{O}((mnd+n^{\omega})\log(1/\epsilon))$.

The follow lemmas are standard in literature.

Suppose $m=\Omega(d\log(n/\delta))$, then with probability $1-\delta$, we have the following

$\|J_{W_{0},x_{i}}\|_{F}=O(1)$, for $i\in[n]$.

Suppose $m=\Omega(\lambda^{-2}n^{2}\log(n/\delta))$, then with probability at least $1-\delta$, we have

When weights do not change very much, we have

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

(2) For the second claim, we have for any $i\in[n]$

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

It is easy to see $A_{i,r}$ happens if and only if $w_{r}(0)^{\top}x_{i}\in[-R/\sqrt{m},R/\sqrt{m}]$. By the anticoncentration of Gaussian (see Lemma A.3), we have $\E[s_{i,r}]=\Pr[A_{i,r}]\leq\frac{4}{5}R/\sqrt{m}$. Thus we have

holds for any $t>0$. The second inequality comes from the Hoeffding bound (see Lemma A.2), the last inequality comes from $R>1$. Taking $t=2\log(n/\delta)$ and using union bound over $i$, with probability $1-\delta$, we have

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

The second inequality follows from $m=\widetilde{\Omega}(R^{2}n^{2})$.

We use induction to prove the following two claims recursively. We take $R\approx n/\lambda$ in the proof.

$\|w_{r}(t)-w_{r}(0)\|_{2}\leq R/\sqrt{m}$ holds for any $r\in[m]$ and $t\geq 0$.

$\|f_{t}-y\|_{2}\leq\frac{1}{2}\|f_{t-1}-y\|_{2}$ holds for any $t\geq 1$.

Suppose the above two claims hold up to $t$, we prove they continue to hold for time $t+1$. The second claim is more delicate, we are going to prove it first and we define

where we denote $g^{\star}=(J_{t}J_{t}^{\top})^{-1}(f_{t}-y)$ to be the optimal solution to Eq. (6). The second step follows from the definiton of $J_{t,t+1}$ and simple calculus. The third step follows from the updating rule of the algorithm.

since $g_{t}$ is an $\epsilon_{0}(\epsilon_{0}\leq\frac{1}{6})$ approximate solution to regression problem (6).

The third step follows from the second claim in Lemma C.4 and the fact that

The second inequality follows from $\sigma_{\min}(J_{t})=\sqrt{\lambda_{\min}(J_{t}^{\top}J_{t})}\geq\sqrt{\lambda/2}$ (see Lemma C.5).

Combining Eq. (19), (20) and (21), we have

since $m=\widetilde{\Omega}(\lambda^{-4}n^{4})$.

The first step comes from $\lambda_{\min}(J_{t}J_{t}^{\top})=\lambda_{\min}(G_{t})\geq\lambda/2$ (see Lemma C.4) and the last step comes from $g_{t}$ is an $\epsilon_{0}$ approximate solution to Eq. (6). The fourth step follows from Eq. (C) and the fact that $\|(J_{t}J_{t}^{\top})^{-1}\|\leq 2/\lambda$. The last step follows from $g_{t}$ is an $\epsilon_{0}$ ($\epsilon_{0}\leq\sqrt{\lambda/n}$) approximate solution to the regression (6).

The second step follows from Eq. (20) and (C) and the fact that $\|J_{t}\|\leq O(\sqrt{n})$ (see Lemma C.4) The last step follows from the $m\geq\Omega(n^{4}\lambda^{-4})$. Combining Eq. (17), (18), (22), and (C), we have proved the second claim, i.e.,

It remains to show that $W_{t}$ does not move far away from $W_{0}$. First, we have

where the third step follows from Eq. (C) and the last step follows from the obvious fact that $1/\sqrt{n\lambda}\leq 1/\lambda$.

Hence, for any $r\in[m]$ and $0\leq k\leq t$, if we use $g_{k,i}$ to denote the $i^{th}$ indice of $g_{k}$, then we have

The first step follows from the updating rule, the second step follows from triangle inequalities and the fact that $a_{r}=\pm 1$, $\|x_{r}\|_{2}=1$. The third step comes from Cauchy-Schwartz inequality, and the fouth step comes from Eq. (26) and Eq. (C). The last inequality comes from the fact that $\|f_{0}-y\|_{2}\leq O(\sqrt{n})$ (see Lemma C.2). Consequently, we have

Thus we also finish the proof of the first claim.

It remains to give an analysis on the running time of our algorithm. In each iteration, besides evaluating function value and doing backpropagation, which generally takes $O(mnd)$ time, we also need to solve the regression problem in (6), which takes $\widetilde{O}(mnd\log(\kappa(J_{t}J_{t}^{\top})/\epsilon_{0})+n^{3})$ time by Lemma B.1. From Lemma C.4, we know $\|J_{t}J_{t}^{\top}\|=\|G_{t}\|\leq O(n)$ and $\lambda_{\min}(J_{t}J_{t}^{\top})=\lambda_{\min}(G_{t})\geq O(\lambda)$. Moreover, we only need to set $\epsilon_{0}=\min\{\sqrt{\lambda/n},1/6\}$. Thus the total computation cost in each iteration is $\widetilde{O}(mnd+n^{3})$, and the total running time to reduce the trainning loss below $\epsilon$ is $\widetilde{O}((mnd+n^{3})\log(1/\epsilon))$. ∎

Appendix D Application: Convex Optimization

We apply our technique to convex optimization problem. We follow the problem formulation in [PW17] and consider the problem

where $f$ is $\gamma$-strongly convex, $\beta$-smooth and its Hessian matrix $\nabla^{2}f(x)$ is $L$ Lipschitz continuous,

The function $f$ is $\gamma$-strongly convex if

The function $f$ is $\beta$-strongly convex if

The Hessian matrix of function $f$ is $L$ Lipschitz continuous is

As in [PW17], we further assume we have access to

For more examples, we refer interested reader to Section 3.3 in [PW17]

Naive implementation of Newton method needs to compute

and it costs $O(nd^{2})$ time. The original analysis of NewtonSketch in [PW17] takes $\widetilde{O}(nd+d^{3})$, but it requires $n\geq d\kappa^{2}$, where $\kappa$ is the condition number defined as $\kappa=\beta/\gamma$. There are many follow up work [XYR+16, YLZ17, BBN19] intending to get rid of the extra dependence on the condition number $\kappa$. We present an alternative approach and improve the running time to $\widetilde{O}((n\log(\kappa)+d^{2})d\log(1/\epsilon))$ by incorporating the “fast regression solver” introduced in this paper.

Our algorithm is shown in Algorithm 3. Formally, we have

Suppose function $f$ is $\gamma$-strongly convex, $\beta$-smooth and its Hessian is $L$ Lipschitz continuous. Given an initialization point $x_{0}$ satisfying $\|x_{0}-x^{\star}\|_{2}\leq\gamma/(2L)$, there is an algorithm (procedure FastNewtonUpdate in Algorithm 3) achieves

Consequently, in order to find an $\epsilon$ approxmate optimal solution, the running time is

Using fast matrix multiplication, the running time can be further reduced to $\widetilde{O}\left((nd\log(\kappa)+d^{\omega})\log(1/\epsilon)\right)$.

We first analyze the correctness, and then give an analysis on the running time. Denote

The first step follows from the definition of $\widetilde{g}_{t}$ in Eq. (30), the second step follows from $\nabla f(x^{\star})=0$. If the Hessian is $L$ Lipschitz continuous, we have For the second term

The first step follows from Eq. (30), the second step holds since $g_{t}$ is an $1/(4\kappa)$ approximate solution to Eq. (28). The third step follows from the smoothness of $f$. Consequently, we have

The first step follows from the convexity of $f$. The second step follows from Eq. (31), (32), and (33). Thus we prove the correctness of Eq. (29). Since we know $\kappa(\nabla^{2}f(x_{t})^{\frac{1}{2}})=\sqrt{\kappa}$, the running time per iteration is $\widetilde{O}(nd\log(\kappa)+d^{3})$ by Lemma B.1. Thus we conclude the proof. ∎

References