Learning-Augmented Sketches for Hessians 

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 & Richtá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 leverage scores of the Hessian . Our focus is on sketching techniques , and in particular , we consider the framework of ( Pilanci & Wainwright , 2016 ; 2017 ) which introduces the iterative Hessian sketch and the 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 COUNT-SKETCH matrices ; the latter was done in ( Cormode & Dickens , 2019 ) . The Newton sketch was proposed by 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-thin , 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 ‖SAx‖2 = ( 1± ) ‖Ax‖2 for all x simultaneously . 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 a single non-zero entry per column , then SA can be computed in O ( nnz ( A ) ) time , where nnz ( A ) denotes the number of nonzeros in A . This is also referred to as input-sparsity running time . Each iteration of a second order method often involves solving an equation of the formA > Ax = A > b , where A > A is the Hessian and b is the gradient . For a number of problems , one has access to a matrix A ∈ 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 & Nguyê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 data-oblivious 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 smaller errors for low rank approximation . Dong et al . ( 2020 ) made significant improvements to nearest neighbor search using learned sketches . More recently , Liu et al . ( 2020 ) extended learned sketches to several problems in numerical linear algebra , including least-squares 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 develop the first framework of learned sketching that applies to a wide number of problems in convex optimization . Namely , we apply learned sketches to constrained least-squares problems , including LASSO and matrix regression with nuclear norm constraints . We show empirically that learned sketches demonstrate superior accuracy over classical oblivious random sketches for each of these problems . All of our learned sketches S are extremely sparse , meaning that they contain a single non-zero entry per column and that they can be applied in input-sparsity time . For such sketches , there are two things to learn : the position of the non-zero entry in each column and the value of the non-zero entry . 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 ( the value is no longer limited to −1 and 1 ) . Here we consider a new learning objective , that is , we optimize the subspace embedding property of the sketching matrix instead of optimizing the error in the objective function of the optimization problem we are trying to solve . This demonstrates a significant advantage over non-learned sketches , and has a fast training time . Our experiments show that the convergence rate is reduced by 44 % over the nonlearned COUNT-SKETCH ( a classical extremely sparse sketch ) for the LASSO problem on a real-world dataset . Recall that a smaller convergence rate means a faster convergence . We prove theoretically that S can take fewer rows , with optimized positions of nonzero entries , when the input matrix A has a small number of rows of heavy leverage score . More specifically , COUNT-SKETCH takes O ( d2/ ( δ 2 ) ) rows with failure probability δ , while our S requires only O ( ( dpolylog ( 1/ ) + log ( 1/δ ) ) / 2 ) rows if A has at most dpolylog ( 1/ ) / 2 rows of leverage score at least /d . This is a quadratic improvement in d and an exponential improvement in δ . Applying S to A runs in input-sparsity time and the resulting SA may remain sparse if A is sparse . In practice , it is not necessary to calculate the leverage scores . Instead , we show in our experiments that the indices of the rows of heavy leverage score can be learned and the induced S achieves a comparable accuracy for the abovementioned LASSO problem to classical dense sketches such as Gaussian matrices . Combining both aspects , the value of the nonzero entry and the indices of the rows of heavy leverage score , we obtain even better learned sketches . For the same LASSO problem , we show empirically that such learned sketches reduce the convergence rate by a larger 79.9 % to 84.6 % over non-learned sketches . 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 . 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 . 2 PRELIMINARIES Algorithm 1 LEARN-SKETCH : Gradient descent algorithm for learning the sketch values Require : Atrain = { Ai } Ni=1 ( Ai ∈ Rn×d ) , learn- ing rate α 1 : Randomly initialize p , v for a COUNT- SKETCH-type sketch as described in the text 2 : for t = 0 to step do 3 : Form S using p , v 4 : Sample batch Abatch from Atrain 5 : v ← v − α∂L ( S , Abatch ) ∂v Notation . We denote by Sn−1 the unit sphere in the n-dimensional Euclidean space Rn . For a matrix A ∈ Rm×n we denote by ‖A‖op its operator norm , which is defined as ‖A‖op = supx∈Sn−1 ‖Ax‖2 . We also denote by σmax ( A ) and σmin ( A ) the largest and smallest singular values of A , respectively , and by colsp ( A ) the column space of A . The condition number of A is defined to be κ ( A ) = σmax ( A ) /σmin ( A ) . Leverage Scores . We only consider matrices of full column rank1 . Suppose thatA ∈ Rm×n ( m ≥ n ) has full column rank . It has m leverage scores , denoted by τ1 ( A ) , . . . , τm ( A ) , which are defined as τi ( A ) = ‖e > i A ( A > A ) −1A > ‖22 , where { e1 , . . . , em } is the canonical basis of Rm . Equivalently , letting A = UΣV > be the singular value decomposition of A , where U ∈ Rm×n , Σ , V ∈ Rn×n , we can also write τi ( A ) = ‖e > i UU > ‖22 = ‖e > i U‖22 , which is the squared ` 2 norm of the i-th row of U . Classical Sketches . Below we review several classical sketches that have been used for solving optimization problems . • Gaussian sketch : S = 1√ m G , where G ∈ Rm×n with i.i.d . N ( 0 , 1 ) entries . • COUNT-SKETCH : Each column of S has only a single non-zero entry . The position of the non-zero entry is chosen uniformly over the m entries in the column and the value of the entry is either +1 or −1 , each with probability 1/2 . Further , the columns are chosen independently . • Sparse Johnson-Lindenstrauss Transform ( SJLT ) : S is the vertical concatenation of s independent COUNT-SKETCH matrices , each of dimension m/s× n. 1This can be assumed w.l.o.g . by adding artbirarily small random noise to the input , or one can first quickly use sketching to find a subset of columns of maximum rank , and replace the inut with that subset of columns . COUNT-SKETCH-type Sketch . A COUNT-SKETCH-type sketch is characterized by a tuple ( m , n , p , v ) , where m , n are positive integers and p , v are n-dimensional real vectors , defined as follows . The sketching matrix S has dimensions m × n and Spi , i = vi for all 1 ≤ i ≤ n , while all the other entries of S are 0 . When m and n are clear from context , we may characterize such a sketching matrix by ( p , v ) only . Subspace Embeddings . For a matrix A ∈ Rn×d , we say a matrix S ∈ Rm×n is a ( 1± ) -subspace embedding for the column span of A if ( 1− ) ‖Ax‖2 ≤ ‖SAx‖2 ≤ ( 1 + ) ‖Ax‖2 for all x ∈ Rd . The classical sketches above , with appropriate parameters , are all subspace embedding matrices with probability at least 1− δ ; our focus is on COUNT-SKETCH which can be applied in input sparsity running time . We summarize the parameters needed for a subspace embedding below : • Gaussian sketch : m = O ( ( d + log ( 1/δ ) ) / 2 ) . It is a dense matrix and computing SA costs O ( m · nnz ( A ) ) = O ( nnz ( A ) ( d+ log ( 1/δ ) ) / 2 ) time . • COUNT-SKETCH : m = O ( d2/ ( δ 2 ) ) ( Clarkson & Woodruff , 2017 ) . Though the number of rows is quadratic in d/ , the matrix S is sparse and computing SA takes only O ( nnz ( A ) ) time . • SJLT : m = O ( d log ( dδ ) / 2 ) and has s = O ( log ( dδ ) / ) non-zeros per column ( Nelson & Nguyên , 2013 ; Cohen , 2016 ) . Computing SA takes O ( snnz ( A ) ) = O ( nnz ( A ) log ( dδ ) / ) time . Iterative Hessian Sketch . The Iterative Hessian Sketching ( IHS ) method ( Pilanci & Wainwright , 2016 ) solves the constrained least-squares problem ( 1 ) by iteratively performing the update xt+1 = arg min x∈C { 1 2 ‖St+1A ( x− xt ) ‖22 − 〈A > ( b−Axt ) , x− xt〉 } , ( 2 ) where St+1 is a sketching matrix . It is not difficult to see that for the unsketched version ( St+1 is the identity matrix ) of the minimization above , the optimal solution xt+1 coincides with the optimal solution to the constrained least squares problem ( 1 ) . The IHS approximates the Hessian A > A by a sketched version ( St+1A ) > ( St+1A ) to improve runtime , as St+1A typically has very few rows . Unconstrained Convex Optimization . Consider an unconstrained convex optimization problem minx f ( x ) , where f is smooth and strongly convex , and its Hessian ∇2f is Lipschitz continuous . This problem can be solved by Newton ’ s method , which iteratively performs the update xt+1 = xt − arg min z ∥∥∥ ( ∇2f ( xt ) 1/2 ) > ( ∇2f ( xt ) 1/2 ) z −∇f ( xt ) ∥∥∥ 2 , ( 3 ) provided it is given a good initial point x0 . In each step , it requires solving a regression problem of the form minz ∥∥A > Az − y∥∥ 2 , which , with access to A , can be solved with a fast regression solver in ( van den Brand et al. , 2020 ) . The regression solver first computes a preconditioner R via a QR decomposition such that SAR has orthonormal columns , where S is a sketching matrix , then solves ẑ = arg minz′ ∥∥ ( AR ) > ( AR ) z′ − y∥∥ 2 by gradient descent and returns Rẑ in the end . Here , the point of sketching is that the QR decomposition of SA can be computed much more efficiently than the QR decomposition of A , since S has only a small number of rows . Learning a Sketch . We use the same learning algorithm in ( Liu et al. , 2020 ) , given in Algorithm 1 . The algorithm aims to minimize the mean loss function L ( S , A ) = 1N ∑N i=1 L ( S , Ai ) , where S is the learned sketch , L ( S , A ) is the loss function of S applied to a data matrix A , and A = { A1 , . . . , AN } is a ( random ) subset of training data .

Summary: In this paper, the authors extend a line of work focused on sketching the Hessian for convex problems to help accelerate second order optimization methods. In particular, they present an algorithm for learning weights of a sketching matrix (with one non-zero entry per columns chosen uniformly) by using gradient descent; the Hessians used for training are treated as draws from a distribution of Hessians. They also discuss how leverage scores can be used to improve convergence rates by ensuring that "heavy" rows are sampled with probability 1.   The authors show empirically that the learned sketches improve convergence rates and reduce the number of iterations for several problems. They also provide theoretical results stating a reduction in rows required in the sketch when heavy leverage scores rows are known and error/time complexity bounds for the Hessian sketch/regression problems.