L EARNING -A UGMENTED S KETCHES FOR H ESSIANS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 speciﬁc 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.
I NTRODUCTION
Large-scale optimization problems are abundant and solving them efﬁciently requires powerful tools to make the computation practical.
This is especially true of second order methods which often are less practical than ﬁrst 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, ﬁrst 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 ´ arik, 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 reﬁnement given in (van den Brand et al., 2020).
If one were to run Newton’s method to ﬁnd a point where the gradient is zero, in each iteration one needs to solve an equation involving the current Hessian and gradient to ﬁnd the update direction.
When the Hessian can be decomposed as A (cid:62) A for an n × d matrix A with n (cid:29) 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 C OUNTS KETCH 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 (cid:15) 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 .
If A 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 (cid:107) SAx (cid:107)
= (1 ± (cid:15) )
(cid:107) Ax (cid:107) 2 .
An observation exploited in (Cormode & Dickens, 2019), building off of the C OUNT -S KETCH 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 (cid:62)
Ax =
A (cid:62) b , where A (cid:62) A is the Hessian and b is the gradient.
For a number of problems, one has access to a matrix A ∈ R n × d with n (cid:29) 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:
[Equation]
where C is a convex constraint set in R d .
For the unconstrained case ( C = R d ), various classical sketches that attain the subspace embedding property can provably yield high-accuracy approximate solutions (see, e.g., (Sarlos, 2006; Nelson & Nguy ˆ en, 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 ﬁrst used in the data stream context for ﬁnding 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 signiﬁcantly small errors for low rank approximation.
In (Dong et al., 2020), signiﬁcant improvements to nearest neighbor search were obtained via learned sketches.
More recently, Liu et al.
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 ﬁrst 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.
Speciﬁcally, compared with three classical sketches (Gaussian, C OUNT -S KETCH and Sparse Johnson-Lindenstrauss Transforms; see deﬁnitions in Section 2), the learned sketches in each of the ﬁrst few iterations
• improve the LASSO error f ( x ) − f ( x ∗ ) by 80% to 87% in two real-world datasets, where f ( x ) = 12 (cid:107)
Ax − b (cid:107) 22 + (cid:107) x (cid:107) 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 )
= (cid:107) Bx (cid:107) 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 ) = (cid:107) AX − B (cid:107) 2 F .
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 signiﬁcant 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.
P RELIMINARIES
Classical Sketches
• Gaussian sketch: S = 1 √
m G , where G is an m × n Gaussian random matrix.
• C OUNT -S KETCH :
Each column of S has only a single non-zero entry.
The position of the nonzero 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 C OUNT -S KETCH matrices, each of dimension m/s × n .
C OUNT -S KETCH -type Sketch.
A C OUNT -S KETCH -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, deﬁned as follows.
The sketching matrix S has dimensions m × n and S p
i ,i = v i 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 ∈ R n × d , we say a matrix S ∈ R m × n is
a (1 ± (cid:15) ) -subspace embedding for the column span of A if (1 − (cid:15) )
(cid:107) Ax (cid:107) 2 ≤ (cid:107) SAx (cid:107) 2 ≤ (1 + (cid:15) )
(cid:107) Ax (cid:107) 2 for all x ∈ R d .
The classical sketches above, with appropriate parameters, are all subspace embedding matrices with at least a constant probability; our focus is on C OUNT -S KETCH which can be applied in input sparsity running time.
We summarize the parameters needed for a subspace embedding below:
• Gaussian sketch: m =
O ( d/(cid:15) 2 ) .
It is a dense matrix and computing SA costs O ( m · nnz( A ))
=
O (nnz( A ) d/(cid:15) 2 ) time.
• C OUNT -S KETCH : m =
O ( d 2 /(cid:15) 2 )
(Clarkson & Woodruff, 2017).
Although the number of rows is quadratic in d/(cid:15) , the sketch matrix S is sparse and computing SA takes only O (nnz( A )) time.
• SJLT: m =
O ( d/(cid:15) 2 ) and has s = O (1 /(cid:15) ) non-zeros per column (Nelson & Nguyˆen, 2013; Cohen, 2016).
Computing SA takes O ( s nnz( A ))
=
O (nnz( A ) /(cid:15) ) 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
[Equation]
where S t +1 is a sketching matrix.
It is not difﬁcult to see that for the unsketched version ( S t +1 is the identity matrix) of the minimization above, the optimal solution x t +1 coincides with the optimal solution to the constrained least square problem (1).
The IHS approximates the Hessian A (cid:62) A by a sketched version ( S t +1 A ) (cid:62) ( S t +1 A ) to improve runtime, as S t +1 A typically has very few rows.
Unconstrained Convex Optimization.
Consider an unconstrained convex optimization problem min x f ( x ) , where f is smooth and strongly convex, and its Hessian ∇ 2 f is Lipschitz continuous.
This problem can be solved by Newton’s method, which iteratively performs the update
[Equation]
provided it is given a good initial point x 0 .
In each step, it requires solving a regression problem of the form min z (cid:13)(cid:13)
A (cid:62) Az − y
(cid:13)(cid:13) 2 , which, with access to A , can be solved with a fast regression solver in (van den Brand et al., 2020).
The regression solver ﬁrst computes a preconditioner R via a QR decomposition such that SAR has orthonormal columns, where S is a sketching matrix, then solves (cid:98) z =
arg min
z (cid:48) (cid:13)(cid:13) ( AR ) (cid:62) ( AR ) z
(cid:48) − y (cid:13)(cid:13) 2 by gradient descent and returns R (cid:98) z in the end.
Here, the point of sketching is that the QR decomposition of SA can be computed much more efﬁciently than the QR decomposition of A since S has only a small number of rows.
Require:
A train =
{ A 1 , ..., A N } ( A i ∈ R n × d ), learning rate α 1: Randomly initialize p, v for a Count-Sketchtype sketch 2: for t = 0 to steps do 3:
Form S using p, v 4:
Sample batch A batch from A train 5: v ← v − α ∂ L ( S, A batch ) ∂v 6: end for
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 )
= 1 N (cid:80) Ni
=1 L ( S, A i ) , where S is the learned sketch, L ( S, A ) is the loss function of S applied to a data matrix A , and
A =
{ A 1 , . . .
, A N } is a (random) subset of training data.
H ESSIAN S KETCH
Algorithm 1: S 1 ← learned sketch, S 2 ← random sketch 2: ( (cid:98) Z
i, 1 , (cid:98) Z i, 2 ) ← E STIMATE ( S i , A ) , i = 1 ,3: if (cid:98) Z
, 2 / (cid:98) Z 1 , 1 < (cid:98)
Z 2 , 2 / (cid:98) Z 2 , 1 then 4: (cid:98) x ← solution of (3) with S = S5: else 6: (cid:98) x ← solution of (3) with S = S7: end if 8: return (cid:98) x 9: function E STIMATE ( S, A ) 10:
T ← sparse (1 ± η ) -subspace embedding matrix for d -dimensional subspaces 11: ( Q, R ) ← QR ( TA ) 12: (cid:98)
Z 1 ← σ min ( SAR − 1 ) 13: (cid:98) Z 2 ← (1 ± η ) -approximation to (cid:13)(cid:13) ( SAR − 1 ) (cid:62) ( SAR − 1 ) − I (cid:13)(cid:13) op 14: return ( (cid:98) Z 1 , (cid:98) Z 2 ) 15: end function
In this section, we consider the minimization problem
[Equation]
which is used as a subroutine for the IHS (cf.
We present an algorithm with the learned sketch in Algorithm 2.
To analyze its performance, we let R be the column space of
A ∈ R m × n and deﬁne the following quantities (corresponding exactly to the unconstrained case in Pilanci & Wainwright (2016))
[Equation]
where S n − 1 denotes the Euclidean unit sphere in R n .
The following is the estimation guarantee of (cid:98) Zand (cid:98) Z 2 .
The proof is postponed to Appendix A.
Lemma 3.1.
Suppose that η ∈ (0 , 1 3 ) is a small constant, A is of full rank and S has O ( d 2 ) rows.
The function E STIMATE ( S, A ) returns in O ((nnz( A ) log(1 /η )
+poly( d/η )) time (cid:98) Z 1 , (cid:98) Z 2 which with probability at least 0 .
99 satisfy that Z 1 ( S ) 1+ η ≤ (cid:98) Z 1 ≤ Z 1 ( S ) 1 − η and Z 2 ( S ) ([-S )+
-]
[Table]
Note for a matrix A , (cid:107) A (cid:107) op =
{+sup x (cid:54) =
+}
(cid:107) Ax (cid:107) 2 (cid:107) x (cid:107) 2 is its operator norm.
Similar to (Pilanci & Wainwright, 2016, Proposition 1), we have the following guarantee.
The proof is postponed to Appendix B.
Theorem 3.2.
Let η ∈ (0 , 13 ) be a small constant.
Suppose that A is of full rank and S 1 and S 2 are both C OUNT -S KETCH -type sketches with O ( d 2 ) rows.
Algorithm 2 returns a solution (cid:98) x which, with probability at least 0 . 98 , satisﬁes that (cid:107) A ( (cid:98) x − x ∗ ) (cid:107) 2 ≤ (1 + η ) 4 (cid:16) min (cid:110) (cid:98) Z 1 , 2 (cid:98) Z 1 , 1 , (cid:98) Z 2 , 2 (cid:98) Z 2 , 1 (cid:111) + 4 η (cid:17) (cid:107)
Ax ∗ (cid:107)in O (nnz( A ) log( 1 η ) + poly( dη )) time, where x ∗ = arg min x ∈C (cid:107)
Ax − b (cid:107) 2 is the least-squares solution.
H ESSIAN R EGRESSION
1: S S 2: ( Q i , R i ) ← QR ( S i A ) , i = 1 ,3: ( σ i , σ (cid:48) i ) ← E IG ( AR − 1 i ) , i = 1 ,(cid:46) E IG ( B ) returns the max and min singular values of B 4: if σ 1 /σ (cid:48)
1 < σ 2 /σ (cid:48) 2 then 5:
P ← R − 1 1 , η ← 1 / ( σ 21 + ( σ (cid:48) 1 ) 2 ) 6: else 7:
P ← R − 1 2 , η ← 1 / ( σ 22 + ( σ (cid:48) 2 ) 2 ) 8: end if 9: z 0 ←10: while (cid:13)(cid:13) A (cid:62) APz t − y (cid:13)(cid:13) 2 ≥ (cid:15) (cid:107) y
(cid:107) 2 do 11:
z t +1 ← z t − η
( P (cid:62) A (cid:62) AP )
( P (cid:62) A (cid:62) APz t − P (cid:62) y ) 12: end while 13: return Pz t
In this section, we consider the minimization problem
[Equation]
which is used as a subroutine for the unconstrained convex optimization problem min x f ( x ) with A (cid:62) A being the Hessian matrix ∇ 2 f ( x ) (See Section 2).
Here A ∈ R n × d , y ∈ R d , and we have access to A .
We incorporate a learned sketch into the fast regression solver in (van den Brand et al., 2020) and present the algorithm in
Algorithm 3.
Here the subroutine E IG ( B ) applies a (1 + η ) -subspace embedding sketch T to B for some small constant η and returns the maximum and the minimum singular values of TB .
Since B admits the form of AR , the sketched matrix TB can be calculated as ( TA ) R and thus can be computed in
O (nnz( A ) + poly( d )) time if T is a C OUNT -S
KETCH matrix of O ( d 2 ) rows.
The extreme singular values of TB can be found by SVD or the Lanczos’s algorithm.
Similar to Lemma 4.2 in (van den Brand et al., 2020), we have the following guarantee of Algorithm 3.
The proof parallels the proof in (van den Brand et al., 2020) and is postponed to Appendix C.
Theorem 4.1.
Suppose that S 1 and S 2 are both C OUNT -S KETCH -type sketches with O ( d 2 ) rows.
Algorithm 3 returns a solution x (cid:48) such that (cid:13)(cid:13)
A (cid:62) Ax (cid:48) − [Equation] ≤ [Equation]
(cid:107) y (cid:107) 2 in O (nnz( A ))
+ (cid:101) O ( nd · (min { σ 1 /σ (cid:48) 1 , σ 2 /σ (cid:48) 2 } ) 2 · log( κ ( A ) /(cid:15) )
+ poly( d )) time.
Remark 4.2.
In Algorithm 3, S 2 can be chosen to be a subspace embedding matrix for ddimensional subspaces, in which case, AR − 1 2 has condition number close to 1 (see, e.g., (Woodruff, 2014, p38)) and the full algorithm would run faster than the trivial O ( nd 2 ) -time solver to (4).
Remark 4.3. For the original unconstrained convex optimization problem min x f ( x ) , one can run the entire optimization procedure with learned sketches versus the entire optimization procedure with random sketches, compare the objective values at the end, and choose the better of the two.
For least-squares f ( x ) = 12 (cid:107)
Ax − b (cid:107) 22 , the value of f ( x ) can be approximated efﬁciently by a sparse subspace embedding matrix in O (nnz( A )
+ nnz( b ) + poly( d )) time.
IHS E XPERIMENTS
Training.
We learn the sketching matrix in each iteration (also called a round) separately.
Recall that in the ( t +1)
-st round of IHS, the optimization problem we need to solve (2) depends on x t .
An issue is that we do not know x t , which is needed to generate the training data for the ( t +1)
-st round.
Our solution is to use the sketch matrix in the previous round to obtain x t by solving (2), and then use it to generate the training data in the next round.
That is, in the ﬁrst round, we train the sketch matrix S 1 to solve the problem x 1 = arg min x ∈C 12 (cid:107) S 1 Ax (cid:107) 22 −(cid:104)
A (cid:62) b, x (cid:105) and use x 1 to generate the training data for the optimization problem for x 2 , and so on.
The loss function we use is the unsketched objective function in the ( t +1) -st iteration, i.e., L ( S t +1 , A )
= 12 (cid:107)
A ( x t +1 − x t ) (cid:107) 22 −(cid:104)
A T ( b − Ax t ) , x t +1 − x t (cid:105) , where x t +1 is the solution to (2) and thus depends on S t +1 .
Comparison.
We compare the learned sketch against three classical sketches: Gaussian, C OUNTS KETCH , and SJLT (see Section 2) in all experiments.
The quantity we compare is a certain error, deﬁned individually for each problem, in each round of the iteration of the IHS or as a function of the runtime of the algorithm.
All of our experiments are conducted on a laptop with 1.90GHz CPU and 16GB RAM.
5. LASSO
We deﬁne an instance of LASSO regression as
[Equation] [Equation]
) where λ is a parameter.
We use two real-world datasets:
• CO emission 1 : the dataset contains 9 sensor measures aggregated over one hour (by means of average or sum), which can help to predict the CO emission.
We divide the raw data into 120 ( A i , b i ) such that A i ∈ R 300 × 9 , b i ∈ R 300 × 1 .
The data in each matrix is sorted in chronological order.
| ( A, b ) train | = 96 , | ( A, b ) test | = 24 .
• Greenhouse gas 2 : time series of measured greenhouse gas concentrations in the California atmosphere.
Each ( A, b ) corresponds to a different measurement location
A i ∈ R 327 × 14 , b i ∈ R 327 × 1 , and | ( A, b ) train | = 400 , | ( A, b ) test | = 100 .
(This dataset was also used in (Liu et al., 2020).)
We choose λ = 1 in the LASSO regression (5).
We choose m = 5 d and m = 3 d for the CO emission dataset and m = 6 d and m = 3 . 5 d for the greenhouse gas dataset.
We consider the error ( 12 (cid:107)
Ax − b (cid:107) 22 + (cid:107) x (cid:107) 1 ) − ( 12 (cid:107)
Ax ∗ − b (cid:107)
+ (cid:107) x ∗ (cid:107) 1 ) and take an average over ﬁve independent trials.
We plot in (natural) logarithmic scale the mean errors of the two datasets in Figuresand 3.
We can observe that the learned sketches consistently outperform the classical random sketches in all cases. Note that our sketch size is much smaller than the theoretical bound. For a smaller sketch size, classical random sketches converge slowly or do not yield stable results, while the learned sketches can converge to a small error quickly. For larger sketch sizes, all three classical sketches have approximately the same order of error, and the learned sketches reduce the error by a 5 / 6 to 7 /factor in all iterations.
5. S UPPORT V ECTOR M ACHINE
In the context of binary classiﬁcation, a labeled sample is a pair ( a i , z i ) , where a i ∈ R n is a vector representing a collection of features and z i ∈ {− 1 , +1 } is the associated class label.
Given a set of labeled patterns { ( a i , z i ) } di =1 , the support vector machine (SVM) estimates the weight vector w ∗ by minimizing the function
[Equation]
where C is a parameter.
Here we use the squared hinge loss g ( z i , (cid:104) w, a i (cid:105) ) :
= (1 − z i (cid:104) w, a i (cid:105) ) 2+ .
The dual of this problem can be written as a constrained minimization problem (see, e.g., (Li et al., 2009; Pilanci & Wainwright, 2015)),
[Equation]
over the domain ∆ d = { x ∈ R d : x ≥ 0 and (cid:107) x (cid:107) 1 = 1 } , the positive simplex in R d .
Here
B =
[( AD ) (cid:62) 1 √ C I d ] (cid:62) ∈
R ( n + d ) × d , where A is an n × d matrix with a i ∈ R n as its i -th column and D = diag
( z ) is a d × d diagonal matrix.
We conduct experiments on the following three datasets:
• Random Gaussian (synthetic):
We follow the same construction in Pilanci & Wainwright (2016).
We generate a two-component Gaussian mixture model, based on the component distributions N ( µ 0 , I ) and N ( µ 1 , I ) , where µ 0 and µ 1 are uniformly distributed in [ − 3 , 3] .
Placing equal weights on each component, we draw d samples from this mixture distribution.
• Swarm behavior 3 : Each instance in the dataset has n = 2400 features and the task is to predict whether the instance is ﬂocking or not ﬂocking.
We use only the ﬁrst 6000 instances of the raw data, and divide them into 200 smaller groups of instances.
Each group contains d =instances, corresponding to a B i of size 2430 × 30 .
The training data consists of 160 groups and the test data consist of 40 groups.
• Gisette 4 : Gisette is a handwritten digit recognition problem which asks to to separate the highly confusable digits ‘4’ and ‘9’.
Each instance has n = 5000 features.
The raw data is divided into 200 smaller groups, where each contains d = 30 instances and corresponds to a B i of size 5030 × 30 .
The training data consists of 160 groups and the test data consists of 40 groups.
We choose m = 10 d in all experiments and deﬁne the error as (cid:107) Bx (cid:107) 22 − (cid:107)
Bx ∗ (cid:107) 22 .
For random sketches, we take the average error over ﬁve independent trials, and for learned sketches, over three independent trials.
For the Gisette dataset, we use the learned sketch in all rounds.
We plot in a (natural) logarithmic scale the mean errors of the three datasets in Figures 5 to 7.
For Gisette and random Gaussian datasets, using the learned sketches reduced the error by 10%–30%, and for the
Swarm Behavior dataset, the learned sketches reduce the error by about 30%–40%.
5. M ATRIX E STIMATION WITH N UCLEAR N ORM C ONSTRAINT
In many applications, for the problem
[Equation]
it is reasonable to model the matrix X ∗ as having low rank.
Similar to (cid:96) 1 -minimization for compressive sensing, a standard relaxation of the rank constraint is to minimize the nuclear norm of X , deﬁned as (cid:107) X (cid:107) ∗ :
= (cid:80) min { d 1 ,d 2 } j = σ j ( X ) , where σ j ( X ) is the j -th largest singular value of X .
Hence, the matrix estimation problem we consider here is
[Equation]
such that (cid:107)
[Equation]
≤ ρ. where ρ > 0 is a user-deﬁned radius as a regularization parameter.
We conduct experiments on the following two datasets:
• Synthetic Dataset : We generate the pair ( A i , B i ) as B i =
A i X ∗ i
+ W i , where A i ∈ R n × d 1 with i.i.d N (0 , 1) entries.
X ∗ i ∈ R
d 1 ×
d 2 is a matrix with rank at most r .
W i is noise with i.i.d N (0 , σ 2 ) entries.
Here we set n = 500 , d 1 = d 2 = 7 , r = 3 , ρ = 30 .
| ( A, B ) | train = 270 , | ( A, B ) | test = 30 .
• Tunnel 5 : The data set is a time series of gas concentrations measured by eight sensors in a wind tunnel.
Each ( A, B ) corresponds to a different data collection trial.
A i ∈ R 13530 × 5 , B i ∈ R 13530 × 6 , | ( A, B ) | train =
144 , | ( A, B ) | test = 36 .
The same dataset and parameters were also used in (Liu et al., 2020) for regression tasks.
In our nuclear norm constraint, we set ρ = 10 .
We choose m = 50 for both datasets and deﬁne the error to be 12 ( (cid:107) AX − B (cid:107) 2 F − (cid:107) AX ∗ − B (cid:107) 2 F ) .
For each data point, we take the average error of three independent tri - als The mean errors of the two datasets
are
plotted in a (natural logarithmic scale in Figuresand 9. We observe
that the classical sketches yield approximately the same order of error, while the learned sketches improve the
error by
at
least 30% for
the
synthetic dataset
and
surprisingly
by at
least 95% for the Tunnel dataset.
The huge improvement
on the
Tunnel
dataset may be
due
to
the fact that the matrices A i have many duplicate
rows 5. R UNTIME OF L EARNED S KETCHES [Figure] As stated in Section 2, our learned sketch matrices S are all
COUNT - SKETCH -type matrices (each column contains a single nonzero entry), the matrix product SA can thus be computed in O (nnz( A )) time and the overall algorithm is expected to be fast.
To verify this, we plot in an error-versus-runtime plot for the SVM and matrix estimation with nuclear norm constraint tasks in Figures 10 and 11 (corresponding to the datasets in
Figures 7 and 9).
The runtime consists only of the time for sketching and solving the optimization problem and does not include the time for loading the data.
We run the same experiment three times.
Each time we take an average over all test data.
From the plot we can observe that the learned sketch and
C OUNT -S KETCH have the fastest runtimes, which are slightly faster than that of the SJLT and signiﬁcantly faster than that of the Gaussian sketch.
F AST R EGRESSION E XPERIMENT
We consider the unconstrained least-squares problem, i.e., (5) with λ = 0 , using the CO emission, greenhouse gas datasets, and the following Census dataset:
• Census data 6 : this dataset consists of annual salary and related features on people who reported that they worked 40 or more weeks in the previous year and worked 35 or more hours per week.
We randomly sample 5000 instances to create ( A i , b i ) , where A ∈ R 5000 × 11 and b ∈ R 5000 × 1 , | ( A, B ) | train = 160 , | ( A, B ) | test = 40 .
Training.
We optimize the learned sketch S 1 by gradient descent (Algorithm 1), where L ( S, A )
=
κ ( AR − 1 1 ) , where R 1 is computed as in Algorithm 3 and κ ( M ) denotes the condition number of a matrix M .
Next we discuss how to generate the training data.
Since we use Newton’s method to solve an unconstrained convex optimization problem (see Section 2), in the t -round, we need to solve a regression problem
[Equation]
Reformulating it as min z (cid:13)(cid:13)
A (cid:62) Az − y
(cid:13)(cid:13) 2 , we see that A and y depend on the previous solution x t .
Hence, we take x t to be the solution obtained from Algorithm 3 using the learned sketch S t , and this generates A and y for the ( t + 1) -st round.
Experiment.
For the CO emission dataset, we set m = 70 , and for the Census dataset, we set m = 500 .
For the η in the Algorithm 3, we set η = 1 for the ﬁrst round and η = 0 .
2 for the next rounds for the CO emission dataset, and η = 1 for all rounds for the Census dataset. We leave the setting and results on the greenhouse gas dataset to Appendix D.
We examine the accuracy of the subproblem (4) and deﬁne the error to be (cid:13)(cid:13)
A (cid:62) ARz t − y (cid:13)(cid:13) 2 / (cid:107) y (cid:107) 2 .
We run the ﬁrst three subroutines of solving the subproblem for the CO emission dataset and the Census dataset.
The average error of three independent trials is plotted in Figures 12, 13 and 14.
We observe that for the CO emission dataset, the classical sketches have a similar performance and the learned sketches lead to a fast convergence in the subroutine with the ﬁrst-round error at least 80% smaller; for the Census dataset, the learned sketch achieves the smallest error for all three rounds, where we reduce about 60% error in the ﬁrst round and about 50% error in the third round.
Note that the learned sketch always considerably outperforms C OUNT -S
KETCH in all cases.
C ONCLUSION
We demonstrated the superiority of using learned sketches, over classical random sketches, in the Iterative Hessian Sketching method and fast regression solvers for unconstrained least-squares.
Compared with random sketches, our learned sketches of the same size can considerably reduce the error in the loss function (i.e., f ( x ) − f ( x ∗ ) , where x is the output of the sketched algorithm and x ∗ the optimal solution to the unsketched problem) for a given threshold of maximum number of iterations or maximum runtime.
Learned sketches also admit a smaller sketch size.
When the sketch size is small, the algorithm with random sketches may fail to converge, or converge slowly, while the algorithm with learned sketches converges quickly.
A P ROOF OF L EMMA 3.
Suppose that AR − 1 = UW , where U ∈ R n × d has orthonormal columns, which form an orthonormal basis of the column space of A .
Since T is a subspace embedding of the column space of A with probability 0 . 99 , it holds for all x ∈ R d that
[Equation]
Since
[Equation]
and
[Equation]
(6) we have that 1 +
[Equation]
It is easy to see that
[Equation]
and thus,
[Equation]
Recall that SUW = SAR − 1 .
We see that
[Equation]
By deﬁnition,
[Equation]
It follows from (7) that
[Equation]
and from (7), (6) and (Vershynin, 2012, Lemma 5.36) that
[Equation]
Since
[Equation]
and (cid:13)(cid:13)
[Equation]
it follows that
[Equation]
We have so far proved the correctness of the approximation and we shall analyze the runtime below.
Since S and T are sparse, computing SA and TA takes O (nnz( A )) time.
The QR decomposition of TA , which is a matrix of size poly( d/η ) × d , can be computed in poly( d/η ) time.
The matrix
SAR − 1 can be computed in poly( d ) time.
Since it has size poly( d ) × d , its smallest singular value can be computed in poly( d ) time.
To approximate Z 2 ( S ) , we can use the power method to estimate (cid:13)(cid:13) ( SAR − 1 ) T SAR − 1 − I (cid:13)(cid:13) op up to a (1 ± η ) -factor in O ((nnz( A ) + poly( d )) log(1 /η )) time
[Equation]
[-+ poly( d )) log(
-]
B P ROOF OF T HEOREM 3.
In Lemma 3.1, we have {+with probability at least 0 . 
 that+}
[Equation]
When S is random subspace embedding, it holds with probability at least 0 . 99 that Z 1 ( S ) ≥ 3 /and so, by a union bound, it holds with probability at least 0 . 98 that
[Equation]
or,
[Equation]
The correctness of our claim then follows from (Pilanci & Wainwright, 2016, Proposition 1), together with the fact that S 2 is a random subspace embedding.
The runtime follows from Lemma 3.
and (Cormode & Dickens, 2019, Theorem 2.2).
C P ROOF OF T HEOREM 4.
The proof follows almost an identical argument as that in (van den Brand et al., 2020, Lemma
B.1).
In (van den Brand et al., 2020), it is assumed (in our notation) that 3 / 4 ≤ σ min ( AP ) ≤ σ max ( AP ) ≤ 5 / 4 and thus one can set η = 1 in Algorithm 3 and achieve a linear convergence.
The only difference is that here we estimate σ min ( AP ) and σ max ( AP ) and set the step size η in the gradient descent algorithm accordingly.
By standard bounds for gradient descent (see, e.g., (Boyd & Vandenberghe, 2004, p468)), with a choice of step size η = 2 / ( σ 2max ( AP ) + σ 2min ( AP )) , after O (( σ max ( AP ) /σ min ( AP ))
2 log(1 /(cid:15) )) iterations, we can ﬁnd z t such that
[Equation]
where z ∗ = arg min
z (cid:13)(cid:13) P (cid:62) A (cid:62) APz − P (cid:62) y
(cid:13)(cid:13) 2 is the optimal least-squares solution.
This establishes
Eq.
(11) in the proof in (van den Brand et al., 2020), and the rest of the proof follows as in there.
D F AST R EGRESSION E XPERIMENT : G REENHOUSE G AS
We set m = 100 , η = 0 .
2 in the ﬁrst round and η = 1 in the second round.
The mean errors of the ﬁrst two calls of the fast regression subroutine are plotted in Figure 14.
We can observe that
Gaussian and sparse JL sketches have a signiﬁcantly better performance than C OUNT -S KETCH , and the learned sketch shows again a signiﬁcant reduction of more than 50% in the ﬁrst-round error.