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