Newton-LESS: Sparsification without Trade-offs for the Sketched Newton Update
Year: 2022

Abstract:
In second-order optimization, a potential bottleneck can be computing the Hessian matrix of the optimized function at every iteration. Randomized sketching has emerged as a powerful technique for constructing estimates of the Hessian which can be used to perform approximate Newton steps. This involves multiplication by a random sketching matrix, which introduces a trade-off between the computational cost of sketching and the convergence rate of the optimization algorithm. A theoretically desirable but practically much too expensive choice is to use a dense Gaussian sketching matrix, which produces unbiased estimates of the exact Newton step and which offers strong problem-independent convergence guarantees. We show that the Gaussian sketching matrix can be drastically sparsified, significantly reducing the computational cost of sketching, without substantially affecting its convergence properties. This approach, called Newton-LESS, is based on a recently introduced sketching technique: LEverage Score Sparsified (LESS) embeddings. We prove that Newton-LESS enjoys nearly the same problem-independent local convergence rate as Gaussian embeddings, not just up to constant factors but even down to lower order terms, for a large class of optimization tasks. In particular, this leads to a new state-of-the-art convergence result for an iterative least squares solver. Finally, we extend LESS embeddings to include uniformly sparsified random sign matrices which can be implemented efficiently and which perform well in numerical experiments.

heading: 1 Introduction
text: Consider the task of minimizing a twice-differentiable convex function f : Rd → R:
find x∗ = argmin x∈Rd f(x).
One of the most classical iterative algorithms for solving this task is the Newton’s method, which takes steps of the form xt+1 = xt − µt∇2f(xt)−1∇f(xt), and which leverages second-order information in the d×d Hessian matrix∇2f(xt) to achieve rapid convergence, especially locally as it approaches the optimum x∗. However, in many settings, the cost of forming the exact Hessian is prohibitively expensive, particularly when the function f is given as a sum of n d components, i.e., f(x) =∑n i=1 fi(x). This commonly arises in machine learning when f represents the training loss over a
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
dataset of n elements, as well as in solving semi-definite programs, portfolio optimization, and other tasks. In these contexts, we can represent the Hessian via a decomposition∇2f(x) = Af (x)>Af (x), where Af (x) is a tall n × d matrix, which can be easily formed, and the main bottleneck is the matrix multiplication which takes O(nd2) arithmetic operations. To avoid this bottleneck, many randomized second-order methods have been proposed which use a Hessian estimate in place of the exact Hessian (e.g., [BCNN11, EM15, ABH17, RKM19]). This naturally leads to a trade-off between the per-iteration cost of the method and the number of iterations needed to reach convergence. We develop Newton-LESS, a randomized second-order method which eliminates the computational bottleneck while minimizing the convergence trade-offs.
An important family of approximate second-order methods is known as the Newton Sketch [PW17]: x̃t+1 = x̃t − µt ( Af (x̃t) > S>t StAf (x̃t) )−1∇f(x̃t), (1) where µt is the step size, and St is a random m × n sketching matrix, with m n, that is used to reduce Af (x̃t) to a small m × d sketch StAf (x̃t). This brings the complexity of forming the Hessian down to O(md2) time plus the cost of forming the sketch.
Naturally, sketching methods vary in their computational cost and they can affect the convergence rate, so the right choice of St depends on the computation-convergence trade-off. On one end of this spectrum are the so-called Sub-Sampled Newton methods [RKM19, XRKM17, YXRKM18], where St simply selects a random sample of m rows of Af (x) (e.g., a sample of data points in a training set) to form the sketch. Here the sketching cost is negligible, since St is extremely sparse, but the convergence rate can be highly variable and problem-dependent. On the other end, we have what we will refer to as the Gaussian Newton Sketch, where St is a dense matrix with i.i.d. scaled Gaussian entries (a.k.a. a Gaussian embedding). While the O(mnd) cost of performing this sketch limits its practical appeal, Gaussian Newton Sketch has a number of unique and desirable properties [LP19]: it enjoys strong problem-independent convergence rates; it produces unbiased estimates of the exact Newton update (useful in distributed settings); and it admits analytic expressions for the optimal step size.
A natural way to interpolate between these two extremes is to vary the sparsity s of the sketching matrix St, from s = 1 non-zero element per row (Sub-Sampling) to s = n non-zero elements (Gaussian embedding), with the sketching complexity O(mds).1 Motivated by this, we ask:
Can we sparsify the Gaussian embedding, making its sparsity closer to that of Sub-Sampling, without suffering any convergence trade-offs?
In this paper, we provide an affirmative answer to this question. We show that it is possible to drastically sparsify the Gaussian embedding so that two key statistics of the sketches, namely first
1For a more detailed discussion of other sketching techniques that may not fit this taxonomy, such as the Subsampled Randomized Hadamard Transform and the CountSketch, see Section 1.2.
and second inverse moments of the sketched Hessian, are nearly preserved in a very strong sense. Namely, the two inverse moments of the sparse sketches can be upper and lower bounded by the corresponding quantities for the dense Gaussian embeddings, where the upper/lower bounds are matching not just up to constant factors, but down to lower order terms (see Theorem 6). We use this to show that the Gaussian Newton Sketch can be sparsified to the point where the cost of sketching is proportional to the cost of other operations, while nearly preserving the convergence rate (again, down to lower order terms; see Theorem 1). This is illustrated conceptually in Figure 1, showing how the sparsity of the sketch affects the per-iteration convergence rate as well as the computational cost of the Newton Sketch. We observe that while the convergence rate improves as we increase the density, it eventually flattens out. On the other hand, the computational cost stays largely flat until some point when it starts increasing at a linear rate. As a result, there is a sparsity regime where we achieve the best of both worlds: the convergence rate and the computational cost are both nearly at their optimal values, thereby avoiding any trade-off.
To establish our results, we build upon a recently introduced sketching technique called LEverage Score Sparsified (LESS) embeddings [DLDM21]. LESS embeddings use leverage score techniques [DMIMW12a] to provide a carefully-constructed (random) sparsification pattern. This is used to produce a sub-Gaussian embedding with d non-zeros per row of St (as opposed to n for a dense matrix), so that the cost of forming the sketch StAf (x̃t) matches the cost of constructing the Hessian estimate, i.e., O(md2) (see Section 2). [DLDM21] analyzed the first inverse moment of the sketch to show that LESS embeddings retain certain unbiasedness properties of Gaussian embeddings. In our setting, this captures the bias of the Newton Sketch, but it does not capture the variance, which is needed to control the convergence rate.
Contributions. In this paper, we analyze both the bias and the variance of Newton Sketch with LESS embeddings (Newton-LESS; see Definition 2 and Lemma 7), resulting in a comprehensive convergence analysis. The following are our key contributions:
1. Characterization of the second inverse moment of the sketched Hessian for a class of sketches including sub-Gaussian matrices and LESS embeddings;
2. Precise problem-independent local convergence rates for Newton-LESS, matching the Gaussian Newton Sketch down to lower order terms;
3. Extension of Newton-LESS to regularized minimization tasks, with improved dimensionindependent guarantees for the sketch sparsity and convergence rate;
4. Notable corollary: Best known global convergence rate for an iterative least squares solver, which translates to state-of-the-art numerical performance.

