Year: 2021

Abstract:
Sketching is a dimensionality reduction technique where one compresses a matrix by linear combinations that are chosen at random. 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. We prove that a smaller sketching dimension of the column space of a tall matrix is possible, given an oracle that can predict the indices of the rows of large leverage score. We design such an oracle for various datasets, and this leads to a faster convergence of the well-studied iterative Hessian sketch procedure, which applies to a wide range of problems in convex optimization. We show empirically that learned sketches, compared with their “non-learned” counterparts, do improve the approximation accuracy for important problems, including LASSO and matrix estimation with nuclear norm constraints.

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 & 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.

heading: 3 LEARNING-AUGMENTED SUBSPACE EMBEDDINGS
text: In this section we explain two ways to optimize the subspace embedding property of the sketching matrix. One is to optimize the non-zero positions of the COUNT-SKETCH-type sketch, based on a trained oracle to identify a superset of the rows of large leverage score. The other is to optimize the values of the nonzero entries, which may no longer be −1 or 1, via a learning algorithm based on gradient descent. As we shall see in Section 4 and 5, a better subspace embedding implies a better convergence rate in the IHS, as well as for the subroutine in unconstrained convex optimization.

heading: 3.1 SKETCHED LEARNING: OPTIMIZING THE POSITIONS
text: In this section we consider the problem of embedding the column space of a matrix A ∈ Rn×d, provided that A has a few rows of large leverage score, as well as access to an oracle which reveals a
superset of the indices of such rows. Formally, let τi(A) denote the leverage score of the i-th row of A and let I∗ = {i : τi(A) ≥ ν} be the set of rows with large leverage score. Suppose that a superset I ⊇ I∗ is known to the algorithm. In the experiments we train an oracle to predict such rows. We can maintain all rows in I explicitly and apply a COUNT-SKETCH to the remaining rows, i.e., the rows in [n] \ I . Up to permutation of the rows, we can write
A = ( AI AIc ) and S = ( I 0 0 S′ ) , (4)
where S′ is a random COUNT-SKETCH matrix of m rows. Clearly S has a single non-zero entry per column. We have the following theorem, whose proof is postponed to Section A. Intuitively, the proof for COUNT-SKETCH in (Clarkson & Woodruff, 2017) handles rows of large leverage score and rows of small leverage score separately. The rows of large leverage score are to be perfectly hashed while the rows of small leverage score will concentrate in the sketch by the Hanson-Wright inequality.
Theorem 3.1. Let ν = /d. Suppose that m = O((d/ 2)(polylog(1/ ) + log(1/δ))), δ ∈ (0, 1/m] and d = Ω((1/ ) polylog(1/ ) log2(1/δ)). Then, there exists a distribution on S of the form in (4)
with m+ |I| rows such that Pr { ∀x ∈ colsp(A), ∣∣∣‖Sx‖22 − ‖x‖22∣∣∣ > ‖x‖22 } ≤ δ . Hence, if there happen to be at most dpolylog(1/ )/ 2 rows of leverage score at least /d, the overall sketch length for embedding colsp(A) can be reduced to O((dpolylog(1/ ) + log(1/δ))/ 2), a quadratic improvement in d and an exponential improvement in δ over the original sketch length of O(d2/( 2δ)) for COUNT-SKETCH. In the worst case there could be O(d2/ ) such rows, though empirically we do not observe this. The following is an immediate corollary, by setting δ = 1/m.
Corollary 3.2. Suppose that d = Ω((1/ ) polylog(1/ )) and |I| = O((d/ 2) polylog(d/ )) with ν = /d. There exists a distribution on S of the form in (4) with O((d/ 2) polylog(d/ )) rows such
that Pr { ∀x ∈ colsp(A), ∣∣∣‖Sx‖22 − ‖x‖22∣∣∣ > ‖x‖22 } ≤ 3 . We remark that our S is of the COUNT-SKETCH type, which has a twofold benefit. First, SA can be applied in O(nnz(A)) time. This is faster than a chained subspace embedding of the form S2S1A, where S1 is a COUNT-SKETCH matrix of O(d2/ 2) rows and S2 is a subspace embedding matrix of O(d/ 2) rows. Computing S1A takes O(nnz(A)) time but computing S2(S1A) will take an additional time of poly(d/ ) or O(nnz(S1A) log(d)/ ). The latter terms can be quite large and even comparable to n if say, n is close to d2. Second, our S allows the sketched matrix SA to be sparse when A is sparse, while the other designs such as Subsampled Randomized Hadamard Transforms and Sparse Johnson-Lindentrauss Transforms either would not guarantee that SA is sparse, or would yield a worse sparsity than a matrix of the COUNT-SKETCH type. The sparsity of SA is also important for solving regression problems involving B = SA in intermediate steps, as algorithms such as conjugate gradient, which use matrix-vector products, become more efficient.
We note that approximate leverages scores of all rows can be found in time O(nnz(A) log n + poly(d/ )) (Clarkson & Woodruff, 2017). Hence, one can approximate the leverage score of every row in a preprocessing step before running the IHS. This time will be amortized by the IHS iterations, because the matrix A remains the same throughout the process. Moreover, in Section 6, we show that for a number of real-world datasets, it is possible to learn the indices of the heavy rows. In practice, one can shrink the size of the superset I by restricting I to the rows with large `2 norms in AI . We shall demonstrate in Section 6 that this heuristic works well on some real-world datasets.

heading: 3.2 SKETCHED LEARNING: OPTIMIZING THE VALUES
text: As mentioned in Section 2, when we fix the positions of the non-zero entries, we aim to optimize the values by gradient descent. We propose the following objective loss function for the learning algorithm L(S,Ai) = ‖(AiRi)>AiRi − I‖F , over all the training data, where Ri comes from the QR-decomposition of SAi = QiR−1i . We found empirically that not squaring this loss function works better than squaring it. We think one of the reasons is that the version without squaring may be less sensitive to outliers. The intuition for this loss function is given by the lemma below, whose proof is deferred to Section B.
Lemma 3.3. Suppose that ∈ (0, 12 ), S ∈ R m×n, A ∈ Rn×d has full column rank, and SA = QR is the QR-decomposition of SA. If ‖(AR−1)>AR−1 − I‖op ≤ , then S is a (1 ± )-subspace embedding of the column space of A.
Lemma 3.3 implies that if the loss function overAtrain is small and the distribution ofAtest is similar to Atrain, it is reasonable to expect that S is a good subspace embedding of Atest. Here we use the Frobenius norm rather than operator norm in the loss function because it will make the optimization problem easier to solve, and our empirical results also show that the performance of the Frobenius norm is better than that of the operator norm.
4 HESSIAN SKETCH
Algorithm 2 Solver for (5) 1: S1 ← learned sketch, S2 ← random sketch 2: (Ẑi,1, Ẑi,2)← ESTIMATE(Si, A), i = 1, 2 3: i∗ ← arg mini=1,2(Ẑi,2/Ẑi,1) 4: x̂← solution of (5) with S = Si∗ 5: return x̂ 6: function ESTIMATE(S,A) 7: T ← sparse (1±η)-subspace embedding
matrix for d-dimensional subspaces 8: (Q,R)← QR(TA) 9: Ẑ1 ← σmin(SAR−1)
10: Ẑ2 ← (1± η)-approximation to∥∥(SAR−1)>(SAR−1)− I∥∥ op
11: return (Ẑ1, Ẑ2)
In this section, we consider the minimization problem
min x∈C
{ 1
2 ‖SAx‖22 − 〈A
>y, x〉 } , (5)
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 define the following quantities (corresponding exactly to the unconstrained case in (Pilanci & Wainwright, 2016))
Z1(S) = inf v∈colsp(A)∩Sn−1 ‖Sv‖22 , Z2(S) = sup u,v∈colsp(A)∩Sn−1
〈 u, (S>S − In)v 〉 .
When S is a (1 + )-subspace embedding of colsp(A), we have Z1(S) ≥ 1− and Z2(S) ≤ 2 .
For a general sketching matrix S, the following is the approximation guarantee of Ẑ1 and Ẑ2, which are estimates of Z1(S) and Z2(S), respectively. The proof is postponed to Appendix C. The main idea is that AR−1 is well-conditioned, where R is as calculated in Algorithm 2. Lemma 4.1. Suppose that η ∈ (0, 13 ) is a small constant, A is of full rank and S has poly(d/η) 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η.
Similar to Proposition 1 of (Pilanci & Wainwright, 2016), we have the following guarantee. The proof is postponed to Appendix D. Theorem 4.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 poly(d/η) 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 leastsquares solution.
Theorem 4.2 suggests the following. If the ratio of the learned sketch Z2(S1)/Z1(S1) is a constant smaller than that of the random sketch Z2(S2)/Z2(S2) and η is a constant fraction of the ratio gap, then Ẑ1,2/Ẑ1,1 is a constant smaller than Ẑ2,2/Ẑ2,1, which means that the procedure of IHS will converge faster with the learned sketch. In particular, if Si is a (1 + i)-subspace embedding matrix for colsp(A) with i < 1/3 and η < γmin{| 1 − 2|, 1, 2} for some small constant γ > 0, we have Z2(Si)/Z1(Si) ≤ 3 i and the guarantee in Theorem 4.2 becomes ‖A(x̂− x∗)‖2 ≤ O(min{ 1, 2}) ‖Ax∗‖2, that is, a better subspace embedding can lead to a faster convergence. Hence, if the learned sketch is a better subspace embedding than a random sketch, theoretically we can obtain a better convergence by setting η small enough; in practice we shall observe this.
Furthermore, if we know the indices of the rows of large leverage scores of A and the assumptions in Corollary 3.2 are satisfied, we can use O(d2/ 2) rows to obtain a ( 1 + O( √ d/ polylog(d/ ) ) ) - subspace embedding using Corollary 3.2, which is almost a √ d-factor better than the usual guarantee
of a random COUNT-SKETCH matrix of the same dimension, leading to an algorithm of faster convergence.
5 HESSIAN REGRESSION
Algorithm 3 Fast Regression Solver for (6) 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
estimates of σmax(B) and σmin(B) 4: i∗ ← mini=1,2(σi/σ′i) 5: P ← R−1i∗ 6: η ← 1/(σ2i∗ + (σ′i∗)2) 7: z0 ← 0 8: while
∥∥A>APzt − y∥∥2 ≥ ‖y‖2 do 9: zt+1 ← zt − η(P>A>AP )(P>A>APzt−P>y)
10: return Pzt
In this section, we consider the minimization problem
min z ∥∥A>Az − y∥∥ 2 , (6)
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 (3)). 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 σmax(TB) and σmin(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 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 E. Theorem 5.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 with probability at least 0.97. The runtime is O(nnz(A)) + Õ(nd · (min{σ1/σ′1, σ2/σ′2})2 · log(κ(A)/ ) + poly(d)). Remark 5.2. In Algorithm 3, S2 can be chosen to be a subspace embedding matrix for d-dimensional subspaces, in which case, AR−12 has condition number close to 1 (see, e.g., p38 of (Woodruff, 2014)) and the full algorithm would run faster than the trivial O(nd2)-time solver to (6). Remark 5.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, and 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: 6 EXPERIMENTS
text: 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 iteration of the IHS or the internal regression problem in fast regression. All of our experiments are conducted on a laptop with a 1.90GHz CPU and 16GB RAM. The offline training is done separately and the training of a single sketch matrix in our dataset can be finished within 5 minutes using a single GPU. For the learned sketches with learned values of nonzero entries, we take an average over three independent trials; for all other sketches, we take an average over five independent trials. The details of the implementation are deferred to Appendix H.
We elaborate on the reason that the horizontal axes in the plots are in terms of iterations rather than in terms of runtime. The learned matrix S is trained offline only once using the training data. It is not computed while solving the optimization problem on the test data. Hence, no additional computational cost is incurred in generating S other than solving the iteration step using COUNT-SKETCH. Since Gaussian matrices and sparse JL transforms are denser than COUNTSKETCH matrices, they will be considerably slower in each round. Since we want to understand the convergence behavior, an iteration count is more revealing than an overall time bound. If our learned sketch performs no worse with respect to the total number of rounds (which our experiments show), then it has an even greater advantage in runtime. To substantiate this claim, we show in Appendix F an error-versus-runtime plot for the task of matrix estimation with nuclear norm constraints.

heading: 6.1 IHS EXPERIMENTS: LASSO
text: We define an instance of LASSO regression to be:
x∗ = arg min ‖x‖1≤λ
1 2 ‖Ax− b‖22 , (7)
where λ is a parameter. We use two real-world datasets:
• Electric2: residential electric load measurements. Each row of the matrix corresponds to a different residence. Matrix columns are consecutive measurements from different times. Ai ∈ R370×9, bi ∈ R370×1, and |(A, b)train| = 320, |(A, b)test| = 80. We set λ = 15. • Greenhouse gas (GHG)3: 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. We set λ = 30.
Experiment Setting: We choose m = 6d, 8d, 10d for both datasets. We consider the error 1 2 ‖Ax− b‖ 2 2 − 1 2 ‖Ax
∗ − b‖22. For the two datasets, we use both the methods proposed in Section 3. For the heavy-row Count-Sketch, we allocate 30% of the sketch space to the rows of heavy leverage score. For the Electric dataset, each row represents a specific residence and the indices of the heavy rows do not vary much across the matrices in the training data. We select the heavy rows according to the number of times each row is heavy in the training data for the heavy rows. We also consider optimizing the non-zero values after identifying the heavy rows. For the GHG dataset, each row represents a specific time point and the heavy rows are not very concentrated. Nevertheless, we can find a superset of about 30% of the rows that contains most of the heavy rows, based on the counts on the training data. Then we prune the superset by selecting the rows with the largest `2 norms, subject to the dimension budget. This will incur an additional computational cost, but the time is almost the same as the time to read the sub-matrix of these rows, and it can be used in all iterations, so the time of this step is negligible compared to the total runtime. We might lose a small fraction of heavy rows, but it only negligibly affects the experiments. The distribution on the indices of the heavy rows over the dataset is discussed in Appendix G.
Experimental Result: We plot in a logarithmic scale the mean errors of the two datasets in Figures 1 and 2. We see all methods display linear convergence, that is, letting ek denote the error in the k-th iteration, we have ek ≈ ρke1 for some convergence rate ρ. A smaller convergence rate implies a faster convergence.
We calculate an estimated rate of convergence ρ = (ek/e1)1/k with k = 10 for the GHG dataset, and with k = 7 for the Electric dataset. For the GHG dataset, we can see that when the sketch size is small (m = 6d), the gradient-based learned sketch has a rate of convergence that is 56% of that of COUNT-SKETCH, and the heavy-rows sketch has a convergence rate that is 86.9%. When the sketch size is large (m = 10d), the gradient-based learned sketch has a convergence rate that is 63.7%, and the heavy-rows sketch is 82.1%. For the Electric dataset, both sketches, especially the heavy-rows sketch, show significant improvements. When the sketch size is small, the combined-learned sketch has a convergence rate that is just 21.1% of that of sparse JL, and when the sketch size is large, the combined-learned sketch has a smaller convergence rate that is just 15.4%.
We also conducted IHS experiments for the matrix estimation problem with a nuclear norm constraint in Appendix F.
2https://archive.ics.uci.edu/ml/datasets/ElectricityLoadDiagrams20112014 3https://archive.ics.uci.edu/ml/datasets/Greenhouse+Gas+Observing+Network
6.2 FAST REGRESSION EXPERIMENTS
We consider the unconstrained least squares problem minx f(x) with f(x) = 12 ‖Ax− b‖ 2 2 using the Electric dataset.
Training: Note that ∇2f(x) = A>A, independent of x. In the t-th round of Newton’s method, by (3), we need to solve a regression problem minz ∥∥A>Az − y∥∥2 2
with y = ∇f(xt). Hence, we can use the same two methods in the preceding subsection to optimize the learned sketch Si. For a general
problem where ∇2f(x) depends on x, one can take xt to be the solution obtained from Algorithm 3 using the learned sketch St to generate A and y for the (t+ 1)-st round, train a learned sketch St+1, and repeat this process.
Setup for Experiments: For the Electric dataset, we setm = 10d = 90. We compare the heavy-rows COUNT-SKETCH matrix with the three classical random sketches, COUNT-SKETCH, Gaussian and Sparse-JL. For the parameter η in Algorithm 3, we set η = 1 in all iterations for heavy-rows sketches. For the classical random sketches, we set η in the following two ways: (a) η = 1 in all iterations and (b) η = 1 in the first iteration and η = 0.2 in all subsequent iterations. Experimental Results: We examine the accuracy of the subproblem (6) and define the error to be ∥∥A>ARzt − y∥∥2 / ‖y‖2. We consider the subproblems in the first three iterations of the global Newton method. The results are plotted in Figure 3. In this task, the COUNT-SKETCH causes a terrible divergence of the subroutine and is thus omitted in the plots. Still, we observe that in setting (a) of η, the other two classical sketches cause the subroutine to diverge. In setting (b) of η, the other two classical sketches lead to convergence but their error is significantly larger than that of the heavy-rows sketches, in each of the first three calls to the subroutine. The error of the heavy-rows sketch is less than 0.01 in all iterations of all three subroutine calls, in both setting (a) and (b) of η.
We also plot a figure on the convergence of the global Newton method. Here, for each subroutine, we only run one iteration, and plot the error of the original least squares problem. The result is shown in Figure 4, which clearly displays a significantly faster decay with heavy-rows sketches. The rate of convergence using heavy-rows sketches is 80.6% of that using Gaussian or sparse JL sketches.
CONCLUSION. We demonstrated the superiority of using learned sketches over classical random sketches, for the Iterative Hessian Sketch method which is used for a number of problems in convex optimization. Compared with random sketches, our learned sketches of the same size yield considerably faster convergence. We also provably show a better subspace embedding property of a sketch of the same size given an oracle for predicting a superset of rows with large leverage score. Our experiments show the construction of such an oracle is possible for real data sets, and they demonstrate a significant advantage over non-learned sketches for problems in convex optimization.

