Learning-Augmented Sketches for Hessians
Authors: Yi Li
Honghao Lin
David Woodruff
Venue:  (2021)

Abstract: We study learning-based sketching for Hessians, which is known to provide considerable speedups to second order optimization. A number of works have shown how to sketch or subsample the Hessian to speed up each iteration, but such sketches are usually specific to the matrix at hand, rather than being learned from a distribution. We extend such schemes to learned sketches, where we learn different potentially different sketches for the different iterations, and show empirically that learned sketches, compared with their "non-learned" counterparts, improve the approximation accuracy for a large number of important problems, including LASSO, SVM, and matrix estimation with nuclear norm constraints. 

Under review as a conference paper at ICLR 2021
LEARNING-AUGMENTE SKETCHES FOR HESSIANS
Anonymous authors Paper under double-blind review
ABSTRACT
Sketching is dimensionality reduction technique where compresses a matrix by often random linear combinations. A line of work has shown how to sketch the Hessian speed up each iteration in 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 learn a distribution on sketching matrices that is optimized for the specific distribution of input matrices. We show how to design learned sketches for the Hessian in the context of second order methods, where we learn potentially different sketches for the different iter- ations of an optimization procedure. We show empirically that learned sketches, compared with their "non-learned" counterparts, improve the approximation ac- curacy for important problems, including LASSO, SVM, and matrix estimation with nuclear norm constraints. Several of our schemes can be proven to perform no worse than their unlearned counterparts.
INTRODUCTION
Large-scale optimization problems are abundant and solving them efficiently requires powerful tools to make the computation practical. This is especially true of second order methods which often 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 al. [2000] Li al. 2020). Our focus is on skete king techniques, which often consist of multiplying the Hessian by a random matrix chosen independently of the Hessian. Sketching has a long history theoretical computer science (see, e.g., (Woodruff) 2014) for a survey), and we describe such methods more below. A special case of sketching is sampling. which in practice is often uniform sampling, and hence oblivious to properties of the actual matrix. Other times the sampling is non-uniform, and based on squared norms of submatrices of the Hessian or on the so-called leverage scores of the Hessian.
Our focus is sketching techniques, and in particular, Wainwright 2016 2017) which introduce the iterative Hessian sketch and Newton sketch, as well as the high accuracy refinement given in (yan den Brand 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 decomposed as A 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 S which could be i.i.d. Gaussian or drawn from more structured family of random matrices such as the Subsampled Randomized Hadamard Transforms or COUNT- SKETCHI matrices; the latter was done in (Cormode & Dickens 2019). The Newton sketch was proposed in (Pilanci & Wainwright 2017). which extended sketching methods beyond constrained least-squares problems to any twice differentiable function subject 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
follow the framework of (Pilanci &
find the update direction. When the Hessian can be

---

Under review as a conference paper at ICLR 2021
preconditioner, 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 random matrix S, drawn from certain family of random matrices, and computes SA. If A is -and-nnn then is short-and-fan and thus SAA is small, roughly square matrix. Moreover, SA preserves important properties of A. One typically desired property is that S is subspace embedding, meaning that simultaneously for all x, one has || - (1 + e) Ax/l2- An observation exploited in (Cormode & Dickens 2019), building off of the
random matrices S introduced in randomized linear algebra in (Clarkson & Woodruff 2017). that contains ((1) non-zero entries per column, then SA can be computed in O (Inz(AA) time, where nnz( A) denotes the number of nonzeros in A. This is sometimes referred to as input sparsity running time.
Each iteration of a second order method often involves solving an equation of the form A Ar A A where A A the Hessian and is the gradient. For a number of problems, one has access to matrix A E Redd with >> d, which is also an assumption made in Pilaacii & Wainwright (2117). Therefore, the solution æ is the minimizer to a constrained least squares regression problem:
min xEC
|| 6ll2
where C is convex constraint set in R. 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 (0166) 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 future sample B, also drawn from D. the learned sketching matrix S will be such that S B 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 number of other problems on large data. For exam- ple, Indyk et al.] (2019) showed that learned sketches yield significantly small errors for low rank approximation. In (Dong et al., 2020). significant improvements to nearest neighbor search were obtained via learned sketches More recently, Liu et al. (2020) extended learned sketches to sev- eral problems in numerical linear algebra, including least-squares and robust regression, as well as k-means clustering.
Despite the number of problems that learned sketches have been applied to, they have not been applied to convex optimization in general. Given that such methods often require solving a large overdetermined least squares problem in each iteration, iteration using learned sketches. However, number of natural questions arise: (1) how should we learn the sketch? (2) should we apply the same learned sketch each iteration, or learn il in the next iteration by training on a data set involving previously learned sketches from prior iterations?
hopeful that one can improve each
Our Contributions. In this work we answer the above questions and derive the first learned sketches for wide number of problems in convex optimization.
Namely, we apply learned sketches to constrained least-squares problems, including LASSO, sup- port vector machines (SVM), and matrix regression with nuclear norm constraints. We show empir- ically that learned sketches demonstrate superior accuracy over random oblivious sketches for each of these problems. Specifically, compared with three classical sketches (Gaussian, COUNT-SKETOR and Sparse SainsonnLindesstrauus Transforms; see definitions in Section the learned sketches in each of the first few iterations
2

---

Under review as a conference paper at ICLR 2021
improve the LASSO error f(x) f(x²) by 80% to 87% in two real-world datasets, where f(x) || - bll2 + ||=||;
improve the dual SVM error f(z) - f(x*) by 10-30% for synthetic and a real-world dataset, as well as by 30%-40% for another real- real-world dataset, where f(x ||
improve the matrix estimation error f(X) - f(X*) by at least 30% for a synthetic dataset and at least 95% for a real-world data set, where f(x) = ||AX - B
Therefore, the learned sketches attain a smaller error within the same number of iterations, and in fact, within the same limit on the maximum runtime, since our sketches are extremely sparse (see below).
We also study the general framework of convex optimization in (van den Brand et al. 2020), and show that also for sketching-based preconditioning, learned sketches demonstrate considerable ad- vantages. More precisely. by using 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.
All of our learned sketches S are extremely sparse, meaning that they contain single non-zero entry per column. Following the previous work of (Indyk et al. (019), we choose the position of the nonzero entry in each column to be uniformly random, while the value of the nonzero entry is learned. This already demonstrates a significant advantage over non-learned sketches, and has fast training time. Importantly, because of such sparsity, our sketches can be applied in input sparsity time given a new optimization problem.
We also provide several theoretical results, showing how to algorithmically use learned sketches in conjunction with random sketches so as to do no worse than random sketches.
2
PRELIMINARIES
Classical Sketches. Below we review several classical sketches that have been used for solving optimization problems.
Gaussian sketch: S - where G is an x Gaussian random matrix.
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 -1, each with probability 1/2. Further, the columns are chosen independently. Sparse Johnson-Lindenstrann Transform (SJLT): S is the vertical concatenation of indepen- dent COUNT-SKETCH matrices, each of dimension m/s n.
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 m and Sp,, = v2 for all 1 < i S n while all the other entries of S are O. When m and n are clear from context, we may characterize such a sketching matrix by (p, v) only.
COUNT-SKETCH-type Sketch. A COUNT-SKETCH-O
Subspace Embeddings. For a matrix A E Baxd we say matrix S E Rmxa is (1 = sububppace embedding for the column span of if e) || S ||AAI|| S [)++)]|Ar)) for all x Rd. The classical sketches above, with appropriate parameters, are all subspace embedding matrices with at least constant probability; our focus is COUNT-SKETCH which can be applied in input sparsity running time. We summarize the parameters needed for subspace embedding below:
Gaussian sketch: m (d/ce). It is a dense matrix and computing SA costs O(m mnz(A)) (unz(A)d/e2) time.
COUNT-SKETCH: m Odd//e2 CCarkson & Woodruff 2017). Although the number of rows is quadratic in d/e, the sketch matrix S is sparse and computing SA takes only O (A) time.
((1/e) non-zeros per column (Nelson Nguyen 2013 Cohen 2016). Computing SA takes ((s nnz (mz( = O (nnz( time.
SJLT: m Odde2) and has
Iterative Hessian Sketch. The Iterative Hessian Sketching (IHS) method (Pilanci & Wainwright) 2016) solves the constrained least-squares problem 1 by iteratively performing the update
1 It+1 = arg min ー || 15411((x 2 (AT (b -Az;), -xt),
(2)
#EC 2
3

---

Under review as a conference paper at ICLR 2021
where S1+1 is sketching matrix. It is not difficult to that for the unsketched version (S++1 is the identity matrix) of the minimization above, the optimal solution x²+1 coincides with the optimal solution to the constrained least square problem ( The IHS approximates the Hessian A A by sketched version (S++1A) (S++1 A) to improve runtime, as St+1 A typically has very few rows.
Unconstrained Convex Optimization. Consider unconstrained convex optimization problem min, fzx), where f is smooth and strongly convex, and its Hessian f is Lipschitz continuous. This problem can be solved by Newton's 's method, which iteratively performs the update
T++1 I2 arg min || (V> (V>)))))))))
V v
provided it is given good initial point zo. In each step, it requires solving regression problem of the form ming A A= y y|اا, which, with access to A, can solved with fast regression solver in (van den Brand et al.] 2020). The regression solver first computes preconditioner R via QR decomposition such that SAR has orthonormal columns, where Š is a sketching matrix, then solves T arg ming || (AR) (AR)z' 2/l2 by gradient descent and returns R2 in the end. Here, the point of sketching is that the QR decomposition of S A can be computed much more efficiently than the QR decomposition of A since has only small number of rows.
Learning a Sketch. We use the same learn- ing algorithm in (Liu et al. 2020), given in Algorithm I The algorithm aims to min- imize the mean loss function A)
Algorithm LEARN-SKETCH: SKETCH: Gradient descent algorithm for learning the sketch values
Require: Atrain {Ai AN} (A, E BBxd),
learning rate a
1 EN L(S, At). where S is the learned sketch, L(S, )) is the loss function of S applied to a data matrix A, and A = {AI,..., AN) is (random) subset of training data.
1: Randomly initialize p, for a Count-Sketch- type sketch
2: for 0 to steps do
3: Form S using
4: Sample batch Anatch from Atrain
HESSIAN SKETCH
5: e f- 1 - a De
6: end for
In this section, consider the minimization problem
Algorithm 2 Solver for (3)
1: S + learned sketch, S₂ t random sketch
{ ||AAT| (A v. uA 2
2: 24.1,2.22 + ESTIMATE(SS. A), = 1,2
min ZEC
(3)
3: if 2 2,,/21,1 < 2 /2,, then
which is used as a subroutine for the IHS (ef. ()). We present an algorithm with the learned sketch in Algorithm 2 To analyze its performance, let R be the column space of A E R and define the following quantities (corresponding exactly to the unconstrained case in Pilanci & Wainwright (2016)
4:
f + solution of with with S S
5: else
6:
T + solution of 3 with S = S2
7: end if
8: return T
9: function ESTIMATE(SSA)
10:
T sparse (1 + -)ususpaace embedding
matrix for d-dimensional subspaces
ZS( = inf
eERnS" ||
11:
(Q, R) + QR(TA)
2, x- Omin (SAR-1)
(z(S)
sup =
(u, S- -m))))))))))
13:
Z2 x- (1 + y)-approximmtion to
((SAR (ASAARR)T(SAAA (SAR -llop
where g--1 denotes the Euclidean unit sphere R"
14: return (2,, 2z)
end function
The following is the estimation guarantee of Z and Zzz The proof is postponed to Appendix A
Lemma 3.1. Suppose that E (0, } is small constant, A is of full rank and S has O(d2) rows. The function ESTIMATE(S, SS, A) returns in log(//)) poly (d/n)) time 2, 2, which with
probability at least 0.99 satisfy that 2,(5) < Z < I-n and 2,(5)
(1(5)
(1+9)) 31 < Z. < 11-1)
Zi(S)
3n
1+
Note for matrix A, ||A]|op sup
|z is its operator norm. Similar to PPllanci & Wainwright T
2016 Proposition 1), have the following guarantee. The proof is postponed to Appendix
4

---

Under review as a conference paper at ICLR 2021
Theorem 3.2. Let E (0, (1) be small constant. Suppose that A is of of rank and S1 and S2 are both COUNT-SKETCH--pp sketches with O(d²) rows. Algorithm 2 returns solution F which, with
probability at least 0.98, satisfies that ||A(ā x²)1l2 - (1 n) min
Z1.2 22.2
+ ||Ax" ll2
%1,1 12
((nz(()) (A) log(;)) + poly( time, where arg minzec solution.
|| b||2 is the least-squares
Algorithm 3 Fast Regression Solver for 4
4 HESSIAN REGRESSION
1: S +- learned sketch, S2 <- random sketch
(Q:, R₂) < QR A), = 1,2
In this section, we consider the minimiza- tion problem
3: (0 (AR, ") 1,2
EIG returns the max and min singular values B
min A Az / 2112. 1
(4)
4: if 0 / 02. 02/02 then
5:
P R1
1/(of (o1)²)
which used as a subroutine for the un- constrained convex optimization problem min, f(x) with A being the Hessian matrix f(x(() (See Section Here A Raxd E R2 and we have access to A. We incorporate learned sketch into the fast regression solver in (van den Brand al. 2020) and present the algorithm in Algorithm
6: clse
P R + 1/ +
8: end if
9: 20 <
10: while
|| AP 1112 2 <1lll12 do
11: 2++1 < 2t - (p (PP
12: end while
13: return P2
Here the subroutine EIG(B) applies a (1 + -)-subppace embedding sketch T to B for some small constant 1) and returns the maximum and the minimum singular values of TB Since B admits the form of AR, the sketched matrix TB can be calculated as (TA) and thus can be computed in ((nzz(A A) + poly time if T is COUNT-SKETCH matrix of (dd²) rows. The extreme singular values of TB can be found by SVD or the Lanezos's algorithm.
Similar to Lemma 4.2 in (van den Brand et al. 2020). we have the following guarantee of Algo- rithm The proof parallels the proof (van den Brand et al.. 2020) and is postponed to Ap- pendix
Theorem 4.1. Suppose that S₁ and S2 are both COUNT-SKETCHH-yp sketches with O(d)) rows. Algorithm returns solution x' such that || Ax' |A''''' -yl2 S c <llyll2 in O(nnz(A) (A)) O(nd
(min(ar/on (102//)))2 (oo/oo))).oo(((((A)) + poly(d) time.
Remark 4.2. Algorithm 3 52 can be chosen to be a subspace embedding matrix for d- dimensional subspaces, in which case, AR has condition number close to (see, e.g., (Woodruff 2014 p38)) and the full algorithm would run faster than the trivial O(nd2)-tima solver to
Remark 4.3. For the original unconstrained convex optimization problem min, fxx), 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) = 1/2 Ax - bllz- the value of (x) can be approximated efficiently by a sparse subspace embedding matrix in O(nz(()) + nnz(b) + poly(d) time.
IHS EXPERIMENTS
Training. We learn the sketching matrix in each iteration (also called round) separately. Recall that in the (++1)-st round of IHS, the optimization problem we need to solve (2) depends on 2 An issue is that we do not know IG, which is needed to generate the training data for the (t ( 1)-st round. Our solution use the sketch matrix in the previous round obtain x; by solving 2 and then use to generate the training data in the next round. That is, in the first round, we train the sketch matrix S to || Ax|]
solve the problem # arg mingee for the optimization problem for Ø2, and so on. The loss function we use is the unsketched objective function in the (t+1)-+ iteration, i.c., [(S A A) = 1/ ||A( - xx11 (AT (b-Ax,)) Ze+1 - =Zx), where 2 is the solution to (2) and thus depends S
(A b, b) and use 1 to generate the training data
5

---

Under review as a conference paper at ICLR 2021
Comparison. We compare the learned sketch against three classical sketches: Gaussian, COUNT- SKETCH, and SJLT (see Section in all experiments. The quantity we compare a certain error, defined individually for each problem, in each round of the iteration of the IHS or function of the runtime the algorithm. All of our experiments are conducted on a laptop with 1.90GHz CPU and 16GB RAM.
5.1 LASSO
We define an instance of LASSO regression as
puidis
x
arg min x
|
bll2
||=|||
(5)
where À is a parameter. We use two real-world datasets:
CO emissior the dataset contains 9 sensor measures ag- gregated over one hour (by means of average or sum), which can help to predict the CO emission. We divide the raw data into 120 (A,, b,) such that A, E R20009 b, 3000x The data in each matrix is sorted in chrono- logical order. (A, traail = 96, [(A, 6)asal = 24.
Figure 1: Test eπror of LASSO CO emissions dataset, m = 5d
Greenhouse ga time series of measured greenhouse
gas concentrations in the California atmosphere. Each
(A, b) corresponds to a different measurement location.
A; E R222x111 b, E R327x1, and [(A, b)rrinn
400.
Figure 2: Test efror of LASSO on
(A. rest = 100. (This dataset was also used in
(Luu et al,
CO emissions dataset, m = 3d
2020
1 in the LASSO regression 5). We choose 5d and = 3d for the CO emission dataset and = 6d and m = 3.5d for the greenhouse gas dataset. We consider the
We choose
error ( Ar - bll2 + +|==|||]=--))-00000> [| Alx - +|="+1 and take an average over five independent trials. We plot (natural) log- arithmic scale the mcan errors of the two datasets in Figures to
Figure 3: Test error of LASSO on greenhouse gas dataset, m = 6d
We can observe that the learned sketches consistently outper- form the classical random sketches in all cases. Note that our sketch size much smaller than the theoretical bound. For a smaller sketch size, classical random sketches converge slowly or do not yield stable results, while the learned sketches can converge to small eтror quickly. For larger sketch sizes, all three classical sketches have approximately the same order of error, and the learned sketches reduce the error by a 5/6 7/8 factor in all iterations.
Figure 4: Test error of LASSO
greenhouse gas dataset, m = 3.5d
5.2 SUPPORT VECTOR MACHINE
the context of binary classification, a labeled sample a pair (a), where E R. is vector representing a collection of features and z; E {-1,+1} is the associated class label. Given a set of labeled patterns {(a,, +)) the support vector machine (SVM) estimates the weight vector w" by minimizing the function
+1]
[c Dg g(z,, (wi, (()),a))) + |/0
1
"" arg min WEBn
2
where a parameter. Here we use the squared hinge loss g(z,, a.)) = (1 - 2, a:))

---

Under review as a conference paper at ICLR 2021
lerned dourtidettt
(smmed
larred
cont-shatch
grusiann
gaussion sparte I.
sparne
sparse
Ixpienent
Figure 5: Test error of SVM on Figure 6: Test error of SVM on Figure 7: Test error of SVM on random Gaussian dataset
swarm behavior dataset
Gisette dataset
The dual of this problem can be written as constrained minimization problem (see, e.g., LLi al, 2009 Pilanci & Wainwright 2015)).
arg min |r 2EA4
the domain A {s Rd 2 0 and | 1}, the positive simplex in Rd. Here 1 [a]_ g(++d)xd, where an matrix with a; R2 its i-th column vC [u]
B [AD)
and = diag(z) 2) is d x diagonal matrix.
We conduct experiments the following three datasets:
Random Gaussian (synthetic): We follow the same construction in Pilanci Wainwright (2016). We generate a two-component Gaussian mixture model, based the component distri- butions N N(o, /) and N(,, where Pin and li are uniformly distributed in 33. Placing equal weights on each component, we draw d samples from this mixture distribution.
2400 features and the task is to predict whether the instance flocking or not flocking. We use only the first 6000 instances of the raw data, and divide them into 200 smaller groups of instances. Each group contains d = 30 instances, corresponding to B, of size 2430 30. The training data consists of 160 groups and the test data consist of 40 groups.
Swarm behavior Each instance in the dataset has
Gisette") Gisette is handwritten digit recognition problem which asks to to separate the highly confusable digits '4' and '9'. Each instance has 5000 features. The raw data is divided into 200 smaller groups, where each contains = 30 instances and corresponds to a B, of size 5030 30. The training data consists of 160 groups and the test data consists of 40 groups
10d in all experiments and define the error || [xll Bx 1| For random sketches, we take the average error over five independent trials, and for learned sketches, over three independent trials. For the Gisette dataset, we use the learned sketch in all rounds. We plot in (natural) logarithmic scale the mean errors of the three datasets in Figures to 7 For Gisette and random Gaussian datasets, using the learned sketches reduced the error by 10%-30% and for the Swarm Behavior dataset, the learned sketches reduce the error by about 30%-400%
We choose
5.3
MATRIX ESTIMATION WITH NUCLEAR NORM CONSTRAINT
In many applications, for the problem
X = arg min ||AX BI XgR41x42
is reasonable model the matrix X as having low rank. Similar to & mmmimmiizationn for com- pressive sensing, standard relaxation of the rank constraint is to minimize the nuclear norm of X, defined as X ΣØ= Emin(di d2) a,(X), where or (X) is the j-th largest singular value of X.
Hence, the matrix estimation problem we consider here is
X arg min || - BIK such that ||]||, S p. X Xig
where p > 0 is user-defined radius as regularization parameter.

---

Under review as a conference paper at ICLR 2021
We conduct experiments on the following two datasets:
Synthetic Dataset: We generate the pair (A,, B,) as B, = A,X; + W,, where A, Ruxdi with iid xd2 (0) entries. X E gdixda is a matrix with rank at most r. W; is noise with i.idd (0, o²) entries. Here we set n 500, d1 = d2 7, r 3, - 30.
Figure 8: Test error of matrix estimation on synthetic data, m 50
(AA, B) train = 270, [AA, B),,esttss = 30.
Tunne The data set is time series of gas concen- trations measured by eight sensors in a wind tunnel. Each (A, B) corresponds to different data collection B, E R 13530x6 [(A, B = 144, I(A, B) teat 36. The same dataset and param- eters were also used in (Liu et al. 2020) for regression tasks. In our nuclear norm constraint, we set p = 10.
kined (illia)
trial. A; E IR 135330x5
Figure 9: Test error of matrix estimation on Tunnel dataset, m = 50
We choose m = 40, 50 for the synthetic dataset and m = 10, 50 for the Tunnel dataset and define the eπror to be 1 [I BI 2 || - BlI²)). ). For each data point, we take the average error of five independent trials. The mean errors of the two datasets when m 50 are plotted in a (natural) logarithmic scale in Figures and 9 We observe that the classical sketches yield approximately the same order of error, while the learned sketches improve the error by at least 30% for the synthetic dataset and surprisingly by at least 95% for the Tunnel dataset. The huge improvement on the Tunnel dataset may be due to the fact that the matrices A, have many duplicate rows. We defer the results when m = 10 or 40 to Appendix D which show that the leameed sketches yield much smaller errors than the random sketches and the random sketches could converge significantly more slowly with considerably larger errors in the first several rounds.
2
tames
5.4 RUNTIME OF LEARNED SKETCHES
As stated in Section 2 our learned sketch matrices S are all COUNT-SKETCH-yye matrices (each column contains single nonzero entry). the matrix product SA can thus be computed in (nnz(A)) time and the overall algorithm is expected to be plot for
67
14
fast. To verify this, we plot in an cetoo--erus--uutninn the SVM and matrix estimation with nuclear norm constraint tasks in Figures 10 and [ (corresponding to the datasets in Figures 7 and
Figure 10: Test error of SVM on Gisette dataset
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 exper- iment three times. Each time we take an average over all test data. From the plot we can observe that the learned sketch and COUNT- HOUNT-SKETH SKETCH have the fastest runtimes, which are slightly faster than that of the SJLT and significantly faster than that of the Gaussian sketch.
Figure 11: Test error of matrix es- timation on Tunnel dataset
6 FAST REGRESSION EXPERIMENT
We consider the unconstrained least-squares problem, i.e., (5) with À = O. using the CO emission, greenhouse gas datasets, and the following Census dataset:
Census data this dataset consists of annual salary and related features people who reported that they worked 40 or more weeks in the previous year and worked 35 or more hours per week. We randomly sample 5000 instances to create bz), where A E R5000x11 and b 13000x1 [A,, B) [rrin = 160, (A, beest 40.
8

---

Under review as a conference paper at ICLR 2021
01-
boreel pirste
04-
glaike
02
01
Figure 12: Test error of fast regression for the CO emission dataset, first three calls in solving an unconstrained least-squares
Training. We optimize the learned sketch S by gradient descent (Algorithm D. where C(S, A) = R(AR, )) where R2 is computed as in Algorithm and k(()) denotes the condition number of a matrix M.
Next we discuss how to generate the training data. Since we use Newton's method to solve an unconstrained convex optimization problem (see Section 2 in the t-round, we need to solve regression problem min: [ (xx)) (xx][
()(=))]l2 Reformulating it as ming | A= - 21l2. we see that A and y depend on the previous solution zz. Hence, we take T to be the solution obtained from Algorithm using the learned sketch S,, and this generates A and y for the (t + s)t round.
Experiment. For the CO emission dataset, set m = 500. For the 17 in the Algorithm 3 we set 17 = for the first round and 1j = 0.2 for the next rounds for the CO emission dataset, and 17 1 for all rounds for the Census dataset. We leave the setting and results on the greenhouse gas dataset to Appendix E
70, and for the Census dataset, we set
We examine the accuracy of the subproblem 4 and define the error to be | AR y ||yll2--
We run the first three subroutines of solving the subproblem for the CO emission dataset and the Census dataset. The average error of three independent trials is plotted in Figures [2 [3] and [6 We observe that for the CO emission dataset, the classical sketches have similar performance and the learned sketches lead to fast convergence in the subroutine with the first-round error least 80% smaller; for the Census dataset, the learned sketch achieves the smallest error for all three rounds, where we reduce about 60% error in the first round and about 50% error in the third round. Note that the learned sketch always considerably outperforms COUNT-SKETCH in all cases.
kerred
oour torstch
deth guaras
co4 derth
gacis
0.3
0.5
12
0.1
00
0.
meratorrourd
Figure 13: Test error of fast regression for the Census dataset, first three calls in solving an uncon- strained least-squares
7 CONCLUSION
We demonstrated the superiority of using learned sketches, over classical random sketches, in the It- erative Hessian Sketching method and fast regression solvers for unconstrained least-squares. Com- pared with random sketches, our learned sketches of the same size can considerably reduce the error in the loss function (i.e., f(x f(x*), where x is the output of the sketched algorithm and x² the optimal solution to the unsketched problem) for a given threshold of maximum number of iterations or maximum runtime. Learned sketches also admit smaller sketch size. When the sketch size is small, the algorithm with random sketches may fail to converge, or converge slowly, while the algorithm with learned sketches converges quickly.
9

---

Under review as a conference paper at ICLR 2021
REFERENCES
Marcin Andrychowicz, 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. In Advances in neural information processing systems, pp. 3981-3989, 2016.
Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, 2004.
Kenneth L. Clarkson and David Woodruff. Low-rank approximation and regression in input sparsity time. J. ACM, 63(6), January 2017. ISSN 0004-5411. doi: 10.1145/3019134, URL
https///do..or/1001114/311111111
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. CoRR, abs/1910. 1.111.14166 14166, 2019. URL http:://arxxv.org/b///////////// 14166)
Nikita Doikov and Peter Richtárik. Randomized block cubic Newton method. In Proceedings of the 35th International Conference on Machine Learning, ICML Stockholmsmüssann 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 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 Informatinn Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 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, NeurIPS 2019, 8-14 December 2019, Vancouver, BC, Canada, pp. 614-623, 2019.
Chen-Ye Yu Hsu, Piotr Indyk, Dina Katabi, and Ali Vakilian. Learning- based frequency estimation algorithms. In 7th International Conference on Learning Representations, ICLR 2019, New Or- leans, LA, USA, May 6-9, 2019, 2019.
Piotr Indyk, Ali Vakilian, and Yang Yuan. searning-based low-rank approximations. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Process- ing Systems 2019, NeurIPS 2019, 8-14 December 2019, Vancouver, BC, Canada, pp. 7400-7410, 2019.
Sudhir B. Kylasa, Fred (Farbod) Roosta, Michael W. Mahoney, and Ananth Grama. GPU accelerated sub-sampled Newton'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.
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 Shirty-ecoond 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.
10

---

Under review as a conference paper at ICLR 2021
Yu- Yu-Feng Li, Ivor W. Tsang. James Tin- Yau Kwok, and Zhi-Hua Zhou. Tighter and convex maximum margin clustering. In David A. Van Dyk and Max Welling (eds.), Proceedings of the Twelfth Inter- national Conference Artificial Intelligence and Statistics, AISTATS 2009, Clearwater Beach, Florida, USA, April 16-18, 2009, volume 5 of JMLR Proceedings, pp. 344-351. A.L........ JMLR.org, 2009.
Simin Liu, Tianrui Liu, Ali Vakilian, Yulin Wan, and David P. Woodruff. On learned sketches for randomized numerical linear algebra. arXiv:2007 20X0:2007.098900 .09890 [cs.LGG], 2020. URL. https://arxiv.org) abs/2007.09890
J. Nelson and H. I.. 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.
Mert Pilanci and Martin J. Wainwright. Randomized sketches of convex programs with sharp guar- antees, IEEE Trans. Inf. Theory, 610)):509-55115 2015.
Mert Pilanci and Martin J. Wainwright. Iterative Hessian sketch: Fast and accurate solution approx- imation for constrained least-squares. J. Mach. Learn. Res., 17:53:1-53:38, 2016.
Mert Pilanci and Martin Wainwright. Newton sketch: A near linear-time optimization algorithm with lincar-quadratic convergence. SIAM J. Optim., 27(1):205-245, 2017.
Farbod Roosta-Khorasee and Michael W. Mahoney. Sub-sampled Newton methods. Math. Pro-
gram., 174(1-2):293-326, 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 (overparametrizee)) neural networksin near-linear time. arXiv:2006. 11648 [cs.LG], 2020.
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices, pp. 210-268. .006.
Cambridge University Press, 2012. doi: 10.1017/CB009880
David P. Woodruff. Sketching as a tool for numerical linear algebra. 10(1-2):1-177, October 2014. ISSN 1551-305X. doi: 1040000000000000000 URL httpps://doiorgg/// http://0000000/00000000000000000
Peng Xu, Jiyan Yang, Farbod Khorasain Christopher Ré, and Michael W. Mahoney. Sub- sampled Newton methods with non-uniform sampling. Advances in Neural Information Pro- cessing Systems 29; Annual Conference on Neural Information Processing Systems 2016, De- cember 5-10, 2016, Barcelona, Spain, pp. 3000-3008, 2016.
Peng Xu, Fred Roosta, and Michael W. Mahoney. Second-order optimization for non-convex ma- chine learning: an empirical study. In Proceedings of the 2020 SIAM International Conference on Data Mining, SDM 2020, Cincinnati, Ohio, USA, May 7-9, 2020, pp. 199-207, 2020.
A PROOF OF LEMMA 3.1
Suppose that AR-1 UW, where U E RRxd has orthonormal columns, which form an orthonor- mal 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 E Rd that
TAR #/12- < ||R 'x|12- < 1- TAR 20112
1+
Since
|TA "2|12 1|Qx112 ||>|||
and
WWrll |[ ||
11

---

Under review as a conference paper at ICLR 2021
have that
1
+ 1) |0|| S [a S
1 1-1 ||
x E EER..
It is easy to see that
|SUWall
21(5) = min ZES--1
||SUx|1 - min y=0
[ ull
and thus,
|SUWall vll2 - (S) S min(1 + 17)
||SWWbl2
min(1 11) V=0
1ly12
840
1/112
Recall that SUW SAR-!
We see that
(i) - (1+ m) (SAR ) (1 77) (SAR s
By definition,
=2(5)=11'2 - - S-nnulll
follows from that
(1 m)² -ST op < Zz(S) S (1+ (1+ (1"""" wn)ull,,
and from
7 (6) and (Vershynin 2012 Lemma 5.36) that
- ol - 31.
Since
= ||AR ()S L-,R op
and
11,0 -((AR-IT(AR-)) 7
< (TATR))((S'S -
< (|A STR - 71..p +[(AR-1) ((AR-)) I op'
follows that
(1- 1)² ( SAR - 711 -3(1 - =
$ (2)
$ <1 + 7) [SARR SAR - op + + n)
We have so far proved the correctness of the approximation and we shall analyze the runtime below.
Since S and T sparse, computing SAA and TAA takes ((nx(A)) time. The QR decomposition of A, which a matrix of size (dly(d/n) x 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) x d, its smallest singular value can be computed in (d) time. To approximate Z2(S), we can use the power method to estimate [(AAR SAR Ilop up to a (1 + -fctor in ((nnz(() + poly(d)) (og(1/))) time.
B PROOF OF THEOREM 3.2
In Lemma we have with probability at least 0.99 that
22
1++)
(2) - 3n
1- Zz(S)
3n
2
1.))
(1+ 1)2 (S)
Zi(S)
When S' is random subspace embedding, it holds with probability at least 0.99 that (i 2 3/4 and so, by a union bound, holds with probability least 0.98 that
2 N
1
Zz(S)
4m-
2
(1 + 17 (i(S)
12

---

Under review as a conference paper at ICLR 2021
or,
Zz(5) (i(
Z2 21
S (1+ n)*
+ 417
The correctness of our claim then follows from (Pilanci Wainwright 2016 Proposition 1), gether with the fact that S2 is random subspace embedding. The runtime Tollows from Lemma 3.1 and (Cormode & Dickens 2019 Theorem 2.2).
C PROOF OF THEOREM
The proof follows almost an identical argument as that in (van den Brand al. 2020 Lemma
B.1). In (van den Brand et al. 2020), it is assumed (in our notation) that 3/4 S T (AP) <
Tmax (AP) < 5/4 and thus one can set 1/ = Algorithm 3 and achieve linear convergence. The only difference is that here we estimate minn (AP) and max (AP) m
and set the step size in the gradient descent algorithm accordingly. By standard bounds for gradient descent (see, c.g., (Boyd
& Vandenberghe 2004 P468)), with a choice of step size
(AP) + (AP)) after 2/ (o (AP)
O(( (AP)/ 0 miin P))) log(1/e)) iterations, we can find 2, such that
|| |P A AP - 2²)112 S S A
APPAATAPLE AP(=0 x²) 22)1
where arg ming || || A AP= PTy, y is the optimal least-squares solution. This establishes Eq. (11) in the proof in (van den Brand et al. 2020), and the rest of the proof follows as in there.
D IHS EXPERIMENT: MATRIX ESTIMATION WITH NUCLEAR NORM CONSTRAINT
stated Section 5.3 the mean errors of the two datasets when m 40 the synthetic dataset and 10 for the Tunnel dataset are plotted in (natural) logarithmic scale in Figures 14 and 15
sparee II
kamed cortistanch
0-
Figure 14: Test error of matrix estimation on synthetic dataset, 40
Figure 15: Test error of matrix estimation on Tunnel dataset, m = 10
E FAST REGRESSION EXPERIMENT: GREENHOUSE GAS
We 100, 1 0.2 in the first round and 1 in the second round. The mean errors of the first two calls of the fast regression subroutine are plotted in Figure [6 We can observe that Gaussian and sparse JL sketches have significantly better performance than COUNT-SKETCN the learned sketch shows again a significant reduction of more than 50% in the first-round error.
and
13

---

Under review as a conference paper at ICLR 2021
08-
giLetinn
hanee
0.4
05
02-
Deratonn ruund
Figure 16: Test error of fast regression for the greenhouse gas emission dataset, first two calls in solving an unconstrained least-squares
14


Reviews:
[{"review_id": "dgR_plzbfIE", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 5, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposed a learned variant of the well-known iterative Hessian sketch (IHS) method of Pilanci and Wainwright, for efficiently solving least-squares regression. The proposed method is essentially a learned variant of the count-sketch, where the positions of the non-zero entries are random while the value is learned. While getting a learned variant for IHS is an interesting direction, the current theoretical contribution of this paper is only incremental, and most importantly, the reviewer is unconvinced for the practicality of the current approach.\n\nThe learned sketching matrix S^t is computed at each iteration, which should introduce a significant computational overhead. Then the plots comparing the learned IHS with unlearned IHS is unfair since it is regarding the iteration count. The iteration count does not reflect such overhead, and the reviewer believe that the wall-clock time comparision would be more sensible.\n\nThe reviewer found that the authors are over-selling their numerical results. For example in Figures 1, 2, 6, 7, that the learned IHS has the same linear convergence rate as the unlearn IHS, and actually only slightly faster overall (only 1 iteration faster if we view the plots horizontally), but the authors' claims are like “We can observe that the classical sketches have approximately the same order of error and the learned sketches reduce the error by a 5/6 to 7/8 factor in all iterations”, “We observe that the classical sketches yield approximately the same order of error, while the learned sketches improve the error by at least 30% for the synthetic dataset and surprisingly by at least 95% for the Tunnel dataset.”, which the authors seem to purposefully report from vertical views to exaggerate the results. This is very misleading to the readers which are not familiar with the related works.\n\n********* Update ********\n\nThe reviewer appreciates the efforts authors have made. However, the response does not fully address the issues. In the new experiments the authors used small sketch size to show clearer advantages of the learned IHS (Figure 2 & 4), but at that regime all of the methods including learned-IHS are converging only slowly and not practical compare to the slightly larger sketch size choices. The reviewer believes that such a comparison is not meaningful. In order to truly demonstrate the benefits of learned-IHS (which seems to be robust to small sketch sizes), the authors should choose for each algorithm the best sketch size and then compare them in run time, at least between learned-IHS and Count-sketch IHS.\n\nBeside the flaws in numerical experiments, the reviewer found that the theoretical contribution of the current version is incremental. From the reviewer's point of view, the meaningful analysis for learned IHS should certainly be how the converge relates to  the statistics of the training data, to show the benefit of the learned sketch theoretically (i.e. to show how much better the learned sketch SA compare to the unlearned ones in terms of Z_1 and Z_2). The authors have avoided such type of analysis and the main results are \"safe-guarded\" by the concentration of the random sketch, which are easy to derive.\n\nOverall, the idea of combining IHS with learned sparse sketch is interesting, but the reviewer believes that the current version needs significant amount of rework to be publishable in top conferences.\n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "8n6vaMJ_Cqb", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nPrevious work has shown that sketching the Hessian can improve the running time of second-order optimization methods. However, these prior sketching techniques were data-oblivious. This paper leverages learned sketches to improve the convergence rate of second-order optimization methods.\n\nStrengths: \n+ The learned sketch generally converges better than the other baseline.\n\nWeaknesses:\n+ The technical novelty of the paper seems incremental. The learned sketch algorithm from (Liu et al., 2020) serves as a replacement for the data-oblivious sketch in these second-order optimization methods.\n\nQuestions:\n1) In experiments, is it possible to plot the convergence rate of the non-sketched, full-sized Hessian? I would like to compare the best convergence rate against the other baselines.\n2) What is the total running time of the learned sketch against the other baselines?\n3) The paper asks whether we should apply the same learned sketch each iteration or train the sketch continuously? However, this question was not thoroughly investigated in the paper. In the experiments, the sketching matrix is learned in each iteration. Is there evidence that supports this choice?\n\nReferences:\n1) Simin Liu, Tianrui Liu, Ali Vakilian, Yulin Wan, and David P. Woodruff. On learned sketches for randomized numerical linear algebra. arXiv:2007.09890 [cs.LG], 2020. URL https://arxiv.org/ abs/2007.09890.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "2Vhz0tDT4M3", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "SUMMARY:\n\nSketching is a popular technique in numerical linear algebra for achieving various desirable properties (e.g., lower complexity, one pass methods). The present paper considers a particular kind of sketch for which the sketch matrix is learned from data. It shows how such learned sketches can be used in two types of problems: Hessian sketching (Sec. 3) and Hessian regression (Sec. 4). The authors give both algorithms and provide theoretical guarantees. They also apply these techniques to a number of both synthetic and real datasets in the experiments. For the most part, the experiments indicate that the proposed methods give a consistent, but not necessarily very large, improvement.\n\nI think the idea of using learned sketches is interesting, and it seems like it has not been applied to the problems considered in this paper before. I think the paper strikes a good balance between algorithms, theory and experiments. I like the paper, but there are a few issues that must be addressed before it can be published. Most importantly, there seems to be errors in both Lemma 3.1 and Theorem 3.2. These should hopefully be easy to fix though. \n\nIn the current state, I vote to reject the paper. But if my concerns, especially the issues with Lemma 3.1 and Theorem 3.2, can be addressed, I'll be happy to increase the rating. \n\n\nADVANTAGES:\n\n- Learned sketching is a fairly new and interesting idea.\n- It seems like learned sketching has not been used for Hessian sketching and regression before.\n- Good balance of algorithms, theory and experiments.\n\n\nCONCERNS/QUESTIONS:\n\n- In the list that describes performance improvements under \"Our Contributions\", it is not clear what the quantities $x^*$ and $X^*$ represent. In Sections 5 and 6, these quantities are used to represent the optimal solution to the unsketched problems, but here they seem to represent the solution for the sketched problem using one of the random sketches. If it's the latter, are $x^*$ and $X^*$ the best performing solutions produced by the competing methods? It would be helpful if you clarified this. \n\n- In Algorithm 1, how is $\\alpha$ chosen? Is $\\alpha$ updated adaptively, or is it just a fixed small number? Is there any rule of thumb for how to do this? \n\n- Lemma 3.1: Is seems like the bound on $\\hat{Z}_2$ stated in the lemma is wrong. Based on the upper and lower bounds on $Z_2(S)$ towards the end of Section A, it seems like this bound should read\n$$\n\\frac{Z_2(S)}{(1+\\eta)^2} - 3 \\eta \\leq \\hat{Z}_2 \\leq \\frac{Z_2(S)}{(1-\\eta)^2} + 3 \\eta.\n$$\n\n- Lemma 3.1: Since you use that $\\max(\\eta,\\eta^2) = \\eta$ when applying the result of Vershynin (2012) in the proof, you should add the condition that $\\eta \\leq 1$ in the lemma statement.\n\n- Proof of Lemma 3.1: In Section A, 2nd sentence, you say \"Since $T$ is a subspace embedding of the column space of $A$...\". I think you should add that this is true with probability at least 0.99, to clarify where the 0.99 probability of success in the lemma statement comes from. \n\n- Proof of Lemma 3.1: In Section A, you use the inverse of $R$ in multiple places. However, it seems like $R$ won't be invertible unless $A$ is of full rank. Can the proof be adapted for a case when $A$ is rank deficient? If not, you should add that $A$ is assumed to be full rank in the lemma statement.\n\n- Proof of Lemma 3.1: This is related to the previous point. In the second to last equation on page 10, it seems like the equation\n$$\n\\min_{x \\in S^{d-1}} || S U x || = \\min_{y \\neq 0} \\frac{|| S U W y ||}{|| W y ||}\n$$\nonly will hold if W is full rank, which requires A to be full rank (since $AR^{-1} = UW$ and $U$ has $d$ columns and is full rank). Can the proof be adapted for a case when $A$ is rank deficient? If not, you should add that $A$ is assumed to be full rank in the lemma statement.\n\n- Theorem 3.2: Given the error in Lemma 3.1, the bound in Theorem 3.2 needs to be updated accordingly. Also, in the proof in Section B, in the first inequality you use\n$$\n\\frac{1}{\\hat{Z}_1} \\geq \\frac{1}{\\frac{1}{1+\\eta} Z_1(S)}.\n$$\nThere's a sign error here; it should read\n$$\n\\frac{1}{\\hat{Z}_1} \\geq \\frac{1}{\\frac{1}{1-\\eta} Z_1(S)}.\n$$\n\n- Proof of Theorem 3.2: In Section B, 2nd sentence, it would be helpful for the reader if you said \"holds with probability 0.99\" instead of \"holds with high probability\", and then also clarify that you do a union bound with the 0.99 probability from Lemma 3.1 to get a success probability of at least 0.98 in Theorem 3.2. This may be obvious to readers familiar with the area, but being clear with these things would make the paper accessible to a wider audience.\n\n- Solving (4) deterministically would cost $O(nd^2)$. For Alg. 3 to be worthwhile, it therefore seems like the number $\\min(\\sigma_1/\\sigma_1', \\sigma_2/\\sigma_2')$ in Theorem 4.1 must be smaller than $d$. Could you say something about why we expect this $\\min$ to be small? Is it the case that the $\\min$ may be large with some small probability?\n\n- In the experiments, $m$ is chosen as $m = kd$ for some integer $k$. This is different from the $m = O(d^2)$ required for CountSketch to be a subspace embedding. Have you found that this choice $m = O(d)$ always works well in practice, or have you encountered datasets where such a choice has proven to be too small?\n\n- For the experiments in Section 5.2, is there any reason why you only use learned sketches in the first round for the Gaussian and Swarm behavior datasets, rather than use them all rounds as for the other datasets and experiments?\n\n- In Section 6, is there a reason why you choose $\\eta = 1$ and $\\eta = 0.2$ rather than using the rule for setting $\\eta$ in Alg. 3?\n\n- In Section C, you refer to \"standard bounds for gradient descent\". Can you please provide a reference for those that are unfamiliar with the literature?\n\n- The paper ends abruptly with no conclusion.\n\n\nMINOR CONCERNS/QUESTIONS:\n\n- On page 2, in the 2nd paragraph, 3rd sentence, you say \"If $A$ is tall-and-skinny, then $S$ is wide-and-fat...\". Should it be \"short-and-fat\" rather than \"wide-and-fat\"?\n\n- The usage of R is a bit confusing. It is used to mean the R matrix in a QR factorization in Alg. 2, the inverse of the R matrix from a QR factorization in Alg. 3, and the upper bound on the nuclear norm in Section 5.3. It would be less confusing if a the inverses in Alg. 3 and the nuclear norm bounds were called something other than R.\n\n- In Section 5.2, for the Swarm behavior and Gisette datasets, you say that $B_i$ is of size $2430 \\times 30$ and $5030 \\times 30$, respectively. Should this be $2400 \\times 30$ and $5000 \\times 30$ since $n$ is 2400 and 5000 respectively for the two datasets?\n\n#######################\n\nUpdate:\n\nThe authors have addressed my questions adequately. In particular, my main concerns with the theoretical results and proofs have been fixed. I have updated the score from 3 to 6 for now.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper considers convex optimization problems whose solutions involve the solution of linear systems defined in terms of the Hessian. It presents algorithms that reduce the runtime of standard iterative approaches to solving these problems by iteratively sketching the Hessian; the novelty lies in the fact that the authors use the idea of learned sketches which have been used prior for problems in data mining. In particular, the authors use the approach to learning sketches of Liu et al., 2020 to learn the entries in sparse sketching matrices for the Hessian, and propose using the Iterative Hessian Sketch algorithms of Pilanci and Wainright, 2016 to iteratively solve the concerned optimization problem. The advantages of learned sketches are that they may allow using smaller sketch sizes while making progress on the problem, as they are learned to work well on the distribution of Hessians from which the problem instance is drawn.

The consensus of the reviews is that the idea of using learned sketches for convex optimization seems to be novel, and this paper is an interesting attempt, but falls short of the level of contribution required for publication in ICLR. The main concern is that the theory provided for the use of the learned sketches is incremental: the analysis does not reflect the fact that the sketches are learned; instead, the algorithm builds in a safeguard by using both a random sketch and a learned sketch, and the analysis uses the properties of the random sketch to proceed. The empirical results are suggestive, but the convergence rates of the learned and random sketches do not vary much, so the benefits seem marginal for most of the problems considered (with the exception of a standard least squares problem, for which we know learned sketching performs well).

 The paper is recommended to be rejected, as the theory is weak, and the empirical results are borderline. 
Decision: Reject