heading: 1.1 Main results
text: As our main contribution, we show that, under standard assumptions on the function f(x), NewtonLESS achieves the same problem-independent local convergence rate as the Gaussian Newton Sketch, despite drastically smaller per-iteration cost. Theorem 1. Assume that f(x) is (a) self-concordant, or (b) has a Lipschitz continuous Hessian. Also, let H = ∇2f(x∗) be positive definite. There is a neighborhood U containing x∗ such that if x̃0 ∈ U , then Newton-LESS with sketch size m ≥ Cd log(dT/δ) and step size µt = 1− dm satisfies:(
Eδ ‖x̃T − x∗‖2H ‖x̃0 − x∗‖2H
)1/T ≈ d
m for = O ( 1√ d ) ,
where EδX is expectation conditioned on an event that holds with a 1 − δ probability, ‖v‖M =√ v>Mv, and a ≈ b means that |a− b| ≤ b.
Remark 2. The same guarantee holds for the Gaussian Newton Sketch, but it is not known for any fast sketching method other than Newton-LESS (see Section 1.2). The alternative assumptions of self-concordance and Lipschitz continuous Hessian are standard in the local convergence analysis of the classical Newton’s method, and they only affect the size of the neighborhood U (see Section 4). Global convergence of Newton-LESS follows from existing analysis of the Newton Sketch [PW17].
The notion of expectation Eδ allows us to accurately capture the average behavior of a randomized algorithm over a moderate (i.e., polynomial in d) number of trials even when the true expectation
is not well behaved. Here, this guards against the (very unlikely, but non-zero) possibility that the Hessian estimate produced by a sparse sketch will be ill-conditioned.
To illustrate this result in a special case (of obvious independent interest), we provide a simple corollary for the least squares regression task, i.e., f(x) = 12‖Ax − b‖
2. Importantly, here the convergence rate of ( dm ) T holds globally. Also, for this task we have 12‖x− x ∗‖2H = f(x)− f(x∗), so the convergence can be stated in terms of the excess function value. To our knowledge, this is the best known convergence guarantee for a fast iterative least squares solver. Corollary 3. Let f(x) = 12‖Ax − b‖
2 for A ∈ Rn×d and b ∈ Rn. Then, given any x̃0 ∈ Rd, Newton-LESS with sketch size m ≥ Cd log(dT/δ) and step size µt = 1− dm satisfies:(
Eδ f(x̃T )− f(x∗) f(x̃0)− f(x∗)
)1/T ≈ d
m for = O ( 1√ d ) .
Prior to this work, a convergence rate of ( dm ) T was known only for dense Gaussian embeddings, and only for the least squares task [LP19]. On the other hand, our results apply as generally as the standard local convergence analysis of the Newton’s method, and they include a broad class of sketches. In Section 3, we provide general structural conditions on a randomized sketching matrix that are needed to enable our analysis. These conditions are satisfied by a wide range of sketching methods, including all sub-Gaussian embeddings (e.g., using random sign entries instead of Gaussians), the original LESS embeddings, and other choices of sparse random matrices (see Lemma 7). Moreover, we develop an improved local convergence analysis of the Newton Sketch, which allows us to recover the precise convergence rate and derive the optimal step size. In Appendix D, we also discuss a distributed variant of Newton-LESS, which takes advantage of the near-unbiasedness properties of LESS embeddings, extending the results of [DLDM21].
The performance of Newton-LESS can be further improved for regularized minimization tasks. Namely, suppose that function f can be decomposed as follows: f(x) = f0(x) + g(x), where g(x) has a Hessian that is easy to evaluate (e.g., l2-regularization, g(x) = λ2 ‖x‖
2). In this case, a modified variant of the Newton Sketch has been considered, where only the f0 component is sketched:
x̃t+1 = x̃t − µt ( Af0(x̃t) > S>t StAf0(x̃t) +∇2g(x̃t) )−1∇f(x̃t), (2) where, again, we let Af0(x) be an n× d matrix that encodes the second-order information in f0 at x. For example, in the case of regularized least squares, f(x) = 12‖Ax− b‖ 2 + λ2 ‖x‖ 2, we have Af0(x) = A and ∇2g(x) = λI for all x. We show that the convergence rate of both Newton-LESS and the Gaussian Newton Sketch can be improved in the presence of regularization, by replacing the dimension d with an effective dimension deff. This can be significantly smaller than d when the Hessian of f0 at the optimum exhibits rapid spectral decay or is approximately low-rank:
deff = tr ( ∇2f0(x∗)∇2f(x∗)−1 ) ≤ d.
Theorem 4. Assume that f0 and f are (a) self-concordant, or (b) have a Lipschitz continuous Hessian. Also, let ∇2f0(x∗) be positive definite and let ∇2g(x∗) be positive semidefinite, with H = ∇2f(x∗). There is a neighborhood U containing x∗ such that if x̃0 ∈ U , then Regularized Newton-LESS (2), with sketch size m ≥ Cdeff log(deffT/δ) and step size µt = 1− deffm , satisfies:(
Eδ ‖x̃T − x∗‖2H ‖x̃0 − x∗‖2H
)1/T ≤ deff
m · (1 + ) for = O ( 1√ deff ) .
Remark 5. The same guarantee holds for the Gaussian Newton Sketch. Unlike in Theorem 1, here we can only obtain an upper-bound on the local convergence rate, because the exact rate may depend on the starting point x̃0 (see Section 4). For regularized least squares, f(x) = 12‖Ax− b‖ 2 + λ2 ‖x‖ 2, the above convergence guarantee holds globally, i.e., U = Rd. Note that deff can be efficiently estimated using sketching-based trace estimators [ACW16, CEM+15].
Finally, our numerical results show that Newton-LESS can be implemented very efficiently on modern hardware platforms, improving on the optimization cost over not only dense Gaussian embeddings, but also state-of-the-art sketching methods such as the Subsampled Randomized Hadamard Transform, as well as other first-order and second-order methods. Moreover, we demonstrate that our theoretical predictions for the optimal sparsity level and convergence rate are extremely accurate in practice.