heading: A PROOF OF THEOREM 3.1
text: First we prove the following lemma. Lemma A.1. Let δ ∈ (0, 1/m]. It holds with probability at least 1− δ that
sup x∈colsp(A) ∣∣∣‖Sx‖22 − ‖x‖22∣∣∣ ≤ ‖x‖22 , provided that
m & −2((d+ logm) min{log2(d/ ), log2m}+ d log(1/δ)), 1 & −2ν((logm) min{log2(d/ ), log2m}+ log(1/δ)) log(1/δ).
Proof. We shall adapt the proof of Theorem 5 in (Bourgain et al., 2015) to our setting. Let T denote the unit sphere in colsp(A) and set the sparsity parameter s = 1. Observe that ‖Sx‖22 = ‖xI‖22 + ‖SxIc‖ 2 2, and so it suffices to show that
Pr {∣∣∣‖S′xIc‖22 − ‖xIc‖22∣∣∣ > } ≤ δ
for x ∈ T . We make the following definition, as in (2.6) of Bourgain et al. (2015):
Aδ,x := m∑ i=1 ∑ j∈Ic δijxjei ⊗ ej ,
and thus, S′xIc = Aδ,xσ. Also by E ‖S′xIc‖22 = ‖xIc‖ 2 2, one has
sup x∈T ∣∣∣‖S′xIc‖22 − ‖xIc‖22∣∣∣ = sup x∈T ∣∣∣‖Aδ,xσ‖22 − E ‖Aδ,xσ‖22∣∣∣ . (8) Now, in (2.7) of Bourgain et al. (2015) we instead define a seminorm
‖x‖δ = max1≤i≤m ∑ j∈Ic δijx 2 j 1/2 . Then (2.8) continues to hold, and (2.9) as well as (2.10) continue to hold if the supremum in the left-hand side is replaced with the left-hand side of (8). At the beginning of Theorem 5, we define U (i) to be U , but each row j ∈ Ic is multiplied by δij and each row j ∈ I is zeroed out. Then we have in the first step of (4.5) that
∑ j∈Ic δij ∣∣∣∣∣ d∑ k=1 gk〈fk, ej〉 ∣∣∣∣∣ 2 ≤ ∥∥∥U (i)g∥∥∥2 2 ,
instead of equality. One can verify that the rest of (4.5) goes through. It remains true that ‖·‖δ ≤ (1/ √ s) ‖·‖2, and thus (4.6) holds. One can verify that the rest of the proof of Theorem 5 in Bourgain
et al. (2015) continues to hold if we replace ∑n j=1 with ∑ j∈Ic and max1≤j≤n with maxj∈Ic , noting
that E ∑ j∈Ic δij ‖PEej‖22 = s m ∑ j∈Ic 〈PEej , ej〉 ≤ s m d
and E(U (i))∗U (i) = ∑ j∈Ic (Eδij)uju∗j 1 m . Thus, the symmetrization inequalities on∥∥∥∥∥∥ ∑ j∈Ic δij ‖PEej‖22 ∥∥∥∥∥∥ Lpδ and ∥∥∥∥∥∥ ∑ j∈Ic δijuju ∗ j ∥∥∥∥∥∥ Lpδ
continue to hold. The result then follows, observing that maxj∈Ic ‖PEej‖2 ≤ ν.
The subspace embedding guarantee now follows as a corollary.
Theorem 3.1. Let ν = /d. Suppose that m = Ω((d/ 2)(polylog(1/ ) + log(1/δ))), δ ∈ (0, 1/m) and d = Ω((1/ ) polylog(1/ ) log2(1/δ)). Then, there exists a distribution on S with m+ |I| rows such that
Pr { ∀x ∈ colsp(A), ∣∣∣‖Sx‖22 − ‖x‖22∣∣∣ > ‖x‖22} ≤ δ. Proof. One can verify that the two conditions in Lemma A.1 are satisfied if
m & d
2
( polylog( d ) + log 1
δ
) ,
d & 1 ( log 1
δ
)( polylog( d ) + log 1
δ
) .
The last condition is satisfied if
d & 1 ( log2 1
δ
) polylog ( 1 ) .

