Learning-Augmented Sketches for Hessians

Sketching is a dimensionality reduction technique where one compresses a matrix by often random linear combinations . A line of work has shown how to sketch the Hessian to speed up each iteration in a second order method , but such sketches usually depend only on the matrix at hand , and in a number of cases are even oblivious to the input matrix . One could instead hope to learn a distribution on sketching matrices that is optimized for the specific distribution of input matrices . We show how to design learned sketches for the Hessian in the context of second order methods , where we learn potentially different sketches for the different iterations of an optimization procedure . We show empirically that learned sketches , compared with their “ non-learned ” counterparts , improve the approximation accuracy for important problems , including LASSO , SVM , and matrix estimation with nuclear norm constraints . Several of our schemes can be proven to perform no worse than their unlearned counterparts . 1 INTRODUCTION . Large-scale optimization problems are abundant and solving them efficiently requires powerful tools to make the computation practical . This is especially true of second order methods which often are less practical than first order ones . Although second order methods may have many fewer iterations , each iteration could involve inverting a large Hessian , which is cubic time ; in contrast , first order methods such as stochastic gradient descent are linear time per iteration . In order to make second order methods faster in each iteration , a large body of work has looked at dimensionality reduction techniques , such as sampling , sketching , or approximating the Hessian by a low rank matrix . See , for example , ( Gower et al. , 2016 ; Xu et al. , 2016 ; Pilanci & Wainwright , 2016 ; 2017 ; Doikov & Richtárik , 2018 ; Gower et al. , 2018 ; Roosta-Khorasani & Mahoney , 2019 ; Gower et al. , 2019 ; Kylasa et al. , 2019 ; Xu et al. , 2020 ; Li et al. , 2020 ) . Our focus is on sketching techniques , which often consist of multiplying the Hessian by a random matrix chosen independently of the Hessian . Sketching has a long history in theoretical computer science ( see , e.g. , ( Woodruff , 2014 ) for a survey ) , and we describe such methods more below . A special case of sketching is sampling , which in practice is often uniform sampling , and hence oblivious to properties of the actual matrix . Other times the sampling is non-uniform , and based on squared norms of submatrices of the Hessian or on the so-called leverage scores of the Hessian . Our focus is on sketching techniques , and in particular , we follow the framework of ( Pilanci & Wainwright , 2016 ; 2017 ) which introduce the iterative Hessian sketch and Newton sketch , as well as the high accuracy refinement given in ( van den Brand et al. , 2020 ) . If one were to run Newton ’ s method to find a point where the gradient is zero , in each iteration one needs to solve an equation involving the current Hessian and gradient to find the update direction . When the Hessian can be decomposed as A > A for an n × d matrix A with n d , then sketching is particularly suitable . The iterative Hessian sketch was proposed in ( Pilanci & Wainwright , 2016 ) , where A is replaced with S · A , for a random matrix S which could be i.i.d . Gaussian or drawn from a more structured family of random matrices such as the Subsampled Randomized Hadamard Transforms or COUNTSKETCH matrices ; the latter was done in ( Cormode & Dickens , 2019 ) . The Newton sketch was proposed in ( Pilanci & Wainwright , 2017 ) , which extended sketching methods beyond constrained least-squares problems to any twice differentiable function subject to a closed convex constraint set . Using this sketch inside of interior point updates has led to much faster algorithms for an extensive body of convex optimization problems Pilanci & Wainwright ( 2017 ) . By instead using sketching as a preconditioner , an application of the work of ( van den Brand et al. , 2020 ) ( see Appendix E ) was able to improve the dependence on the accuracy parameter to logarithmic . In general , the idea behind sketching is the following . One chooses a random matrix S , drawn from a certain family of random matrices , and computes SA . IfA is tall-and-skinny , then S is short-and-fat , and thus SA is a small , roughly square matrix . Moreover , SA preserves important properties of A . One typically desired property is that S is a subspace embedding , meaning that simultaneously for all x , one has ‖SAx‖2 = ( 1± ) ‖Ax‖2 . An observation exploited in ( Cormode & Dickens , 2019 ) , building off of the COUNT-SKETCH random matrices S introduced in randomized linear algebra in ( Clarkson & Woodruff , 2017 ) , is that if S contains O ( 1 ) non-zero entries per column , then SA can be computed in O ( nnz ( A ) ) time , where nnz ( A ) denotes the number of nonzeros in A . This is sometimes referred to as input sparsity running time . Each iteration of a second order method often involves solving an equation of the form A > Ax = A > b , where A > A is the Hessian and b is the gradient . For a number of problems , one has access to a matrixA ∈ Rn×d with n d , which is also an assumption made in Pilanci & Wainwright ( 2017 ) . Therefore , the solution x is the minimizer to a constrained least squares regression problem : min x∈C 1 2 ‖Ax− b‖22 , ( 1 ) where C is a convex constraint set in Rd . For the unconstrained case ( C = Rd ) , various classical sketches that attain the subspace embedding property can provably yield high-accuracy approximate solutions ( see , e.g. , ( Sarlos , 2006 ; Nelson & Nguyên , 2013 ; Cohen , 2016 ; Clarkson & Woodruff , 2017 ) ) ; for the general constrained case , the Iterative Hessian Sketch ( IHS ) was proposed by Pilanci & Wainwright ( 2016 ) as an effective approach and Cormode & Dickens ( 2019 ) employed sparse sketches to achieve input-sparsity running time for IHS . All sketches used in these results are dataoblivious random sketches . Learned Sketching . In the last few years , an exciting new notion of learned sketching has emerged . Here the idea is that one often sees independent samples of matrices A from a distribution D , and can train a model to learn the entries in a sketching matrix S on these samples . When given a future sample B , also drawn from D , the learned sketching matrix S will be such that S · B is a much more accurate compression of B than if S had the same number of rows and were instead drawn without knowledge of D. Moreover , the learned sketch S is often sparse , therefore allowing S · B to be applied very quickly . For large datasets B this is particularly important , and distinguishes this approach from other transfer learning approaches , e.g. , ( Andrychowicz et al. , 2016 ) , which can be considerably slower in this context . Learned sketches were first used in the data stream context for finding frequent items ( Hsu et al. , 2019 ) and have subsequently been applied to a number of other problems on large data . For example , Indyk et al . ( 2019 ) showed that learned sketches yield significantly small errors for low rank approximation . In ( Dong et al. , 2020 ) , significant improvements to nearest neighbor search were obtained via learned sketches . More recently , Liu et al . ( 2020 ) extended learned sketches to several problems in numerical linear algebra , including least-squares and robust regression , as well as k-means clustering . Despite the number of problems that learned sketches have been applied to , they have not been applied to convex optimization in general . Given that such methods often require solving a large overdetermined least squares problem in each iteration , it is hopeful that one can improve each iteration using learned sketches . However , a number of natural questions arise : ( 1 ) how should we learn the sketch ? ( 2 ) should we apply the same learned sketch in each iteration , or learn it in the next iteration by training on a data set involving previously learned sketches from prior iterations ? Our Contributions . In this work we answer the above questions and derive the first learned sketches for a wide number of problems in convex optimization . Namely , we apply learned sketches to constrained least-squares problems , including LASSO , support vector machines ( SVM ) , and matrix regression with nuclear norm constraints . We show empirically that learned sketches demonstrate superior accuracy over random oblivious sketches for each of these problems . Specifically , compared with three classical sketches ( Gaussian , COUNT-SKETCH and Sparse Johnson-Lindenstrauss Transforms ; see definitions in Section 2 ) , the learned sketches in each of the first few iterations • improve the LASSO error f ( x ) − f ( x∗ ) by 80 % to 87 % in two real-world datasets , where f ( x ) = 12 ‖Ax− b‖ 2 2 + ‖x‖1 ; • improve the dual SVM error f ( x ) − f ( x∗ ) by 10–30 % for a synthetic and a real-world dataset , as well as by 30 % –40 % for another real-world dataset , where f ( x ) = ‖Bx‖22 ; • improve the matrix estimation error f ( X ) − f ( X∗ ) by at least 30 % for a synthetic dataset and at least 95 % for a real-world data set , where f ( X ) = ‖AX −B‖2F . Therefore , the learned sketches attain a smaller error within the same number of iterations , and in fact , within the same limit on the maximum runtime , since our sketches are extremely sparse ( see below ) . We also study the general framework of convex optimization in ( van den Brand et al. , 2020 ) , and show that also for sketching-based preconditioning , learned sketches demonstrate considerable advantages . More precisely , by using a learned sketch with the same number of rows as an oblivious sketch , we are able to obtain a much better preconditioner with the same overall running time . All of our learned sketches S are extremely sparse , meaning that they contain a single non-zero entry per column . Following the previous work of ( Indyk et al. , 2019 ) , we choose the position of the nonzero entry in each column to be uniformly random , while the value of the nonzero entry is learned . This already demonstrates a significant advantage over non-learned sketches , and has a fast training time . Importantly , because of such sparsity , our sketches can be applied in input sparsity time given a new optimization problem . We also provide several theoretical results , showing how to algorithmically use learned sketches in conjunction with random sketches so as to do no worse than random sketches .

Summary: The paper proposed a learned variant of the well-known iterative Hessian sketch (IHS) method of Pilanci and Wainwright, for efficiently solving least-squares regression. The proposed method is essentially a learned variant of the count-sketch, where the positions of the non-zero entries are random while the value is learned. While getting a learned variant for IHS is an interesting direction, the current theoretical contribution of this paper is only incremental, and most importantly, the reviewer is unconvinced for the practicality of the current approach.