heading: 1.2 Related work
text: LEverage Score Sparsified (LESS) embeddings were proposed by [DLDM21] as a way of addressing the phenomenon of inversion bias, which arises in distributed second-order methods [WRKXM18, DM19, DBPM20, GGD+21]. Their results only establish the near-unbiasedness of Newton-LESS iterates (i.e., that E[x̃t+1] ≈ xt+1), but they did not provide any improved guarantees on the convergence rate. Also, their notion of LESS embeddings is much narrower than ours, and so it does not capture Regularized Newton-LESS or uniformly sparsified sketches (LESS-uniform).
Convergence analysis of the Newton Sketch [PW17, LP19, LLDP20] and other randomized secondorder methods [BCNN11, BCNW12, EM15, RKM19] has been extensively studied in the machine learning community, often using techniques from Randomized Numerical Linear Algebra (RandNLA) [DM16, DM21]. Some of the popular RandNLA methods include the Subsampled Randomized Hadamard Transform (SRHT, [AC09]) and several variants of sparse sketches, such as the CountSketch [CW17, MM13] and OSNAP [NN13, Coh16]. Also, row sampling based on Leverage Scores [DMM06, AM15] and Determinantal Point Processes [Der19, DCV19, DM21] has been used for sketching. Note that CountSketch and OSNAP sparse sketches differ from LESS embeddings in several ways, and in particular, they use a fixed number of non-zeros per column of the sketching matrix (as opposed to per row), so unlike LESS, their rows are not independent. While all of the mentioned methods, when used in conjunction with the Newton Sketch, exhibit similar per-iteration complexity as LESS embeddings (see Section 2), their existing convergence analysis is fundamentally limited: The best known rate is (C log d · dm )
T , which is worse than our result of ( dm )
T , by a factor of C log d, where C > 1 is a non-negligible constant that arises in the measure concentration analysis.
In the specific context of least squares regression where f(x) = 12‖Ax − b‖ 2, the Hessian A>A remains constant, and an alternative strategy is to keep the random sketch SA fixed at every iteration. Many efficient randomized iterative solvers are based on this precondition-and-solve approach [RT08, AMT10, MSM14]: form the sketch SA, compute an easy-to-invert square-root matrix H̃ 1 2 of A>S>SA and apply an iterative least squares solver to the preconditioned objective minz 1 2‖AH̃
− 12 z− b‖2, e.g., Chebyshev iterations or the preconditioned conjugate gradient. In contrast to the Newton Sketch, these methods do not naturally extend to more generic convex objectives for which the Hessian matrix changes at every iteration. Also, similarly as the Newton Sketch, their convergence guarantees are limited to (C log d · dm )
T when used in conjunction with fast sketching methods such as SRHT, OSNAP, or leverage score sampling.