heading: B PROOF OF LEMMA 3.3
text: Proof. On the one hand, since Q = SAR is an orthogonal matrix, we have
‖x‖2 = ‖Qx‖2 = ‖SARx‖2 . (9)
On the other hand, the assumption implies that∥∥(ARx)T (ARx)− xTx∥∥ 2 ≤ ‖x‖22 ,
that is, (1− ) ‖x‖22 ≤ ‖ARx‖ 2 2 ≤ (1 + ) ‖x‖ 2 2 . (10)
Combining both (9) and (10) leads to √
1− ‖SARx‖2 ≤ ‖ARx‖2 ≤ √ 1 + ‖SARx‖2 , ∀x ∈ R d
Equivalently, it can be written as
1√ 1 + ‖SAy‖2 ≤ ‖Ay‖2 ≤ 1√ 1− ‖SAy‖2 , ∀y ∈ R d.
The claimed result follows from the fact that 1/ √ 1 + ≥ 1 − and 1/ √
1− ≤ 1 + whenever ∈ (0, √ 5−1 2 ].

heading: C PROOF OF LEMMA 4.1
text: Suppose thatAR−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 (11)
we have that 1
1 + η ‖x‖2 ≤ ‖Wx‖2 ≤
1
1− η ‖x‖2 , x ∈ R d. (12)
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 (12) that
(1− η)2 ∥∥WTUT (STS − In)UW∥∥op ≤ Z2(S) ≤ (1 + η)2 ∥∥WTUT (STS − In)UW∥∥op .
and from (12), (11) and Lemma 5.36 of Vershynin (2012) 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: D PROOF OF THEOREM 4.2
text: In Lemma 4.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 Proposition 1 of Pilanci & Wainwright (2016), together with the fact that S2 is a random subspace embedding. The runtime follows from Lemma 4.1 and Theorem 2.2 of Cormode & Dickens (2019).

