Year: 2020

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

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

heading: 1 INTRODUCTION
text: 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.

heading: 2 PRELIMINARIES
text: Classical Sketches. Below we review several classical sketches that have been used for solving optimization problems.
• Gaussian sketch: S = 1√ m G, where G is an m× n Gaussian random matrix. • COUNT-SKETCH: 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 COUNT-SKETCH matrices, each of dimension m/s× n.
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 at least a constant probability; 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/ 2). It is a dense matrix and computing SA costs O(m · nnz(A)) = O(nnz(A)d/ 2) time. • COUNT-SKETCH: m = O(d2/ 2) (Clarkson & Woodruff, 2017). Although the number of rows is quadratic in d/ , the sketch matrix S is sparse and computing SA takes only O(nnz(A)) time. • SJLT: m = O(d/ 2) and has s = O(1/ ) non-zeros per column (Nelson & Nguyên, 2013; Cohen, 2016). Computing SA takes O(snnz(A)) = O(nnz(A)/ ) 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 square 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 ,
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.
Algorithm 1 LEARN-SKETCH: Gradient descent algorithm for learning the sketch values Require: Atrain = {A1, ..., AN} (Ai ∈ Rn×d),
learning rate α 1: Randomly initialize p, v for a Count-Sketch-
type sketch 2: for t = 0 to steps do 3: Form S using p, v 4: Sample batch Abatch from Atrain 5: v ← v − α∂L(S,Abatch)∂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 ∑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.
3 HESSIAN SKETCH
Algorithm 2 Solver for (3) 1: S1 ← learned sketch, S2 ← random sketch 2: (Ẑi,1, Ẑi,2)← ESTIMATE(Si, A), i = 1, 2 3: if Ẑ1,2/Ẑ1,1 < Ẑ2,2/Ẑ2,1 then 4: x̂← solution of (3) with S = S1 5: else 6: x̂← solution of (3) with S = S2 7: end if 8: return x̂ 9: function ESTIMATE(S,A) 10: T ← sparse (1± η)-subspace embedding matrix for d-dimensional subspaces 11: (Q,R)← QR(TA) 12: Ẑ1 ← σmin(SAR−1) 13: Ẑ2 ← (1± η)-approximation to∥∥(SAR−1)>(SAR−1)− I∥∥
op
14: return (Ẑ1, Ẑ2) 15: end function
In this section, we consider the minimization problem
min x∈C
{ 1
2 ‖SAx‖22 − 〈A
>y, x〉 } , (3)
which is used as a subroutine for the IHS (cf. (2)). We present an algorithm with the learned sketch in Algorithm 2. To analyze its performance, we let R be the column space of A ∈ Rm×n and define the following quantities (corresponding exactly to the unconstrained case in Pilanci & Wainwright (2016))
Z1(S) = inf v∈R∩Sn−1
‖Sv‖22 ,
Z2(S) = sup u,v∈R∩Sn−1
〈 u, (S>S − In)v 〉 ,
where Sn−1 denotes the Euclidean unit sphere in Rn.
The following is the estimation guarantee of Ẑ1 and Ẑ2. The proof is postponed to Appendix A. Lemma 3.1. Suppose that η ∈ (0, 13 ) is a small constant, A is of full rank and S has O(d
2) rows. The function ESTIMATE(S,A) returns inO((nnz(A) log(1/η)+poly(d/η)) time Ẑ1, Ẑ2 which with probability at least 0.99 satisfy that Z1(S)1+η ≤ Ẑ1 ≤ Z1(S) 1−η and Z2(S) (1+η)2 − 3η ≤ Ẑ2 ≤ Z2(S) (1−η)2 + 3η.
Note for a matrixA, ‖A‖op = supx 6=0 ‖Ax‖2 ‖x‖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 S1 and S2 are both COUNT-SKETCH-type sketches with O(d2) rows. Algorithm 2 returns a solution x̂ which, with
probability at least 0.98, satisfies that ‖A(x̂− x∗)‖2 ≤ (1 + η)4 ( min { Ẑ1,2
Ẑ1,1 , Ẑ2,2 Ẑ2,1
} + 4η ) ‖Ax∗‖2
in O(nnz(A) log( 1η ) + poly( d η )) time, where x ∗ = arg minx∈C ‖Ax− b‖2 is the least-squares solution.
4 HESSIAN REGRESSION Algorithm 3 Fast Regression Solver for (4)
1: S1 ← learned sketch, S2 ← random sketch 2: (Qi, Ri)← QR(SiA), i = 1, 2 3: (σi, σ′i)← EIG(AR −1 i ), i = 1, 2
. EIG(B) returns the max and min singular values of B
4: if σ1/σ′1 < σ2/σ′2 then 5: P ← R−11 , η ← 1/(σ21 + (σ′1)2) 6: else 7: P ← R−12 , η ← 1/(σ22 + (σ′2)2) 8: end if 9: z0 ← 0
10: while ∥∥A>APzt − y∥∥2 ≥ ‖y‖2 do 11: zt+1 ← zt − η(P>A>AP )(P>A>APzt−P>y) 12: end while 13: return Pzt
In this section, we consider the minimization problem
min z ∥∥A>Az − y∥∥ 2 , (4)
which is used as a subroutine for the unconstrained convex optimization problem minx f(x) with A>A being the Hessian matrix∇2f(x) (See Section 2). Here A ∈ Rn×d, y ∈ Rd, 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 EIG(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 COUNT-SKETCH matrix of O(d2) 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 S1 and S2 are both COUNT-SKETCH-type sketches with O(d2) rows. Algorithm 3 returns a solution x′ such that ∥∥A>Ax′ − y∥∥ 2 ≤ ‖y‖2 in O(nnz(A)) + Õ(nd · (min{σ1/σ′1, σ2/σ′2})2 · log(κ(A)/ ) + poly(d)) time. Remark 4.2. In Algorithm 3, S2 can be chosen to be a subspace embedding matrix for ddimensional subspaces, in which case, AR−12 has condition number close to 1 (see, e.g., (Woodruff, 2014, p38)) and the full algorithm would run faster than the trivial O(nd2)-time solver to (4).
Remark 4.3. For the original unconstrained convex optimization problem minx 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 ‖Ax− b‖ 2 2, the value of f(x) can be approximated efficiently by a sparse subspace embedding matrix in O(nnz(A) + nnz(b) + poly(d)) time.

heading: 5 IHS EXPERIMENTS
text: 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 xt. An issue is that we do not know xt, 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 xt by solving (2), and then use it to generate the training data in the next round. That is, in the first round, we train the sketch matrix S1 to solve the problem x1 = arg minx∈C 1 2 ‖S1Ax‖ 2 2−〈A>b, x〉 and use x1 to generate the training data for the optimization problem for x2, and so on. The loss function we use is the unsketched objective function in the (t+1)-st iteration, i.e.,L(St+1, A) = 12 ‖A(xt+1 − xt)‖ 2 2−〈AT (b−Axt), xt+1−xt〉, where xt+1 is the solution to (2) and thus depends on St+1.
Comparison. We compare the learned sketch against three classical sketches: Gaussian, COUNTSKETCH, and SJLT (see Section 2) in all experiments. The quantity we compare is a certain error, defined 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.

heading: 5.1 LASSO
text: Figure 2: Test error of LASSO on CO emissions dataset, m = 3d
Figure 4: Test error of LASSO on greenhouse gas dataset,m = 3.5d
gregated over one hour (by means of average or sum), which can help to predict the CO emission. We divide the raw data into 120 (Ai, bi) such that Ai ∈ R300×9, bi ∈ R300×1. The data in each matrix is sorted in chronological order.|(A, b)train| = 96, |(A, b)test| = 24. • Greenhouse gas2: time series of measured greenhouse gas concentrations in the California atmosphere. Each (A, b) corresponds to a different measurement location. Ai ∈ R327×14, bi ∈ R327×1, and |(A, b)train| = 400, |(A, b)test| = 100. (This dataset was also used in (Liu et al.,
5.2 SUPPORT VECTOR MACHINE
In the context of binary classification, a labeled sample is a pair (ai, zi), where ai ∈ Rn is a vector representing a collection of features and zi ∈ {−1,+1} is the associated class label. Given a set of labeled patterns {(ai, zi)}di=1, the support vector machine (SVM) estimates the weight vector w∗ by minimizing the function
w∗ = arg min w∈Rn
{ C
2 d∑ i=1 g(zi, 〈wi, ai〉) + 1 2 ‖w‖22
} ,
where C is a parameter. Here we use the squared hinge loss g(zi, 〈w, ai〉) := (1− zi〈w, ai〉)2+.
1https://archive.ics.uci.edu/ml/datasets/Gas+Turbine+CO+and+NOx+Emission+Data+Set 2https://archive.ics.uci.edu/ml/datasets/Greenhouse+Gas+Observing+Network
The dual of this problem can be written as a constrained minimization problem (see, e.g., (Li et al., 2009; Pilanci & Wainwright, 2015)),
x∗ := arg min x∈∆d ‖Bx‖22 ,
over the domain ∆d = {x ∈ Rd : x ≥ 0 and ‖x‖1 = 1}, the positive simplex in Rd. Here B = [(AD)> 1√
C Id] > ∈ R(n+d)×d, where A is an n × d matrix with ai ∈ Rn 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 behavior3: Each instance in the dataset has n = 2400 features and the task is to predict whether the instance is flocking or not flocking. We use only the first 6000 instances of the raw data, and divide them into 200 smaller groups of instances. Each group contains d = 30 instances, corresponding to a Bi of size 2430× 30. The training data consists of 160 groups and the test data consist of 40 groups. • Gisette4: 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 Bi of size 5030× 30. The training data consists of 160 groups and the test data consists of 40 groups.
We choose m = 10d in all experiments and define the error as ‖Bx‖22 − ‖Bx∗‖ 2 2. For random sketches, we take the average error over five 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%.

heading: 5.3 MATRIX ESTIMATION WITH NUCLEAR NORM CONSTRAINT
text: In many applications, for the problem
X∗ := arg min X∈Rd1×d2 ‖AX −B‖2F ,
it is reasonable to model the matrix X∗ as having low rank. Similar to `1-minimization for compressive sensing, a standard relaxation of the rank constraint is to minimize the nuclear norm of X , defined as ‖X‖∗ := ∑min{d1,d2} j=1 σj(X), where σj(X) is the j-th largest singular value of X .
Hence, the matrix estimation problem we consider here is
X∗ := arg min X∈Rd1×d2 ‖AX −B‖2F such that ‖X‖∗ ≤ ρ.
where ρ > 0 is a user-defined radius as a regularization parameter. 3https://archive.ics.uci.edu/ml/datasets/Swarm+Behaviour 4https://www.csie.ntu.edu.tw/∼cjlin/libsvmtools/datasets/binary.html#gisette
Figure 8: Test error of matrix estimation on synthetic data, m = 50
Figure 9: Test error of matrix estimation on Tunnel dataset, m = 50
We conduct experiments on the following two datasets:
• Synthetic Dataset: We generate the pair (Ai, Bi) as Bi = AiX ∗ i + Wi, where Ai ∈ Rn×d1 with i.i.d
N(0, 1) entries. X∗i ∈ Rd1×d2 is a matrix with rank at most r. Wi is noise with i.i.d N(0, σ2) entries. Here we set n = 500, d1 = d2 = 7, r = 3, ρ = 30. |(A,B)|train = 270, |(A,B)|test = 30. • Tunnel5: 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. Ai ∈ R13530×5, Bi ∈ R13530×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 = 40, 50 for the synthetic dataset and m = 10, 50 for the Tunnel dataset and define the error to be 1 2 (‖AX −B‖ 2 F − ‖AX∗ −B‖ 2 F ). For each data point,
we take the average error of five independent trials. The mean errors of the two datasets when m = 50 are plotted in a (natural) logarithmic scale in Figures 8 and 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 Ai have many duplicate rows. We defer the results when m = 10 or 40 to Appendix D, which show that the learned sketches yield much smaller errors than the random sketches and the random sketches could converge significantly more slowly with considerably larger errors in the first several rounds.
5.4 RUNTIME OF LEARNED SKETCHES
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 COUNT-SKETCH have the fastest runtimes, which are slightly faster than that of the SJLT and significantly faster than that of the Gaussian sketch.

heading: 6 FAST REGRESSION EXPERIMENT
text: 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 data6: 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 (Ai, bi), where A ∈ R5000×11 and b ∈ R5000×1, |(A,B)|train = 160, |(A,B)|test = 40.
5https://archive.ics.uci.edu/ml/datasets/Gas+sensor+array+exposed+to+turbulent+gas+mixtures 6https://github.com/chocjy/randomized-quantile-regression-solvers/tree/master/matlab/data
Training. We optimize the learned sketch S1 by gradient descent (Algorithm 1), where L(S,A) = κ(AR−11 ), where R1 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 minz ∥∥(∇2f(xt)1/2)>(∇2f(xt)1/2)z −∇f(xt)∥∥2. Reformulating it as minz ∥∥A>Az − y∥∥ 2 , we see that A and y depend on the previous solution xt. Hence, we take xt to be the solution obtained from Algorithm 3 using the learned sketch St, 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 first 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 E.
We examine the accuracy of the subproblem (4) and define the error to be ∥∥A>ARzt − y∥∥2 / ‖y‖2. We run the first 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 16. 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 first-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 first round and about 50% error in the third round. Note that the learned sketch always considerably outperforms COUNT-SKETCH in all cases.

heading: 7 CONCLUSION
text: 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.

heading: A PROOF OF LEMMA 3.1
text: Suppose that AR−1 = UW , where U ∈ Rn×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 ∈ Rd that
1
1 + η
∥∥TAR−1x∥∥ 2 ≤ ∥∥AR−1x∥∥ 2 ≤ 1 1− η ∥∥TAR−1x∥∥ 2 .
Since ∥∥TAR−1x∥∥ 2
= ‖Qx‖2 = ‖x‖2 and
‖Wx‖2 = ‖UWx‖2 = ∥∥AR−1x∥∥ 2 (6)
we have that 1
1 + η ‖x‖2 ≤ ‖Wx‖2 ≤
1
1− η ‖x‖2 , x ∈ R d. (7)
It is easy to see that
Z1(S) = min x∈Sd−1 ‖SUx‖2 = min y 6=0 ‖SUWy‖2 ‖Wy‖2 ,
and thus,
min y 6=0 (1− η) ‖SUWy‖2 ‖y‖2 ≤ Z1(S) ≤ min y 6=0 (1 + η) ‖SUWy‖2 ‖y‖2 .
Recall that SUW = SAR−1. We see that
(1− η)σmin(SAR−1) ≤ Z1(S) ≤ (1 + η)σmin(SAR−1).
By definition, Z2(S) = ∥∥UT (S>S − In)U∥∥op . It follows from (7) that
(1− η)2 ∥∥WTUT (STS − In)UW∥∥op ≤ Z2(S) ≤ (1 + η)2 ∥∥WTUT (STS − In)UW∥∥op .
and from (7), (6) and (Vershynin, 2012, Lemma 5.36) that∥∥(AR−1)>(AR−1)− I∥∥ op ≤ 3η. Since ∥∥WTUT (STS − In)UW∥∥op = ∥∥(AR−1)>(STS − In)AR−1∥∥op and ∥∥(AR−1)>STSAR−1 − I∥∥ op − ∥∥(AR−1)>(AR−1)− I∥∥ op
≤ ∥∥(AR−1)>(STS − In)AR−1∥∥op
≤ ∥∥(AR−1)>STSAR−1 − I∥∥ op + ∥∥(AR−1)>(AR−1)− I∥∥ op ,
it follows that
(1− η)2 ∥∥(SAR−1)>SAR−1 − I∥∥
op − 3(1− η)2η
≤ Z2(S) ≤ (1 + η)2 ∥∥(SAR−1)>SAR−1 − I∥∥ op + 3(1 + η)2η.
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 Z2(S), we can use the power method to estimate∥∥(SAR−1)TSAR−1 − I∥∥
op up to a (1± η)-factor in O((nnz(A) + poly(d)) log(1/η)) time.

heading: B PROOF OF THEOREM 3.2
text: In Lemma 3.1, we have with probability at least 0.99 that
Ẑ2 Ẑ1 ≥
1 (1+η)2Z2(S)− 3η
1 1−ηZ1(S)
≥ 1− η (1 + η)2 Z2(S) Z1(S) − 3η Z1(S) .
When S is random subspace embedding, it holds with probability at least 0.99 that Z1(S) ≥ 3/4 and so, by a union bound, it holds with probability at least 0.98 that
Ẑ2 Ẑ1 ≥ 1 (1 + η)4 Z2(S) Z1(S) − 4η,
or, Z2(S)
Z1(S) ≤ (1 + η)4
( Ẑ2
Ẑ1 + 4η
) .
The correctness of our claim then follows from (Pilanci & Wainwright, 2016, Proposition 1), together with the fact that S2 is a random subspace embedding. The runtime follows from Lemma 3.1 and (Cormode & Dickens, 2019, Theorem 2.2).

heading: C PROOF OF THEOREM 4.1
text: 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 ) + σ 2 min(AP )), after O((σmax(AP )/σmin(AP )) 2 log(1/ )) iterations, we can find zt such that∥∥P>A>AP (zt − z∗)∥∥2 ≤ ∥∥P>A>AP (z0 − z∗)∥∥2 ,
where z∗ = arg minz ∥∥P>A>APz − P>y∥∥ 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 IHS EXPERIMENT: MATRIX ESTIMATION WITH NUCLEAR NORM CONSTRAINT
As stated in Section 5.3, the mean errors of the two datasets when m = 40 for the synthetic dataset and m = 10 for the Tunnel dataset are plotted in a (natural) logarithmic scale in Figures 14 and 15.

heading: E FAST REGRESSION EXPERIMENT: GREENHOUSE GAS
text: We set m = 100, η = 0.2 in the first round and η = 1 in the second round. The mean errors of the first two calls of the fast regression subroutine are plotted in Figure 16. We can observe that Gaussian and sparse JL sketches have a significantly better performance than COUNT-SKETCH, and the learned sketch shows again a significant reduction of more than 50% in the first-round error.

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

Sketching is a popular technique in numerical linear algebra for achieving various desirable properties (e.g., lower complexity, one pass methods). The present paper considers a particular kind of sketch for which the sketch matrix is learned from data. It shows how such learned sketches can be used in two types of problems: Hessian sketching (Sec. 3) and Hessian regression (Sec. 4). The authors give both algorithms and provide theoretical guarantees. They also apply these techniques to a number of both synthetic and real datasets in the experiments. For the most part, the experiments indicate that the proposed methods give a consistent, but not necessarily very large, improvement.