heading: 2 Preliminaries
text: Notation. We let ‖v‖M = √
v>Mv. We define a ≈ b to mean |a− b| ≤ b, whereas a = b± means that |a− b| ≤ , and C denotes a large absolute constant. We use EE to denote expectation conditioned on E , and for a δ ∈ (0, 1), we use Eδ as a short-hand for: “There is an event E with probability at least 1− δ s.t. EE ...”. Let pd and psd mean positive definite and positive semidefinite. Random variable X is sub-Gaussian if Pr{|X| ≥ t} ≤ exp(−ct2) for all t ≥ 0 and some c = Ω(1). We next introduce some concepts related to LEverage Score Sparsified (LESS) embeddings. We start with the notion of statistical leverage scores [DMIMW12a], which are importance weights assigned to the rows of a matrix A ∈ Rn×d. The definition below for leverage scores is somewhat more general than standard definitions, because it allows for a regularized leverage score, where the regularization depends on a d× d matrix C. When C is a scaled identity matrix, this matches the definition of ridge leverage scores [AM15]. Definition 1 (Leverage scores). For given matrices A ∈ Rn×d and psd C ∈ Rd×d, we define the ith leverage score li(A,C) as the squared norm of the ith row of U = AH− 1 2 , where H = A>A+C is
assumed to be invertible. The effective dimension of A (given C) is defined as deff = ∑ i li(A,C) = tr(U>U), whereas the coherence of A (given C) is τ = ndeff maxi li(A,C) ∈ [1, n deff ].
Next, we define a class of sparsified sub-Gaussian sketching matrices which will be used in our results. This captures LESS embeddings, as well as other sketching matrices that are supported by the analysis. To that end, we define what we call a sparsifier, which is an n-dimensional random vector ξ that specifies the sparsification pattern for one row of the m× n sketching matrix S.
Definition 2 (LESS embeddings). Let t1, ..., ts be sampled i.i.d. from a distribution p = (p1, ..., pn). Then, the random vector ξ> = (√ b1 sp1 , ..., √ bn spn ) , where bi = ∑s j=1 1[tj=i], is a (p, s)-sparsifier. A (p, s)-sparsified sub-Gaussian sketch is a random matrix S that consists of i.i.d. row vectors distributed as c · (x ◦ ξ)>, where ◦ denotes the entry-wise product, x has i.i.d. mean zero, unit variance and sub-Gaussian entries, and c is some constant. For given matrices A ∈ Rn×d and psd C ∈ Rd×d, we focus on two variants of Leverage Score Sparsified embeddings:
1. LESS. We assume that pi ≈1/2 li(A,C)/deff, and let s ≈1/2 deff. For C = 0 and s = d, we recover the LESS embeddings proposed by [DLDM21].
2. LESS-uniform. We simply let pi = 1/n (denoted as p = unif). This avoids the preprocessing needed for approximating the li(A,C), but we may need larger s to recover the theory.
Computational cost. To implement LESS, we must first approximate (i.e., there is no need to compute exactly [DMIMW12a]) the leverage scores of A. This can be done in timeO(nnz(A) log n+ d3 log d) by using standard RandNLA techniques [DMIMW12a, CW17], where nnz(A) is the number of non-zero entries in A and it is bounded by nd. Since the prescribed sparsity for LESS satisfies s = O(d), the sketching cost is at most O(md2). Thus, the total cost of constructing the sketched Hessian with LESS is O(nnz(A) log n+md2), which up to logarithmic factors matches other sparse sketching methods such as leverage score sampling (when implemented with approximate leverage scores [DMIMW12a]), CountSketch, and OSNAP. In comparison, using the SRHT leads to O(nd logm+md2) complexity, since this method does not take advantage of data sparsity. Note that, in practice, the computational trade-offs between sketching methods are quite different, and significantly hardware-dependent (see Section 5). In particular, the cost of approximating the leverage scores in LESS embeddings can be entirely avoided by using LESS-uniform. Here, the total cost of sketching is O(mds), but the sparsity of the sketch that is needed for the theory depends on A and C. Yet, in Section 5, we show empirically that this approach works well even for s = d.

