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

1 Introduction . 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 . 1.1 Main results . 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 .

Summary: 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).