heading: E PROOF OF THEOREM 5.1
text: The proof follows a similar argument to 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., p468 of Boyd & Vandenberghe (2004)), 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.
We use three subspace embeddings here, S1, S2 and one used in the EIG subrountine. Each subspace embedding uses O(d2) rows with a constant distortion parameter and a failure probability of 0.01. The overall failure probability is thus 0.03.
F IHS EXPERIMENTS: MATRIX ESTIMATION WITH NUCLEAR NORM CONSTRAINT
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 the `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.
We conduct experiments on the following datasets:
• Tunnel4: 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. In our nuclear norm constraint, we set ρ = 10.
Experiment Setting: We choose m = 7d, 10d for the Tunnel dataset. We consider the error 1 2 ‖AX −B‖ 2 2 − 1 2 ‖AX
∗ −B‖22. The leverage scores of this dataset are very uniform. Hence, for this experiment we only consider optimizing the values of the non-zero entries.
Results of Our Experiments: We plot in a logarithmic scale the mean errors of the two datasets in Figures 5. We can see that when m = 7d, the gradient-based sketch, based on the first 6 iterations, has a rate of convergence that is 48% of the random sketch, and when m = 10d, the gradient-based sketch has a rate of convergence that is 29% of the random sketch.
Runtime of Learned Sketch. As stated in Section 2, our learned sketch matrices S are all COUNTSKETCH-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 matrix estimation with nuclear norm constraint tasks with m = 10d in Figures 6 (corresponding to the datasets in Figure 5). 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.
4https://archive.ics.uci.edu/ml/datasets/Gas+sensor+array+exposed+to+turbulent+gas+mixtures

heading: G HEAVY LEVERAGE SCORE ROWS DISTRIBUTION OVER THE DATASET
text: In our experiments, we hypothesize that in real-world data that there may be an underlying pattern which can help us identify the heavy rows. In the Electric dataset, each row of the matrix corresponds to a specific residence and the heavy rows are always concentrated on some specific rows; in the GHG data set, each row corresponds to a specific time point and we can select some specific time points to be a superset of the heavy rows and then select the heavy rows based on their `2-norm in this superset.
To exemplify this, we study the heavy leverage score rows distribution over the Electirc dataset. For a row i ∈ [370], let fi denote the times that row i is heavy out of 320 training data points from the Electric dataset, where we say row i is heavy if `i ≥ 5d/n. Below we list all 74 pairs (i, fi) with fi > 0.
(195,320), (278,320), (361,320), (207,317), (227,285), (240,284), (219,270), (275,232), (156,214), (322,213), (193,196), (190,192), (160,191), (350,181), (63,176), (42,168), (162,148), (356,129), (363,110), (362,105), (338,95), (215,94), (234,93), (289,81), (97,80), (146,70), (102,67), (98,58), (48,57), (349,53), (165,46), (101,41), (352,40), (293,34), (344,29), (268,21), (206,20), (217,20), (327,20), (340,19), (230,18), (359,18), (297,14), (357,14), (161,13), (245,10), (100,8), (85,6), (212,6), (313,6), (129,5), (130,5), (366,5), (103,4), (204,4), (246,4), (306,4), (138,3), (199,3), (222,3), (360,3), (87,2), (154,2), (209,2), (123,1), (189,1), (208,1), (214,1), (221,1), (224,1), (228,1), (309,1), (337,1), (343,1)
Observe that the heavy rows are concentrated on a set of specific row indices. There are only 30 rows i with fi ≥ 50. We view this as strong evidence for our hypothesis.
H IMPLEMENTATION DETAILS
As we state in Section 3.2, when we fix the positions of the non-zero entries (uniformly chosen in each column or sampling according to the heavy leverage score distribution), we aim to optimize the values by gradient descent mentioned in Algorithm 1. Here the loss function is given in Section 3.2. In our implementation, we use PyTorch (Paszke et al. (2019)), which can compute the gradient automatically (here we can use torch.qr() and torch.svd() to define our loss function). For a more
nuanced loss function, which may be beneficial, one can use the package released in Agrawal et al. (2019), where the authors studied the problem of computing the gradient of functions which involve the solution to certain convex optimization problem.
As mentioned in Section 2, each column of the sketch matrix S has exact one non-zero entry. Hence, the i-th coordinate of p can be seen as the non-zero position of the i-th column of S. In the implementation, to sample p randomly, we can sample a random integer in {1, . . . ,m} for each coordinate of p. For the heavy rows mentioned in Section 3.1, we can allocate positions 1, . . . , k to the k heavy rows, and for the other rows, we randomly sample an integer in {k+ 1, . . . ,m}. We note that once the vector p, which contains the information of the nonzero position in each column of S, is chosen, it will not be changed during the optimization process in Algorithm 1.
Next, we introduce some parameters for our experiments.
• bs: batch size, the number of training samples used in one iteration. • lr: learning rate of the gradient descent(the α in Algorithm 1). • iter: the number of iteration for Algorithm 1.
In our experiments, we set bs = 20, iter = 1000 for all dataset. We set lr = 10 for the Green House Gas dataset and lr = 0.1 for the Electric dataset.

heading: I ADDITIONAL EXPERIMENTS FOR LASSO
text: In this section, we consider the IHS experiments for LASSO on data of a larger size. The experiment setting is the same as that in Section 6. We conduct our experiments on the following dataset:
• Gas Sensor.5 A chemical detection platform composed of 8 chemoresistive gas sensors was exposed to turbulent gas mixtures generated naturally in a wind tunnel. Each matrix represents the measurements at dense time points during a short time period. Ai ∈ R95000×19, bi ∈ R95000×1, and |(A, b)train| = 30, |(A, b)test| = 9. We set λ = 10.
The results are shown in Table 1. Here we choose m = 300. The leverage scores of the rows on this dataset is very uniform hence we choose random positions for the nonzero entries and only optimize the values in the learned sketch. For a matrix of such size, the gaussian sketching matrix is extremely slow, hence, we only consider the Count-Sketch matrix and the Sparse-JL matrix. From the table below we can see that the gradient-based learned sketch has a converge rate that is 74.6% of that of the random sketch.

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

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. 

This paper studies matrix sketching algorithms where the sketching matrix is `learned', or generated adaptively. Theoretically, this is done by favoring more the rows of A with high leverage scores, while practically this is done by associating that set with rows of large norms. It evaluates this algorithm on both L_1 and L_2 regression, and observes improved test errors.