heading: 3 Equivalence between LESS and Gaussian Embeddings
text: In this section, we derive the basic quantities that determine the convergence properties of the Newton Sketch, namely, the first and second moments of the normalized sketched Hessian inverse. Our key technical contribution is a new analysis of the second moment for a wide class of sketching matrices that includes LESS embeddings and sub-Gaussian sketches.
Consider the Newton Sketch update as in (2), and let ∆̃t = x̃t − x∗. Denoting Ht = ∇2f(x̃t), gt = ∇f(x̃), and using pt = −µtH−1t gt to denote the exact Newton direction with step size µt, a simple calculation shows that:
‖∆̃t+1‖2Ht − ‖∆̃t‖ 2 Ht = 2∆̃ > t H 1 2 t Q̃H 1 2 t pt + p > t H 1 2 t Q̃ 2H 1 2 t pt, (3)
where Q̃ = H 1 2 t (Af0(x̃t) > S>t StAf0(x̃t) +∇2g(x̃t))−1H 1 2 t . From this, we have that the expected decrease in the optimization error is determined by the first two moments of the matrix Q̃, i.e., E[Q̃] and E[Q̃2]. In the unregularized case, i.e., g(x) = 0, these moments can be derived exactly for the Gaussian embedding. For instance if we let St be an m× n matrix with i.i.d. standard normal entries scaled by 1√
m−d−1 , then we obtain that:
E[Q̃] = I, E[Q̃2] = (m−1)(m−d−1)(m−d)(m−d−3) · I ≈ m m−d · I,
for = O(1/d). This choice of scaling for the Gaussian Newton Sketch ensures that each iterate x̃t+1 is an unbiased estimate of the corresponding exact Newton update with the same step size, i.e., that E[x̃t+1] = x̃t + pt. For most other sketching techniques, neither of the two moments is analytically tractable because of the bias coming from matrix inversion. Moreover, if we allow for regularization, e.g., g(x) = λ2 ‖x‖
2, then even the Gaussian embedding does not enjoy tractable formulas for the moments of Q̃. However, using ideas from asymptotic random matrix theory, [DLDM21] showed that in the unregularized case, the exact Gaussian formula for the first moment holds approximately for sub-Gaussian sketches and LESS embeddings: Eδ[Q̃] ≈ I. This implies near-unbiasedness of
the unregularized Newton-LESS iterates relative to the exact Newton step, but it is not sufficient to ensure any convergence guarantees.
In this work, we develop a general characterization of the first and second moments of Q̃ for a wide class of sketching matrices, both in the unregularized and in regularized settings. For the sake of generality, we will simplify the notation here, and analyze the first and second moment of Q = H 1 2 (A>S>SA + C)−1H 1 2 for some matrices A ∈ Rn×d and C ∈ Rd×d such that A>A + C = H. In the context of Newton Sketch (2), these quantities correspond to A = Af0(x̃t) and C = ∇2g(x̃t). Also, as a shorthand, we will define the normalized version of matrix A as U = AH− 1 2 . The following are the two structural conditions that need to be satisfied by a sketching matrix to enable our analysis.
The first condition is standard in the sketching literature. Essentially, it implies that the sketching matrix S produces a useful approximation of the Hessian with high probability (although this guarantee is still far too coarse by itself to obtain our results). Condition 1 (Property of random matrix S). Given U ∈ Rn×d, the m×n random matrix S satisfies ‖U>S>SU−U>U‖ ≤ η with probability 1− δ.
This property is known as the subspace embedding property. Subspace embeddings were first used by [DMM06], where they were used in a data-aware context to obtain relative-error approximations for `2 regression and low-rank matrix approximation [DMM08]. Subsequently, data-oblivious subspace embeddings were used by [Sar06] and popularized by [Woo14]. Both data-aware and dataoblivious subspace embeddings can be used to derive bounds for the accuracy of various algorithms [DM16, DM18].
For our analysis, it is important to assume that S has i.i.d. row vectors cs>i , where c is an appropriate scaling constant. The second condition is defined as a property of those row vectors, which makes them sufficiently similar to Gaussian vectors. This is a relaxation of the Restricted Bai-Silverstein condition, proposed by [DLDM21], which leads to significant improvements in the sparsity guarantee for LESS embeddings when the Newton Sketch is regularized. Condition 2 (Property of random vector s). Given U ∈ Rn×d, the n-dimensional random vector s satisfies Var [ s>UBU>s ] ≤ α · tr(UB2U>) for all p.s.d. matrices B and some α = O(1).
Given these two conditions, we are ready to derive precise non-asymptotic analytic expressions for the first two moments of the regularized sketched inverse matrix, which is the main technical contribution of this work (proof in Appendix A).
Theorem 6. Fix A and assume that C is psd. Define H = A>A + C and U = AH− 12 . Let S consist of m i.i.d. rows distributed as 1√
m−deff s>, where E[ss>] = In and deff = tr(U>U). Also,
let d̃eff = tr((U>U)2). Suppose that the matrix consisting of the first m/3 rows of S scaled by √
3 satisfies Condition 1 w.r.t. U, for η ≤ 1/2 and probability 1− δ/3, where δ ≤ 1/m3. Suppose also that s satisfies Condition 2 w.r.t. U. If m ≥ O(deff), then, conditioned on event E that holds with probability 1− δ, matrix Q = H 12 (A>S>SA + C)−1H 12 satisfies ‖Q− I‖ ≤ O(η) and:∥∥EE [Q]− I∥∥ ≤ O(√deffm ), ∥∥EE [Q2]− (I + deffm−d̃eff U>U)∥∥ ≤ O(√deffm ). Theorem 6 shows that, for a wide class of sketching matrices, we can approximately write E[Q] ≈ I and E[Q2] ≈ I+ deff
m−d̃eff U>U, with the error term scaling asO( √ deff m ). In the case of the first moment,
this is a relatively straightforward generalization of the unregularized formula for the Gaussian case. However, for the second moment this expression is considerably more complicated, including not one but two notions of effective dimension, d̃eff 6= deff. To put this in context, in the unregularized case, i.e., C = 0, we have deff = d̃eff = d and U>U = I, so we get I + deffm−d̃eff U
>U = mm−dI, which matches the second moment for the Gaussian sketch (up to lower order terms).
In the case of the first moment, the proof of Theorem 6 follows along the same lines as in [DLDM21], using a decomposition of E[Q] that is based on the Sherman-Morrison rank-one update of the inverse. This approach was originally inspired by the analysis of Stieltjes transforms that are used to establish the limiting spectral distribution in asymptotic random matrix theory (e.g., [BS10, CD11]), and applied to sketching by [DLLM20, DLDM21]. Our key contribution lies in deriving the bound for the second moment, which requires a substantially more elaborate decomposition of E[Q2].
In the following lemma, we establish that the assumptions of Theorem 6 are satisfied not only by Gaussian, but also sub-Gaussian and LESS embedding matrices (proof in Appendix C). Lemma 7. Fix A and assume that C is psd. Let S be a sketching matrix withm i.i.d. rows distributed as 1√
m−deff s>. Then, S satisfies Conditions 1 and 2 as long as one of the following holds:
1. Sub-Gaussian: s is an i.i.d. sub-Gaussian random vector and m ≥ C(deff + log(1/δ))/η2;
2. LESS: s is a (p, deff)-sparsified i.i.d. sub-Gaussian random vector with pi ≈1/2 li(A,C)/deff, and m ≥ Cdeff log(deff/δ)/η2;
3. LESS-uniform: s is a (unif, τdeff)-sparsified i.i.d. sub-Gaussian random vector, where τ is the coherence of A, i.e., ndeff maxi li(A,C), and m ≥ Cdeff log(deff/δ)/η 2.

heading: 4 Convergence Analysis for Newton-LESS
text: In this section, we demonstrate how our main technical results can be used to provide improved convergence guarantees for Newton-LESS (and, more broadly, any sketching methods that satisfy the conditions of Theorem 6). Here, we will focus on the more general regularized setting (2), where we can only show an upper bound on the convergence rate (Theorem 4). The unregularized result (Theorem 1) with matching upper/lower bounds follows similarly.
To start, we introduce the standard assumptions on the function f , which are needed to ensure strong local convergence guarantees for the classical Newton’s method [BV04]. Assumption 8. Function f : Rd → R has a Lipschitz continuous Hessian with constant L, i.e., ‖∇2f(x)−∇2f(x′)‖ ≤ L ‖x− x′‖ for all x,x′ ∈ Rd. Assumption 9. Function f : Rd → R is self-concordant, i.e., for all x,x′ ∈ Rd, the function φ(t) = f(x + tx′) satisfies: |φ′′′(t)| ≤ 2(φ′′(t))3/2.
Only one of those two assumptions needs to be satisfied for our analysis to go through, and the choice of the assumption only affects the size of the neighborhood around the optimum x∗ for which our local convergence guarantee is satisfied. To clarify this, below we give an expanded version of Theorem 4 (proof in Appendix B). Theorem 10 (Expanded Theorem 4). Let H0 = ∇2f0(x∗) be pd and C = ∇2g(x∗) be psd. Define deff = tr(H0H −1) and d̃eff = tr((H0H−1)2) for H = H0 + C. Assume one of the following:
1. f0 and f satisfy Assumption 8, and U = {x : ‖x − x∗‖H < √ deff m (λmin)
3/2/L}, where λmin is the smallest eigenvalue of H0;
2. f0 and f satisfy Assumption 9, and U = {x : ‖x− x∗‖H < √ deff m }.
Then, Newton Sketch (2) starting from x̃0 ∈ U , using any St from Lemma 7 (i.e., sub-Gaussian, LESS, or LESS-uniform) with δ replaced by δ/T , and step size µt = 1− deffm+deff−d̃eff , satisfies:(
Eδ ‖x̃T − x∗‖2H ‖x̃0 − x∗‖2H
)1/T ≤ deff · (1 + )
m+ deff − d̃eff ≤ deff m · (1 + ), for = O ( 1√ deff ) .
Note that, compared to Theorem 4, here we present a slightly sharper bound which uses both types of effective dimension, d̃eff ≤ deff, that are present in Theorem 6. The statement from Theorem 4 is recovered by replacing d̃eff with deff in the step size and in the bound. The key step in the proof of the result is the following lemma, which uses Theorem 6 to characterize the Newton Sketch iterate x̃t+1 in terms of the corresponding Newton iterate xt+1. Note that this result holds globally for arbitrary x̃t and without the smoothness assumptions on f . Recall that we let ∆̃t = x̃t − x∗ denote the error residual at step t. Lemma 11. Fix Ht = ∇2f(x̃t) and let x̃t+1 be the Newton Sketch iterate with St as in Lemma 7. If the exact Newton step xt+1 = x̃t − µtH−1t gt is a descent direction, i.e., ‖∆t+1‖Ht ≤ ‖∆̃t‖Ht where ∆t+1 = xt+1 − x∗, then
Eδ ‖∆̃t+1‖2Ht = ‖∆t+1‖ 2 Ht + deff(x̃t) m−d̃eff(x̃t) ‖xt+1 − x̃t‖2∇2f0(x̃t) ±O (√ deff m ) ‖∆̃t‖2Ht .
Importantly, the second term on the right-hand side uses the norm ‖ · ‖∇2f0(x̃t), which is different than the norm ‖ ·‖Ht used for the remaining terms. As a result, in the regularized setting, it is possible that the second term will be much smaller than the last term (the approximation error). This prevents us from obtaining a matching lower-bound for the convergence rate of Regularized Newton-LESS. On the other hand, when g(x) = 0, then f0 = f and we obtain matching upper/lower bounds.
The remainder of the proof of Theorem 10 is essentially a variant of the local convergence analysis of the Newton’s method. Here, note that typically, we would set the step size to µt = 1 and we would expect superlinear (specifically, quadratic) convergence rate. However, for Newton Sketch, convergence is a mixture of linear rate and superlinear rate, where the linear part is due to the approximation error in sketching the Hessian. Sufficiently close to the optimum, the linear rate will dominate, and so this is what we focus on in our local convergence analysis. The key novelty here is that, unlike prior work, we strive to describe the linear rate precisely, down to lower order terms. As a key step, we observe that the convergence of exact Newton with step size µt < 1, letting ∆t+1 = xt+1 − x∗, is given by:
‖∆t+1‖2Ht = (1− µt) 2‖∆̃t‖2Ht︸ ︷︷ ︸
linear rate
+µt(∆t+1 + (1− µt)∆̃t)>(Ht∆̃t − gt)︸ ︷︷ ︸ superlinear rate , (4)
where recall that ∆̃t = x̃t−x∗ corresponds to the previous iterate, and gt = ∇f(x̃t). Here, (1−µt)2 represents the linear convergence rate. The superlinear term vanishes near the optimum x∗, because of the presence of Ht∆̃t − gt, which (under the smoothness assumptions on f ) vanishes at the same rate as ‖∆̃t‖2Ht . Interestingly, with this precise analysis, the superlinear term does not have a quadratic rate, but rather a 3/2 rate. Entering (4) into the guarantee from Lemma 11, we obtain that the local rate of Newton Sketch can be expressed as (1− µt)2 + deffm−d̃effµ 2 t , and minimizing this expression over µt we obtain the desired quantities from Theorem 10. Finally, we note that while the norms and effective dimensions in the above exact calculations are stated with respect to the Hessian Ht at the current iterate x̃t, these can all be approximated by the corresponding quantities computed using the Hessian at the optimum x∗ (as in Theorem 10), relying on smoothness of f .

heading: 5 Numerical Experiments
text: We evaluated our theory on a range of different problems, and we have found that the more precise analysis that our theory provides describes well the convergence behavior for a range of optimization problems. In this section, we present numerical simulations illustrating this for regularized logistic regression and least squares regression, with different datasets ranging from medium to large scale: the CIFAR-10 dataset, the Musk dataset, and WESAD [SRD+18]. Data preprocessing and implementation details, as well as additional numerical results for least squares and regularized least squares, can be found in Appendix E.
We investigate first the effect of the row density of a LESS-uniform embedding on the Newton Sketch convergence rate and on the time for computing the sketch SA for least squares regression. In Figure 2, we report these two performance measures versus the row density. (This is the empirical analog for real data of Figure 1). Note that here we also consider a synthetic data matrix with highcoherence, which aims to be more challenging for a uniform sparsifier. Remarkably, our prescribed row density of d non-zeros per row offers an excellent empirical trade-off: except for the CIFAR-10
dataset, for which random row sampling performs equally well to Gaussian embeddings, we observe that one can drastically decrease the row density without significantly impairing the convergence rate.
Next, in Figure 3, we investigate a minimization task for the regularized logistic regression loss. Namely, given a data matrix A ∈ Rn×d with rows ai, a target vector b ∈ {±1}n and a regularization parameter λ > 0, the goal is to solve:
min x∈Rd
1
n n∑ i=1 log(1 + exp(−bia>i x)) + λ 2 ‖x‖22 .
For each dataset, we choose the value of λ among {10−j | j = 0, . . . , 8} that minimizes the error on a hold out validation set. For CIFAR-10 and Musk, we pick λ = 10−4. For WESAD, we pick λ = 10−5. We plot the error versus wall-clock time for the Newton Sketch with LESS-uniform, Gaussian, Subsampled Randomized Hadamard Transform (SRHT) and Random Row Sampling (RRS) matrices, and we compare it with two standard second-order optimization baselines: Newton’s method and BFGS. We also included three first-order baselines: Gradient Descent (GD), Accelerated GD (AGD), and Stochastic GD (SGD), observing much worse performance. Based on the top plots in Figure 3, we conclude that Newton-LESS (i.e., NS LESS uniform) offers significant time speed-ups over all baselines.
Finally, we compare wall-clock time on different hardwares (CPU versus GPU) for the Newton Sketch to reach a 10−6-accurate solution (see Appendix E for hardware details). From Figure 3 (bottom plots), we conclude the following: first, when switching from CPU to GPU for WESAD and CIFAR-10, Gaussian embeddings become more efficient than SRHT, despite a worse time complexity, by taking better advantage of the massively parallel architecture; second, Random Row Sampling performs better than either of them, despite having much weaker theoretical guarantees; and third, LESS-uniform is more efficient than all three other methods, on both CPU and GPU hardware platforms, observing a significant speed-up when switching to the parallel GPU architecture.

heading: 6 Conclusions
text: We showed that, when constructing randomized Hessian estimates for second-order optimization, we can get the best of both worlds: the efficiency of Sub-Sampling methods and the precision of Gaussian embeddings, by using sparse sketching matrices known as LEverage Score Sparsified (LESS) embeddings. Our algorithm, called Newton-LESS, enjoys both strong theoretical convergence guarantees and excellent empirical performance on a number of hardware platforms. An important future direction is to explain the surprising effectiveness of the simpler LESS-uniform method, particularly on high-coherence matrices, which goes beyond the predictions of our current theory.

heading: Acknowledgements
text: We would like to acknowledge DARPA, IARPA, NSF, and ONR via its BRC on RandNLA for providing partial support of this work. MP and JL were partially supported by the National Science Foundation under grants IIS-1838179, ECCS-2037304 and the Army Research Office. Our conclusions do not necessarily reflect the position or the policy of our sponsors, and no official endorsement should be inferred.

Review summaries:
The paper considers using sketching to accelerate Newton's optimization method. It is argued in the paper that there exist some cases where the Hessian of the function to be optimized can be represented as A^\top A for a tall and skinny matrix A of size N \times d. The authors argue that if the cost of computing the matrix A is small and the complexity is dominated by the matrix multiplication time for computing A^\top A then approximating the Hessian with A^\top S^\top S A, for some sparse sketch matrix S, can accelerate Newton's method. The paper analyzes the convergence of the sketched Newton's method for a recently proposed sketching method known as LESS and proves a linear convergence rate (in expectation). 

This paper analyzes a recently proposed sketching transform (LESS) that has been suggested for Newton Sketch. LESS (originally proposed in [DLDM20]) is aimed to be between full Gaussian sketch and row sampling. It achieves that by sampling the Gaussian sketch, with a sampling scheme based on the leverage scores. In the context of Newton Sketch, Gaussian sketch has been shown to work well, but using is mostly useless since the sketch operation is as expensive as computing the Hessian. Other sketches (e.g. SRHT) reduce per iteration cost, but increase the number of iterations. With LESS, the authors show we can avoid this tradeoff for certain types of functions.   The convergence rate the authors prove is linear (and not superlinear as one can hope from a second order method). But crucially, the rate is independent of the condition number.  The authors also consider the special case of ridge regression. 

This work proposes a Sketched Newton method based on sparsified random Gaussian matrix. The authors provide a precious convergence analysis by a careful characterization of the first and second moment of the inverse sketched Hessian. The authors also provide empirical studies to validate the efficiency of the proposed method.