Learning the Positions in CountSketch
Authors: Yi Li
Honghao Lin
Simin Liu
Ali Vakilian
David Woodruff
Venue: ICLR 2023 notable top 25% (2023)
Keywords: learning-augmented sketches
count-sketch
low-rank approximation
iterative Hessian sketch

Abstract: We consider sketching algorithms which first compress data by multiplication with a random sketch matrix, and then apply the sketch to quickly solve an optimization problem, e.g., low-rank approximation and regression. In the learning-based sketching paradigm proposed by Indyk et al., the sketch matrix is found by choosing a random sparse matrix, e.g., CountSketch, and then the values of its non-zero entries are updated by running gradient descent on a training data set. Despite the growing body of work on this paradigm, a noticeable omission is that the locations of the non-zero entries of previous algorithms were fixed, and only their values were learned.
In this work, we propose the first learning-based algorithms that also optimize the locations of the non-zero entries. Our first proposed algorithm is based on a greedy algorithm. However, one drawback of the greedy algorithm is its slower training time. We fix this issue and propose approaches for learning a sketching matrix for both low-rank approximation and Hessian approximation for second-order optimization. The latter is helpful for a range of constrained optimization problems, such as LASSO and matrix estimation with a nuclear norm constraint. Both approaches achieve good accuracy with a fast running time.  Moreover, our experiments suggest that our algorithm can still reduce the error significantly even if we only have a very limited number of training matrices.

