Learning-Augmented Sketches for Hessians 
Authors: Yi Li
Honghao Lin
David Woodruff
Venue: ICLR 2022 Submitted (2022)
Keywords: least squares
convex optimization
iterative Hessian sketch
subspace embedding
learning-augmented sketch

Abstract: Sketching is a dimensionality reduction technique where one compresses a matrix by linear combinations that are typically 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, assuming the knowledge of the indices of the rows of large leverage scores. This would lead to faster convergence of the iterative Hessian sketch procedure. We also design a new objective to learn the sketch, whereby we optimize the subspace embedding property of the sketch. 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.

Under review as a conference paper at ICLR 2022
LEARNING-AUGMENTE SKETCHES FOR HESSIANS
Anonymous authors Paper under double-blind review
ABSTRACT
Sketching is dimensionality reduction technique where one compresses a matrix by linear combinations that are chosen at random. A line of work has shown how 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 number of cases are even oblivious to the input matrix. One could instead hope 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 smaller sketching dimension of the column space a tall matrix is possible, given an oracle that can predict the indices of the rows large leverage score. We design such 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.
INTRODUCTION
Large-scale optimization problems 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 [2066]
2017 Doikov Richtárik] [2018] Gower et al. 2018 KoostaaKKorassee
& | [209]] Gower
Li et al. (2020). Our focus is sketching techniques, which often consist of multiplying ying the Hessian by random matrix chosen independently of the
al. 2019 Kylasa et al.) 2019 Xu et al. 2020
Hessian. Sketching has 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 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 (017) 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 's method find point where the gradient zero, in each iteration one needs solve an equation involving the current Hessian and gradient to find the update direction. When the Hessian can be decomposed as A for an n x 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 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 (017). which extended sketching methods beyond constrained least- squstesquares problems to any twice differentiable function subject to closed convex constraint set. Using this sketch inside of interior point updates has led to much faster algorithms

---

Under review as a conference paper at ICLR 2022
for 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. (020) (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. If A is tall-and-thin, then S is short-and-fa -att and thus SA is a small, roughly square matrix. Moreover, S preserves important properties of A. One typically desired property is that S is subspace embedding, meaning that || Az) 1 (1 + e) Ax|l2 for all x simultancously. An observation exploited in Cormode & Dickens (2199, building off of the COUNT-SKETCH random matrices introodueed 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[nz 4)) time, where nnz( A) denotes the number of nonzeros in A. This also 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 where A A is the Hessian and b is the gradient. For number of problems, one has access to matrix A E Raxdd with n >> d, which is also an assumption made in Pilanci & Wainwright (0117).
Therefore, the solution is the minimizer to a constrained least squares regression problem:
min |r - bll2 ZEC 2
(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 & Nguyen| 2013) Cohen 2016: Clarkson & Woodruif 2017)): for the general constrained case, the Iterative Hessian Sketch (IHS) was proposed by Pillanci &
Wainwright (2016) as an effective approach Cormode & Dickens (2019) employed sparse sketches
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 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 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 frequeett items (Hsu et al., 2019) and have subsequently been applied to a number of other problems on large data. For example, Indyk 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 the next iteration by training on a data set involving previously learned sketches from prior iterations?
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 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.
Our Contributions.
2

---

