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: This paper considers how learned CountSketches can be used in a variety of optimization tasks. The authors propose a method for predicting rows with high leverage score based on the training data. These rows are then sampled deterministically alongside either a standard CountSketch, or the learned CountSketch introduced by Indyk et al. (2019). Some novel theory is presented for this new sketch.