Published as a conference paper at ICLR 2023
LEARNING THE POSITIONS IN COUNTSKETCH
Li
Honghao Lin, Simin Liu
Nanyang Technological University edu.sg
Carnegie Mellon University (honghaol, siminliu) (awdiiueaannw...mm cmu. .dd:
Ali Vakilian
David P. Woodruff
Toyota Technological Institute Chicago vakilian&ter edu
Carnegie Mellon University
@roodruf@anndze cmu, edu
ABSTRACT
We consider sketching algorithms which first compress data by multiplication with random sketch matrix, and then apply the sketch to quickly solve optimization problem, e.g., low-rank approximation and regression. ing paradigm proposed by Indyk, Vakilian, and Yuan (2019), the sketch matrix is found by choosing a random sparse matrix, e.g., CountSketch, and then the values of its non-zero entries are updated by running gradient descent on a training data set. Despite the growing body of work on this paradigm, noticeable omission is that the locations of the non-zero entries previous algorithms were fixed, and only their values were learned. In this work, we propose the first learning- based al- gorithms that also optimize the locations of the greedy algorithm. However, one drawback of the greedy
the learning-based sketch-
entries. Our first proposed
algorithm is based algorithm is its slower training time. We fix this issue and propose approaches for learning sketching matrix for both low-rank approximation and Hessian ap- proximation for second order optimization. The latter is helpful for a range constrained optimization problems, such LASSO and matrix estimation with a nuclear norm constraint. Both approaches achieve good accuracy with a fast run- ning time. Moreover, our experiments suggest that our algorithm can still reduce the eтror significantly even if we only have very limited number of training matrices.
INTRODUCTION
The work of (Indyk et al., 2019) investigated learning-based sketching algorithms for low-rank approximation. A sketching algorithm is method of constructing approximate solutions for opti- mization problems via summarizing the data. In particular, linear sketching algorithms compress data by multiplication with sparse "sketch matrix" and then use just the compressed data to find an approximate solution. Generally, this technique results in much faster or more space-efficient algorithms for fixed approximation error. The pioneering work of Indyk et al. (2019) shows is possible to learn sketch matrices for low-rank approximation (LRA) with better average performance than classical sketches.
In this model, we assume inputs come from an unknown distribution and learn a sketch matrix with strong expected performance over the distribution. This distributional assumption is often realistic there are many situations where sketching algorithm is applied to large batch of related data. For example, genomies researchers might sketch DNA from different individuals, which is known to exhibit strong commonalities. The tighhpeeformancc computing industry also sketching, e.g.. researchers at NVIDIA have created standard implementations sketching algorithms for CUDA, a widely used GPU library. They investigated the (classical) sketched singular value decomposition (SVD), but found that the solutions were not accurate enough across a spectrum of inputs (Chien & Bernabeu, 2019). This is precisely the issue addressed by the learned sketch paradigm where optimize for "good" average performance across range of inputs.
"All authors contributed equally.

---

Published as a conference paper at ICLR 2023
While promising results have been shown using previous learned sketching techniques, notable gaps remain. In particular, all previous methods work by initializing the sketching matrix with a random sparse matrix, e.g., each column of the sketching matrix has single non-zero value chosen at uniformly random position. Then, the values of the non-zero entries are updated by running gradient descent training data set, or via other methods. However, the locations of the non-zero entries are held fixed throughout the entire training process.
Clearly this is sub-optimal. Indeed, suppose the input matrix A is an equal to the d x d identity matrix, and remaining rows equal to O. A random sketching matrix S with a single non-zero per column known require m = (2(d²) rows in order for S-A to preserve the rank of A (Nelson & Nguyên, 2014); this follows by a birthday paradox argument. On the other hand, it is clear that if is a d n matrix with first d rows equal to the identity matrix, then ||S AzIl2 = |xll for all vectors and so S preserves not only the rank of A but all important spectral properties. A random matrix would be very unlikely to choose the non-zero entries in the first d columns of S so perfectly, whereas algorithm trained to optimize the locations of the non-zero entries would notice and correct for this. This is precisely the gap in our understanding that we seek to fill.
matrix with first d rows
Learned CountSketch Paradigm of Indyk et al. (2019). Throughout the paper, we assume our data A E R is sampled from an unknown distribution D. Specifically, we have a training set Tr = {A...... A E DD. The generic form of our optimization problems is minx f (A, X), where A Radd is the input matrix. For a given optimization problem and set S of sketching matrices, define ALG(S, A) to be the output of the classical sketching algorithm resulting from using S; this uses the sketching matrices in S to map the given input A and construct an approximate solution X. We remark that the number of sketches used by an algorithm can vary and in its simplest case, S is a single sketch, but in more complicated sketching approaches we may need to apply sketching more than once hence may also denote a set of more than one sketching matrix.
The learned sketch framework has two parts: (1) offline sketch learning and (2) "online" sketching (i.e., applying the learned sketch and some sketching algorithm to possibly unseen data). In offline sketch learning, the goal is to construct CountSketch matrix (abbreviated as CS matrix) with the minimum expected error for the problem of interest. Formally, that is,
arg mincs s EAET (A, ALG(S, A)) - [(A, X*) = arg mincs s E ALG(S, A)),
where X* denotes the optimal solution. Moreover, the minimum is taken over all possible construc- tions of CS. We remark that when ALG needs more than one CS to be learned (e.g., in the sketching algorithm we consider for LRA), we optimize each CS independently using a surrogate loss function.
In the second part of the learned sketch paradigm, take the sketch from part and use it within sketching algorithm. This learned sketch and sketching algorithm can be applied, again and again, to different inputs. Finally, we augment the sketching algorithm provide worst-case guarantees when used with learned sketches. The goal is to have good performance on A E DD while the worst-case performance on A & DD remains comparable to the guarantees of classical sketches. We remark that the learned matrix S is trained offline only once using the training data. Hence, no additional computational cost is incurred when solving the optimization problem on the test data.
Our Results. In this work, in addition learning the values of the non-zero entries, we learn the locations of the non-zero entries. Namely, we propose three algorithms that learn the locations of the non-zero entries in CountSketch. Our first algorithm (Section 4) based greedy search. The empirical result shows that this approach can achieve a good performance. Further, we show that the greedy algorithm is provably beneficial for LRA when inputs follow a certain input distribution (Section F). However, one drawback of the greedy algorithm is its much slower training time. We then fix this issue and propose two specific approaches for optimizing the positions for the sketches for low-rank approximation and second-order optimization, which run much faster than all previous algorithms while achieving better performance.
For low-rank approximation, our approach is based on first sampling small set of rows based on their ridge leverage scores, assigning each of these sampled rows to a unique hash bucket, and then placing each non-sampled remaining row in the hash bucket containing the sampled row for which it is most similar to, i.c., for which it has the largest dot product with. We also show that the worst-case guarantee of this approach is strictly better than that of the classical Count-Sketch (see Section 5).
2

---

Published as a conference paper at ICLR 2023
For sketch-based second-order optimization where we focus the case that >> d, observe that the actual property of the sketch matrix we need is the subspace embedding property. We next optimize this property of the sketch matrix. We provably show that the sketch matrix S needs fewer rows, with optimized positions of the non-zero entries, when the input matrix A has a small number of rows with heavy leverage score. More precisely, while CountSketch takes O(d²/ (5e2) rows with failure probability 5, in our construction, S requires only ((d polylog(1/ee) + log( [1)))/=²) rows if A has at most volylog(L/e)// rows with leverage score at least e/d. This is quadratic improvement in d and an exponential improvement in 6. In practice, it is not necessary to calculate the leverage scores. Instead, we show in our experiments that the indices of the rows of heavy leverage score can be learned and the induced S' is still accurate. We also consider a new learning objective, that is, we directly optimize the subspace embedding property of the sketching matrix instead of optimizing the error in the objective function of the optimization problem in hand. This demonstrates a significant advantage over non-lcarned sketches, and has a fast training time (Section 6).
We show strong empirical results for real-world datasets. For low-rank approximation, our methods reduce the errors by 70% than classical sketches under the same sketch size, while we reduce the errors by 30% than previous learning- ceaninngbbased sketches. For second-orde optimization, we show that the convergence rate can be reduced by 87% over the non-learned CountSketch for the LASSO problem on a real-world dataset. We also evaluate our approaches in the few-shot learning setting where we only have a limited amount of training data (I (Indyk et al., 2021). We show our approach reduces the error significantly even if we only have one training matrix (Sections 7 and 8). This approach clearly runs faster than all previous methods.
Additional Related Work. In the last few years, there has been much work leveraging machine learning techniques to improve classical algorithms. We only mention few examples here which are based on learned sketches. One related body of work is data-dependent dimensionality reduction, such as an approach for pair-wise/multiwiis similarity preservation for indexing big data (Wang et al., 2017), learned sketching for streaming problems (Indyk et al., 2019: Aamand et al., 2019; Jiang et al., 2020; Cohen et al., 2020; Eden et al., 2021; Indyk et al., 2021), learned algorithms for nearest neighbor search (Dong et al., 2020), and method for learning linear projections for general applications (Hegde et al., 2015). While we also learn linear embeddings, our embeddings are optimized for the specific application of low rank approximation. In fact, one of our central challenges is that the theory and practice of learned sketches generally needs to be tailored to each application. Our work builds off of (Indyk et al., 2019), which introduced gradient descent optimization for LRA, but a major difference is that we also optimize the locations of the non-zero entries.
2 PRELIMINARIES
Notation. Denote the canonical basis vectors of R. by c1,..., Pu-- Suppose that A has singular value decomposition (SVD) A = UEVT. Define [A] = U, V. to be the optimal rank-k approximation to A, computed by the truncated SVD. Also, define the Moore-Penrose pseudo-inverse of A to be A = VI-'U" where - is constructed by inverting the non-zero diagonal entries. Let row(A) and col(A) be the row space and the column space of A, respectively.
CountSketch. We define Se E Ix as classical CountSketch (abbreviated as CS). It is a sparse matrix with one nonzero entry from {+1} per column. The position and value of this nonzero entry are chosen uniformly at random. CountSketch matrices can be succinetly represented by two vectors. We define p E [m]". , E IR" as the positions and values of the nonzero entries, respectively. Further, we let CS(p, v) be the CountSketch constructed from vectors and v.
Below we define the objective function f-))) and a classical sketching algorithm ALG(S, A) for each individual problem.
Low-rank approximation (LRA). LRA, we find a rank-k approximation of our data that mini- mizes the Frobenius norm of the approximation error. For A E Raxd, min B fLRA (A, B) min x ||A- BII Usually, instead of outputting the a whole B Raxd the algorithm outputs two factors Y E Rank and X E Rkxd such that B Y for efficiency.
3

---

Published as a conference paper at ICLR 2023
Indyk al. (2019) considered Algorithm which only compresses one side of the input matrix A. However, in practice often both dimensions of the matrix A are large. Hence, in this work we consider Algorithm 2 that compresses both sides of A.
Constrained regression. Given a vector E IR", a matrix A E Ruxd (n > d) and a convex set C, want to find x: to minimize the squared error
minzec fREG ((A b), x) minzed || - -b1
(2.1)
Iterative Hessian Sketch. The Iterative Hessian Sketching (IHS) method (Pilanci & Wainwright, 2016) solves the constrained least-squares problem by iteratively performing the update
{ |S - =,)112 (AT Ar;) mi)
T++1 = arg min
|
(2.2)
#EC
where S1+1 is a sketching matrix. is not difficult to see that for the unsketched version (S1+1 the identity matrix) of (2.2), the optimal solution x²+1 coincides with the optimal solution to the original constrained regression problem (2.1). The IHS approximates the Hessian A A by a sketched version (S)+1A)T (5 AA) to improve runtime, A typically has very few rows.
Algorithm 1 Rank-k approximation of A using a sketch S (see (Clarkson & Woodruff, 2009, Sec. 4.1.1))
Learning-Based Algorithms in the Few-Shot Setting. Recently, Indyk et al. (2021) studied learning-based algorithms for LRA in the setting where we have access to limited data or comput- ing resources. We provide brief explanation of learning-basee algorithms in the Few-Shot setting in Appendix A.3.
Input: A E Baxd S Rmxx
1: V. + COMPACTSVD(SAA)
(r rSak(SAA) E Baxr V E Rx
2: Return: [AV]&V
Leverage Scores and Ridge Leverage Scores. Given a matrix A, the leverage score of the i-th row of A is defined to be T; == a,(A' A)taT which is the squared 42-norm of the i-th row of U, where A = UΣV" is the singular value decomposition of A. Given a regularization pa- rameter X, the ridge leverage score of the i-th row a, of is defined to be т, := (A XT)ta. Our learning-based algorithms ploys the ridge leverage score sampling tech- nique proposed in (Cohen et al., 2017), which shows that sampling proportional to ridge lever- age scores gives a good solution to LRA.
Algorithm 2 ALGLRA LownrCH Sar- los (2006); Clarkson & Woodruff (2017); Avron et al. (2017).
Input: A E Raxd, S E Rm R
Rmaxd
VE E RR Xn W E IR xd
T TD
1: Uc [To Tel
VART,
SAWT with Uc, Un orthogonal
2: G G VAWT. Zi.2'r + 10/ GUDl
T Z
3: Zz + -[Z((T)))
7 +-
0
4: Z + Zr.ZR_
5: return: AR ZSA in form Paxks Qkxd
3 DESCRIPTION OF OUR APPROACH
We describe our contributions to the Sarning-based sketching paradigm which, as mentioned, is to learn the locations of the non-zero values in the sketch matrix. To learn a CountSketch for the given training data set, locally optimize the following in two stages:
ming EAED (A, (4G(S,,
(3.1)
(I) compute the positions of the non-zere entries, then (2) fix the positions and optimize their values.
Stage 1: Optimizing Positions. In Section 4, we provide greedy search algorithm for this stage, as our starting point. In Section 5 and 6, we provide our specific approaches for optimizing the positions for the sketches for low-rank approximation and second-order optimization.
Stage 2: Optimizing Values. This stage is similar to the approach of Indyk et al. (2019). However,
instead of the power method, 2019), and we pass it our objective min,er" EAED (A, ALG v), A))], implemented as a chain of differentiable operations. It will automatically compute the gradient using the chain rule. We
automatic differentiation package, PyTorch (Paszke et al.,
4

---

Published as a conference paper at ICLR 2023
also consider new approaches to optimize the values for LRA (proposed in Indyk et al. (2021), see Appendix A.3 for details) and recond-order optimization (proposed in Section 6).
Worst-Cases Guarantees. In Appendix D, we show that both of approaches for the above two problems can perform no worse than a classical sketching matrix when A does not follow the distribution D. In particular, for LRA, show that the sketch monotonicity property holds for the time-optimal sketching algorithm for low rank approximation. For second-order optimization, we propose an algorithm which runs in input-sparsity time and can test for and use the better of random sketch and a learned sketch.
4 SKETCH LEARNING: GREEDY SEARCH
When S is CountSketch, computing S4 amounts to hashing the n rows of A into the m < rows of SA. The optimization is combinatorial optimization problem with an em- pirical risk minimization (ERM) objective. The naive solution is to compute the objective value of the exponentially many (m") possible place- ments, but this is clearly intractable. Instead, we iteratively construct full placement in greedy fashion. We start with S as zero matrix. Then, iterate through the columns of in an order determined by the algorithm, adding a nonzero entry to each. The best position in each column is the one that minimizes Eq. (3.1) if an entry were to be added there. For each column, we evaluate Eq. (3.1) (m) times, once for each prospective half-buil sketch.
Algorithm 3 POSITION OPTIMIZATION: GREEDY
SEARCH
Input: f. ALAGG, Tr
(Ai,... AN E Raxd];
sketch dimension m
1: initialize Sr. = 1 m
2: for i = to n do
3: j j+ arg min Σ
ALG(S)
= =eee.
jE(ml AETr
4: S2 - Sr + (eze.)
5: end for
6: return p for S CS(p, v)
While this greedy strategy is simple to state, additional tactics are required for each problem to make more tractable. Usually the objective evaluation (Algorithm 3, line 3) is too slow, so we must leverage our insight into their sketching algorithms to pick proxy objective. Note that we can reuse these proxies for value optimization, since they may make gradient computation faster too.
Proxy objective for LRA. For the two-sided sketching algorithm, X, Y has the form Y AR Y and X XSA. where S and R are both CS matrices, so we optimize the positions in both and R. We cannot use f(A, ALG R, A)) as our objective because then we would have to consider combinations of placements between S and R. To find proxy, note that prerequisite for good performance is for row (S) and col( (AAR T) to both contain a good rank-k & ap- proximation to A (see proof of Lemma C.5). Thus, we decouple the optimization of S and R. The proxy objective for S is [AV]& V AAl where SA = UΣVT, In this expression, X = [AV] VT is the best rank- approximation to in row(SA) The proxy objective for is defined analogously.
assume that the two factors
In Appendix F, show the greedy algorithm is provably beneficial for LRA when inputs follow the spiked covariance or the Zipfian distribution. Despite the good empirical performance we present in Section 7, one drawback is its much slower training time. Also, for the iterative sketching method for second-order optimization, it is non-trivial to find a proxy objective because the input of the i-th iteration depends
the solution the (i 1)-th iteration, for which the greedy approach sometimes does not give a good solution. In the next section, we will propose our specific approach for optimizing the positions of the sketches for low-rank approximation and second-order optimization, both of which achieve a very high accuracy and can finish in a very short amount of time.
5 SKETCH LEARNING: LOW-RANK APPROXIMATION
Now present conceptually new algorithm which runs much faster and empirically achieves similar error bounds as the greedy search approach. Moreover, we show that this algorithm has strictly better guarantees than the classical Count-Sketch.
To achieve this, we need a more careful analysis. To provide some intuition, rank(S = A and SA UEVT, T, then the rank-k approximation cost is exactly |VV A p the projection cost
5

---

Published as a conference paper at ICLR 2023
onto (V). Minimizing it is equivalent maximizing the sum of squared projection coefficients:
arg min A AVV = arg min
Σ (E(n)
(1A11
Σ vz) = arg max
Σ Σ((A,)
(A. vy)²
JEIM
S
KE(M]JE[
As mentioned, computing SA actually amounts hashing the rows of A to the m rows of SA. Hence, intuitively, if we can put similar rows into the same bucket, we may get smaller error.
Algorithm 4 Position optimization: Inner Product
Our algorithm is given in Algorithm 4. Suppose that we want to form matrix S with m rows. At the beginning of the algorithm. we sample rows according to the ridge leverage scores of A. By the property of the ridge leverage score, the subspace spanned by this set of sampled rows contains an approximately optimal solution to the low rank approximation problem. Hence, we map these rows separate "buckets" of S.A. Then, need to decide the locations of the remaining rows (i.e., the non-sampled rows). Ideally, want similar rows to be mapped into the same bucket. To achieve this, we use the sampled rows as reference points and assign each (non- sampled) row A, to the p-th bucket in SA the normalized row A; and Cp have the largest inner product (among all possible buckets).
Input: A E Rnxd, average of Tr; sketch dim. m 1: initialize Si, S2 = ① m
2: Sample set C = {C1 ... -m } of rows using ridge leverage score sampling (see Section 2).
3: for i to n do
Cp (Tc, "AAla
4: Pi, + arg max pE)))
5: Si < 1/j
6: end for
7: for = 1to m do
8:
1; +--
9:
A +- restriction of A to rows in I,
10: u; < the top left singular vector of (0)
11: S Ti] 4 u
12: end for
for i ===1=t0 to m do
14: f index such that C, is the g-tth row of A
Once the locations of the non-zero entries are fixed, the next step is to determine the values of these entries. We follow the same idea pro- posed in (Indyk et al., 2021); for each block A one natural approach is to choose the unit vector s; R that preserves as much of the Frobenius norm of A as possible, i.e., to maximize ss A 1| Hence, we set s; to be the top left singular vector of A In our experiments, we observe that this step reduces the error of downstream value optimizations performed by SGD.
92 +
16: end for
17: return S₁ or (3,
To obtain a worst-case guarantee, we show that w.h.p.. the row span of the sampled rows C, is good subspace. We set the matrix S2 to be the sampling matrix that samples C. The final output of our algorithm is the vertical concatenation of S1 and S2. Here S₁ performs well empirically, while S₃ has worst-case guarantee for any input. Combining Lemma E.2 and the sketch monotonicity for low rank approximation in Section D, we get that Olk log + k/e) rows is enough for (1 + e)- approximation for the input matrix A induced by Tr, which is better than the (((A2) rows required of non-leerned Count-Sketch, even if its non-zero values have been further improved by the previous learning-base based algorithms in (Indyk et al., 2019; 2021). 1). As a result, under the assumption of the input data, may expect that S will still be good for the test data. We defer the proof to Appendix E.
In Appendix A, we shall show that the assumptions we make in Theorem 5.1 are reasonable. We also provide empirical comparison between Algorithm 4 and some of its variants, well as some adaptive sketching methods on the training sample. The evaluation result shows that only our algorithm has a significant improvement for the test data, which suggests that both ridge leverage score sampling and row bucketing are essential.
Theorem 5.1. Let S 22mxn be given by concatenating the sketching matrices S1, S2 computed by Algorithm 4 with input A induced by Tr and let B E R Then with probability at least 1 - 8, we have minrank-k (rew(W)Srow((eeeee
B - X X < (1 + +) B Bxl
p if one of the following holds:
1. m = O(B (k log ls + k/c)), 5 = 0.1, and (B) 2 (A) for all i [n]].
2. Ol log tt + k/e), 0.1 + 1.1 and the total variation distance dev (p, 4) B, where (BB
p, q are sampling probabilities defined as p; Σ. [(A]
(((A) and q:
- T(B
Time Complexity, As mentioned, an advantage of our second approach is that it significantly reduces the training time. We now discuss the training times of different algorithms. For the value-learning

---

Published as a conference paper at ICLR 2023
algorithms in (Indyk al., 2019), each iteration requires computing differentiable SVD to perform gradient descent, hence the runtime is at least a(n : T), where na is the number of iterations (usually set 500) and T is the time to compute an SVD. For the greedy algorithm, there are m choices for each column, hence the runtime is at least [mmm T). For our second approach, the most complicated step is compute the ridge leverage scores of A and then the SVD of each submatrix. Hence, the total runtime is at most (TT). We note that the time complexities discussed here are all for training time. There is no additional runtime cost for the test data.
6 SKETCH LEARNING: SECOND-ORDER OPTIMIZATION
In this section, we consider optimizing the sketch matrix in the context of second-order methods. The key observation is that for many sketching-baser second-order methods, the crucial property of the sketching matrix is the so-called subspace embedding property: for matrix A E Rnxdd we say a matrix S E Rm*x is (1 + c)-subspace embedding for the column space of A if (1 c) ||Az||| 1 ||SAz||2 S (1 + e) || for all E Rd. For example, consider the iterative Hessian sketch, which performs the update (2.2) to compute {x+),. Pilanci & Wainwright (2016) showed that if -.....,,++, are (1 + ((p))-subspace embeddings of A, then ||A((² - x² Il2 S p² ||Ax" Thus, if S; is good subspace embedding of A and we will have a good convergence guarantee. Therefore, unlike (Indyk et al., 2019), which treats the training objective in black-box manner, we shall optimize the subspace embedding property of the matrix A.
Optimizing positions. We consider the case that A has a few rows of large leverage score, as well as access to an oracle which reveals a superset of the indices of such rows. Formally, let (A4) be the leverage score of the i-th row of A and I = {i : (x(A) 2 1/3 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 explicitly and apply Count-Sketer to the remaining rows, i.e., the rows in [n]] \ L. Up to permutation of the rows, we can write
A₁ Aye
and S = 5-66))
A =
(6.1)
where S' is a random Count-Sketch matrix of m rows. Clearly S has a single non-zero entry per column. We have the following theorem, whose proof is postponed to Section F.2. 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 inequality,
Theorem 6.1. Let = c/d. Suppose that m = ((d/e² ((((dd/e))oo))) (polylog(1 (/c) + log(1/6))), 8 E (0, 1/m] and d= = ((11ee) polylog(1// /e)
Then, there exists a distribution on S of the form in (6. I) with m + [7] rows such that Pr (Vr E col(A), || || - 1|=1|| > c [|= } I 5. In particular, when 8 = 1/m, the sketching matrix S has O((d/e)) polylog(d/e) rows.
Hence, if there happen to be at most d polylog(1/e))ee rows of leverage score at least e/d, the overall sketch length for embedding cAlsp(A) can be reduced to O(d pol)logL1/e)++ + ())()))))/²²). quadratic improvement in d and exponential improvement in 8 over the original sketch length of Odd² (cc()) for Count-Sketch. In the worst case there could be O such rows, though empirically we do not observe this. In Section 8, we shall show it is possible learn the indices of the heavy rows for real-world data.
Optimizing values. When we fix the positions of the non-zero entries, aim to optimize the values by gradient descent. Rather than the previous black-box way in (Indyk et al., 2019) that minimizes Σ, f(A, ALG(S, A)), we propose the following objective loss function for the learning algorithm L(S. E A.EA R.) A,R, Ille. over all the training data, where R, comes from the QR decomposition of S.A,- = QiR,1 The intuition for this loss function is given by the lemma below, whose proof is deferred to Section E.3.
Lemma 6.2. Suppose that e E (0, }), S E RmX", A E Ruxd full column rank, and SA = QR is the -R-decomposition of SA. If (AR-1) -TTTRRRR AR- - -llopp I e, then S is a (1 + e)-subspace
embedding of col(A).
7

---

Published as a conference paper at ICLR 2023
Lemma 6.2 implies that if the loss function over Atrain is small and the distribution of Atest is similar to Aurain is reasonable to expect that S is a good subspace embedding of Atest Here we use the Frobenius norm rather than operator norm in the loss function because it will make the optimization problem easier to solve, and our empirical results also show that the performance of the Frobenius norm is better than that of the operator norm.
EXPERIMENTS: LOW-RANK APPROXIMATION
In this section, we evaluate the empirical performance of our learning-based approach for LRA on three datasets. For each, we fix the sketch size and compare the approximation error || - X ||A - Axl1 averaged over 10 trials. In order to make position optimization more efficient, in line 3 of Algorithm 3). instead of computing many rank-1 SVD updates, we use formulas for fast rank-1 SVD updates (Brand, 2006). For the greedy method, used several Nvidia GeForce GTX 1080 Ti machines. For the maximum inner product method, the experiments are conducted on a laptop with a 1.90GHz CPU and 16GB RAM.
Datasets. We use the three datasets from (Indyk et al., 2019): (1, 2) Friends, Logo (image): frames from a short video of the TV show Friends and of logo being painted; (3) Hyper (image): hyperspectral images from natural scenes, Additional details are in Table A.I.
Baselines. We compare our approach the following baselines. Classical CS: random Count- Sketch. 19 a sparse sketch with learned values. and random positions for the non-zero entries. Ours (greedy): a sparse sketch where both the values and positions of the non-zero entries are learned. The positions are learned by Algorithm 3. The values are learned similarly to (Indyk et al., 2019). Ours (inner product): sparse sketch where both the values and the positions of the non-zero entries are learned. The positions are learned by S in Algorithm 4. IVY19 and greedy algorithm use the full training set and our Algorithm takes the input as the average over the entire training matrix.
We also give sensitivity analysis for our algorithm, where we compare our algorithm with the following variants: Only row sampling (perform projection by ridge leverage score sampling). & sampling (Replace leverage score sampling with fn-norm downstream step), and Randomly Grouping (Use ridge leverage score sampling but randomly distribute the remaining rows). The result shows none of these variants outperforms non-learned sketching. We defer the results of this part to Appendix A.
sampling and maintain the same
Result Summary. Our empirical results are provided in Table 7.1 for both Algorithm 2 and Algo- rithm 1, where the errors take an average over 10 trials. We use the average of all training matrices from Tr, as the input to the algorithm 4. We note that all the steps of our training algorithms done on the training data. Hence, no additional computational cost is incurred for the sketching algorithm on the test data. Experimental parameters (i.c., learning rate for gradient descent) can be found in Appendix G. For both sketching algorithms, Ours are always the best of the four sketches. is significantly better than Classical CS, obtaining improvements of around 70%. It also obtains a roughly 30% improvement over IVY19.
Offinne learning
Online solving
Wall-Clock Times. The offline learning runtime is in Ta- ble 7.2, which is the time to train a sketch on Atrain We can see that although the greedy method will take much
Uurs cinner product)
0.I6l
Ours (gready)
6300 (L.75)
1VY19
193 (3min)
0.168
Classical
longer (Ih 45min), our second approach is much faster (5 seconds) than the previous algorithm in (Indyk et al., 2019) (3 min) and can still achieve similar error the greedy
Table 7.2: Runtime (in seconds) of LRA on Logo with A- 30, 60
Sketch
Logo
Friends
Hyper
Sketch
Logo
Fricads
Hyper
20, 40. Classical CS
2.371
4.073
6.344
20,40, Classical CS
0.930
1.542
2.971
20, 40, IVY19
0.687
1.048
3.764
20, 40. IVY19
0.255
0.723
1.273
20,
(gready)
0.500
20, 40, (gready)
20. 40. Ours tinner product)
0.532
0.733
2.975
20, 40, Ours tinner product)
0.205
0.407
1.223
301, 61,, Class CS
1.642
2.683
5,390
30, 60,, Classical CS
0.650
1.0575
23315
30, 61, 19
0.734
30, 60,,
0.290
30, 60. Ours (gready)
0.492
0.794
2.492
30, 60, Ours(groody)
0.197
0.106
0.717
30,60, Ours (inner product)
0.436
0.733
2,409
30, 60, Ourstinner product)
0.201
0.340
0.943
Table 7.1: Test errors for LRA. (Left: two-side sketch. Right: one-side sketch)
8

---

Published as a conference paper at ICLR 2023
Coal cotic
Figure 7.1: Test error of LASSO in Electric dataset.
algorithm. The reason is that Algorithm 4 only needs to compute the ridge leverage scores the training matrix once, which is actually much cheaper than IVY19 which needs to compute a differentiable SVD many times during gradient descent.
In Section A.4, we also study the performance of our approach in the few-shot learning setting, which has been studied in Indyk et al. (2021).
8 EXPERIMENTS: SECOND-ORDEN OPTIMIZATION
In this section, we consider the IHS on the following instance of LASSO regression:
arg min |=ll, SA f1) arg |all, SA ||rr bl(
(8.1)
where Å is a parameter. We also study the performance of the sketches on the matrix estimation with nuclear norm constraint problem, the fast regression solver (van den Brand et al. (2021)), well as the use of sketches for first-order methods. The results can be found in Appendix B. All of our experiments are conducted on a laptop with 1.90GHz CPU and 16GB RAM. The offline training is done separately using a single GPU. The details of the implementation are deferred to Appendix G.
Dataset. We the Electric dataset of residential electric load measurements. Each row of the matrix corresponds to a different residence. Matrix columns are consecutive measurements at different times. Here A E p370x9 6² | and (A,, b)train| 320, ((A, bestt 80. We set 15.
Experiment Setting, We compare the learned sketch against the classical Count-Sketch". We choose m 6d, Sd, 10d and consider the error (x) - f(=)). For the heavy-row Count-Sketch, we allocate 30% of the sketch space to the rows of the heavy row candidates. For this dataset, each row represents a specific residence and hence there is strong pattern of the distribution of the heavy rows. We select the heavy rows according to the number of times each row is heavy in the training data. We give detailed discussion about this in Appendix I. We highlight that it is still possible to recognize the pattern of the rows even if the row orders of the test data permuted. We also consider optimizing the non-zero values after identifying the heavy rows, using our new approach in Section
Results. We plot in Figures 7.1 the mean errors on logarithmic scale. The average offline training time is 3.67s to find a superset of the heavy rows over the training data and 66s to optimize the values when parameters. Note that the learned matrix trained offiine only once using the training data. Hence, no additional computational cost incurred when solving the optimization problem on the test data.
10d, which are both faster than the runtime of Indyk et al. (2019) with the same
We see all methods display linear convergence, that is, letting e denote the error in the k-th iteration, have C p²ez for some convergence rate p. A smaller convergence rate implies faster convergence. We calculate estimated rate of convergence p (et/ez)² ///////)// with k 7. We both sketches, especially the sketch that optimizes both the positions and values, show significant improvements. When the sketch size is small (6d), this sketch has convergence rate that just 13.2% of that of the classical Count-Sketch, and when the sketch size is large (10d), this sketch has a smaller convergence rate that is just 12.1%. 1%.
ittppo //archive. fos.uct. wtt.u//mmmommmmmmmmmmmmmm edd//m) /dat Rleettti of tyt oaddi agwams20111000 12014
Thhe framework of Indyk al. (2019) does not apply to the iterative sketching methods in straightforward manner, here we only compare with the classical CountSketch. For more details, please refer to Section B.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
Yi Li would like to thank for the partial support from the Singapore Ministry of Education under Tier 1 grant RG75/21. Honghao Lin and David Woodruff were supported in part by an Office of Naval Research (ONR) grant N00014-18-1-2555 Ali Vakilian was supported by NSF award CCF-1934843.
REFERENCES
Anders Aamand, Piotr Indyk, and Ali Vakilian. (learned) frequency estimation algorithms under zipfian distribution. arXiv preprini arXiv: 198, 2019.
Akshay Agrawal, Brandon Amos, Shane T. Barratt, Stephen P. Boyd, Steven Diamond, and J. Zico Kolter. Differentiable convex optimization layers. In Advances Neural Information Processing Systems, pp. 9558-9570, 2019.
Haim Avron, Kenneth L.. Clarkson, and David P. Woodruff. Sharper bounds for regularized data fitting. In Approximation, Randomization, and Combinatorial Optimization. Algorithms and Techniques, (APPROX/RANDOM) pp. 27:1-27:22, 2017.
Jean Bourgain, Sjoerd Dirksen, and Jelani Nelson. Toward unified theory of sparse dimensionality reduction in Euclidean space. Geometric and Functional Analysis, pp. 1009-1088, 2015.
Matthew Brand. Fast low-rank modifications of the thin singular value decomposition. Linear
Algebra and its Applications, 415.1, 2006.
Lung-Sheng Chien and Samuel Rodriguez Bernabeu. Fast singular value decomposition gpu. NVIDIA presentation at GPU Technology Conference, 2019. URL httpp:://eveelppp download, nvidia com http://wwwweeegpttttttooo//t//// s
on gpus v2...s-v2.pp....
Kenneth L. Clarkson and David P Woodruff. Numerical linear algebra the streaming model. In
Proceedings of the 2009.
forty- forty-first annual symposium on Theory of computing (STOC), pp. 205-214,
Kenneth I. Clarkson and David Woodruff. Low-rank approximation and regression in input sparsity time. Journal of the ACM (JACM), 63(6):54, 2017.
Edith Cohen, Ofir Geri, and Rasmus Pagh. Composable sketches for functions of frreueences: Beyond the worst case. In International Conference on Machine Learning, pp. 2057-2067. PMLR, 2020.
Michael B. Cohen, Cameron Musco, and Christopher Musco. Input sparsity time low-rank approxi- mation via ridge leverage score sampling. In Proceedings of the Twenty-Eightt Annual ACM-SIAN Symposium on Discrete Algorithms, (SODA), pp. 1758-1777, 2017.
Graham Cormode and Charlie Dickens. Iterative hessian sketch in input sparsity time. In Proceedings of 33rd Conference Neural Information Processing Systems (NeurIPS), Vancouver, Canada, 2019.
Yihe Dong, Piotr Indyk, Ilya Razenshteyn, and Tal Wagner. Learning sublinear-time indexing for nearest neighbor search. In International Conference on Learning Representations, 2020.
Talya Eden, Piotr Indyk, Shyam Narayanan, Ronitt Rubinfeld, Sandeep Silwal, and Tal Wagner. Learning-based support estimation in sublinear time. In International Conference on Learning Representations, 2021.
Chinmay Hegde, Aswin C. Sankaranarayanan, Wotao Yin, and Richard G. Baraniuk. Numax: A convex approach for learning near-isometric linear embeddings. In IEEE Transactions on Signal Processing, pp. 6109-6121, 2015.
Piotr Indyk, Ali Vakilian, and Yang Yuan. Learning-based low-rank approximations. In Advances in Neural Information Processing Systems, pp. 7400-7410, 2019.
Piotr Indyk, Tal Wagner, and David Woodruff. Few-shot data-driven algorithms for low rank approximation. Advances in Neural Information Processing Systems, 34, 2021.
10

---

Published as a conference paper at ICLR 2023
Tanqiu Jiang, Yi Li, Honghao Lin, Yisong Ruan, and David P. Woodruff. Ceaaning-uugmentt data stream algorithms. In International Conference on Learning Representations, 2020.
Xiangrui Meng and Michael W Mahoney. Low-distortion subspace embeddings in input-sparsity time and applications to robust linear regression. In Proceedings of the forty-fifth annual ACM symposium on Theory of computing, pp. 91-100, 2013.
Jelani Nelson and Huy L Nguyên. Osnap: Faster numerical linear algebra algorithms via sparser subspace embeddings. In Foundations of Computer Science (FOCS), 2013 IEEE 54th Annual Symposium on, pp. 117-126, 2013.
Jelani Nelson and Huy L. Nguyên. Lower bounds for oblivious subspace embeddings. In Automata, Languages, and Programming 4/st International Colloquium (ICALP), pp. 894, 2014.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, and Trevor Killeen. Pytorch: An imperative style, perfhoreffrmmcce deep learning library. 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::8, 2016.
Tamas Sarlos. Improved approximation algorithms for large matrices via random projections. In 47th
Annual IEEE Symposium on Foundations of Computer Science (FOCS), pp. 143-152, 2006.
Jan van den Brand, Binghui Peng, Zhao Song, and Omri Weinstein. Training (overparametrized) neural networks in near-linear time. In James R. Lee (ed.), 12th Innovations in Theoretical Computer Science Conference, ITCS, volume 185, pp. 63:1-63:15 2021.
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/CB0978051 1194499999981199908880 1794308 0006.
Jingdong Wang, Ting Zhang, Nicu Sebe, and Heng Tao ShenWang. A survey on learning to hash. In IEEE Transactions on Pattern Analysis and Machine Intelligence, pp. 769-790, 2017.

---

Published as a conference paper at ICLR 2023
A
ADDITIONAL EXPERIMENTS, LOW-RANK APPROXIMATION
The details (data dimension, Nirain, etc.) are presented in Table A.I.
Name
Type
Dimension
Nirain
Nrest
Friends
Image
5760 1080
400
100
Logo
Image
5760 X 1080
400
100
Hyper
Image
1024 x 768
400
100
Table A.1: Data set descriptions
A.1 SENSITIVITY ANALYSIS OF ALGORITHM 4
Sketch
Logo
Friends
Hyper
Ours(inner product)
0.311
0.470
1.232
l2 Sampling
0.698
0.935
1.293
Only Ridge
0.994
1.493
4.155
Randomly Grouping
0.659
1.069
2.070
Table A.2: Sensitivity analysis for our approach (using Algorithm from Indyk et (2019) with one sketch)
In this section we explore how sensitive the performance of our Algorithm is the ridge leverage score sampling and maximum inner product grouping process. We consider the following baselines:
& norm sampling: sample the rows according their squared length instead of doing ridge leverage score sampling.
Only ridge leverage score sampling: the subspace spanned by only the sampled rows from ridge leverage score sampling.
Randomly grouping: we put the sampled rows into different buckets as before, but randomly divide the non-sampled rows into buckets.
The results are shown in Table A.2. Here set k 30, 60 as example. To show the difference of the initialization method more clearly, we compare the error using the one-sided sketching Algorithm and do not further optimize the non-zeros values. From the table we can see both that ridge leverage score sampling and the downstream grouping process are necessary, otherwise the error will be similar or even worse than that of the classical Count-Sketch.
A.2 TOTAL VARIATION DISTANCE
As have shown in Theorem 5.1, if the total variation distance between the sampling probability distributions p and q is (L)) we have a worst-case guarantee of O log k + k/e), which is strictly better than the s2(k22 lower bound for the random CountSketch, even when its non-zero values have been optimized. We now study the total variation distance between the train and test matrix in our dataset. The result shown in Figure A.1. From the figure we can see that for all the three dataset, the total variation distance is bounded by constant, which suggests that the assumptions are reasonable for real-world data.
23
350
Figure A.l: Total variation distance between train and test matrices. left: Logo, middle: friend, right: Hyper.
12

---

Published as a conference paper at ICLR 2023
A.3 LEARNING- BEARNING-BASED ALGORIHHMS FOR LOW-RANK APPROXIMATION IN THE FEW-SHOT SETTING
In this section, we will give a brief explanation of the two algorithms proposed in Indyk et al. (2021).
Both algorithms aim to optimize the non-zero values of a Count-Sketch matrix under fixed locations of the non-zero entries.
One-shot closed-form algorithm. Given a sparsity pattern of a Count-Sketch matrix S E m it partitions the rows of A into m blocks A ...., A(m) as follows: let I; (j Sij = 1}. The block A E R xd is the sub-matrix of A that contains the rows whose indices are in 4, The goal here is for each block A to choose a (non-sparse) one-dimensional sketching vector E R!4,, The first approach is to set S; to be the top left singular vector of A which is the algorithm Veott2 Another approach is set 8; to be left singular vector of ( chosen randomly and proportional to its squared singular value. The main advantage of the latter approach over the previous one is that it endows the algorithm with provable guarantees on the LRA error. The 1Shot2 Vec algorithm combines both ways, obtaining the benefits of both approaches. The advantage of these two algorithms is that they extract a sketching matrix by an analytic computation, requiring neither GPU access nor auto-grat tutt--rraieet functionality.
Fhw-shot SGD algorithm. In this algorithm, the authors propose a new loss function for LRA, namely,
min E ST. ST SU - -oll
CS TET
where A UEVT is the SVD-decomposition of A and Us E TR denotes the submatrix of that contains its first l columns. To R4xd denotes the result of augmenting the identity matrix of order with 1 additional zero columns on the right. This loss function motivated by the analysis of prior LRA algorithms that use random sketching matrices. It is faster to compute and differentiate than the previous empirical loss in Indyk al. (2019). In the experiments the authors also show that this loss function can achieve smaller error in shorter amount of time, using small number of randomly sampled training matrices, though the final error will be larger than that of the previous algorithm in Indyk et al. (2019) if allow longer training time and access to the whole training set Tr.
A.4 EXPERIMENTS: LRA IN THE FEW-SHOT SETTING
In the rest of this section, we study the performance of our second approach in the few-shot learning setting. We first consider the case where we only have one training matrix randomly sampled from Tr. Here, we compare our method with the 1Shot2Vee method proposed in (Indyk et al., 2021) 1) the same setting (k- 10, m 40) as in their empirical evaluation. The result is shown in Table A.3. Compared to Veeo our method reduces the error by around 50%, and has an even slightly faster runtime.
Indyk et al. (2021) also proposed FewShotSGD algorithm which further improves the non-zero values of the sketches after different initialization methods. We compare the performance this approach for different initialization methods: in all initialization methods, we only use one training matrix and we use three training matrices for the FewShotSGD step. The results are shown in aim to
Table A.4. We report the minimum error of 50 iterations of the FewShotSGD because compare the computational efficiency for different methods. From the table we see that our approach plus the FewShotSGD method can achieve much smaller error, with around 50% improvement upon (Indyk et al., 2021). Moreover, even without further optimization by FewShotSGD, our initialization method for learning the non-zero locations in CountSketch obtains a smaller error than other methods (even when they are optimized with IShotSGD or FewShotSGD learning).
B ADDITIONAL EXPERIMENTS: SECOND-ORDER OPTIMIZATION
As we mentioned in Section 8, despite the number of problems that learned sketches have been applied to, they have not been applied to convex optimization, say, iterative sketching algorithms in general. To demonstrate the difficulty, we consider the Iterative Hessian Sketch (IHS) as an example. In that scheme, suppose that we have l iterations of the algorithm. Then we need independent
13

---

Published as a conference paper at ICLR 2023
Algorithm
Dutaset
Few shot Error
Training Time
Classical CS
Friends
Hyper
Logo
0.524
1,082
x
Ours (Initialzation only)
Sketch
Lego
0.065
0.139
Friends
0.535
Hyper
Vshot2Vee
Fricads
Logo
0.306
0,171
5.680
5.682
Ours + FewSShiSSGD
1Shot Vec only
0.171
0.048
0.306
0.125
0.795
0.443
Hyper
0.795
1.054
[Shot Vec FewShot SGD
0.104
0.229
0.636
Ours (inner product)
Friends
Logo
0.139
4,773
4.515
Classical CS
Classical CS + FewSShot SGD
0.331
0.173
0.524
0.279
0.771
1.082
Hyper
0.535
0.623
Table A.4: Test errors for LRA in the few-
Table A.3: Test emrors and training times for LRA in the one-shot setting (using Alg. with one sketch)
shot setting (using Alg. 1 from Indyk al. (2019) with one sketch)
sketching matrices (otherwise the solution may diverge). A natural way is to follow the method in (Indyk et al., 2019), which is to minimize the following quantity
min E f(A.ALL(S),.......AA S1....g. AED
where the minimization is taken over &s Count-Sketch matrices Si...., St. In this case, however, calculating the gradient with respect to S₁ would involve all iterations and in each iteration we need to solve constrained optimization problem. Hence, it would be difficult and intractable to compute the gradients. An alternative way is to train l sketching matrices sequentially, that is, learn the sketching matrix for the i-th iteration using local loss function for the i-th iteration, and then using the learned matrix in the i-th iteration to generate the training data for the (i 1)-st iteration. However, the empirical results suggest that it works for the first iteration only, because in this case the training data for the (i + 1)-st iteration depends on the solution the i-th iteration and may become farther away from the test data in later iterations. The core problem here is that the method proposed in Indyk et al. (2019) treats the training process in black-box way, which is difficult to extend to iterative methods.
B.1 THE DISTRIBUTION OF THE HEAVY ROWS
In our experiments, we hypothesize that in woal-world data there may be an underlying pattern which can help us identify the heavy rows. In the Electric dataset, each row of the matrix corresponds to a specific residence and the heavy rows are concentrated on some specific rows.
To exemplify this, we study the heavy leverage score rows distribution over the Electric dataset. For a row i E [370], let fi denote the number of 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 f. > 0.
(195,320), (278,320), (361,320), (207,317), (227,285), (240,284), (219,270), (275,232), (156,214),
(322,213), 1993, (190,192), (160,191), (850,181), (63.176), (42,168), (162,148), (356, 129), (363,110), 1362,105), (338,95), (215,94), (234,93), (289,81), (97,80), (146,70), (102,67), (98,58), (48.57), (349,53), (165,46), (101,41), (352,40), (293,34), (344,29), (268,21), (206,20), (217,20), (327,20), (340,19), (230,18), (359,18), (297,14), (357,14), (161,13), (245,10), (100,8), (85.6), (212,6),
(313.6), (129,5). (130.5), (366,5), (103,4), (204,4), (246,4), (306,4), (138,3), (199,3), (222,3), (360,3), (87,2), (154,2), (209,2), (123,1), (189,1), (208,1), (214,1), (221,1), (224,1), (228,1), (309,1), (337,1), (343,1)
Observe that the heavy rows are concentrated on a set of specific row indices. There are only 30 rows i with f: 50. We view this strong evidence for our hypothesis.
Heavy Rows Distribution Under Permutation. We note that even though the order of the rows has been changed, we can still recognize the patterns of the rows. We continue to use the Electric dataset as example. To address the concern that permutation may break the sketch, we can measure the similarity between vectors, that is, after processing the training data, we can instead test similarity on the rows of the test matrix and use this to select the heavy rows, rather than an index which may simply be permuted. To illustrate this method, we use the following example on the Electric dataset, using locality sensitive hashing.
After processing the training data, we obtain a set I of indices of heavy rows. For each i E I, pick q independent standard Gaussian vectors 9i, 92,99 33 Rd. and compute f(x))
14

---

Published as a conference paper at ICLR 2023
(9f ri, 92 Ts: rg: r4) R2, where r; takes an average of the i-th rows over all training sets. Let A be the test matrix. For each E I, let j; = argmin; IIf(A,) - f(r))l2. We take the je-th row to be a heavy row in our learned sketch. This method only needs an additional O(1) passes over the entries of A and hence, the extra time cost is negligible. To test the performance of the method, randomly pick a matrix from the test set and permute its rows. The result shows that when is small, can roughly recover 70% of the top-k heavy rows, and we plot below the regression error using the learned Count-Sketch matrix generated this way, where we set m = 90 and k: = 0.3m = 27. We can see that the learned method still obtains significant improvement.
counteicetch karmeelposkisn-only
Iteration round
Figure B.1: Test error of LASSO on Electric dataset
B.2 MATRIX NORM ESTIMATION WITH A NUCLEAR NORM CONSTRAINT
In many applications, for the problem
X
arg min X X Xd2
||AX
BI
it is reasonable to model the matrix X as having low rank. Similar to E <mminimizationn for com- pressive sensing, a standard relaxation of the rank constraint minimize the nuclear norm of X, defined as | |l == > j=1
minn d2) (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 XeR4 Xd2
BIß such that |X S P.
where p 0 is a user-defined radius as a regularization parameter.
We conduct Iterative Hessian Sketch (IHS) experiments on the following dataset:
Tunnel': The data set a time series of gas concentrations measured by eight sensors a wind tunnel. Each (A, B) corresponds to a different data collection trial. A' e 13333005 B e
R 13330066 |(A, B) train = 144, (AA, best = 36. In our nuclear norm constraint, we set
p 10.
Experiment Setting: We choose m 7d, 10d for the Tunnel dataset. We consider the error ||AX - B12 |X'* B12 The leverage scores of this dataset are very uniform. Hence, for this experiment we only consider optimizing the values of the non-zero entries.
Results of Our Experiments: We plot on logarithmic scale the mean errors of the dataset in Figures B.2, We can see that when m 7d, the gradient- based sketch, based on the first itcrations, has rate of convergence that is 48% of the random sketch, and when = 10d, the gradient-based sketch has a rate of convergence that 29% the random sketch.
B.3 FAST REGRESSION SOLVER
Consider an unconstrained convex optimization problem min. f(x), where f is smooth and strongly convex, and its Hessian 2 is Lipschitz continuous. This problem can be solved by Newton's
Phttps: //arcchi///ccveeee uus.uci. turbulent tasttt ittures xtures
15

---

Published as a conference paper at ICLR 2023
0-
vgg 20.2000mm0)
10
Itrrtton round
famition round
Figure B.2: Test error of matrix estimation with nuclear norm constraint on the Tunnel dataset
12-
Seavin 221 IEETALie 32 (1
10-
Grnsorkita 11 quaadii Driksaed CS
64
64
02
02
00
188
125
1.3
2.7
651
325
2.59
325
2.5
Figure B.3: Test error of the subroutine in fast regression on Electric dataset.
Seaiser Oomiitamed
Figure B.4: Test error of fast regression on Electric dataset
method, which iteratively performs the update
I++1 Zt 1---an -ra min
(V V
(B.1)
provided it is given good initial point zo. In each step. requires solving regression problem of the form min, || Az y which, with access to A, can be solved with fast regression solver in (van den Brand et al., 2021). The regression solver first computes a preconditioner R via a QR decomposition such that SAR has orthonormal columns, where S is sketching matrix, then solves = arg min. [A()) T (AR))' - vll, by gradient descent and returns It 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 a small number of rows.
In this section, We consider the unconstrained least squares problem min, f(x) with f(x) || Ax b1l2 using the Electric dataset, using the above fast regression solver.
Training: Note that f(x) = A A, independent of æ. In the t-th round of Newton's method, by (B.1). we need to solve regression problem min; || - -
with y = fxxx). Hence, we can use the same methods in the preceding subsection to optimize the learned sketch S,. For a general problem where fz) depends on ær one can take to be the solution obtained from the learned sketch S, to generate A and y for the (t + 1)-st round, train a learned sketch S4+1. and repeat this process.
Experiment Setting: For the Electric dataset, we set m 10d = 90. We observe that the classical Count-Sketch matrix makes the solution diverge terribly in this setting. To make clearer comparison, we consider the following sketch matrix:
Gaussian sketch: S 1 G, where G Rmxn with i.id. (O, 1) entries.
16

---

Published as a conference paper at ICLR 2023
Transform (SJLT): S is the vertical concatenation of inde- pendent Count-Sketer Seent-kkethh matrices, each of dimension m/s X n.
Sparse
We note that the above sketching matrices require more time to compute SA but need fewer rows to be a subspace embedding than the classical Sount-Sketchh matrix.
For the step length in gradient descent, we set 1 in all iterations of the learned sketches. For classical random sketches, we set the following two ways: (a) 1) = all iterations and (b)
7 1 in the first iteration and 17 0.2 in all subsequent iterations.
Experimental Results: We examine the accuracy of the subproblem min, | = and define the error to be || ARz, =||2 / /lvl22 We consider the subproblems in the first three iterations of the global Newton method. The results are plotted in Figure B.3. Note that Count-Sketch causes terrible divergence of the subroutine and is thus omitted in the plots. Still, we observe that in setting (a) of 17. the other two classical sketches cause the subroutine to diverge. setting (b) of 17, the other two classical sketches lead to convergence but their error is significantly larger than that of the learned sketches, in each of the first three calls to the subroutine. The error of the learned sketch is less than 0.01 in all iterations of all three subroutine calls, in both settings (a) and (b) of 17.
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 B.4, which clearly displays a significantly faster decay with learned sketches. The rate of convergence using heavy-rows sketches is 80.6% of that using Gaussian sparse JL sketches.
B.4 FIRST-ORDER OPTIMIZATION
In this section, we study the use of the sketch in first-order methods. Particularly, let QR S SA be the QR decomposiition for SA,, where S is sketch matrix. We use R as an (approximate) pre- conditioner and use gradient descent to solve the problem min ARr bllg- Here we use the Electric dataset where A is 370 x 9 and set S have 90 rows. The result shown in the following table, where the time includes the time to compute R. We can see that if we use a learned sketch matrix, the error converges very fast when we set the learning rate to be and 0.1, while the classical Count-Sketch will lead to divergence.
Iteration
10
100
500
Error (learned, lr 1)
2.73
I.5e-7
Error (learned, IT 0.1)
4056
605
4.04e-6
Error (learned, Ir 0.01)
4897
4085
667
0.217
Error (random, lr = 1)
N.A
N.A
N.A
N.A
Error (random, Ir 0.1)
Error (random, Ir 0.01)
4881
3790
685
1.52
Time
0,00048
0.00068
0.0029
0.0013
Table B.1: Test Error for Gradient Descent
C PRELIMINARIES: THEOREMS AND ADDITIONAL ALGORITHMS
In this section, we provide the full description of the time-optimal sketching algorithm for LRA in Algorithm 2. We also provide several definitions and lemmas that are used in the proofs of our results for LRA.
Definition C.1 (Affine Embedding). Given a pair of matrices A and B, a matrix S is affine e-embedding if for all X of the appropriate shape, ||S(AX B)1 (1 =) ||X -BI Lemma C.2 (Clarkson & Woodruff (2017); Lemma 40). Lei A be an n x matrix and let S E ) xm be a CountSketch matrix. Then with constant probability, ||SA|]] = (1 +)
The following result is shown in Clarkson & Woodruff (2017) and sharpened with Nelson & Nguyên (2013); Meng & Mahoney (2013).
Lemma C.3. Given matrices A, B with rows, a CountSketch with Orankk( (rrank(A)//a) rows is an affine e-embedding matrix with constant probability. Moreover, the matrix product SA can be computed in ((nz((A) time, where nnz( A) denotes the number of non-zero entries of matrix A.
17

---

Published as a conference paper at ICLR 2023
Lemma C.4 (Sarlos (2006); Clarkson & Woodruff (2017)). Suppose that A Raxd and B E Ruxd' rank(A)²
Let E mRxx be a CountSketch with m = 2² Let X = arg min X |AX - SBI Then,
|AX AAX - =1 S (1++)
With constant probability,
millran-k ||AX BI In other words, in
(nnz(A)) + nnz(B) + m(d++ d))) time, we can reduce the problem to a smaller (multi-resppol response whose optimal solution is (1 -appproxximate solution to the
regression) problem with original instance.
2. The (1 + <)-approximate solution X can be computed in time (uuz((A) nnz(B) mdd' min(m² d, md)))
Now we turn our attention to the time-optimal sketching algorithm for LRA. The next lemma is known, though we include it for completeness Avron et al. (2017):
Lemma C.S. Suppose that S E Rmsss x and R E Rmaxdd are sparse affine e-embedding matrices for (As, A) and ((SA) ),,, respectively. Then,
|| |SARAAA= XSA AlI S (1++) A
min rank-k
Proof. Consider the following multiple-response regression problem:
min rank-k X
|| - A
(C.1)
Note that since X = 1g is a feasible solution to Eq. (C.1), mingank. ||, - All 2 = ||A, - Alp Let S Rmsx be sketching matrix that satisfies the condition of Lemma C.4 (Item for A = Ak
and A. By the normal equations, the rank-k minimizer of ||SAXX SA is (SAA) SSA Hence,
SA All (1+ A
(C.2)
which in particular shows that + appproxxmma rank-k approximation of A exists the row space of SA. In other words,
min rank-k
IXSA- AII s (1+ Al
(C.3)
Next, let R E Ruaxd xd be a sketching matrix which satisfies the condition of Lemma C.4 (Item 1)
for == ( and B = AT. Let Y denote the rank-k minimizer of R(SA)) X -RA Hence,
[[SA)) AT S 61+
min I|XSA - AAll rank- X
Lemma C.4 (Item 1)
< + O(e)) ||A: All
Eq. (C.3)
(C.4)
Note that by the normal equations, again rowsp ( T) C rowsp RAT) and we can write where rank(2] (Z) k. Thus,
= AR Z
min ||AR XSA A < |R ZSA - A ||(SA)) AT
D = AR
rank-k
|A - AAl
D Eq. (C.4)
Lemma C.6 (Avron al. (2017); Lemma 27). For C E Rpxm D E m G E RPxp' the following problem
min 1CZD -G|2 rauk-k 2
(C.5)
can be solved in Opm''ro + ppmmp + pp'(ro + TC)) time, where TO = rank(C") & min(m', .p) and rD rank(D) & min(m,p')
18

---

Published as a conference paper at ICLR 2023
Lemma C.7. Let S E Rmsxx R Raax be CountSketch (CS) matrices such that
min rank-k X
AR
ARAXSA= XSA Alle S (1+ A AlE
(C.6)
Let V E
and W E R xd be CS matrices. Then, Algorithm 2 gives a (1 + (33 + 7))- mm))m)) + k(nms + dmm)) with constant 84
approximation in time nnz A) + O ms probubility.
g +
Proof. The approximation guarantee follows from Eq. (C.6) and the fact that V and W are affine 3-embedding matrices of AR and SA, respectively (see Lemma C.3).
The algorithm first computes C VAR", D SAWT, G time ((nnz(A)) As an example, we bound the time to compute C = VAR. Note that since V is CS, VA can be computed in ((nn((A) time and the number of non-zero entries in the resulting matrix is at most unz(A)) Hence, since R is a CS as well, C' can be computed in time O(nnz A) + nnz V A)) = (A)z)A))) Then, it takes an extra O + m + m/ time to store C, D and G in matrix form. Next, we showed in Lemma C.6, the time to compute Z in m (ms (m)))
VAWT which be done in
Algorithm 2 O 3r
Finally, takes ((nnz(A) A) + k(nms + dmR)) time to compute Q ART Z and P SASA and to return the solution in the form of Paxa Qxxd- Hence, the total runtime
+
3a
m
m
m (ms + mR) + k(nms + dmn))
O A) +
+
+
32
B
D ATTAINING WORST-CASE GUARANTEES
D.1 LOW-RANK APPROXIMATION
We shall provide the following two methods to achieve worst case guarantees: MixedSketch-whon guarantee is via the sketch monotonicity property, and approximate comparison method (a.k.a. ApproxCheck), which just approximately evaluates the cost of two solutions and takes the better one. These methods asymptotically achieve the worst-case guarantee. However, for any input matrix A and any pair of sketches S, T, the performance of the MixedSketch method on (A, S, T) is never worse than the performance of its corresponding ApproxCheck method on (A, S, T), and can be much better.
Remark D.I. Let A diag(2, 2, V2, /2), and suppose the goal is to find rank-2 approximation of A. Consider two sketches and T such that SA and T A capture span (e1, ex) and span(e2, 64), respectively. Then for both SA and TA, the best solution in the subspace of one of these two spaces (3) >approximmaiion: || A2l 4 and || Ps l1 || Pr Il? 6 where PSA and PA respectively denote the best approximation of A in the space spanned by SA and T'A.
However, if find the best rank-2 approximation of A, Z, inside the span of the union of SA and TA, then ||A4 2
Since Approx Check just chooses the better of SA and A by evaluating their costs, it misses out on the opportunity to do as well as MixedSketch.
Here, we show the sketch monotonicity property for LRA.
Theorem D.2. Let A E Raxd be an inpul matrix, V and W be n-affine embeddings, and S E R xn, R₁ E m be arbitrary matrices. Consider arbitrary extensions S1, RG 3, R (e.g.. S is a concatenation of S₁ with an arbitrary matrix with the same number of columns). Then, |A ALGLRA((S, V. W), A))I S (1 n) ||A ALGLRA((S,, R2 V. W,,
Proof. We have ||A ALGERA ( (3,R, R, V, W), A) S (1+ 7) min ARX3A All²
(1 17) min rank-& (STTM | Al which is in turn at most (1
7) mankeakkk S...rro(SS || Al (1 minrank- |AR XS, Al S <(1
19

---

Published as a conference paper at ICLR 2023
1| ALGLRA R1, V, V), where use the fact the V. W affine n-embeddings (Definition C.1), well as the fact that (col(AAR)) row(S14) _ (col(AA) row(SA)) L
ApproxCheck for LRA. We give the pseudocode for the ApproxCheck method and prove that the runtime of this method for LRA of the same order as the classical time-optimal sketching algorithm
LfA.
Algorithm 5 LRA APPROXCHECK
Input: learned sketches St. R2, VL. WL; classical sketches Sc, Re, Vc. We; 3: A E R Pr. Qr ALG (S Rr,, Vr, Wt. PcQc + ALG Re, Vc, Wc,
2: Let S E (((115))xn E (0(133))x be classical CountSketch matrices
3: As < (P2QL - R Ac ||s (PoQc A)R R 1|
4: if ÂL S. AC then
5: return PLQL
6: end if
7: return PoQc
Theorem D.3. Assume we have data A E Raxd, learned sketches Sr E R RL E
g Vz (roly(")x WL E ( which attain (1 + (())) approximaainnn classical sketches of the same size, Sc, Vc, We, which attain a + (()))(())pppp)l) approxiiatoooionn and a tradeoff parameter B. Then, Algorithm 5 attains a (1 + + min(7, apprrxxinatto in Ounzz(A) (n d) (=
1² (poly())) time.
Proof. Let (PL, QL). (Pc, Qc) be the approximate rank-k approximations of A in factored form using (Sr., Rr.) and (So, Ro). Then, clearly,
min PtQ2 All: PP - = (1+ O(min(,, 7))) ||A - Al
(D.I)
Let L = PLQL - A, C = PcQc A and M = arg mig F ||ST Rll Then, > by Lemma C.2
| s (1 (1+) Rl
(1+ 2 |/
$ (1 + 0 mmn((.)))))) ||A All
D Eq. (D.1)
Runtime analysis. By Lemma C.7, Algorithm 2 computes PL. QL and Po, Qo in time O(nnz(A) + +2)
R ga + 2 (n+ dh?
Next, once we have PL, Pt,Qr. and Pe, Qc it takes O(nnz( + gr) time
&
to compute Ar and Ac.
O(nnz(A)+ 1682 + 2
A3
dk
(n
= O(nn((A) + (n + d
B poly )
24
2
B
To interpret the above theorem, note that when e > k(n d) we can set 3 - O(K(n d)-a) so that Algorithm 5 has the same asymptotic runtime the best (1 + a)aapooxximation algorithm for LRA with the classical CountSketch. Moreover, Algorithm 5 (1 + approxaapooiiiaation when the learned sketch outperforms classical sketches, = o(e). On the other hand, when the learned sketches perform poorly, 7 = (e)), the worst-case guarantee of Algorithm 5 remains (1 + O())).
D.2 SECOND-ORDER OPTIMIZATION
For the sketches for second-order optimization, the monotonicity property does not hold. Below provide an input-sparsity algorithm which can test for and use the better of a random sketch and a learned sketch. Our theorem is as follows.
Theorem D.4. Let c E (0, 0.09) be constant and S a learned Count-Sketch matrix. Suppose that A of full rank. There is algorithm whose output is solution & which, with probability at least 0.98, satisfies that ||A( 222 S ( (min ((().)
e}) || lll where x arg min ||Ar bll, the least-squares solution. Furthermore, the algorithm runs in O log(( + po)(( time.
(2(S.)
20

---

Published as a conference paper at ICLR 2023
Algorithm 6 Solver for (D.2)
S f- learned sketch, S2 random sketch with (dd rows
2: (2,,1, 2,,2) < ESTIMATR(S; A), 1,2
3: 1* < arr min / 24.1)
4: & solution of (D.2) with S S/-
5: return &
6: function ESTIMATE(S, A)
7: 7 < sparse (1 + -ubsspce embedding matrix for d-dimensional subspaces
8: (Q. R) + (R(AA)
9: Z1 < mmin (SAAR
10: 22 + (1 + 17) r)aapppoximation to [ -llop
11: return (2 22)
Consider the minimization problem
| - (A y, v. ..0)
min DEC
(D.2)
which is used as subroutine for the IHS (ef. (2.2)). We note that in this subroutine if we let x 4 x x² x'-1 + 6 b- Ax--1, - c x'=1, we would get the guarantee of the i-th iteration of the original IHS. To analyze the performance of the learned sketch, define the following quantities (corresponding exactly to the unconstrained case in (Pilanci & Wainwright, 2016))
inf Zi( wEcolsp(4)ne--
||
Z2(S) = sup u,p
(u, S - S-
When S' is (1 + e)ssusspace embedding of colsp(A), we have Z1(S) 2 1 - e and Z2(S) S 2e.
For general sketching matrix S. the following is the approximation guarantee of 21 and 22, which are estimates of (5) and (5), respectively. The main idea that AR-1 is well-conditioned, where R is calculated in Algorithm 6.
Lemma D.5. Suppose that 17 E (0, })) is a small constant, A is of full rank and S has poly(d/n)
rows. The function ESTIMATE(S, A) returns in O((nz(()) log 1 + poly(d)) time 21, Zz which with n
2,(S) probability least 0.99 satisfy that (1 2, I- (1+n) 31 < 22 < (1-1)² 31).
(1(5) and (S)
I+"
Proof. Suppose that AR UW, where U E Rnxd has orthonormal columns, which form an orthonormal basis of the column space of A. Since T is subspace embedding of the column space of A with probability 0.99, it holds for all Rd that
1 TAR-'an 2112 < ||AR ||AR-'r/|
|T 20122
I+
1 1-n
Since
[TAR 'x1/2 = =100012 = ||0|||
and
||all, = |UWx|12 = || A=112
(D.3)
have that
| < |W < 1- ||0|| a E Rdd
(D.4)
+
It is easy to see that
|SUWyl vll
Zi(S) = min
SU
= min
xERI-1
0//
|
and thus,
min(1 11) | دl1
|UWWll
< (i( S min(1 + 17)
v=0
11112
3
1 1
21

---

Published as a conference paper at ICLR 2023
Recall that SUW SAR-! We see that
(1 - m)) (SAR - Zi(5) & (1+ o) ( (SAR )
By definition,
Z2(S) - (S -1,))ullpp
It follows from (D.4) that
(1 (1-1) n)² - s Z2()s S (1+ (1=""" wn)WWlopp and from (D.4), (D.3) and Lemma 5.36 of Vershynin (2012) that
((AR-)) - 1llop S 31.
Since
-S = =|((ARt-)) -
and
$ ( 11ap -I(AA-))) 1llop
< 5 1l.pp 1|((ARR (AR-1) ) -
-llopp>
follows that
(1- (1-1) ( - -1lop -3(1 - -
$ (S)
$ (1 + n)² ||(SAR SAR -1 + + n)
We have so far proved the correctness of the approximation and we next analyze the runtime below.
Since S and T are sparse, computing SA and TA takes ((nz(AA) time. The QR decomposition of T A, which is a matrix of size poly(d/n) x d, can be computed in poly(d/m) time. The matrix SAR-1 can be computed in poly(d) time. Since has size poly(d/n) (d/n) d, its smallest singular value can be computed in poly(d/n) time. To approximate (S),), we can use the power method to estimate ( SAR-1 - 7 up to a (1 = -f)-fo (((nnz()) )) + poly(d/n)) log(1/1)) time.
Now we are ready to prove Theorem D.4.
Proof of Theorem D.4. In Lemma D.5, we have with probability least 0.99 that
2
[+t]] (2( 3e
2 1 Zz(S)
3e(1 - e)
2
T- (i)
(1 (1+ (S)
(i
and similarly,
Zg
[1 [1=e)2
1
72( - 3c
< 1+c Zz(S)
+ 3c(1 (c) Zi(S)
<
2,
Itc (S)
(1 - c)2 (I)
Note that since S2 an c-subspace embedding with probability at least 0.99, we have that Z1 2 1 - and Zz(S))ZZ (S2) < 2.2c. Consider (S,).
First, we consider the case where Z, (S,) 1/2. Observe that Z2( 2 1 (S). We have in this case / /2, > 1/5 > 2.2c 2 Ze ( In this case our algorithm will choose S2 correctly.
Next, assume that Zi (5)) 2 1/2. Now we have with probability at least 0.98 that
(1 - 3c) Zz(S,)
2,2
Zz(S,)
3e
<
(1+ 4c)
+ 4c, i= 1,2.
(S,))
2,1
(i.)))
22

---

Published as a conference paper at ICLR 2023
Therefore, when Zz(S:) {Z))))) for some small absolute constant will have 2 < 22.2/Z,, and our algorithm will choose S1 correctly. If Zz (2(5))/2 > Cie for some absolute constant C1 > 0, our algorithm will choose S2 correctly. In the remaining case, both ratios (z(5))/Z1(5 and (z(S))/2((1) are at most max (C2,3/C,, 3)c, and the guarantee of the theorem holds automatically.
The correctness of our claim then follows from Proposition of Pilanci & Wainwright (2016), together with the fact that S2 is random subspace embedding. The runtime follows from Lemma D.5 and Theorem 2.2 of Cormode & Dickens (2019).
E SKETCH LEARNING: OMITTED PROOFS
E. PROOF OF THEOREM 5.1
We need the following lemmas for the ridge leverage score sampling in (Cohen al., 2017).
Lemma E.1 ((Cohen et al., 2017, Lemma 4)). Let = ||A Axllz /k. Then we have Ee T((A) S
2k.
Lemma E.2 ((Cohen et al., 2017, Theorem 7)). Let À ||A Aulle /k and 7; 2 T((A) be overestimate to the i-th ridge leverage score of A. Let p; = E. IfC matrix that is constructed by sampling t O(log k + los(1/6) Ps. then with probability at least 8 we have
Σ; T.) rows of A, each set as with probability
min rank-k (X)om(X))))
|| X I
S (1 -
Recall that the sketch monotonicity for low-rank approximation says that concatenating two sketching matrices S1 and S₂ will not increase the error compared the single sketch matrix S1 or S2, Now matter how S1 and S2 are constructed. (see Section D.I and Section 4 in (Indyk et al., 2019))
Proof. We first consider the first condition. From the condition that TB 2 (A) we know that if sample 0(3 (I: log & + k/c)) rows according (A). The actual probability that the i-th row B gets sampled
(1 - (A))) 0(m T. (A)) (() lol h/cc (B)))
From Ee < 2k and Lemma F.2 (recall the sketch monotonicity property for LRA), have that with probability at least 9/10, S2 is matrix such that
min rank-k X rrow( (ooow(XSSmr (S=B
XI| s (1 (1+6) -B
Hence, since S = -[S] from the the sketch monotonicity property for LRA we have that
min
| X < (1+c) B-B
Bx1
rank-k X (row(
Now consider the second condition. Suppose that {Xihism and {ii) are sequence of m OOk & + k/e) samples from [n]] according to the sampling probability distribution p and q. r.(B) where = Σ; 7:(A] and = Σ, (B) Let S be the set of index i such that X, + Y, From the
T((A)
property of the total variation distance, we get that
Pr Pr|X, + Yi| S (p. = B.
and
B|Ss|| Σ Pr| Yi < 3m.
From Markov's 's inequality get that with probability least - 1.18. < 1/((11)) 9m
10 m.
Let T be the set of index such that X; = Y. We have that with probability at least 11
1.18,
[T] >
1imm m (2k log k/e). Because that (Yi)ee is i.i.d samples according g and the
23

---

Published as a conference paper at ICLR 2023
actual sample we take is (Xi)er. From Lemma E.2 get that with probability at least 9/10, the row space Br satisfies
| XI| - (1 + B B
min
rank-& icrow X) Crow Br)
Similarly, from the the sketch monotonicity property we have that with probability at least 0.9 1.13
rank-k min
| X|| / S + 6) B B&
E.2 PROOF THEOREM 6.1
First we prove the following lemma.
Lemma F.3. Let E (0, 1/m]. It holds with probability at least 1 that
||5x112 - || S |
sup =Ecolap(A)
provided that
2 ((d log m) (dde)))"dd//),dal m) + dllo(1/6)))
1 12 E - vv((og m) minglog? (d/c), m] + log (og(1/))) log(1/6)
Proof. We shall adapt the proof of Theorem 5 in (Bourgain al., 2015) to our setting. Let T denote the unit sphere in colsp(A) A) and set the sparsity parameter Observe that |Sx ? | + || l1 and so suffices to show that
Pr P1[1IS'a11112 ||
<
for x E T. We make the following definition, as in (2.6) of (Bourgain et al., 2015):
m7
A
Σ
dygrees @ e,,
and thus, S'x1- = A Also by E S x1=
12 = || 2 one has
sup
S 1
|z
sup ZET
As. -E o|
(E.1)
#ET
Now, in (2.7) of (Bourgain et al., 2015) we instead define semi-norm
1/2
||>|| = max
Σ jE
8.y
Then (2.8) continues to hold, and (2.9) as well (2. continue hold if the supremum in the left-hand side is replaced with the left- side of (E.1). At the beginning of Theorem 5, we define (+ to be U, but each row j E is multiplied by 5.j and each row j E is zeroed out. Then we have in the first step of (4.5) that
Σ 8ig Σ JE
(fa ei)
1000 2
instead of equality. One can verify that the rest of (4.5) goes through. remains true that ||-11, < (1/vs) Il-llg- and thus (4.6) holds. One can verify that the rest the proof of Theorem 5 (Bourgain al., 2015) continues hold if we replace Σ with Ejere and maxigisn with max noting that
E Σ bú ||Pee JEL
Σ (PEcjee)) S 8 d
m Jel
m
24

---

Published as a conference paper at ICLR 2023
and
)(U()))(())
=
Σ
(E
(Edy)uu =
JEI
Thus, the symmetrization inequalities on
Σ JET-
8ig ||Pee
and
Σ
JE
continue to hold. The result then follows, observing that maxjere
||EEe]1 S v.
The subspace embedding guarantee now follows as a corollary.
Theorem 6.1. Let verv=da cc/d. Suppose that m (dd(e() (polylog)))))) + log(1/6))), E (0. 1/m) and d = ((11/c) polylog(1/c) log6(1))),, Then, there exists distribution on S with m + [ rows such that
Pr Pr(ve colsp((), || ||
>
>
= S 8.
Proof. One can verify that the two conditions in Lemma E.3 are satisfied
d
[polylog( -) log 5)) (polylog( d 2) + log 1 5))
m z
2 log 1o8,)
! ;)
The last condition satisfied if
1 2 log" 5/ polylog (
E.3 PROOF LEMMA 6.2
Proof. On the one hand, since SAR is orthogonal matrix, have
||ا|| = ||@=11 = ||SAR=
(E.2)
On the other hand, the assumption implies that
(ARR)) ((ARR))(AARR)) - I.T z0/22 S ||
that is,
(1- - c) ||=112 5 ||RR (1 (1+ (1
(E.3)
Combining both (E.2) and (E.3) leads to
vI - SARc < |ARx| & VI+e SARr
Va Rd.
Equivalently, it can written as
SAylla S Aylla S -1-c |SAyl12
Vy VyeR" Rdd
VI+C
The claimed result follows from the fact that 1/ - and 111-6511 whenever
(0, v3-..
n
25

---

Published as a conference paper at ICLR 2023
F LOCATION OPTIMIZATION IN COUNTSKETCH: GREEDY SEARCH
While the position optimization idea is simple, one particularly interesting aspect is that it is provably better than random placement in some scenarios (Theorem. F.1). Specifically, it is provably beneficial for LRA when inputs follow the spiked covariance model or Zipfian distributions, which are common for real data.
Spiked covariance model. Every matrix A E Raxd from the distribution Asp(s, E) has 8 < "heavy" rows A.,.... .A, of norm > The indices of the heavy rows can be arbitrary, but must be the same for all members of Asp(s, (s, f) and are unknown to the algorithm. The remaining ("light")
rows have unit norm. In other words, let = .vr.....,,.. For all rows A,, E [n], A; = l- 02 i E R and A, = otherwise, where 2/ is uniformly random unit vector.
Zipfian on squared row norms, Every A E Raxd Asipf has rows which are uniformly random and orthogonal. Each A has 2 rows of squared norm n² 22/22 for E [I...., ((lg(n))) We also assume that each row has the same squared norm for all members of Agipf-
Theorem F1. Consider matrix A from either the spiked covariance model Zipfian distribution. Let St. denote a CountSketch constructed by Algorithm 3 that optimizes the positions of the non-zero values with respect to A. Let Se denote a CountSketch matrix. Then there is a fixed 17 > 0 such that, minkaak-k X (Erowsp(Ss.A)
|| All s (1 (1-1) milrank-k Xe
|| All
Remark F.2. Note that the above theorem implicitly provides an upper bound on the generalization error of the greedy placement method on the two distributions that we considered in this paper. More precisely, for each of these two distributions, if II is learned via our greedy approach over a set of sampled training matrices, the solution returned by the sketching algorithm using II over any (test)
matrix A sampled from the distribution has error most 17) minrank-k XErowep(Sc
|X All?
A key structural property of the matrices from these two distributions that is crucial in our analysis is the c-almost orthogonality of their rows (i.e., (normalized) pairwise inner products are at most c). Hence, we can find QR- UR-factorization of the matrix of such vectors where the upper diagonal matrix R has diagonal entries close to and entries above the diagonal are close to 0.
To state our result, we first provide an interpretation of the location optimization task selection of hash function for the rows of A. Note that left-multiplying A by CountSketch S E TR is equivalent to hashing the rows of A proceeds through the rows of A (in some order) and decides which bin to hash to, denoting this by adding an entry to S. The intuition is that our greedy approach separates heavy-norm rows (which are important "directions" in the row space) into different bins.
bins with coefficients in (+1). The greedy algorithm
Proof Sketch of Theorem F.l The first step is to observe that in the greedy algorithm, when rows are examined according to a non-decreasing order of squared norms, the algorithm will isolate rows into their singleton bins until all bins are filled. particular, this means that the heavy norm rows will all be isolated-ea, e.g., for the spiked covariance model, Lemma F.8 presents the formal statement.
Next, we show that none of the rows left to be processed (all light rows) will be assigned to the same bin heavy row. The main proof idea is to compare the cost of "colliding" with a heavy row to the cost of "avoiding" the heavy rows This is the main place we use the properties of the aforementioned distributions and the fact that each heavy row is already mapped to singleton bin. Overall, we show that at the end of the algorithm no light row will be assigned to the bins that contain heavy rows the formal statement and proof for the spiked covariance model is in Lemma F.12.
Finally, we can interpret the randomized construction of CountSketch as a "balls and bins" experiment. In particular, considering the heavy rows, compute the expected number of bins (i.e., rows in A) that contain a heavy row. Note that the expected number of rows in S that do not contain any heavy row is (1 (1- 2 e - Hence, the number of rows in S that contain a heavy row of A is at most 6(1-e- Thus, at least k(1 - e = heavy rows are not mapped to an isolated bin (i.e., they collide with some other heavy rows). Then, it is straightforward to show that the squared loss of the solution corresponding to Sc is larger than the squared loss of the solution corresponding to S,., the CountSketch constructed by Algorithm 3 -please see Lemma F. 14 for the formal statement of its proof.
26

---

Published as a conference paper at ICLR 2023
Preliminaries and notation. Left-multplying A by CountSketch S E R is equivalent to hashing the rows of A to m bins with coefficients in {+1}. The greedy algorithm proceeds through the rows of A (in some order) and decides which bin to hash which entry to S. We will denote the bins by b, and their summed contents by wt
think of as adding an
F.1 SPIKED COVARIANCE MODEL WITH SPARSE LEFT SINGULAR VECTORS.
To recap. every matrix A Raxd from the distribution Asp(s, E) has s < k "heavy" rows (Ar. A,,) of norm > 1. The indices of the heavy rows can be arbitrary. but must be the same for all members of the distribution and are unknown the algorithm. The remaining rows (called "light" rows) have unit norm.
In other words: let (ri....,ra) For all rows A, E [n]:
A; -=1 L L. 0
if R O.W.
where U; a uniformly random unit vector.
We also assume that Sy Rxxn and that the greedy algorithm proceeds in non non-increasing row norm order.
Proof sketch. First, we show that the greedy algorithm using non-increasing row norm ordering will isolate heavy rows (i.e., each is alone in bin). Then, we conclude by showing that this yields better k-rank approximation error when d is sufficiently large compared to n.
We begin with some preliminary observations that will be of use later.
It is well-known that set of uniformly random vectors e-almost orthogonal (i.e., the magnitudes of their pairwise inner products are at most e).
Observation F.3. Let " Un I be a set of random unit vectors. Then with probability
1 1/ s/lpooyy)) we have Ke,, uy>س 2 log V vicjsnn
We define E 2/ log
Observation F.4. Let u1, ue be a set of vectors such that for each pair of < j S t,
e, and gi..... By E 11,,1]. Then,
Tuyl
1u 2c
Σ i<jSt
|uslll |uyll2 S
Σ gous
11:1 2 2c
Σ i<j
(F.1)
Next, a straightforward consequence of e-almost orthogonality is that we can find a QR-factorization of the matrix of such vectors where R (an upper diagonal matrix) has diagonal entries close 1 and entries above the diagonal are close to
Lemma F.S. Let u1, .tt E RRd be a set of unit vectors such that for any pair of < <jSS t, (u,, uz)l S where 6 = Ot==2). There exists orthonormal basis .., ,t for the subspace
spanned by u.. U4 such that for each S t, us Σ)=1 asgey where ai 2 1 Ei-1 and for each j < i, at < f2.
Proof. We follow the Seham- process to construct the orthonormal basis ... et of the space spanned by "1,".. Mr by first setting e and then processing " Me, one-by-one.
..........
The proof is by induction. We show that once the first vectors u1, ... (y are processed, the statement of the lemma holds for these vectors. Note that the base case of the induction trivially holds as 11 ε1. Next, suppose that the induction hypothesis holds for the first vectors "I,... Ug.
Claim F.6. For each j < e, abt2 222
Proof. The proof of the claim is itself by induction. Note that, for j = 1 and using the fact that (un, ue+1)) 1)] I e, the statement holds and a 41, < e2. Next, suppose that the statement holds for all
27

---

Published as a conference paper at ICLR 2023
j S i < L. Then using that Kus+1, ue+1)l S e,
00+1,8+1| S (1(u+1,, ui41 + Σ
Jett1ol (i441ll)/a444ae
S (c+ Σ /2))/l1-1sst1
b by the induction hypothesis on ae+1,j for j S
S ((c Σ f2)((1 - Σ jr 22/22
> by the induction hypothesis on 0,41,141
1=1
1=1
( Σ (22) (1 Σ 2)1/ (1+2) Σ j2
j=1
j=1
y=1
((c Σ (1+2 ΣS7)
j=l
j=l
S <<(( fc) () 4c +1) j=1
j=1 S s((i+)) > b 2)
Finally, since 11 1, 12111-11 1, 21
Corollary F.7. Suppose that (). There exists an orthonormal basis e,,-.- , for the space spanned by the randomly picked vectors U1, Ut, of unit norm, so that for each Us = Ej=ı aiiee where a²,, 1
for each j i, ai S 5².
Proof. The proof follows from Lemma F.5 and the fact that the set of vectors U1,... Ut is e-almost orthogonal (by Observation F.3).
The first main step is to show that the greedy algorithm (with non-increasing row norm ordering) will isolate rows into their own bins until all bins are filled. In particular, this means that the heavy rows (the first to be processed) will all be isolated.
We note that because we set rank(SA) = k. the k-rank approximation cost is the simplified expression ||AVVT All p. where UEVT SA, rather than ||VV|&V - || This is just the projection cost onto row(S A). Also, observe that minimizing this projection cost is the same as maximizing the sum of squared projection coefficients:
arg min || AA-AV
||
arg min S
||A, - ((A,, +...+ + (As, vu)u
sE(m]
arg min
Σ
(ILA,112
(An,)))
S
sE[m]
JE(N]
(A,, uy)
arg max
S
HE
JE
In the following sections, we will prove that our greedy algorithm makes certain choices by showing that these choices maximize the sum of squared projection coefficients.
Lemma E.8. For any matrix A or batch of matrices A, the end of iteration k. the learned CountSketch matrix S maps each row to an isolated bin. In particular, heavy rows are mapped to isolated bins.
28

---

Published as a conference paper at ICLR 2023
Proof. For any iteration S k, we consider the choice of assigning A, to empty bin versus occupied bin. Without loss of generality, let this occupied bin be b5-1, which already contains A;-1-
We consider the difference in cost for empty versus occupied. We will do this cost comparison for A, with S 2 j 2 + 1, and finally, j E {i 1,
First, let fei....,e:! be orthonormal basis for fAI,..., A,} such that for each S i, A, = Σj= ar,je; where ar, > 0. This exists by Lemma F.5. Let fei,..., orthonormal basis for {AI...., A142, A, 1 A.). Now, E = coes-1 + cie; for some co, C1 because (A₂-1 + A.) projfer (A,-1 = =A,)) E span(e)-1, ej). We note that cj + of because we let è be unit vector. We can find Co,Ci be:
an
01 + d1.i-1
O5
Cu =
C1
(ai-1.-1 + 01,1-1²² aZ
(ai-l,i-1 ai,i-1)2 4 + aZ
1. j S 1-2--2 The cost zero for both cases because A, span(fe)...... e(-2))
j 2 i + We compare the rewards (sum of squared projection coefficients) and find that f1.....,i-2,. is no better than (c)).
() (cn(Aj, 6-11 + (A es))
S (cf (A), + (Aj, E))
D> Cauchy-Schward inequality
(A,, e(11)² (A,, e,)²
3. j E 18-1-111 1, We compute the sum of squared projection coefficients of A;-1 and A, onto ē:
1
(ai-1,i-1 di,i-1)² + (ai,i- (ai-1,i-1 diji-)) + ds,)0;;))
14-1,1-1 12,i-1)2
1
)
a0, (a) +
(a:-1,i-1 -i,i-1
- Ci
+ + + 20,,1-10 20,,--10, + ai,s-1))
(0i-1,i-1
(F2)
On the other hand, the sum of squared projection coefficients of A,-1 and A; onto ei-1 Ue, is:
(4--15511 + 0i.i )2 az (1-1-5-1 0i,i 12 a?,
(af (af- + ai + a²
(F33
Hence, the difference between the sum of squared projections of A,-1 and A4 onto ē and e; Ue is ((F3)- (F2))
df + 01ii-1)² + a;-1,5-1 + ař,1-1 -a,,ii- ( + di,i-1))
(ai-1,i-1 a1,1-13 32 + to,
0-1,1-11111 >0 0
(0:-1,1-1 + 01,i-1)
07
Thus, we find that fei..... e a strictly better basis than Ci-,,, Z). This means the greedy algorithm will choose place A; in empty bin.
Next, we show that none of the rows left to be processed (all light rows) will be assigned to the same bin as heavy row. The main proof idea is to compare the cost of "colliding" with a heavy row to the cost of "avoiding" the heavy rows. Specifically, we compare the decreuse (before and after bin assignment of a light row) in sum of squared projection coefficients, lower-bounding it in the former case and upper-bounding it in the latter.
We introduce some results that will be used in Lemma F.12.
Claim F.9. Let Axtr. r E [l.... 1 ki be light row not yet processed by the greedy algorithm. Let (ei....,ey. be the Gram-Schmidt basis for the current {wi.... .... ...,)) wk). Let B O(n -k-3) upper bound the inner products of the normalized Agtr, WT.... Wg. Then, for any bin i, (e,, Actr) (22 2.2
29

---

Published as a conference paper at ICLR 2023
Proof. This is straightforward application of Lemma F.5. From that, we have (Aktr, e B2. for E [I,...,k], kl,, which means (Ak+r: e1)² 1 k²
Claim F.10. Let Ag+r be a light row that has been processed by the greedy algorithm. Let fei...., cx} be the Gram-Schmidt basis for the current (WI,..., wk }. If Ak-r is assigned to bin - (w.llo.g.). the squared projection coefficient of Ag+r onto e,, o k
most 4,2² k2
where B = O(n - 5) upper bounds the inner products of normalized Ak-r, W1,
Wk.
Proof. Without loss of generality, it suffices to bound the squared projection Aktr onto the direction of wi that is orthogonal the subspace spanned by W1, Wk-1- Let e Eg be an orthonormal basis of W1,-
guaranteed by Lemma F.5. Next, we expand the orthonormal Σj=1
basis to include 6x+1 so that we can write the normalized vector of Autr as uktr By a similar approach to the proof of Lemma F.5, for each j S k 2, by S f2 Next, since (wee, Ck+r) $ 8,
Eh+ byeg-
[b&| v
(1KKk uk+r)l ++ + Σ
1by (ukk ej)l)
Kekk ex)l
k-2
1
13+ +
B² /2 + - (1-1)
> >
Σ .
1=1
B
2-2
B2 j2
(k 1)3
Σ/=1 82 .²
(*k 1)²
< 2(k - 1)8
similar to the proof of Lemma F.5
V 1 k-1 f2, j2
28-
Hence, the squared projection of Astr onto en is at most 4,8² 4,82. 1; hence, the squared projection Auur onto E₁ is at most 4,3², 1.2.
-22|||A+rr11 We assumed ||ARtrll
Claim F.1 We assume that the absolute values of the inner products of vectors "1 Un are at most < 1/(n2 EA.e [|A.ll)) and the absolute values of the inner products of the normalized vectors of W1, ... , Wh are at most 3 = 0 +). Suppose that bin b contains the row Aktr- Then, the squared projection of Autr onto the direction of orthogonal span( (on((w,.".. wk) \ {we is
at most || [|]
+ O(n and is at least |A2+-ll O(n 2).
|
Proof. Without loss of generality, we assume that A&tr is mapped to dy; = wk- First, we provide an upper and a lower bound for Kektr, wxei where for each S k; we let w. = [[ l normalized vector of wi. Recall that by definition uker
wy
denote the
Autr
TAxtella
| EA.eb, 7 E||A/l1
(wk, uktrr) 5
||wkll2
| + n
-2
by Ex ||Aill2
||
A
<
+ -2
> ||wxll2 21
(F.4)
|wxll22
||
Ex,eb, ||A:ll2-
(w,, Uk4r)
|
|| 11
(F.5)
[w
30

---

Published as a conference paper at ICLR 2023
Now, let fei,... , c be orthonormal basis for the subspace spanned by (wi,... wy J guaranteed by Lemma F.5. Next, we expand the orthonormal basis to include ex+1 so that we can write
k+1 byez. By a similar approach to the proof of Lemma F.5, we can show that for each Uktr = 2j=1
j S - 1, b3 S f2 Moreover,
|bx| 5
Kewk ex}l (KuK, Uk+T)
Σ 1=1
|bj (wRR e))l
1
(KWwk Uk+r) + Σ .8...
by Lemma E.5
j== B2.
1
|| (n= +
k-1
Σ
B² (3)
by (F.4)
V1 ΣØ1 j= .2
|100 l
5=1
|| 112
<B.k+ + V1 /2 (n-2+
similar to the proof of Lemma F.5
| 2
||Aktr l|
On + (1 (1+ O(n
by O(n --)
|wx
| + =On 1
|| s
||w/||2
and,
1b
(K Uk+r)
Σ
1bj
(T e,)l
K ex)
2 (Tk, Uk+rr)
Σ 32. j²
D since (WE, S
==1
||Att1llz n2 -2 2
k-1 82 .²
D> by (F:5)
|wmk l
j=1
| llz
On 22)
> by 8 = On-33k- -3))
|xlll2
||44 11 | + +O -2) and is at least
| 111
Hence, the squared projection of Astr onto E7 is at most O(n 2)..
Now, we show that at the end of the algorithm no light row will be assigned to the bins that contain heavy rows.
Lemma F.12. We assume that the absolute values of the inner products of vectors in U1. are at most < min(n-2p 2 3, (n EA.e - At iteration k + r, the greedy algorithm will assign the light row Aktr to a bin that does not contain a heavy row.
Proof. The proof is by induction. Lemma F.8 implies that light row has been mapped bin that contains a heavy row for the first I iterations. Next, we assume that this holds for the first 1: + r iterations and show that is also must hold for the (A: r))th iteration.
To this end, we compare the sum of squared projection coefficients when A avoids and collides with heavy row.
31

---

Published as a conference paper at ICLR 2023
First, we upper bound 3 = m (w,, wy)l/[[[:ll1 l|w,)12)= Let c, and ej respectively denote the number of rows assigned to b; and by.
Kwe, w3y1
Ci c
8 = max cy ||wll2 |00 12
> Observation F.4
v
200
V 200
< 16E /ccj
DE < n 2k-
S 5n
5E < n n 2k
1. If Agtr is assigned to bin that contains e light rows and no heavy rows. In this case, the projection loss of the heavy rows A.. Ag onto rSw(SA) remains zero Thus, we only need to bound the change in the sum of squared projection coefficients of the light rows before and after iteration k + r. Without loss of generality, let wa denote the bin that contains As+r. Since Sk-1 = span( { (w.... between projecting onto the component of wx Ag- orthogonal to S and the component of wkk orthogonal S respectively denoted as ej and Ex.
w)-1)) has not changed, we only need to bound the difference in cost
By Claim F.9, for the light rows that are not yet processed (i.e., A, for j > + r), the squared projection of each onto En is at most g2 Hence, the total decrease in the squared projection is at most (n - k- r) R²2
2. By Claim F.10, for the processed light rows that are not mapped to the last bin, the squared projection of each onto eg is at most 4,3 Hence, the total decrease in the squared projection cost is at most (r- 1). k.
3. For each row A, + Aatr that is mapped to the last bin, by Claim F.11 and the fact ||A: ||A.1
||A:112 1, the squared projection of A; onto is at most Tw 1ww--Ake A14--144 + Onn and the
[ || 1 O(n-))
squared projection of A; onto é is least
Moreover, the squared projection of Astr onto eg compared to eg increases by least ( [| On 2)) On ?)
[w
|| O(n
Hence, the total squared projection of the rows in the bin by decreases by at least:
||4rl1
| - 0(
Σ
+ )))
|w
Ar+al1
[w 1l
A.Ews / [(rr+)]
A;Ewg
12 - Ar+ 12 ++(nn1))
| 11(2-111 O(-
> by Observation F.4
1w Ar+al12
| 1
<1(n-1)
Hence, summing up the bounds in items 1 to 3 above, the total decrease in the sum of squared projection cocfficients is at most O(n
2. If Astr is assigned to a bin that contains heavy row. Without loss of generality, we can assume that AR+r is mapped to by that contains the heavy row A,. In this case, the distance of heavy rows A,, A.-1 onto the space spanned by the rows of SA is zero. Next, we bound the amount of change in the squared distance of A, and light rows onto the space spanned by the rows of SA.
Note that the (A= - -immenssiaaa space corresponding Wi, ..., Wi-1 has not changed. Hence, we only need to bound the decrease in the projection distance of A&tr onto ε compared to eg (where Eg Eg are defined similarly as in the last part).
1. For the light rows other than Astr. the squared projection of each onto e₁ is at most k²2, Hence, the total increase in the squared projection of light rows onto ex is at most (n (n-k).,32 k)2 O(n-1).
32

---

Published as a conference paper at ICLR 2023
2. By Claim F.11, the sum of squared projections of A, and Astr onto eg decreases by least
||A.1| 2 + ||Ax+ Onn
||4.112
||A, + AA+*112
|| + ||
||4, 11
2 n11 + +0(n
by Observation F.4
||A + || Ar
| 1/2 (1,A,112 2 - |
(2) - ||A,112 O(n On
||A.112 + || l1
O(n-1)
| x 11 (1I., 112 | 112) ||Aa O(n o(n
||A 11 + || 1
A 1 (II,A.11 2 A
-O(n
||A,112 + ||Ar+ 12
A (1 - <1IAA++ 1/2 [1A,
-n)
1 + (1IA,+*112 (1AA)1()
|| O(n
1 1
2 2114++111 122 (1
1 1 2 1 -
||all2
Hence, in this case, the total decrease in the squared projection is at least
| 1) = || ((n
| (1
D | =
A
|
= (1/V1/ 0
>> |Asll2 Ve
Thus, for a sufficiently large value of f, the greedy algorithm will assign Astr to a bin that only contains light rows. This completes the inductive proof and in particular implies that at the end of the algorithm, heavy rows are assigned to isolated bins.
Corollary E13. The approximation loss of the best rank-k approximate solution in the rowspace A for A ~ Asp(s, l). where A E Raxd for d 2( k log n) and Sy is the CountSketch constructed by the greedy algorithm with non-increasing order, is at most n 8.
Proof. First, we need to show that the absolute values of the inner products of vectors in U1.-.. are most min(n min(n-2k- 2, EA, Ew [|lll)) so that we can apply Lemma F.12. To show
this, note that by Observation F.3, ē S 2 logn S n since d = 2nn log The proof follows from Lemma F,8 and Lemma F.12. Since all heavy rows are mapped to isolated bins, the projection loss of the light rows is at most n s.
Next, we bound the Frobenius norm error of the best rank-k -ap-----pproximatio solution constructed by the standard CountSketch with a randomly chosen sparsity pattern.
Lemma E.14. Let s ak where 0.7 < a < The expected squared loss of the best rank-k approximate solution in the rowspace S.A for A E Baxd ~ Asp(s,E) where d = e2) and S, is the sparsity pattern of CountSketch is chosen uniformly at random, is at least n + fk - akk 0(1)
Proof. We can interpret the randomized construction of the CountSketch as "balls and bins" experiment. In particular, considering the heavy rows, we compute the expected number of bins (i.e., rows in S, that contain heavy row. Note that the expected number of rows S,A that do not contain any heavy row is k (1 ))²² 2 k न Hence, the number of rows in S,A that contain a heavy row of A is at most k(1 e F Thus, at least mapped an isolated bin (i.e., they collide with some other heavy rows). Then, it is straightforward to show that the squared loss of each such is at least - n -(()
A(1 e --xTT) heavy rows are not
33

---

Published as a conference paper at ICLR 2023
Claim F.15. Suppose that heavy rows ....... Ar, mapped the same bin via CountSketch S. Then, the total squared distances of these rows from the subspace spanned by SA is at least (c- 1) O(n
Proof. Let b denote the bin that contains the rows Arys" Ar. and suppose that it has e' light rows as well. Note that by Claim F. and Claim F.11, the squared projection of each row Ar, onto the subspace spanned by the k bins is at most
|A, [w|] 12
+
+((
2
82
cl + d - 2e(c+ ae +O(n
&
d -011) +n -01)
by 38-1
1
(el + + O(n-1)
26
{ S- + O(
Hence, the total squared loss of these heavy rows is at least ol c- (() 2 (c 1)e O(n
Thus, the expected total squared loss of the heavy rows is at least:
e. (s - k(1 e =)) ))- - 8- 0(1)
26 k(a - 1 -7) - la O(1)
D s a (k where 0.7 <a <
Ek
2 2e
{ n -(1)
D a 0.7
Ek:
Onn -1
D assuming k >
de
Next, we compute a lower bound on the expected squared loss of the light rows. Note that Claim F.10 and Claim F.11 imply that when light row collides with other rows, its contribution to the total squared loss (note that the loss accounts for the amount it decreases from the squared projection of the other rows in the bin well) is at least O(n-1). Hence, the expected total squared loss of the light rows is at least:
(n s - k)) - O(n 2 (n (1 a)) aa-)))4 O(
Hence, the expected squared loss of a CountSketch whose sparsity is picked at random is at least
fk -( + (1 a - O(n 11) n n
Ek
(1 a:k On
4e
4e
a(k - 1) where 0.7 < a < and let (4e+1)m Let Sn the ak CountSketch whose sparsity pattern is learned over training set drawn from Asp via the greedy approach. Let S, be a CountSketch whose sparsity pattern is picked uniformly at random. Then, for an d matrix A Asp where ((n the expected loss the best rank-k approximation of A returned by S, is worse than the approximation loss of the best rank-k approximation of A returned by Sy by least constant factor.
Corollary F.16. Let s
34

---

Published as a conference paper at ICLR 2023
Proof.
E
min
|X - -A1 2 + - - -1+ + aak ((1)
D Lemma F.14
Se rank-k X Crowsp(S,.A)
4e
(4e + 1 1)
> (1 + (111/a)(nn- s)
D L
ak
(1 1/a)
min
|| All
D> Corollary F.13
rank-k X EErowsp(Sp A)
F.2 ZIPFIAN ON SQUARED ROW NORMS.
Each matrix A E R Asipf has rows which are uniformly random and orthogonal. Each A has 2 rows of squared norm n² /2 for i E [I,..., (llo(()))]- We also assume that each has the same squared norm for all members of Asipf.
In this section, the rows with largest norm are called the heavy rows and the remaining are the light rows, For convenience, we number the heavy rows ,,...,;; s; however, the heavy rows can appear at any indices, as long as any row of given index has the same norm for all members of Asipf- Also, we assume that S k/2 and, for simplicity, = Ei= 2 for some h2 E 2+. That means the minimum squared norm of a heavy row is n² /222 and the maximum squared norm of light row is n² +
The analysis of the greedy algorithm ordered by non- non-increasing row norms on this family of matrices similar our analysis for the spiked covariance model. Here analyze the case in which rows are orthogonal. By continuity, if the rows are close enough to being orthogonal, all decisions made by the greedy algorithm will the same.
As first step, by Lemma F.8, at the end of iteration & the first k rows are assigned to different bins. Then, via similar inductive proof, we show that none of the light rows are mapped to a bin that contains one of the top heavy rows.
Lemma F.17. At each iteration l + r, the greedy algorithm picks the position of the non-zero value the (k r)-th column of the CountSketch matrix S that the light row Aktr mapped to bin that does not contain any of top heavy rows.
Proof. We prove the statement by induction. The base case r r= 0 trivially holds as the first rows assigned to distinct bins. Next assume that in none of the first - 1 iterations light row is assigned to bin that contains heavy row. Now, we consider the following cases:
1. If A is assigned to a bin that only contains light rows. Without loss of generality we can assume that Ak+r is assigned by. Since the vectors are orthogonal, we only need to bound the difference in the projection of A&tr and the light rows that are assigned to by onto the direction of wy before and after adding Actr to bg. In this case, the total squared loss corresponding to rows in by and Agtr before and after adding A&+1 are respectively
ΣAS Ebg || 1
before adding A to byi ||Ap+iI
Σ AyEby
|
12
Σ Edi. || 1|2
| 11 + EA Enjen, ||A,11
after adding A to bei ||
| 1
| + E,A ||A,11
AyEb
35

---

Published as a conference paper at ICLR 2023
Thus, the amount of increase in the squared loss is
EaJebr | (2) ||Att 1 + EAjEbe |A,l12 Enyeb, ||A,112 || 112 + Enyeba |Ail1
| EayEb || || EdS Ebe
|Aj
EA || |A]11](AAktr]])) EnnEo 11,)112)
Enycr AAill
-||A+t 12
En, Aglli
= -| 112
Σ Enyeru || ||A 11
Enyebu ||jl1 ||A
12
$ ||ARtr
Σ Eae |4jl12 + || l|
(F.6)
2. If As+r is assigned to bin that contains a heavy row. Without loss of generality and by the induction hypothesis, we assume that A&+r is assigned to a bin b that only contains a heavy row Aj. Since the rows are orthogonal, we only need to bound the difference in the projection of Ak+r and A5 In this case, the total squared loss corresponding to A, and Agtr before and after adding A&+1 to are respectively
before adding ARtr to by: 1|
||
112 1 + ||4,11/
after adding Autr to by: ||A + 11,,112
| 12 + ||
Thus, the amount of increase in the squared loss is
|| 1/ + ||
||4,112 - || ||A;112 + ||
||
||
(F.7)
||A
12 + 1|11
Then (F.7) is larger than (F6) if ||, 2 E0, 114.112- Next, we show that at every inductive iteration, there exists a bin which only contains light rows and whose squared norm is smaller than the squared norm of any heavy row. For each value m, define hm so that m i= 20m+2 2.
Am 2
Recall that all heavy rows have squared norm at least light rows and has squared norm at most
There must be a bin b that only contains
2
n
2
10122 = Σ 1A12
2
+1)
k -
A,Eb
n²
2n
22(h,+11 +1)
2h (k- s)
n
n
D S k/2 and k >>> 2ha+1
22(h, +1)
22h
n
D h4 2 h +1
21h.+I
||
Hence, the greedy algorithm will map A&tr to a bin that only contains light rows.
Corollary F.18. The squared loss of the best rank-k approximate solution in the rowspace of Sy
for A E Raxd Azipt where A E Raxd and Sy the CountSketch constructed by the greedy n²
algorithm with non-increasing order, is < 2-
Dishu+1 Proof. At the end of iteration k. the total squared loss is Sh 24+1
n² After that, in each
2
iteration i + r, by (F.6), the squared loss increases by most | Hence, the total squared
36

---

Published as a conference paper at ICLR 2023
loss in the solution returned by Sg is most
hn
n2²
hp
4n²
n
2(
-nn
2
2
2h
2
11++1
i=
+1
Next, we bound the squared loss the best rapn--taa solution constructed by the standard CountSketch with randomly chosen sparsity pattern.
Observation F.19. Let us assume that the orthogonal rows , Ar. are mapped to the same bin and for each - c, ||,, 112 2 || 112 Then, the total squared loss of A Are after projecting onto Ar + ... + Ar, is at least || l12 +... 1|
Proof. Note that since Ar,.".. .r. are orthogonal, for each i S c, the squared projection of Ar, onto Ar, . Ar, is Ar, 112 Σj=1 |r, ll Hence, the sum of squared projection coefficients of Ap,. Are onto Ar, - ..... - Ar. is
Ar S || Ar, 11
Σ j Arg l1
Hence, the total projection loss of Ar,, Ar, onto A,, - .... + Ar, is at least
A 12 A 11 = | 12 + ... + Ar
Ø=1
In particular, Observation F.19 implies that whenever two rows are mapped into the same bin, the squared norm of the row with smaller norm fully contributes to the total squared loss of the solution.
Lemma E20. For 20 2, the expected squared loss of the best rank-k approximate solution in the rowspace of S,A for Anxd ~ Asipf. where S, is the sparsity pattern of a CountSkeich chosen uniformly at random, is at least 1.95n2
2
Proof. In light of Observation F.19, we need to compute the expected number of collisions between rows with "large" norm. We can interpret the randomized construction of the CountSketch "balls and bins" experiment.
For each 0 S j S. hu, let A, denote the set of rows with squared norm マークhーーク and let As 2hx -+1
Uscighe Ai Note that for each j. |A,1 = 2ha - and |A>₁l = E 2(2h--j 1). Moreover, note that k 2(2hv
Di=1 2
Σ h
1). Next, for row A, A, (0 S < h2). we
compute the probability that at least one row in A>; collides with Ar.
1 y1assl)
Pr at least one row in As, collides with A,] = (1 - (1
k
2 14p2l
2h6-1-1
= (1 e 2hatt 1)
241-1 -
2 2---2
D since
2-
2h -
Hence, by Observation F.19, the contribution of rows in Aj to the total squared loss is at least
(1 (1-6-2117) |4,1- 22(h =(1 ------
n²
- = (1 c----2--3
n2
n
2
2h
37

---

Published as a conference paper at ICLR 2023
Thus, the contribution of rows with "large" squared norm, i.e., Asu, to the total squared loss is at least"
n²
n² 2
2/-1.
-e- 2 1.095
2h -2
f=0
Corollary F.21. Let Sg be CountSketch whose sparsity pattern is learned over a training set drawn from Asp via the greedy approach. Let S, be CountSketch whose sparsity pattern is picked uniformly at random. Then, for an n d matrix expected loss of the best rank-k approximation of A returned by S, is worse than the approximation loss of the best rank-k approximation of returned by Sg by at least a constant factor.
Azziff for sufficiently large value of k, the
Proof. The proof follows from Lemma F.20 and Corollary F.18.
Remark F:22. We have provided evidence that the greedy algorithm that examines the rows of A according to a non-increasing order of their norms (i.e., greedy with non-increasing order) results in better rank-k solution compared the CountSketch whose sparsity pattern is chosen at random. However, still other implementations of the greedy algorithm may result in a better solution compared to the greedy algorithm with non-increasing order. To give an example, in the following simple instance the greedy algorithm that checks the rows of A in random order (i.e., greedy with random order) achieves rank-k solution whose cost is a constant factor better than the solution returned by the greedy with non-increasing order.
Let A be matrix with four orthogonal rows u, u, v, wo where ||u/|2 and || = ||w||2 1 and suppose that the goal is to compute a rank-2 approximation of A. Note that in the greedy algorithm with non-decrcasing order, and w will be assigned to different bins and by a simple calculation can show that the copies of also will be assigned different bins. Hence, the
(1+e)²
2+(1=e) However, the optimal solution will assign v and w to one bin and the two copies of to the other bin which results in squared loss of + c)² which a constant factor smaller than the solution returned by the greedy algorithm with non-increasing order for sufficiently small values of c.
squared loss in the computed rank-2 solution is +
On the other hand, the greedy algorithm with random order, with a constant probability of ( + )) the computed solution is the same as the optimal solution. Otherwise, the greedy algorithm with random order returns the same solution as the greedy algorithm with a non-increasing order. Hence, in expectation, the solution returned by the greedy with random order is better than the solution returned by the greedy algorithm with non-increasing order by a constant factor.
G EXPERIMENT DETAILS
G. LOW-RANK APPROXIMATION
In this section, we describe the experimental parameters in our experiments. We first introduce some parameters in Stage of our approach proposed in Section 3.
bs: batch size, the number of training samples used in one iteration.
lr: learning rate of gradient descent.
iter: the number of iterations of gradient descent.
Table 7.1: Test errors for LRA (using Algorithm 2 with four sketches)
For a given m, the dimensions of the four sketches were: S mm R E Rmxd S2
R5mxn R2 E Rpmxdd
Parameters of the algorithm: bs 1, 1.0, 10.0 for hyper and video respectively, num_ic 1000.
For our algorithm 4, we use the average of all training matrix as the input to the algorithm.
*The numerical calculation computed using WplframAlpha.
38

---

Published as a conference paper at ICLR 2023
Table 7.1:
Test errors for LRA (using Algorithm 1 with one sketch)
Parameters of the algorithm: bs = 1, Ir = 1.0, 10.0 for hyper and video For our algorithm 4, we use the sum of all training matrix as the imput to the algorithm.
v respectively, num_it __tt 1000.
G.2 SECOND-ORDER OPTIMIZATION
As we state in Section 6, when we fix the positions of the non-zero entries (uniformly chosen in each column or sampled according to the heavy leverage score distribution), we to optimize the values by gradient descent, as mentioned in Section 3. Here the loss function is given in Section 6. In our implementation, we use PyTTrcc (Paszke et al. (2019)), which can compute the gradient automatically (here we can use torch.qr( and torch.svd() to define our loss function). For a more nuanced loss function, which may be beneficial, one can use the package released Agrawal al. (2019), where the authors studied the problem of computing the gradient of functions which involve the solution to certain convex optimization problem.
As mentioned in Section 2, each column of the sketch matrix S has exactly non-zero entry. Hence, the i-th coordinate of p can be seen as the non-zero position the i-th column of S. In the implementation, coordinate of p. For the heavy rows mentioned in Section 6, we can allocate positions 1, 1,..., to the k heavy rows, and for the other rows, we randomly sample integer in { + 1, ... m). We note that once the vector p. which contains the information of the non-zero positions in each column of S, chosen, it will not be changed during the optimization process in Section 3.
sample p randomly,
sample random integer in f1,...,m) m for each
Next, we introduce the parameters for our experiments:
bs: batch size, the number of training samples used in one iteration.
lr: learning rate of gradient descent.
iter: the number of iterations of gradient descent
In our experiments, we set bs 20, iler 1000 for all datasets. We set lr 0.1 for the Electric dataset.
39


Reviews:
[{"review_id": "l39N5e-Bvp", "reviewer": "Reviewer_qwvC", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "This paper is about learning linear maps for dimensionality reduction that are based on count sketch. Existing works either use randomized sketches or learn the values in a count sketch matrix (with the positions drawn randomly). This paper develops on such methods by proposing to learn both the positions and the values of the count sketch matrix. The paper studies this idea in the context of two problems: low-rank approximation and iterative Hessian sketching. The contributions of the paper are:\n1. It shows that learning positions and values instead of just values leads to better performance. This is shown using a greedy search algorithm \n2. In order to overcome the high run time of the greedy search algorithm, the paper proposed algorithms to construct the count sketch matrix using ridge leverage score based sampling. \n3. The paper shows using experiments that the proposed ideas lead to better numerical performance. ", "summary_of_the_review": "The paper studies an interesting problem and presents good ideas. However, the presentation needs to be improved and the experiments should be explained more clearly. I have elaborated further in my comments above. ", "strength_and_weaknesses": "Strengths:\n\n1. The idea of learning both the positions and the values is interesting and warrants a study. This paper does a good job of exploring this (previously unexplored) idea. \n\n2. The paper shows strong improvement over previous algorithms in terms of numerical performance.\n\nWeaknesses: \n\n1. The presentation of the paper can be improved a lot. Some examples of poor presentation are: \n    - Algorithms 1, 2 and 3 are placed in unnatural locations, with presumably liberal usage of \\vspace and \\hspace. The paper is visually \n        very dense and doesn't need to be.\n    - On page 3: \"Hence, in this work we consider the following algorithm that compresses both sides of A...\": which algorithm?\n    - Authors should consider completely moving the few-shot learning setting to the appendix and write more clearly about the two main \n        problems considered. \n    - Placement of the results tables are also not great. For example, Table 7.4 appear at a random location in between text. \n\n2. I have the following questions regarding the experiments:\n    - Shouldn't metrics always be relative errors? For example, in the low-rank approximation setting, instead of reporting ||A-A_k|| - ||A-X*||, \n        it should be ( ||A-A_k|| - ||A-X*||)/ ||A||. This would indicate how large the error is w.r.t the Frobenius norm of A itself and gives a better \n        idea of performance.\n    - Why are the errors in the few-shot learning case smaller than when the full training set is used? Intuitively, it should be the opposite. \n    - The authors mention that during training, only the average of the training matrices are used. However, this is not equivalent to \n        minimizing the expected error over the training set. It is also not clear if the baseline methods (especially IVY19) follow this procedure.  \n        Upon reading the IVY19 paper, it appears that they do minimize the error over all the training matrices. Maybe I am missing something, \n        but there is chance that the baseline methods have worse performance because of this? It would be better to present results (at least \n        on the baseline methods) where the training fully takes advantage of the training set and does not use just the average. \n    - Authors mention that they use torch.qr for backpropagating through QR factorization. Do the authors take any measures to ensure \n       that the first few columns are linearly independent (as required by torch.qr: https://pytorch.org/docs/stable/generated/torch.qr.html)?\n    ", "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": "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": {}}, {"review_id": "ZFuZFimD9O", "reviewer": "Reviewer_ZY3u", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "The paper considers sketching based algorithms for low rank approximation and second-order methods for regression where the sketching matrix is learned. If the underlying data matrix is drawn from some unknown distribution, one can use historical samples of the data to learn better sketch matrices rather than the classical, data independent Count Sketch. While prior learning-based sketches only learn the values of non-zero entries of the sketching matrix (and keep the positions random), the paper proposes a framework to also learn the positions of the non-zero entries of the sketching matrix. The proposed algorithm is fairly natural and outperforms prior learned sketching based approaches.", "summary_of_the_review": "Overall I feel the paper makes an important contribution. The key ideas introduced are fairly natural and many of the proofs follow from standard techniques. Empirical results illustrate that learning positions of non-zeros along with their values significantly improves the approximation of classical CountSketch (at a similar sketch size).", "strength_and_weaknesses": "Strengths:\n\n+ The paper tackles a natural drawback of existing approaches for learning-based sketching algorithms. It provides good intuitions regarding the central premise of learning positions of non-zero entries of the sketching matrix as opposed to just their values (as prior work)\n\n+ The paper is well-written and all claims are well justified. The paper is well-written and includes detailed algorithms for clarity. I feel the presentation of the paper got hurt significantly due to page restrictions, and I encourage the authors to publish a full version where the flow of text is not hurt due to page limits.\n", "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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "QFHJkTn9QzE", "reviewer": "Reviewer_oTN6", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "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": "This paper proposes to learn the positions of the no-zero entries in count-sketches for better approximation accuracy. It improves previous approaches learn only the value of the no-zero entries. It first proposes a greedy solution, which has high complexity, then designs simpler algorithms for low rank approximation and Hessian approximation. Theoretical analysis shows that the algorithms reduce the size of the sketch without harming worst case performance.", "summary_of_the_review": "This paper proposes to learn the positions of non-zero entries in count-sketches. The proposed methods are practical and come with strong theoretical guarantees. The empirical results demonstrate the effectiveness of the proposed algorithms.", "strength_and_weaknesses": "Strength\n1. Learning the position of non-zero entries makes sense and is a good idea.\n2. The two algorithms are simple and thus practical, and come with theoretical guarantee.\n3. The algorithms significantly reduce the approximation errors in the experiments. \n\nWeakness\nNothing that I can think of.\n", "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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "3CeIYKLpHjC", "reviewer": "Reviewer_Akjx", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "The paper focuses on sketching for linear algebraic problems like LRA. They move the state-of-the-art boundary theoretically and support it empirically. The main idea is to learn positions of non-zeros in the countsketch matrix . The approach seems reasonable and as shown in experiments is valuable.  Some comments /questions\n\n1. not a learning heavy paper. Crux of contributing algorithms do not use learning and its target problems as well are linear. So not sure if it fits the call of papers. Something for Meta-Reviewers / Area chairs to decide. \n\n2. The setting is not clear. where there is a set of training matrices {A1, .. An} and the test is supposed to be on an unseen data from the same distribution. This is not very clear. For example when you say A \\in \\mathcal{D}, does it mean the order of rows in A is also important? Or does it mean \"each row of A\" \\in \\mathcal{D}. In the latter case, it is difficult ot imagine how a fixed S will work for unseen A. For example S learned on A will not work on a permutation of A.  Maybe it will help clarifying this better.\n\n3. I have not read much literature on learning of sketching matrices. so maybe its well known. I think it is important to point out the downside of these approaches. For example, the CS matrix is O(1) storage due to use of universal functions, however the storage of learned sketch matrix requires O(n) memory. It is important to point out this trade-off between CS vs learned sketch\n\n4. The proposed approach is compared with sketching baselines namely - vanilla CS, Indyk 2019. However, as the authors also use existing techniques like ridge-leverage score sampling etc. It seems important to evaluate a comparison between sampling approaches (which in some sense are sketching as well). For example, a comparison with cohen 2017 - both theoretical and empirical would be good to have. ", "summary_of_the_review": "I believe the paper is a good contribution to the field. I would like resolutions to some of the issues with writing and experiments. ", "strength_and_weaknesses": "[Strengths] The paper is well written and relatively easy to understand (except for some clarification needed. see above) . Intuitions are provided with associated theory which is very much appreaciated. \n[weakness] comparisions where only pure-sketching  methods are considered and the sampling methods such as leverage score sampling which is used by the algorithm itself is not used as a baseline. \n", "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": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%