Under review as a conference paper at ICLR 2022
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 consider a new learning objective, that is, 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 significant advantage over non-lcarned sketches, and has a fast training time. Our experiments show that the convergence rate is reduced by 44% over the nonlearned COUNT-SKETS for the LASSO problem on a real-world dataset. Recall that a smaller convergence rate means a faster convergence.
classical extremely sparse sketch)
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 2a///(e2 (822) rows with failure probability 8, while our S requires only ((doololo (1/c) + log(1/8))/e²) rows if A has at most polylog(1//c/e rows of leverage score at least e/d. This is quadratic improvement in d and an exponential improvement in 6. Applying S to A runs in input-sparsity time and the resulting S'A may remain sparse if A is sparse. In practice, is not necessary 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 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, show empirically that such learned sketches reduce the convergence rate by larger 79.9% to 84.6% over non-learned sketches. Therefore, the learned sketches attain 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. (202)) and show that also for sketching- baseed 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 much better preconditioner with the same overall running time.
A Algorithm 1 LEARN- -KETCH: Gradient descent algorithm for learning the sketch values
PRELIMINARIES
Require: Anain = {A,)) (A; E RPxd), learn-
ing rate a
Notation. We denote by S the unit sphere in the n-dimensional Euclidean space I² For
1: Randomly initialize p.v for COUNT-
SKETCH-type sketch as described in the text
a matrix A E mm we denote by ||-A||,, op its operator norm, which is defined as ||AAllap =
2: for t= 0 to step do
3: Form S using p,v
supzes- - ||.2 We also denote by (A) 4: Sample batch Anatch from Atroin
and min (A) the largest and smallest singular 5: -- V a Oe
OC Asatcc)))))))
values of A, respectively, and by colsp(A) the column space of A. The condition number of A is defined to be x(A) = max (4)/mmin(4 (A).
Leverage Scores. We only consider matrices of full column rank"] Suppose that A E Ruxx (m 2 n) has full column rank. It has m leverage scores, denoted by Ti Ti),...... Tm (A), which are defined as (A) |e_ A(A A)- AA 12 where {e1,.... cm } is the canonical basis of R Equivalently, letting A = UEV be the singular value decomposition of A, where U E Rmmxn E, V E Raxn, we can also write T, (A) = lee UUT 1 = ||e_ U I2. which is the squared &g 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 - -mm where G E RR with i.i.d. NOO, 1) entries.
COUNT-SKETCH: Each column of S has only 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 11, each with probability 1/2. Further, the columns are chosen independently.
Sparse Johnson-LLiddenttran Lind--niinnssssss Transform (SJLT): is the vertical concatenation of independent COUNT-SKETCH matrices, each of dimension m/s / X n.
Thii can he assumed w.l.o.g. by adding artbirarily small random noise the input, or one can first quickly sketching to find subset of columns of maximum rank, and replace the inut with that subset of columns.
3

---

Under review as a conference paper at ICLR 2022
COUNT-SKETCH-byp Sketch. A COUNT-SKETCH-OE sketch is characterized by 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 x and Sp..: e2 for all 1 S i S 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, 1) only.
Subspace Embeddings. For a matrix A E Raxd we say matrix S E R= is (1 -ussuspacee embedding for the column span of A if (1 -) ||Ax||2 S ||SAz||2 S (1 + e) ||Ar||2 for all E IRd The classical sketches above, with appropriate parameters, are all subspace embedding matrices with probability at least 8: our focus is on COUNT-SKETCH which can be applied input sparsity
running time. We summarize the parameters needed for subspace embedding below:
Gaussian sketch: m ((dd log()/ It is a dense matrix and computing SA costs O(m mnz(44) = (0mnz(4)((
time.
COUNT-SKETCH: m = (dd (ōc²) (Clarkson & Woodrufi 2017). Though the number of rows is quadratic in d/e, the matrix S is sparse and computing SA takes only (mnz(A)) time.
SJLT: m = Ogg )c²) and has = O f)/c) non-zeros per column (Nelson & Nguyên Cohen 2016 Computing SA takes O(s (0snmz(AA) (A)) O(nnz(A) log( log())/6 f)c)) time.
2013
Iterative Hessian Sketch. The Iterative Hessian Sketching (IHS) method (Pilanci & Wainwright
2016) solves the constrained least-squares problem ① by iteratively performing the update
T++1 = arg min
2 1| - x, 222 - (A (b Art), - ma)
(2)
#EC
where S1+1 is sketching matrix. It is not difficult the identity matrix) of the minimization above, the optimal solution x²+1 coincides with the optimal solution to the constrained least squares problem ( The IHS approximates the Hessian A A by sketched version (S)+1A) AS)+1A) to improve runtime, as S A typically has very few rows.
that for the unsketched version (S1+1 is
Unconstrained Convex Optimization. Consider unconstrained convex optimization problem min, fz), where f is smooth and strongly convex, and its Hessian V is Lipschitz continuous. This problem can be solved by Newton's 's method, which iteratively performs the update
arg min || (V f(x)) (V fz.) V f(z)))
F++1 I2-
(3)
provided it is given good initial point To. In each step, it requires solving regression problem of the form min, |AA Aa- y|| 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 preconditioner R via QR decomposition such that 5AR has orthonormal columns, where S is a sketching matrix, then solves 2 arg min. [[(A)) (AR)z' - v|l, by gradient descent and returns RE in the end. Here, the point of sketching is that the QR decomposition of SAA can be computed much more efficiently than the QR decomposition of A, since has only a small number of rows.
Learning a Sketch. We use the same learning algorithm in (Liu et al. 2020 given in Algorithm !
The algorithm aims to minimize the mean loss function L(S, A) n >i L(S, A:), where S is the learned sketch, C(S, A) the loss function applied to data matrix A, and A (A is a (random) subset of training data.
3 LEARNING-AUGMENTED SUBSPACE EMBEDDINGS
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--E 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 I [ better subspace embedding implies a better convergence rate in the IHS, as well as for the subroutine in unconstrained convex optimization.
sketch, based
3.1 SKETCHED LEARNING: OPTIMIZING THE POSITIONS
In this section we consider the problem of embedding the column space of a matrix A E Rexd provided that A has a few rows of large leverage score, well as access an oracle which reveals
4

---

Under review as a conference paper at ICLR 2022
superser of the indices of such rows. Formally, let (A) denote the leverage score of the i-th row of A and let
r" [i: (A) 2
be the set of rows with large leverage score. Suppose that a superset I 2 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]] \ T. Up to permutation of the rows, we can write
A1 A/=
and S = 6 0
(4)
where 5' is random COUNT-SKETEN matrix of m rows. Clearly S has 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- Hanson-Wright Wright inequality.
Theorem 3.1. Let v = e/d. Suppose that O((d/e)) (1e) + log(1/o))), d E (0, 1/m and d = (2(1/e) polylog(1 /e) log (1/8).. Then, there exists distribution on S of the form in with m + I| rows such that Pr {vz E colsp(A), || - | > e =| 8.
Hence, if there happen be at most d polyloog(1/ (1/l)))og((/e//e" rows of leverage score least e/d, the overall sketch length for embedding colsp(A) A) can be reduced to ((d polylog(1/e)) (e))ll))()///oooo log(1 (og(() /e))
quadratic improvement in d and exponential improvement in d over the original sketch length of (d /( for COUNT-SKETCHS In the worst case there could be (dd /e) such rows, though empirically we do not observe this. The following is an immediate corollary, by setting d = 1/m.
Corollary 3.2. Suppose that d (2((1/ (c) polylog(1/ / and |T|= ((dde²) polylog(d/e)) with == e/d. There exists a distribution on S of the form in ( with O((d/e) polylog(d/e) (d/e)) rows such (((d//"))ooylo((()
that Pr {Va E colsp(A)), || - ||=|| > || <
We remark that our S is of the applied in ((nz((A) time. This is faster than chained subspace embedding of the form S2S,A, where S1 is a COUNT-SKETCH matrix of Odd /e²) rows and S2 is a subspace embedding matrix of (dde²2 rows. Computing S,A A takes ((nnz(A) time but computing (S) will take an additional time of poly(d/c) or O(nnz(S)4)) A) (o))))c). The latter terms can be quite large and even comparable to if say, n is close to dZ. 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, would yield a worse sparsity than a matrix of the COUNT- SKECH type. The sparsity of SAA is also important for solving regression problems involving
type, which has twofold benefit. First, SA can be
SA in intermediate steps, as algorithms such as conjugate gradient, which use matrix-vects products, become more efficient.
We note that approximate leverages scores of all rows can be found in time (A))(A) + poly(d/c) [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, is possible to learn the indices of the heavy rows. In practice, one can shrink the size of the superset by restricting / to the rows with large &g norms in Aj. We shall demonstrate in Section that this heuristic works well on some wallworld datasets.
3.2 SKETCHED LEARNING: OPTIMIZING THE VALUES
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, A,) = [[(A,R;) AA, R. Tl over all the training data, where R3 comes from the QRedecomposition of SA, = QRR.! 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
5

---

Under review as a conference paper at ICLR 2022
Lemma 3.3. Suppose that E (0, 2), S R=x", A Rnxd has full column rank, and SA QR is the QR-decomposttion of S.. If [|(AR AR- Tlop S c, then S is (1 + c)-subspace embedding of the column space of A.
Lemma 3.3 implies that if the loss function over Atrain is small and the distribution of Atest is similar to Aueain it is reasonable to expect that S is a good subspace embedding of Arest- Here we use the Frobenius norm rather than operator norm in the loss function because 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.
Algorithm 2 Solver for (5)
S + learned sketch, S. t random sketch
HESSIAN SKETCH
(Z,,1, Zi,2) ESTIMATE(SS, A), i 1,2
3: 1" + arg min;=1,2 (2.2/ /Z.1)
In this section, we consider the minimization prob- lem
4: I solution of (5) with S = S.-
5: return T
{ [ (A v.
function ESTIMATE(S,A) A)
min SEC
(5)
7:
T f- sparse (1 -subsspppecee embedding
matrix for d-dimensional subspaces
which is used as subroutine for the IHS (ef. 2 We present an algorithm with the learned sketch
8:
(Q, R) QR(TA))
9:
Z, mmnn (SAR-1)
in Algorithm 2 To analyze its performance, we define the following quantities (corresponding ex- actly to the unconstrained case in (Pilanci & Wain-
10:
2z x (1 + aaapppppooiiiiaainnnn to
[(SAR-)) -
-llppp
wright 2016))
11: return (2,, Zz)
inf (i) "cnnssp( A)
||
Zz(S)
sup w.
(S S In))).
When S (1 + s)ssusspace embedding of colsp( (A), we have (S) 2 c and Zz(55) S 26.
For a general sketching matrix S, the following is the approximation guarantee of 2, and 2g. which estimates of (S) and (2)),, 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 1 (0, 4) a small constant, A is of full rank and S has poly(d/n) rows. The function ESTIMATE(S, A) returns in (((nz((A) log 7
1 + +-ooo((y (2)) time Z,, Ze which with
probability at least 0.99 satisfy that (1(S) < 2 < 2s(S) and (5(5) (1+n)² 3n S 2g - (1-1)² 3n.
Z2(S)
1+t
1-n
Similar to Proposition 1 of (Pilanci & Wainwright, 2016), we have the following guarantee. The proof is postponed to Appendix[D]
Theorem 4.2. Let v) (0, 1) be small constant. Suppose that A is of full rank and S and S₂ are both which, with probability at least 0.98, satisfies that [|A(T - x²)112 S (1 + 1))a min Z.,- Ze.1
sketches with poly(d/n) rows. Algorithm 2
returns a solution
21,2 Z2.2
4n) |Axl|l2 in logg)()))og())) + poly( squares solution.
time, where x= = arg min nec || - bll is the least-
Theorem 4.2 suggests the following. If the ratio of the learned sketch Ze (5)) is a constant smaller than that of the random sketch Zz (52))/Z2(S)) and n is a constant fraction of the ratio gap, then / 2i,,, a constant smaller than T2.2/Z2.1 which means that the procedure of IHS will converge faster with the learned sketch. In particular, if S, is a (1 <.)-subspace embedding matrix for colsp(A) with 6 < 1/3 and 17 < min][e1 - 62). 61, E2) for some small constant
y > 0, we have (2. /Z < 3e, and the guarantee in Theorem 4.2 becomes || x" x')lla S
(minn(11, (2)) |Ax 112 that is, a better subspace embedding can lead to a faster convergence. Hence, if the learned sketch is a better subspace embedding than random sketch, theoretically can obtain a better convergence by setting 7) 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 Jai polylog(azz)))
subspace embedding using Corollary 3.2 which is almost Vdlfactor better than the usual guarantee

---

Under review as a conference paper at ICLR 2022
of a random COUNT-SKETCH matrix of the same dimension, leading to algorithm of faster Algorithm Fast Regression Solver for (6)
convergence.
1: S < learned sketch, S2 random sketch
2: (Qi, R3) + QR(S,A); 112
5 HESSIAN REGRESSION
(ar, of) + EIG A AR, ) ; 1, 1,2 EIG(B) returns
estimates of max ( and minn (B)
In this section, we consider the minimiza- tion problem
4: 1 + 1
P + R.-
min AT A= 1/ v/12-
(6)
11 1/ + (o(.)))
7: 20 -0
which is used as subroutine for the un- constrained convex optimization problem minz f(x) with A A being the Hessian
while ||ATAP: y11 e 11lllp do
9:
21+1 + 2 - 1 AAP) -
matrix (xx) (see (3). Here A E Raadd
10: return P2t
E IRd and we have access to A. We incorporate learned sketch into the fast regression solver in (van den Brand et al. 2020) and present the algorithm in Algorithm 3
Here the subroutine (I) applies (1 + -ussspaee embedding sketch T to B for some small constant y and returns max (T)) and T (T Since admits the form of AR, the sketched matrix TB can be calculated as (TA) A) R and thus can be computed in ((nz((A) (A) )) + poly(d)) time if
i's a COUNT-SKETCH SVD or the Lanczos algorithm.
matrix of (cd)) rows. The extreme singular values of T B can be found by
Similar to Lemma 4.2 in (van den Brand et al. 2020), we have the following guarantee of Algorithm 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-
sketches with ((d²) rows.
Algorithm returns a solution x' such that |A Ax' vll S c llyll, with probability at least 0.97.
The runtime is (Anz((A) + Ö(dd (min(a1/o), 02/0//)/))) (og(((A)))) - poly(d)).
Remark 5.2. In Algorithm S2 can be chosen to be a subspace embedding matrix for d-dimensional
subspaces, in which case, A has condition number close to 1 (see, e.g., p38 of (Woodruff) 2014)
and the full algorithm would run faster than the trivial ((nd))-time solver to (6).
Remark 5.3. For the original unconstrained convex optimization problem min, ((r), 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 -eastes fz) = Az - blß and the value of f(x) can be approximated efficiently by a sparse subspace embedding matrix in (nnz(A) + nnz(b) + poly(d)) time.
6 EXPERIMENTS
Comparison. We compare the learned sketch against three classical sketches: Gaussian, COUNT- SKETCH, and SJLT (see Section (2 in all experiments. The quantity we compare is 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 1.90GHz CPU and 16GB RAM. The offline training is done separately and the training of 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 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 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 COUNT- SKETCH 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 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 prror-versu--rnning plot for the task of matrix estimation with nuclear norm constraints.
7

---

Under review as a conference paper at ICLR 2022
6.1 IHS EXPERIMENTS: LASSO
We define an instance of LASSO regression to be:
x' arg min ||Az 6112
|
where À is a parameter. We use two real-world datasets:
Electri residential electric load measurements. Each row of the matrix corresponds to a different residence. Matrix columns are consecutive measurements from different times.
A; |R370x9, by E R370x1, and ((rin) = 320, ((A, ()sst 80. We set À 155.
Greenhouse gas (GHG 31 time series of measured greenhouse gas concentrations in the Califor- nia atmosphere. Each b) corresponds to different measurement location. A E 3272 b; E 12 and (AA, b)rrain| = 400, [(A, b) test| = 100. We set À = 30.
Experiment Setting: We choose = 6d, Sd, 10d for both datasets. We consider the error ||Ar 2 For the two datasets, use both the methods proposed in Sec- tion 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 specific time point and the heavy rows not very concentrated. Nevertheless, 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 lz norms, subject to the dimension budget. This will incur an additional computational cost, but the time 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
||Az - 6112
We plot a logarithmic scale the mean errors of the two datasets in Figures [ 2 We all methods display linear convergence, that is, letting eg denote the error in the k-th iteration, we have eg Re p²e1 for some convergence rate p. A smaller convergence rate implies a faster convergence.
Experimental Result:
We calculate an estimated rate of convergence p p= (ice//e)2// with k = 10 for the GHG dataset, and with & = 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 and the heavy-rows sketch has a convergence rate that is 86.9%. When the sketch size is large (m 10d), the brased learned sketch has convergence rate that is 63.7% 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- has a convergence rate that is just 21.1% 1% of that of sparse JL, and when the sketch size is large, the combined-learee sketch has smaller convergence rate that is just 15.4%.
sketch
We also conducted IHS experiments for the matrix estimation problem with a nuclear norm constraint in Appendix F
CUETHINSS
Figure 1: Test error of LASSO in the Green House Gas dataset.
8

---

Under review as a conference paper at ICLR 2022
curvied
Figure 2: Test error of LASSO in Electric dataset.
2818 PTEE (2)
00-
1.5
65
64
02
0.8
106
12
356
175
200
335
250
275
200
1.0
12
150
135
100
335
256
235
300
1.0
150
135
220
Figure 3: Test error of the subroutine in fast regression on Electric dataset.
6.2 FAST REGRESSION EXPERIMENTS
We consider the unconstrained least squares problem ming f(x) with f(x || bll using the Electric dataset.
Training: Note that f=x = AA A, independent of x. In the f-th round of Newton's 's method, by G). we need to solve (xt).
a regression problem min, || Aa - v|l with Hence, we can use the same two methods in the preceding subsection to optimize the learned sketch S. For a general problem where 2 f()) depends on x, one can take de to be the solution obtained from Algorithm using the learned sketch S, to generate A and for the (t + 1)-st round, train a learned sketch S.+1+ and repeat this process.
Figure 4: Test eπror of fast regres- sion on Electric dataset
For the Electric dataset, we set m 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 1 in Algorithm 3 set 77 in all iterations for heavy-rows sketches. For the classical random sketches, we set in the following two ways: (a) ry = in all iterations and (b) 1 in the first iteration and 7 0.2 in all subsequent iterations.
Setup for Experiments:
Experimental Results: We examine the accuracy the subproblem and define the error to be A AR U||, llyll2- 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 terrible divergence of the subroutine and is thus omitted in the plots. Still, observe that in setting (a) of 7, the other two classical sketches cause the subroutine to diverge. In setting (b) of 7,, the other two classical sketches lead to convergence but their error 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 less than 0.01 in all iterations all three subroutine calls, in both setting (a) and (b) of n).
We also plot 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 80.6% of that using Gaussian or sparse JI. sketches.
CONCLUSION. We demonstrated the superiority of using learned sketches over classical random sketches, for the Iterative Hessian Sketch method which 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 oracle for predicting 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 significant advantage over nea-led sketches for problems in convex optimization.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Akshay Agrawal, Brandon Amos, Shane T. Barratt, Stephen P. Boyd, Steven Diamond, and Zico Kolter. Differentiable convex optimization layers. In Advances Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019,
December 8-14, 2019, Vancouver, BC, Canada, pp. 9558-9570, 2019.
Marein Andrychowiez, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. Advances in neural information processing systems, pp. 398 1-3999, 2016.
Jean Bourgain, Sjoerd Dirksen, and Jelani Nelson. Toward a unified theory of sparse dimensionality reduction in Euclidean space. Geometric and Functional Analysis, pp. 1009-1088, 2015.
Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, 2004.
Kenneth L. Clarkson and David P. Woodruff. Low-rank approximation and regression in input sparsity time. ACM, 63(6), January 2017. ISSN 0004-5411. doi: 10. 10.1145/3019134 URL https//0ob.or//1011114//11113331
Michael B. Cohen. Nearly tight oblivious subspace embeddings by trace inequalities. In Proceedings of the Twenty-Seventh Annual ACM-SIAM Symposium on Discrete Algorithms, SODA '16, pp- 278 287, USA, 2016. Society for Industrial and Applied Mathematics. ISBN 9781611974331
Graham Cormode and Charlie Dickens. Iterative hessian sketch in input sparsity time. In Proceedings of 33rd Conference on Neural Information Processing Systems (NeurIPS). Vancouver, Canada, 2019.
Nikita Doikov and Peter Richtárik. Randomized block cubic Newton method. In Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018, pp. 1289-1297 2018
Yihe Dong, Piotr Indyk. Ilya P. Razenshteyn, and Tal Wagner. Learning space partitions for nearest neighbor search. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020, 2020.
Robert M. Gower, Donald Goldfarb, and Peter Richtárik. Stochastic block BFGS: squeezing more curvature out of data. In Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, pp. 1869-1878. 2016.
Robert M. Gower, Filip Hanzely, Peter Richtárik, and Sebastian U. Stich. Accelerated stochastic matrix inversion: General theory and speeding up BFGS rules for faster second-order optimiza- tion. In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, 3-8 December 2018, Montréal, Canada, pp. 1626 1636, 2018
Robert M. Gower, Dmitry Kovalev, Felix Lieder, and Peter Richtárik. RSN: randomized subspace Newton. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NewrIPS 2019, 8-14 December 2019, Vancouver, BC, Canada, pp. 614-623, 2019.
Chen-Yo Yu Hsu, Piotr Indyk, Dina Katabi, and Ali Vakilian. Learning-based frequency estimation algorithms. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, 2019.
Piotr Indyk, Ali Vakilian, and Yang Yuan. Learning-based low-rank approximations. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, 8-14 December 2019, Vancouver, BC, Canada, pp. 7400-74100 2019.
Sudhir B. Kylasa, Fred (Farbod) Roosta, Michael W. Mahoney, and Ananth Grama. GPU accelerated sub-sampled Newton's 's method for convex classification problems. In Proceedings of the 2019 SIAM International Conference on Data Mining, SDM 2019, Calgary, Alberta, Canada, May 2-4, 2019, pp- 702-710, 2019.
10

---

Under review as a conference paper at ICLR 2022
Xiang Li, Shusen Wang, and Zhihua Zhang. Do subsampled newton methods work for high- dimensional data? In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York NY, USA, February 7-12, 2020, pp. 4723 4730. AAAI Press, 2020.
Simin Liu, Tianrui Liu, Ali Vakilian, Yulin Wan, and David Woodruff. On learned sketches for randomized numerical linear algebra. 1xiv:20007009800 [cs.LG], 2020. URL https//arxiv.org/abb//
Nelson and H. L. Nguyên. Osnap: Faster numerical linear algebra algorithms via sparser subspace embeddings. In 2013 IEEE 54th Annual Symposium Foundations of Computer Science, pp. 117-126, 2013.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zachary Vetoto, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performange deep learning library. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 8024 8035, 2019.
Mert Pilanci and Martin J. Wainwright. Iterative Hessian sketch: Fast and accurate solution approxi- mation for constrained least-squares. J. Mach. Learn. Res., 17:53:1-53:38, 2016.
Mert Pilanci and Martin J. Wainwright, Newton sketch: A near linear-time optimization algorithm with linear-quadratic convergence. SIAM J. Optim., 27(1):205-245, 2017.
Farbod
and Michael W. Mahoney. Sub-sampled Newton methods. Math. Program.,
174(1-2):233-36, 2019.
T. Sarlos. Improved approximation algorithms for large matrices via random projections. In 2006 47th Annual IEEE Symposium on Foundations of Computer Science (FOCS'06), pp. 143-152, 2006.
Jan van den Brand, Binghui Peng, Zhao Song, and Omri Weinstein. Training (overparametrizeed) neural networksin near-linear time. arXiv:2006. 1116..166666 1648 [cs.LG], 2020.
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. In Yonina C. Eldar and Gitta Kutyniok (eds.), Compressed Sensing: Theory and Applications, pp. 210-268. Cambridge University Press, 2012. doi: 10. 10.1017/CCB077801111940 .006.
David P. Woodruff. Sketching tool for numerical linear algebra. (0(1-)):1157, October 2014. ISSN 1551-305X. doi: 10. 56//04000000000 URI. https:://ooi.rr////// 100. 1561 /0400000060
Peng Xu, Jiyan Yang, Farbod -hostaKKooaaaai. Christopher Ré, and Michael W. Mahoney. Sub- sampled Newton methods with non-uniform sampling. In Advances in Neural Information Process- ing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pp. 3000-3008, 2016
Peng Xu, Fred Roosta, and Michael W. Mahoney, Second-order optimization for non-convex machine learning: empirical study. In Proceedings of the 2020 SIAM International Conference Data Mining, SDM 2020, Cincinnati, Ohio, USA, May 7-9, 2020, pp. 199-207, 2020.
11

---

Under review as a conference paper at ICLR 2022
A
PROOF OF THEOREM 3.1
First we prove the following lemma.
Lemma A.l. Ler 8 (0, 1/m]. holds with probability least ö that
||Sx||2 - |1 S 1
sup #Ecolsp(A)
provided that
m 2 ((d log m) min (d/E), (in()o log" m) + llog(1/))),
12 E v m) min (d/e), log m) log(1/6)) (1g()
Proof. We shall adapt the proof of Theorem 5 (Bourgain et al. 2015) to our setting. Let T denote the unit sphere in cAlsp(A) and set the sparsity parameter | || l1 and it suffices to show that
1. Observe that ||x1] 2
Pr Pr(lls'' x 1
for x E T. We make the following definition, as in (2.6) of Bourgain al. (0015)
A
dszzze, 8 e,,
jete
and thus, S'x A a. Also by || xr
| 1| one has
sup zET
| 15'21=111 -| ||
sup TET
| E
A
(8)
Now, in (2.7) of Bourgain et al. (2015) we instead define a seminorm
1/2
|
Σ 8e
Then (2.8) continues to hold, and (2.9) as well as (2.10) continue to hold the supremum in the left-hand side is replaced with the left-hand side of (8). At the beginning of Theorem 5, we define (+)) to be U, but each row j E I is multiplied by 5.j and each row j is zeroed out. Then we have the first step (4.5) that
Σ JET
8ij
Σ
ge eg)
instead of equality. One verify that the rest of (4.5) goes through. It remains true that ||=laa < (1/v8) 8) |- and thus (4.6) holds. One can verify that the rest of the proof of Theorem 5 in Bourgain et al. (015) continues to hold if we replace Σj=1 with EjEr= and maxisisn with maxjeze, noting that
Σ 5is |EEe JEZ
8 dd 8 [(PEe,, eg) S m
m JEl=
and
E(10))0 Σ (Ed uuj 1
1 m
JEL
Thus, the symmetrization inequalities on
Σ
8ig |
and
Σ
S
j
DE
13
continue hold. The result then follows, observing that maxjer= |PEE}) v.
12

---

Under review as a conference paper at ICLR 2022
The subspace embedding guarantee now follows as corollary.
Theorem 3.1 Let = e/d. Suppose that (((dd/)) (polylog(1/e) + log(1/o))), 8 E (0. 1/m)
and d = ((1 /e) polylog(1 log (1/6)) Then, there exists distribution on S with m + rows
such that
Pr Pr(ve VV E colsp(A), ||5= | >
||
S 6.
Proof. One can verify that the two conditions in Lemma A. are satisfied if
d
1 [polylog( d + log 8
m
1 2 log 1o8g)) 2) [polylog(?
d - log 1 5)
The last condition is satisfied if
log² 8 ;) polylog (:)
B PROOF OF LEMMA 3.3
Proof. On the one hand, since SAR is orthogonal matrix, have ||ا||| ||Q=122 ||AARll2.
(9)
On the other hand, the assumption implies that
(ARR)) (ARx) - x² 20|| S < <|11=11
that is,
(1-c) ||01| 2 ||ART|| (1++)) |
(10)
Combining both ( and (10) leads to
||ARlll S ||ARx| S VI SVI+C||SARMN Ya Rd
Equivalently, it can be written as
|AAll2 < Ayll2 <
VT |SAu|12, EyERd
V++
The claimed result follows from the fact that 1/vI+t2 1 and 1/VI- + whenever
E (0, -5-1
L
C PROOF LEMMA
Suppose that AR- UW. where U Rexd has orthonormal columns, which form orthonormal basis of the column space of A. Since T' is subspace embedding of the column space of with probability 0.99, holds for all Rd that
TAR 2112 - $|AAR S 1 17 |TA
1 1+n
Since
|TAR 1|Qx112 ||
and
||>al|_ |[Wllly ||AR 1 20112
(11)
have that
1
1
| <
Wa - 1- 1
||ا|ا x E Rd.
(12)
+ n
13

---

Under review as a conference paper at ICLR 2022
It is easy to see that
|SU yll2
Zi(5)= = min
||SUx|1 = min
ZE8
y#0
| Wlll
and thus,
|SUW yl1
min(1 +1) +
||SUWals
min(1 11)
< Z)()
M
1lyl12
840
1//12
Recall that SUW = -AR-- We see that
(1 mi)
(S (SAR (IS) - (1 + 7) (
(SAR
By definition,
Zz(S) = (sT - 1,)
It follows from 12 that
(1 1) W - w.) op S Zz(S) S (1+ 7)
and from and Lemma 5.36 of Vershynin (2012) that
[AR (AR 1 -llop - 3m.
Since
US S op = 1| -
AR - op (AR [[AR
and
TT ST - ((AR)))(((AA)))) 1 op
< [(AAT)) - -11.p
STAR T 11.op + 1] (AR 1llop'
follows that
(1-1)² ( ((SAAR))TSAR"* - 11l -3(1 - -
S (2)
(1 + 7)² (|(SAR SARR-- -
-lpp + + + 7)
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 ((nnz(A)) time. The QR decomposition of T', which a matrix of size poly(d/n) (dn)) d, can be computed in poly(d/n) time. The matrix SAR-1 can be computed in poly(d) time. Since it has size poly(d/n) d, its smallest singular value can be computed in poly(d/n) time. To approximate (2)),, we can use the power method to estimate 1 SAR - -llopp up to a a1 = -facttt O(nnz (A) + poly(d/n)) log(1/1)) time.
D PROOF OF THEOREM 4.2
In Lemma 4.1 we have with probability at least 0.99 that
2
(+n) Zo (2( -3
37
1- Zz(S)
37
2 n1+ (i())
1 (s)
Zi(
When S is random subspace embedding, it holds with probability at least 0.99 that (1) 2 3/4 and so, by a union bound, it holds with probability at least 0.98 that
%,2 1
7z(S)
2
4n-
2, (1 1)) Zi(S)
or,
Z2(S) Zi(S)
[Zz
< (1 + n)
4n
2,
The correctness of our claim then follows from Proposition 1 of Pilanci & Wainwright (2016). together with the fact that S, is a random subspace embedding. The runtime follows from Lemma 4.] and Theorem 2.2 of Cormode & Dickens (2019).
14

---

Under review as a conference paper at ICLR 2022
E PROOF OF THEOREM 5.1
The proof follows a similar argument to that in (van den Brand al. 2020 Lemma B. In van den Brand et al. 2020)), it is assumed (in our notation) that 3/4 S mmin(AP) I max (AP) I 5/4 and thus one can set 7 = in Algorithm 3 and achieve a linear convergence. The only dif- ference is that here estimate m (APP and Fmax (APP and set the step size 1j in the gradi- ent descent algorithm accordingly. By standard bounds for gradient descent (see, e.g., p468 of Boyd & Vandenberghe (2004)), with choice of step size 17 = 2/ (AP)) + o miPn) after O[(xmax AAP) miin A))) (og( iterations, we can find 2, such that
|| A AP( - -=))=2 S < A AP(=0 - ==) 12)
where =" arg ming
PT A AP= pTyl| is the optimal least-squares solution. This establishes Eq. (11) in the proof in van den Brand et al. (020), and the rest of the proof follows as in there.
We use three subspace embeddings here, S Sg and one used in the EIG subrountine. Each subspace embedding uses O(d²) 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 ||AX BIE x42
it is reasonable to model the matrix X as having low rank. Similar to the é-minimmization for compressive sensing, a standard relaxation of the rank constraint is to minimize the nuclear norm of X, defined as |X|,_ == E (d), a, (X), where (X) is the j-th largest singular value of X.
Hence, the matrix estimation problem we consider here is
X = arg min ||AX- BI? such that |X||, S
xd2
where p > 0 is a user-defined radius as a regularization parameter.
We conduct experiments the following datasets:
Tunnel" 4 The data set is a time series of gas concentrations measured by eight sensors in
wind tunnel. Each (A, B) corresponds to a different data collection trial. A, E R 13530xx5 B, E
R 11353066 I(A, ()train train 144, (A,, Iest 36. In our nuclear norm constraint, we set
0 0= 10.
Experiment Setting: We choose ||AX - B12 || Bll The leverage scores of this dataset are very uniform. Hence, for this experiment we only consider optimizing the values of the non-zero entries.
7d, 10d for the Tunnel dataset. We consider the error
Results of Our Experiments: We plot in logarithmic scale the mean errors of the datasets in Figures 5) We can see that when m 7d, the gradient- sketch, based on the first 6 iterations, has rate of convergence that 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 learned sketch matrices S all COUNT- SKETCH-type matrices (each column contains single nonzero entry), the matrix product SA can thus be computed in (nnz(AA) time and the overall algorithm is expected be fast. To verify this, we plot in an veror-versu--une runtimm plot for matrix estimation with nuclear norm constraint tasks with m 10d in Figures/o (correspondn to the datasets 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-SKETCII have the fastest runtimes, which are slightly faster than that of the SJLT and significantly faster than that of the Gaussian sketch.
15

---

Under review as a conference paper at ICLR 2022
guamtton Sparso
Sp
0-
10
haratton round
Daration round
Figure 5: Test error of matrix estimation with nuclear norm constraint on Tunnel dataset
lesmed
count-beitch
goussian
spargs
Inporoo
2-
0.1
0.2
0.3 rantime
0.4
0.5
Figure 6: Test error of matrix estimation with nuclear norm constraint on Tunnel dataset
G HEAVY LEVERAGE SCORE Rows DISTRIBUTION OVER THE DATASET
In our experiments, hypothesize that in rerldworld data that there may be an underlying pattern which can help us identify the heavy rows. In the Electric dataset, each row the matrix corresponds specific rows; in the
to specific residence and the heavy rows are always concentrated GHG data set, each row corresponds to specific time point and we can select some specific time points be superset of the heavy rows and then select the heavy rows based on their f--norm in this superset.
To exemplify this, we study the heavy leverage score rows distribution over the Electire dataset. For a row i E [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 is heavy if & 2 5d/n. Below we list all 74 pairs (i, fi) with ft >>0 O.
(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), 1462, (356,129). (363,100,, (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), 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), 1208,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 f: 2 50. We view this strong evidence for 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[] [] 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() define our loss function). For more
16

---

Under review as a conference paper at ICLR 2022
nuanced loss function, which may be beneficial, one use the package released in Agrawal et al. (2019), where the authors studied the problem of computing the gradient functions which involve the solution to certain convex optimization problem.
As mentioned 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 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 f1...., ,m) for each coordinate of p. For the heavy rows mentioned in Section 3.1 we can allocate positions I,..., k to the k heavy rows, and for the other rows, we randomly sample an integer in {k + ,,....m.. We note that once the vector p. which contains the information the nonzero position in each column of S, is chosen, it will not be changed during the optimization process in Algorithm
Next, we introduce some parameters for our experiments.
bs: batch size, the number of training samples used in one iteration.
Ir: learning rate of the gradient descent(the in Algorithm 1
iler: the number of iteration for Algorithm I
In our experiments, we set bs 20, iter 1000 for dataset. We set Ir 10 for the Green House Gas dataset and Ir 0.1 for the Electric dataset.
I ADDITIONAL EXPERIMENTS FOR LASSO
In this section, consider the IHS experiments for LASSO data of a larger size. The experiment setting is the same as that in Section We conduct our experiments on the following dataset:
Gas Sensor chemical detection platform composed of 8 chemoresistive gas sen- sors was exposed to turbulent gas mixtures generated naturally in wind tunnel. Each matrix represents the measurements dense time points during short time period. A 1895000x 19 b E I895000x1, and (A, b)rrain| 30, (A, b)rsal 9. We set À = 10.
The results are shown in Table Here we choose m = 300. The leverage scores of the rows on this dataset is very uniform hence choose random positions for the nonzero entries and only optimize the values in the learned sketch. For matrix of such size, the gaussian sketching matrix is extremely only consider the Count-Sketch matrix and the Sparse-JL matrix. From the table
slow, hence, below we can see that the gradient-basee -aseed learned sketch has converge rate that is 74.6% of that of the random sketch.
Table I: Error of the Sketch Matrix on Gas Senser data
lteration
5
7
8
10
Learnd(value-onm -value-only)
428.13
23.72
1.72
0.092
0.0060
0.00036
10-5
Count-Sketch
1864.66
122.26
12.30
1.46
0.074
0.013
0.00036
Sparse-JL
1897.32
188.31
9.60
1.30
0.16
0.0098
0.00048
17


Reviews:
[{"review_id": "MSS3GYjGdoG", "reviewer": "Reviewer_zU27", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths:\nThe paper is well written and gives a nice background of existing literature with strong motivation. The technical detail is of good quality and the numerical experiments provided show that the method performs well on selected data sets when comparing number of iterations required. The idea is simple and intuitively appealing. The problems addressed, specifically least squares, are ubiquitous and finding efficient ways to solve high dimensional problems is of perennial interest. In addition, error/complexity bounds are provided when used in Iterative Hessian Sketch and Fast Regression.\n\n\n## Weaknesses\nThe paper claims to provide a framework for learned sketching that applies to a large number of problems in convex optimization but still focuses on least squares problems. It seems that the primary contribution of this paper is to apply learned sketches for least squares problems as introduced in works by Liu et al (_On learned sketches for randomized numerical linear algebra [2020]_ and _Learning the positions in countSketch [2020]_) to iterative hessian sketch and fast regression. Although theoretical bounds are provided, the contribution is more limited than initially stated. \n\n\n\n## Questions and comments:\n1. Is it common that the Hessian can be decomposed as $A^T A$ with $A \\in \\mathbb{R}^{n \\times d}$ and $n \\gg d$ for other convex problems? This is common for LS but is it often observed elsewhere that makes it more useful for non LS problems.\n2. For learning the locations of the non-zero entries as discussed in 3.1, how is the oracle trained to predict these heavy rows? Are the number of occurrences of heavy rows tabulated and the the $k$ most common rows to have leverage scores over some threshold selected? More detail on this front would be helpful. \n3. Although the sketches can be learned offline and finished \"within 5 minutes\", it is unclear how the sketch training time compares with the time to solve the problem. A comparison of training time and time to convergence would be helpful to evaluate its merit.  \n\n\n\n## Minor comments\n- $p \\in \\mathbb{N}^n$ rather than being a real vector\n- The notation $|(A_i, b_i)|$ isn't entirely clear. It clear indicates the sixe ", "summary_of_the_paper": "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. \n\nThe 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.", "summary_of_the_review": "Since the paper primarily focuses on the application of a learned sketch to an iteratively solved LS problem (both components well established elsewhere) the contribution seems marginal. I believe the paper is marginally below the acceptance threshold as is. The paper can be improved by addressing questions and comments above, in particular, evidence that the method is applicable to more general convex problems and a more comprehensive comparison for total time to solve compared to naive sketching.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "KNUuUpBoy5O", "reviewer": "Reviewer_SAjC", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "My background is more in theoretical matrix algorithms, so to me, the idea of adaptively generating sketches based on leverage scores is somehow more fundamental than sketching. The authors do do a good job addressing this related literature, but their results are mostly direct consequences of matrix concentration bounds.\n\nThe experimental results are interesting: I'm not aware of such evaluations of test errors in previous works. They also point to significant gains of these methods over oblivious sketching based methods. On the other hand, these studies mostly take place on small to moderate sized data where the performance gains from sketching is unclear. So I feel they serve mostly as proof of concept of the utility of such sketching, and perhaps also point out that different formulations/objectives are necessary for giving better test errors.", "summary_of_the_paper": "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.", "summary_of_the_review": "I feel this paper is well thought out, and the experimental results are highly interesting. However, I also feel that the studies taken here (adaptively select rows) can be taken one step deeper: in particular, it would be very interesting to see theoretical analyses of why such learned sketching give improved results (over oblivious sketches).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "JsIxqefbCYD", "reviewer": "Reviewer_kP9y", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Strengths\n\nS1. Learning of sketch matrices is an interesting twist to the more standard random sketches out there.\n\nS2. The theory for the learned heavy row sketch is a nice attempt at providing justification for why the proposed learning method might improve results. \n\n# Weaknesses\n\nW1. It is unclear how useful these learned sketches actually are. In applications, it seems like the learned sketches could easily break. To put it in machine learning terms, the heavy row sketch uses row indices as features to predict whether or not rows are heavy. This is somewhat akin to classifying images on a social networking site based on the ID of the user who uploaded them: if the user with that ID uploaded mostly cat pictures in the training data, then classify the newly uploaded picture as a cat picture. This means that if the rows of a design matrix in the testing data are permuted, then the learned heavy row sketch isn’t useful anymore. This could for example happen in the electricity dataset if two residents switch homes. It therefore seems like, in practice, you would need to retrain the model frequently to ensure that the leverage score estimates remain accurate.\n\nW2. The datasets in the experiments are very small scale, and can be solved very easily using deterministic methods. It would have been more interesting to see the performance on larger scale datasets where these kind of techniques would potentially be more useful. \n\nW3. The improvements that the learned sketches yield don't seem that substantial in some of the experiments. For example, in Figures 1 and 4-6, it seems like they just save a few iterations compared to doing standard CountSketch.\n\nW4. Some parts of the paper are unclear; see questions below.\n\n# Questions\n\nQ1. How is line 10 in Alg. 2 computed?\n\nQ2. Below Theorem 4.2, you discuss how “a better subspace embedding can lead to a faster convergence.” But it’s not clear from the discussion why this is the case. In particular, it’s not clear why the discussion around improved accuracy guarantees would lead to improved convergence rates. What am I missing?\n\nQ3. Out of curiosity, how do you implement the various CountSketch matrices? Do you form sparse matrices in Python with the relevant nonzeros and multiply, or do you do something more refined to make the computation faster?\n\nQ4. The datasets are poorly introduced. What does $b$ represent in the Electric and GHG datasets? How is the data in $A$ in the GHG dataset organized---is time along the rows with different columns corresponding to different measurement sites? \n\nQ5. Also, does |$(A,b)$_train|=400 and |$(A,b)$_test|=100  mean that you have a total of 500 pairs of design matrices and corresponding vectors $b$, and that you use 400 for training and the rest for testing? If yes, are all the reported results the average results over all pairs $(A,b)$ in the testing data? \n\nQ6. It’s not clear why some methods are not given in some plots. Why is “learned(combined)” not plotted in Figure 1 when it appears in Figure 2? Why are “learned(value-only)” and “learned(combined)” not included in Figure 3? Why are “learned(value-only)”, “learned(combined)” and CountSketch not included in Figure 4?\n\nQ7. What is the difference between the three plots in Figure 1? Please explain in the caption. Also, the y-axis labels are cut off and not showing properly. \n\nQ8. In Section 6.2, why are you using Newton’s method to solve a least squares problem? Is there any benefit to doing that rather than just doing a single sketch-and-solve on the least squares problem itself?\n\nQ9. In the last paragraph before the conclusion, you say that you “only run one iteration” for each subroutine. What does that mean? It looks like you do multiple iterations in Figure 4?\n\nQ10. One of the questions you mention in the introduction is “(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?” You then say that you will answer this question, but I didn’t see this discussed in the paper. You briefly mention it in the paragraph “Training” in Section 6.2, but I didn’t see it properly addressed anywhere. Did I miss something? \n\n# Other minor things\n\n- In footnote on page 3, “inut” should be “input”.\n\n- The “Conclusion” heading is too big for a paragraph heading. You should put it as a proper section heading.\n\n- You should specify the unit used in the “runtime” label on the x-axis of Figure 6. I’m assuming it’s seconds?\n", "summary_of_the_paper": "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. ", "summary_of_the_review": "The core idea of learning sketch matrices is interesting. The theory that provides a potential justification for why learned heavy row sketching might improve on the oblivious CountSketch is welcome. However, the “learned oracle” for the heavy row sketch seems very easy to break: Simply permuting rows would mean that it has to be relearned. Overall, it’s not clear how useful the learned sketch would be in practice for that reason. The experiments are all done on very small problems that don’t need sketching. Parts of the paper are also unclear as outlined above. For these reasons, I think it’s below the acceptance threshold. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "JTb_6xra10g", "reviewer": "Reviewer_tMi8", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Pros:**\n\n- This paper shows provable speed-ups on the second-order methods introduced in previous works, and further demonstrates with experiments.\n- The previous works are properly surveyed and discussed in the context of the proposed methods.\n\n\n**Cons, and questions to authors:**\n\nDespite the above contributions, the writing is ambiguous and does not properly draw relationships among different parts, especially among Section 2, 3 and 4. I would increase my score if the authors could properly address the doubts. \n\n- How does COUNT-SKETCH introduced in Section 3.1 relate to the methods proposed in Section 4 and 5? It seems both Section 4 and 5 rely on the QR method in Section 3.2; the results in 3.1 are more of standard leverage score sampling results and are irrelevant to the proposed approaches.\n- How can the matrix $T$ be quickly chosen in Algorithm 2 Line 7, and what is the relationship between $T$ and $S_i$, and $T$ and $A$? It is not clear from the writing why such a $T$ is readily available. \n- Cost of each iteration: at the bottom of Page 7, the authors claim that “no additional computational cost is incurred in generating $S$ other than solving the iteration step using COUNT-SKETCH”, and Guassian matrices and sparse JL transforms \"will be considerably slower\". However, the iteration step itself requires running Algorithm 2 or 3, each of which includes QR on the original $A$ matrix (Algorithm 2 Line 8, Algorithm 3 Line 2), which are also expensive. I would suggest the authors compare per-iteration complexities of the approaches shown in Figure 1. Besides, it is not clear how Algorithm 2 Line 10 is implemented.\n- Section 3.2 at the bottom of Page 5 says “We found empirically that not squaring this loss function works better than squaring it”. Why does (not-)squaring matter if $S$ is solved by just minimizing the loss $\\mathcal{L}(S, A_i)$ itself?\n\n\n**And some minor issues:**\n\n- The notations at the bottom of Page 5 and in Lemma 3.3 at the beginning of Page 6 are inconsistent: Page 5 says $SA_i = Q_i R_i^{-1}$, whereas Page 6 says $SA = QR$. Also, it would be better if the authors could expand the $\\mathcal{L}(S, A_i)$ term at the bottom of Page 5 to explicitly show how S goes into the loss. \n- What are the $A_i$, $b_i$s in Section 6.1 below Equation 7, are they just $A$ and $b$ (the features and labels of the corresponding datasets)?\n- It may be nice to show these numerics: the convergence of first-order methods in FIgure 1-4 as a baseline, and how the spectrum of $A$ looks like. ", "summary_of_the_paper": "This paper shows sketching methods to speed up iterative Hessian sketch (IHS) and Hessian regression. With decently accurate subspaces learned from QR decomposition, IHS and Hessian regression achieve faster convergence than data oblivious methods.", "summary_of_the_review": "As mentioned above, despite the contributions, the writing (if not the method itself) is often confusing. I would suggest the authors address these ambiguities in both the author response and the revised version.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a new contribution in the recent literature on learning distributions of sketches. While all reviewers have recognized the overall good quality of the presentation, two factors seem to weight heavily on a negative decision: clarifications on the contribution's scope (presented as a tool for general Hessians in the introduction, but ultimately only applied to least-square errors of linear predictors, to recover an explicit factorization of the Hessian matrix) and links with existing literature; weakness of experiments whose small scale does not justify using sketches in the first place. Since this is a "learning" approach, I am particularly sensitive to the latter point, and therefore am inclined to reject, but I encourage the authors to address these two issues with the current draft.
Decision: Reject