A framework for learned CountSketch
Authors: Simin Liu
Tianrui Liu
Ali Vakilian
Yulin Wan
David Woodruff
Venue:  (2021)
Keywords: Compression
sketching

Abstract: Sketching is a compression technique that can be applied to many problems to solve them quickly and approximately. The matrices used to project data to smaller dimensions are called "sketches". In this work, we consider the problem of optimizing sketches to obtain low approximation error over a data distribution. 

We introduce a general framework for "learning" and applying CountSketch, a type of sparse sketch. The sketch optimization procedure has two stages: one for optimizing the placements of the sketch's non-zero entries and another for optimizing their values. Next, we provide a way to apply learned sketches that has worst-case guarantees for approximation error. 

We instantiate this framework with three sketching applications: least-squares regression, low-rank approximation (LRA), and k-means clustering. Our experiments demonstrate that our approach substantially decreases approximation error compared to classical and naively learned sketches. 

Finally, we investigate the theoretical aspects of our approach. For regression and LRA, we show that our method obtains state-of-the art accuracy for fixed time complexity. For LRA, we prove that it is strictly better to include the first optimization stage for two standard input distributions. For k-means, we derive a more straightforward means of retaining approximation guarantees.

Under review as a conference paper at ICLR 2021
A FRAMEWORK FOR LEARNED KOUCH
Anonymous authors Paper under double-blind review
ABSTRACT
Sketching is compression technique that can be applied many problems solve them quickly and approximately. The matrices used to project data to smaller dimensions are called "sketches". In this work, consider the problem of optimizing sketches to obtain low approximation eπrоr over data distribution.
We introduce general framework for "learning" and applying CountSketch, a type of sparse sketch. The sketch optimization procedure has two stages: one for optimizing the placements of the sketch's 's non-zero entries and another for optimizing their values. Next, we provide a way to apply learned sketches that has worst-case guarantees for approximation error.
We instantiate this framework with three sketching applications: least-squares re- gression, low-rank approximation (LRA), and k-means clustering. Our experi- ments demonstrate that our approach substantially decreases approximation error compared to classical and naively learned sketches
Finally, we investigate the theoretical aspects of our approach. For regression and LRA, we show that our method obtains state-of-the art accuracy for fixed time complexity. For LRA, optimization stage for two standard input distributions. For k-means, we derive a more straightforward means of retaining approximation guarantees.
prove that it is strictly better to include the first
INTRODUCTION
In recent years, we have seen the influence of machine learning extend far beyond the field of artificial intelligence. The underlying paradigm, which assumes that a given algorithm has an input distribution for which algorithm parameters can be optimized, has even been applied to classical algorithms. Examples of classical problems that have benefitted from ML include cache eviction strategies, online algorithms for job scheduling, frequency estimation of data stream elements, and indexing strategies for data structures (Lykouris & Vassilvitskii, 2018; Purohit et al., 2018; Hsu et al., 2019; Kraska et al., 2018). This input distribution assumption is often realistic. For example, many real-world applications use data streaming track things like product purchasing statistics in real time. Consecutively streamed datapoints are usually tightly correlated and closely fit certain distributions.
We are interested in how this distributional paradigm can be applied to sketching, data compres- sion technique. With the dramatic increase in the dimensions of data collected in the past decade, compression methods are more important than ever. Thus, it is of practical interest to improve the accuracy and efficiency of sketching algorithms.
We study a sketching scheme in which the input matrix is compressed by multiplying it with a "sketch" matrix with small dimension. This smaller, sketched input is then used to compute an approximate solution. Typically, the sketch matrix and the approximation algorithm are designed to satisfy worst-case bounds on approximation error for arbitrury inputs. With the MI. perspective in mind, we examine if it is possible construct sketches which also have low error in expectation over an imput distribution. Essentially, we aim for the best of both worlds: good performance in practice with theoretical worst-case guarantees.
Further, we are interested in methods that work for multiple sketching applications. Typically, sketching is very application-speeffe.. The sketch construction and approximation algorithm are tailored to individual applications, like robust regression or clustering (Sarlos, 2006; Clarkson & Woodruff, 2009; 2014; 2017; Cohen al., 2015; Makarychev et al., 2019). Instead, consider

---

Under review as a conference paper at ICLR 2021
three applications once (regression, LRA, k-means) and propose generalizable methods, well as extending previous application-specee work.
Our results. a high level, our work's aim make sketch learning more effective, general, and ultimately, practical. We propose a framework for constructing and using learned CountSketch. We chose CountSkerch because it is a sparse, input-independent sketch (Charikar et al., 2002). Specif- ically, it has one non-zero entry (+:1) per column and does not need to be constructed anew for each input matrix it is applied to. These qualities enable CountSketch to be applied quickly, since sparse matrix multiplication is fast and we reuse the same CountSketch for different inputs. Our "learned" CountSketch will retain this characteristic sparsity pattern and input- "pputtindependence" but its non-zero entries will range in R.
We list our main contributions and follow this with a discussion.
Two-stage sketch optimization: to first place the non-zero entries and then learn their values.
Theoretical worst-case guarantees, two ways: we derived a time-optimal method which applies to MRR, LRA, k-means, and more. We also proved a simpler method works for k-means.
SOTA experimental results: we showed the versatility of our method on 5 data sets with types. Our method dominated on the majority of experiments.
Theoretical analysis on the necessity of two stages: we proved that including the first stage is strictly better for LRA and two common input distributions.
Empirical demonstration of the necessity of two stages: showed that including the first stage gives 12, 20% boost for MRR. LRA.
Our sketch learning algorithm first places the sparse non-zero entries using greedy strategy, and then learns their values using gradient descent. The resulting learned CountSketch is very differ- from the classical CountSketch: the non-zero entries no longer have random positions and +1 values. As result, the usual worst-case guarantees do not hold.
We sought way obtain worst-case guarantees that was fast and reasonably general. Our solution is a fast comparison step which performs an approximate evaluation of learned and classical sketches and takes the better of the two. Importantly, we can run this step before the approximation algorithm without increasing its overall time complexity. As such, this solution is time-optimal and applies to MRR, LRA, i-mmans, and more.
An alternate method was proposed by a previous work, but it was only proved for LRA (Indyk et al., 2019). This "sketch concatenation" method just involves sketching with the concatenation of learned and a classical sketch. Since it is somewhat simpler, we wanted to extend its applicability. In a novel theoretical result, we proved this works for k-means as well.
We also ran diverse set of experiments to demonstrate the versatility and practicality of our ap- proach. We chose five data sets spanning three categories (image, text, and graph) to test our method on three applications (MRR, LRA, k-means). Importantly, these experiments have real-world coun- terparts. For example, LRA and k-means can be used to compress images, applying SVD (LRA) to text data is the basis of a natural language processing technique, and LRA can be used to compute approximate max cuts on graph adjacency matrices. Ultimately, our method dominated on the vast majority of tests, giving 31, 70% improvement over classical CountSketch for MRR, LRA.
Finally, we conducted ablation study of the components of our algorithm. In another novel theoreti- cal result, we proved that including the time-consuming first optimization stage strictly better than not to for LRA and two input distributions (spiked covariance and Zipfian). Empirically, this is case for all 3 applications.
Related work. In the last few years, there has been much work on leveraging ML to improve clas- sical algorithms; we only mention a few examples here. One related body of work is data-dependent
While learned CountSeetch is data-dependent (it is optimized using sample input matrices), it is still con- sidered input-independent because it applied unseen input matrices (test samples).
2

---

Under review as a conference paper at ICLR 2021
dimensionality reduction, such approach for pair- saiit-ise/multi-we wise similarity preservation for indexing big data (Wang et al., 2017) and a method for learning linear projections for general ap- plications (Hegde al., 2015). We note that multiplying sketch is equivalent to hashing its rows to a small number of bins. Thus, we also find connections with the body of work on learned hashes, most of which addresses the nearest neighhor search prob- lem (see Wang et al. for survey). However, in order to obtain approximation guarantees, our "hash function" (sparse sketch) must satisfy properties which these learned hashes usually do not, such as affine e-embedding (Def. A.1).
input matrix the left with sparse
In particular, we build off of the work of Indyk et al. (2019), which introduced gradient descent op- timization for the LRA application. It also gave an LRA-specific method for worst-case guarantees. We have surpassed the sketching performance and breadth of this work. Namely, we introduce a sparsity pattern optimization step which is clearly crucial for sparse sketches. We also provide more general method for worst-case guarantees and extend their method k-means.
PRELIMINARIES
Our learned sketches have the sparsity pattern of the classical CountSketch. The construction of this sketch is described below. We also define affine c-embeddings, which is a class of sketches that includes CountSketch. This c-embedding property is desirable because allows us to prove that certain sketching algorithms give (I + -)aapproximations.
Definition 2.1 (Classical CountSketch). The CountSketch (abbreviated as CS) matrix S has one non-zero entry in each column with random location and random value in {=1}.
Definition 2.2 (A Affine Embedding). Given a pair of matrices A and B, matrix S is affine e-embedding for all X of the appropriate shape, ||S(AX -B)I (1 e) |AX - B
Notation. We denote the singular value decomposition (SVD) of A by A UVVT with orthog- onal U, V and diagonal Σ. Relatedly, the Moore-Penrose pseudo-inverse of A is At = - where Σ is constructed by inverting the non-zero diagonal entries.
FRAMEWORK
We describe framework for learned CountSketch that be adopted by many different applica- tions, including least-squares regression (MRR), low-rank approximation (LRA), and k-means clus- tering. We will return these applications in the next section. In this section, we first describe how to optimize a CountSketch over a set of training samples. Then, we explain how to use the learned CountSketch to achieve good expected performance with worst-case guarantees. By running "fast comparison" step before the approximation algorithm, we can do this in optimal time.
SKETCH OPTIMIZATION
Most applications are optimization problems. That is, they are defined by an objective function, f(-). For example, least-squares regression solves
min |AX
BI
given E Bexd B Radd" Of course, the optimal solution is function of the inputs:
X min AX
B
ABB
However, in the sketching paradigm, we compute an approximately optimal solution as a function of the sketched (compressed) inputs. Taking CountSketch m for m & n, we have:
X (SA)
Our goal is to minimize the expected approximation error with respect to S, which is constrained to the set of CountSketch-sparse matrices (CS). However, this is simply equivalent minimizing the
3

---

Under review as a conference paper at ICLR 2021
objective value of the approximate solution.
S arg min
(<AB((SA)(SB)) [A.((X')]
SEcs (A,B)~D
arg min
SECS (A.B)~D (CA.(((A)((S))
For ease of notation, we will define G(-) as a function which maps a sketch and inputs to ap- proximate solution. G(-) is defined by the application-specifis approximation algorithm. For MRR, G(S, (A, B)) =
More generally, the sketch optimization objective is:
S = arg min E (GA, SECS AND
(3.1)
If the application is regression, we let A be (A, B)-
We will solve this constrained optimization two stages. For both stages, we approximate the expectation in empirical risk minimization (ERM) fashion. That is, we approximate the expectation over the true distribution by averaging over a sampled batch of the training set. Now, in the first optimization stage, we compute positions for the Souttkktch--parss nonzero entries. Then, in the second stage, we fix the positions and optimize the nonzero values.
Stage 1:
Placing the nonzero entries. We want to maintain the sparsity pattern of CS (one nonzero entry per column), but we are free place that nonzero entry wherever we like for each column. A naïve method would be to evaluate the objective for the exponential number of full place- ments. This is clearly intractable, so we consider greedy alternative. essence, we construct the sketch one nonzero entry at a time, and we choose the location of the next entry by minimizing (3.1) over the discrete set of possibilities.
More precisely, we build the sketch S E RR iteratively, placing one nonzero entry at time. For each nonzero entry, we consider m locations and 2 values for each location (+1). We evaluate the sketch optimization objective (3.1) for all 2m incremental updates to S and choose the minimizing update. In the pseudo-code below, we iterate through the n columns S, each which contains non-zero entry. Note that Swij = S + w(Eeet adds single entry w to the i-th column, j-th row of the current, partially-consttuuceed.. S.
Algorithm 1 GREEDY STAGE
Require: Auaan {AI,..., AND with A, E Rexd sketch dimension initialize 1 mxn
2: for i = to do
3: wt j arg
E (G( A,)) where Sw.j = S w(eeet
4: S[ w*
5: end for
For some applications it can be inefficient to evaluate (3. since it requires computing the approx- imate solution. For MRR and LRA, the approximate solution has closed form, but for k-means, it must be computed iteratively. This is prohibitively expensive, since we perform many evaluations. In this case, recommend finding a surrogate L(-) with closed-form solution, as we illustrate in later sections.
Stage 2: Optimizing the nonzero values. We now fix the positions of the nonzero entries and optimize their values using gradient descent. To fix the positions, we represent S as just vector of its nonzero entries, 전 E R" We will denote H (7) : R" + mx as the function which maps this concise representation of S to the full matrix. H.) depends on the positions computed in the last stage, which are fixed.
Now, we simply differentiate
[CA(G(H( A))] (3.1) with respect to A~D
4

---

Under review as a conference paper at ICLR 2021
Algorithm 2 GRADIENT STAGE
Require: Atrain {AI AND with A, E IRuxd; H-)) from Alg. 1; learning rate a
1: for = Nittr do
2: S = mm
sample Abatch from Atrai
4: ūtū-a Σ
8C (C(()(((T))
06
AEAbbatch
5: end for
LEARNED SKETCH WITH WORST-CASE GUARANTEES
We first run fast comparison between our learned sketch and classical one and then run the ap- proximation algorithm with "winner". This allows us to compute an approximate solution (to the applications consider here; i..., MRR and LRA) that does not perform worse than classical CountSketch. In other words, our solution has the same worst-case guarantees as classical CountS- ketch. The benefit of this is that these guarantees hold for arbitrary inputs, from out-of-distue -it-of--istribution inputs as well as in-distribution inputs which were not well-represented in the training data.
method is protected
More precisely, for a given input A, we quickly compute a rough estimate of the approximation errors for learned and classical CountSketches. This rough estimate can be obtained by sketching. We take the sketch with the better approximation error and use it to run the usual approximation algorithm.
The choice to compute rough estimate of the approximation error rather than the exact value is crucial here. It allows to append this fast comparison step without increasing the time complexity of the approximation algorithm. Thus, the whole procedure is still time-optimal.
Though this method is simple, an even simpler method exists for some applications. Indyk et al. proved that "sketch concatenation" (sketching with classical sketch appended to the learned onc) retains worst-case guarantees for LRA. We prove that this also works for k-means clustering (The- orem 4.6).
Algorithm 3 LEARNED- SEAAREED-KEETCC--LLLCOTTTT
Require: learned sketch S,; classical sketch Se; trade-off parameter B3; input data (., 2: S +- CountSketch 3 x R t CountSketch e R xd
1: Define My, such that (G A)) = || II. CA(G(S A)) = ||Mcl
3: Ar. < S R [ Ac < | Rt ||
4: if Ar D then
5: return A)
6: end if
7: return A(Sc,A)
This algorithm can be used for applications that minimize a Frobenius norm. the MRR example,
L CG ||4 ((LL))((SB) B|
so Mr ((S) B. Note that 8 is parameter which trades off the precision of the approximation enror estimate and the runtime of Algorithm 3.
4 INSTANTIATIONS
For each of 3 problems (least-squares regression, low-rank approximation. k-means clustering), define the problem's 's objective, C.).). and the approximation algorithm, G(-), which maps a sketch and inputs (S, A) approximate minimizer of L.(.).

---

Under review as a conference paper at ICLR 2021
4.1 LEAST-SQUARES REGRESSION
We consider generalized version of &2 regression called "multiple-response regression" (MRR). Given a matrix of observations (A Raxd, with n
> d) and matrix of corresponding values
(B E Raxd" with > d'), the goal of MRR is solve
min (X) = min
X AX - B1I
Approximation algorithm. The algorithm is simply to sketch A, B with a sparse sketch S (e.g,, CS) and compute the closed-form solution SA,SB, which are small (Algorithm 4).
Algorithm SKETCH-REGRESSION (Sarlos, 2006; Clarkson & Woodruff, 2017)
Require: A Raxd B E Raxdd S Rmxa I: return: (SB)
Sketch optimization. For both the greedy and gradient stages, we use the objective (G(,, B))) = (SB)S))((B) Bll
Learned sketch algorithm. For fixed accuracy, our learned sketch algorithm achieves state- of-the-art time complexity, besting the classical algorithm. We prove an equivalent statement: the learned sketch algorithm yields better approximation error for fixed time complexity. (See Lemma A.4 for the classical algorithm's 's worst-case guarantee). In the following theorem, we as- sume Alg. uses a learned sketch which an affine 3-embedding and also a classical sketch of the same size which is an affine e-embedding. If 3 < e, the above statement is true.
Theorem 4.1 (Learned sketching with guarantees for MRR). Given a learned sparse sketching
matrix SL E 2 which attains (1 + 7) approoxiimannnn for MRR on A E Raxd B E Rgaxd Alg. 3 gives a (1 + O(3 + min(y,c)) ())())))))appppp appproxxxxiniiooonn for MRR on A in O(nnz(A) + nnz B) + d'd'e + d3-4)) time where 3 is trade-off parameter.
Remark 4.2. By setting the trade-off parameter 8 Od' Alg. 3 has the same asymp- totic runtime as the best (1 + d-aapproximation algorithm of MRR with classical sparse em- bedding matrices. Moreover, Alg. for MRR achieves strictly better approximation bound (1 + O(3 + 7)) (1 + o(e)) when = ofe). On the other hand, in the worst case scenario when the learned sketch performs poorly (i.e., 7 Ω())) the approximation guarantee of Alg. remains (1 + O(e)).
4.2 LOW-RANK APPROXIMATION
Given input matrix A E R and desired rank k. the goal of LRA solve
min X,rank & CA(X) = min X.rank
|X-
- Al
Approximation algorithm. We consider the time-optimal (up low order terms) approximation algorithm by Sarlos; Clarkson & Woodruff; Avron et al. (Algorithm 5) with worst-case guarantees described in Lemma A.8.
Algorithm 5 SKETCH-LOWRANS (Sarlos, 2006; Clarkson & Woodruff, 2017; Avron et al., 2016). Require: A RRndd_ S Rmaxxn R mmaxd S2 E mR xm, R2 Rma, xd
1: UC [Te Tel + SZART
D
U t SARJ with Uc, Up orthogonal
T
2: G++Sa SARR
3: Zn + GUD)
T- T-114 2A
4: 2, [I
0] ,
5: ZıZZR
6: return: AR ZSA in form Paxks Qkxd

---

Under review as a conference paper at ICLR 2021
Sketch optimization. For the greedy stage, optimize sketches S and R individually. However, do not LA(G(((, R, S2, R2). A)) = || All with X from Alg. 5 as the objective. This is because the optimization for one sketch would then depend on the other sketches. Instead, we use proxy objective. We observe that the proof that Alg. is an e-approximation uses the fact that the row space of SAA and the column space of ARt both contain good rank-k approximation to A. Thus, the proxy objectives for S and R are for k-rank approximation in the row/column space of S SAARR, respectively. For example, the proxy objective for S is LA (G(,, A)) = [[[AV] VT - All where V is from the SVD SA UEVT and Fl takes the optimal k-rank approximation using truncated SVD. The proxy objective for R is defined analogously.
For the gradient stage, we optimize all four sketches (8, R, S2, R2) simultancously using
C R, S2. R2), || All with X from Algorithm 5, since it can easily be imple- mented with differentiable functions.
Learned sketch algorithm. Just like for regression (Section 4.1), we can prove that our learned sketch algorithm achieves a better accuracy than the classical one given the same runtime.
Theorem 4.3 (Low-Rank Approximation). Given learned sparse sketching matrix St R RR E g which attain a (1 + -p-approximmtton for LRA on A E Raxd Alg. gives (1 + B + min(ry, e)) approximmaainnn for LRA on A in O(nnz( (A) (n + d) ()ly( + 3r p)ly(5) time where (3 is a trade-off parameter.
Remark 4.4. For > k(n+ + dd)- by setting the trade-off parameter 8 (k )) Alg.
has the same asymptotic runtime as the best (1 + c)-approximation algorithm of LRA with classical sparse embedding matrices. Moreover, Alg. 3 for LRA achieves a strictly better approximation bound (1 + O(3 + 7)) (1 + o(e)) when = o(e). On the other hand, in the worst case scenario when the learned sketches perform poorly (i.e., 7 Ω(e)) the approximation guarantee of Alg. remains (1 + O(e)).
Greedy stage offers strict improvement. Finally, we prove for LRA that including the greedy stage is strictly better than not including it. We can show this for two different natural distributions (spiked covariance and Zipfian). The intuition that the greedy algorithm separates heavy-norm rows (which are important "directions" in the row space) into different bins.
Theorem 4.5. Consider a matrix A from either the spiked covariance or Zipfian distribution. Let Sr. denote sparse sketch that Algorithm has computed by iterating through indices in order of non-increasing row norms of A. Let So denote a CountSketch matrix. Then, there is a fixed 1j >> 0
such that rinkak (Erowsp(S,A) I|X - All S (1 11) miDrank-k X Erross((Se |X - Al
4.3 k-MEANS CLUSTERING
Let A E Raxd represent a of points, A An Rd. In k-means clustering, the goal is to find a partition of A,,... An into le clusters C = (C1, Ck to
min CA(C) = min C
min
Σ
||Ag - -
C
R.RRd
JEC.
where μr denotes the center of cluster C₁.
Approximation algorithm. First, compress A into AV, where V from SVD UEVT and S E RO(22/2)x) is CountSketch. Then, we use an approximation algorithm for k-means, such as Lloyd's algorithm with a k-means++ initialization (Lloyd, 1982; Arthur & Vassilvitskii, 2007). Solving with AV gives a (1 + + have (1 + ()o(lg(()) approximation overall.
(Cohen et al., 2015) and Lloyd's gives an (llo((k)
approximation,
Sketch optimization. h-means is an interesting case study because it presents a challenge for both stages of optimization. For the greedy stage, we observe that k-means does not have closed form solution, which means that evaluating the objective requires a (relatively) cine-consuming itera- tive process. In the gradient stage, it is possible to differentiate through this iterative computation,
7

---

Under review as a conference paper at ICLR 2021
but propagating the gradient through such nested expression is time-consumingg Thus, proxy objective which simple and quick to evaluate would be useful for both stages.
Cohen et al. showed that get good k-means solution if A is projected to approximate top singular vector space. This suggests that a suitable proxy objective would be low-rank approxima- tion: (G(S, A)) = [[[V]< - All like in Section 4.2. We use this objective for both the greedy and gradient stages.
Learned sketch algorithm. For k-means, we can use different method to obtain worst-case guarantees. We prove that by concatenating a classical sketch to our learned sketch, our sketched solution will be a (1 + oppo)-apprxximation.
Theorem 4.6 (Sketch monotonicity property for k-means). For given classical CountSketch Sc E [o(poly(</e)xx sketching with any extension of Sc (i.e., by a learned sparse sketch St) yields better approximate solution for k-means than sketching with Se itself.
EVALUATION
We implemented and compared the following sketches:
Ours: a sparse sketch for which both the values and the positions of the non-zero entries have been optimized
GD only: a sparse sketch with optimized values and random positions for non-zero entries
Random: classical random CountSketch
We also consider two "naively" learned sketches for LRA, which are computed on just one sample from the training set.
Exact SVD: sketch as AVm, where Vn, contains the top m right singular vectors of random sample A, Atrain
Column sampling: sketch AR, where R is computed from randomly sampled A; E Atrain Each column of R contains one entry. The location of this entry is sam- pled according to the squared column norms of A₁; the value of this entry is the inverse of the selected norm.
Data sets. We used high- nigh-dimensional data sets of three different rypes (image, text, and graph) to test the performance and versatility of our method. Note that the regression datasets are formed from the LRA/k-means datasets by splitting each matrix into two parts.
Table 5.1: Data set descriptions
Name
Description
shape LRA, k-meuns) 5760 x 1080
IP) shapes (for MRR)
Nirats
Neese
Friends
Frames from a scene the TV
(5700 x 1079), (5760 1)
100
100
show Friende"
Hyper
Hyperspectral images dopicting outdoor scenes (Imamoglu et al., 2018)
1024 x 768
(1024 x 767), (1024
400
100
Logo
Frames from video logo being painted"
5760 x 1080
(5760 x 1079), (5760 1)
400
Yelp
If-idf (Ramos, 2003) of restaurant reviews, grouped by restaurant
7000 x 800
(7000 6400)
(7000 160)
260
65
Graph
Gruph udjacency matrices of social circles Google+ (Leskovec & McAuley, 2012)
1052 x 1052
(1052 x 842), (1052 210)
147
37
http::://youuu.be/l.at
"https://www,
htttps://www.yelp..coo wttp:/wwwwwloo...mo.....oo..m..a ctwp/.p..poppaaaaaaaa.aeaa
8

---

Under review as a conference paper at ICLR 2021
Evaluation metric. To evaluate the quality of sketch S, we compute the difference between the objective value of the sketched solution and the optimal solution (with no sketching). That is, the values in this section's 's tables denote
As LAG(S, Atest) C'
averaged over 10 trials.
Analysis of results. For least-squares regression and LRA, our method is the best. It is signif- icantly better than random, obtaining improvements of around 31% and 70% for MRR and LRA respectively compared to a classical sparse sketch. For k-means, "Column Sampling" and "Exact SVD" dominated on several data sets. However, we note that our method was always close second and also, more importantly the values for k-means were only trivially different (< 1%) between methods.
Ablation of greedy stage We find empirically that including the greedy optimization stage is al- ways better than not including it (compare the "Ours" vs. "SGD only" methods). For regression, LRA, and k-means, including this step offers around 12%, 20%, and < 1% improvement respec- tively. We should note that for k-means, the values generally do not differ much between methods.
Table 5.2: Average errors for least-squares regression
Parameters
Algorithm
Datasets
M
Friends
Hyper
Logo
Yelp
Graph
20
Ours
0.4360
15.6e-02
0.6242 19.7e-03
0.1942
11.2e-02
0.2734
L1.ce-03
58,6568
1.66c-022
SGD only
0.5442
49.0--0 03
0.6790
0.22--02
0.2547
+-376--03
0,2781
+6630-04
58.7231
4-8.7o-02
Random
0,7792 +2.50
0.8266
4-4.4e-02
02
0.3773
+1.50-02
0,3173
#1.30
59,5535
+1.20-01
40
Ours
0,2485
11.1
02
0.5380
18.3e-
03
0.1607
15.7e-03
0,2639
14.1o-04
04
55,5267
18800-02
SGD only
0.3688
+3.00
08
0.5718
11.26
02
0.1858
1.2.80
03
0,2701
+7.20-
-04
55.6120
+-0.5e-02
Randon
0.5213
11.6
02
0.6414
11.9c-00 02
0,2389
1..00- 003
0.3114
1.266
2266-03
03
55,7896
1.6. 1.6.7e002
Table 5.3; Average error for LRA
Parameters
Algorithm
Datasets
runk K, m
Friends
Hyper
Lugo
Yelp
Gruph
(20, 40)
Ours
0.8998
22-76-02
2,4977
E1.8e-01
0.5009
2e-12
0,1302
#3.5e-01
30,0969
1:133-01
SGD only
1,0484
30
3,7648
44.20-102 172
0.6879
03
0,1316 +8.80-11
30.5237
+1.30-01
Random
4,0730 1.77
01
6.3445
11.80
01
2.3721
18.3e
02
0,1661 1.44 03
33.0651
1.3.4 01
Exact
1,5744 -1.10-01
3,400
+3.70-01
0.7470
0.1394
+3.70-0
31,0017
+3.40-01
Column sampling
5,9837
16.00
9.7126
18.20
001
4,2008
16.00
01
0,1881
1
43.9920
155000 01
(30, 60)
Ours
0.7948
1.5g
02
2,4920
12.60
01
0.4929
-2.3
02
0.1128
1:3.0
04
30.6163
1.88 01
SGD only
1,0722 =1.26-02
3,1488
-1.80-122
0.7348
=7 03
0,1137
30,9279
4:200-00
Rundom
2.6836 17.40
02
3.3934
17.76
022
1,6228
4.30
02
0,1463
1:330
32,7905
122 01
Exact SVD
1,1678
25.26 02
3,0648
18. 50
0.545%
=7
02
0,1326
1:2 3e 00
32,0861
1:5 90-
Column sampling
4.1899
=2 22
8,2314
11.70-01
3,2200
0.1581 3:30
44.2672
15 20
Table 5.4: Average errors for k-means clustering
Algorithm clusters, 2n, rank (20, 40, 20)
Parameters
Data sets
Friends
Hyper
Logo
Yelp
Graph
Ours
0.2034
86-66
0.6110
04
0.1251
04
2.4755
#7.20-0)
2.797 #1100-03
SGD only
0,2935
#1.4
0.6117
=2.00
53.86
16.
2,4566
#7.00
2.8T00
#:9.0
Random
0.203 + 34
0.6123
0.1254
+1330-14
2,4599
4-7 04
2.9844
Exect SVD
0,2034
1.22
0,6113
+2.50
1253
+2.50
04
2,4977
+0.00
19779
Colomn sampling
0,2932 + 7e
0.6129
=3, 04
0.1250
10-70-14
2,4922
43 88-48
3.0137
(30, 60, 30)
Ours
0.2604
0.5H4
1.62
1e-7e-05
2,4455
Se-01
2,8765
1111--03
SGD only
0,2605
1e
08
0.5452
10 04
0. 1063
1.26
04
2.4476
17
2,777
1.
Random
0.260
Se-
0.3153
0. 1062
2: 110
2,4177
1e-10-11
18191
Exact SVD
0,2005
Se
04
0.5446
2.00
0. 1063
1: 06
2.4532
1.166- 04
2.8799
11770
Cokuma sampling
0,2602
2e-
0.5159
23.00-01
0.1059
1e-11
2/1665
18-08
25002
22.00-13
9

---

Under review as a conference paper at ICLR 2021
REFERENCES
David Arthur and Sergei Vassilvitskii. k-means++: The advantages of careful seeding. In Proceed- ings of the 18th Annual ACM SIAM Symposium on Discrete Algorithms, pp. 1027-1035, 2007.
Haim Avron, Kenneth Clarkson, and David Woodruff. Sharper bounds for regularized data fitting. arXiv preprint drXiv:1611.03225,.00 2016.
Matthew Brand. Fast low-rank modifications of the thin singular value decomposition. Linear Algebra and its Applications, 415.1, 2006.
Moses Charikar, Kevin Chen, and Martin Farach-Colton. Finding frequent items in data streams. In International Colloquium on Automata, Languages, and Programming, pp. 693-703. Springer, 2002.
Kenneth L Clarkson and David P Woodruff. Numerical linear algebra in the streaming model. In Proceedings of the forty-first annual symposium on Theory of computing (STOC), pp. 205-214, 2009.
Kenneth L. Clarkson and David Woodruff. Sketching for m-estimators: A unified approach to robust regression. In Proceedings of the 26th Annual ACM SIAM Symposium on Discrete Algorithms, pp. 921-939, 2014.
Kenneth L Clarkson and David Woodruff. Low-rank approximation and regression in input spar- sity time. Journal of the ACM (JACM), 63(6):54, 2017.
Michael B Cohen, Sam Elder, Cameron Musco, Christopher Musco, and Madalina Persu. Dimen- sionality reduction for k-means clustering and low rank approximation. In Proceedings of the forty-seventh annual ACM symposium on Theory of computing, pp. 163-172, 2015.
Dan Feldman, Melanic Schmidt, and Christian Sohler. Turning big data into tiny data: Constant-size coresets for k-means, pea and projective clustering. In Proceedings the 24th annual ACM-SIAM symposium on Discrete Algorithms, pp. 1453, 2013.
Chinmay Hegde, Aswin C. Sankaranarayanann Wotao Yin, and Richard G. Baraniuk. Numax: A convex approach for learning near-isometric linear embeddings. IEEE Transactions on Signal Processing, pp. 6109-6121, 2015.
Chen-Yo Yu Hsu, Piotr Indyk, Dina Katabi, and Ali Vakilian. Learning-base based frequency estimation algorithms. International Conference on Learning Representations, 2019.
Nevrez Imamoglu, Yu Oishi, Xiaoqiang Zhang, Guanqun Ding, Yuming Fang. Toru Kouyama, and Ryosuke Nakamura. Hyperspectral image dataset for benchmarking on salient object detection. In International Conference on Quality of Multimedia Experience, pp. 1-3. IEEE, 2018.
Piotr Indyk, Ali Vakilian, and Yang Yuan. Learning-based low-rank approximations. In Advances in Neural Information Processing Systems, pp. 7400-7410, 2019.
Tim Kraska, Alex Beutel, Ed H Chi, Jeffrey Dean, and Neoklis Polyzotis. The case for learned index structures. In Proceedings of the 2018 International Conference on Management of Data, pp. 489-5044 2018.
Jure Leskovec and Julian McAuley. Learning to discover social circles in ego networks. In Advances in Neural Information Processing Systems, pp. 539-547, 2012.
Stuart Lloyd. Least squares quantization pem. IEEE transactions on information theory, 28(2): 129-137, 1982.
Thodoris Lykouris and Sergei Vassilvitskii. Competitive caching with machine learned advice. In International Conference Machine Learning, pp. 3302-3311, 2018.
Konstantin Makarychev, Yury Makarychev, and Ilya P. Razenshteyn. Performance of johnson- lindenstrauss transform for k means and k-medians clustering. In Proceedings of the 51st Annual ACM SIGACT Symposium Theory of Computing, STOC, pp. 1027-1038, 2019.
10

---

Under review as a conference paper at ICLR 2021
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- frrty-fitt annual ACM symposium 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.
Manish Purohit, Zoya Svitkina, and Ravi Kumar. Improving online algorithms via ml predictions. In Advances Neural Information Processing Systems, pp. 9661-9670, 2018.
Juan Ramos. Using tf-idf to determine word relevance in document queries. In First Instructional Conference on Machine Learning, pp. 133-142, 2003.
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.
Jingdong Wang, Ting Zhang, Nicu Sebe, and Heng Tao ShenWang. A survey on learning to hash.
In IEEE Transactions on Puttern Analysis and Machine Intelligence, pp. 769-790, 2017.
David P Woodruff. Sketching tool for numerical linear algebra. Foundations and Trends in Theoretical Computer Science, 10(1-2):1-157, 2014.
A APPENDIX
CORRESPONDING TO SECTION 4.1: LEAST-SQUARES REGRESSION
Definition A.1 (Subspace Embedding). Given d matrix A, matrix S is affine embedding for (the column space of) A ffr R2, ||SAx]| = (1 +) |x
The following result is shown in Clarkson & Woodruff (2017) and sharpened with Nelson & Nguyên (2013); Meng & Mahoney (2013).
Lemma A.2. Given matrices A, B with n rows, sparse embedding matrix (i.e., random CS) with ((rank(A)2/e) rows is an affine e-embedding matrix with constant probability. Moreover, the matrix product SA can be computed in O(nnz(A) time, where nnz(A) denotes the number of non-zero entries of matrix A.
Lemma A.3 (Clarkson & Woodruff (2017); Lemma 40). Let A be an X d matrix and let S E O 1 be randomly chosen sparse embedding matrix (i.e., random CS). Then with constant probability, ||| (1 + <) ||
Lemma A.4 (Sarlos (2006); Clarkson & Woodruff (2017)). Güven classical CountSketch S Rmxn,
(A S) returns a (1 + c)-approximation in time (nnz(AA) + (Bz(B)= dd' + min dm?)).
Proof: Since S is an affine c-embedding matrix of A, B, then
||SAX SBI (1 + ||AX -
Next, by the normal equations, (SA)+ is minimizer of minx ||SAX SBB and is a + appropxiiee solution of minx ||AX - 2
To bound the runtime, note that since S is CountSketch O(nnz(A) + anz(B) (B) and reduce the problem to an instance of multiple-response regression with m rows. Then, we solve the reduced size problem in time O(dd' m² + min(d dm O0min(dre) m) to compute (SA) and O(dd'm² to compute (SA)+(SB)
compute SA and SB in time
Proof of Theorem 4.1. By Lemma A.2, a random CS So with O ( rows is affine e-embedding matrix of A, B with constant probability.
Next, let X2 and Xo be respectively the solutions returned by SEETCH-REGRESSIONEE (A, B, So). By Lemma A.4, with constant probability,
and
|Aoo B1} < (1 + 3c) min |AX -B12 B
11

---

Under review as a conference paper at ICLR 2021
Together with the assumption that the solution constructed from S2, which is denoted XL, over Auaii is a 7-approximate solution,
mix( - BIi 11AXo - BII;) & (1 min(3e, 7)) min ||AX - -BI
(A.1)
efficiently. Note that it takes 12(n d d') to compute these values exactly. However, for our purpose it suffices to compare (1 + -estimmass of these values and return the minimum the estimates. To achieve
Hence, only remains compute the minimum of ||, BII and |AXx - -Bli
this, we use two applications of Lemma A.3 with Rd 12 ,8 S )xn particular, both Xo and Lt), with constant probability,
For any X' (in
(A.2) Let Tr. and To respectively denote AX, B and AXo B and let T' arg min( ST, R [[S[[)TTT)))))) By Eq. (A.2) and union bound over Xo and Xt, with con- stant probability,
||S R || (AX' B) T (
B)I
|r s (1+ (L+) a RI
by Lemma A.3
5 (1+00()) min[ oll (
by Lemma A.3
$ 10000((++nne)))) ||AX - BIE
by Eq. (A.1)
Runtime Analysis. By Lemma A.4, Xo an X can be computed in O(nnz(A) + mnz( + d time. Next, the time to compute A, and Do is
d
(Ann((A) + unz( X R) + nnz B) + z
O(mm((A) + nnz(B) + dd' + 3z))
where X is either Xo or Xr, and we use that fact that mnz(X) S d' (i.e., the total number of cells in X).
Thus, the total runtime of Algorithm 3 is O(mz((A) + nnz(B) d' +
Theorem A.5 (Least Squares Regression). Suppose there exists a learned, sparse, subspace em- bedding matrix SL computed over Auaan with poly( () rows that attains a B-approximation over Atest: Then, there exists algorithm that runs in time O(nnz(A) + poly(d/c)) that outputs (1 + min cppappproinnation to the least squares regression problem.
Proof: Note that to solve minzer" || |r - bll2- given a sparse embedding matrix $ E mx for the columns of A together with the vector b, the problem can be solved within + e approxximati in time unz(A) + poly(d/e) (e.g., see Theorem 2.14 Woodruff (2014)).
The proof outline is similar to the proof of Theorem 4.1. First, we compute solutions IT, zo Rd to the given instance using respectively learned sketching matrix Sy and a learning-free sketching algorithm. Then, compare || - bll2 and ||Aro bll2 and report the better. Note that since L zo are vectors, unlike the case of MRR, the naive comparison (i.c., without applying any sketch- ing matrices) takes mnz(A). Hence, the algorithm runs in time nnz(A) + poly(d/c) and returns (1 + min(c, 3)) approximmtt solution of the least squares regression problem.
Moreover, we can employ the best known sketching techniques for the least squares regression problem and achieve the dependence d/e² 2 in the running time (e.g., see Section 2.5 of Woodruff (2014)).
A.1 CORRESPONDING TO SECTION 4.2: LOW-RANK APPROXIMATION
Lemma A.6. Suppose that S E Rm and R B sparse affine e-embedding matrices for (AT, A) and ((SA) A.). Then,
min | XSA AIl; < (1+c) A 2 rank-k X
Proof: Consider the following multiple-responss regression problem:
min || AAl rank-k x
(A.3)
12

---

Under review as a conference paper at ICLR 2021
Note that since X Ig is feasible solution to Eq. (A.3), min x || All
||A - Alle- Let S R be a sketching matrix that satisfies the condition of Lemma A.9 for A ;= A& and B := A. By the normal equations, the rank-k minimizer of ||S - SA is (SAR)* SA. Hence,
( All - (1+ 1 A
(A.4)
which in particular shows that a (1 + +) rank-k approxmmaiio of A exists in the row space of SA. In other words,
min |XSA Al S (1 + +) - All
(A.5)
Next, let R E gmaax be a sketching matrix that satisfies the condition of Lemma A.9 for = (SA) and B == AT. Let Y denote the rank-k minimizer of [R(SA) X RR | Hence,
[[(SA) y A 1TI S c1+c)
min IXSA All rank-k X
Lemma A.9
(1+0 | -Alle
> Eq. (A.5)
(A.6)
Note that by the normal equations, again rowsp(Y") $ rowsp(RAA) and we can write Y ART Z where rank(2) ()) = k. Thus,
||R T XSA AllE & |AR ZSA Alle rauk-KX
min
= =I(SA) - -A -A
D = AR
(1+00()) || All
b Eq. (A.6)
Lemma A.7 (Avron et al. (2016); Lemma 27). For C E IRxxmm following problem
D E pmxp' G E Rexp', the
min |CZD - G rank-k 2
(A.7)
can be solved in Opm''r + p'mrD + pp' (rD + +r)) time, where rc = rank(C) S min(m', p} and nn rank(D) D) < min(m.p') p').
Proof: Let Uc and UD be orthogonal bases for colsp(C) and rowsp(D), respectively, that for each Z, CZD U'U for some Z'. Let Pc and PD be the projection matrices onto the subspaces spanned by the rows of CT and DT, respectively: Pe = UcUU and Pp UDUT. Then by the Pythagorean theorem,
JCZDDGT G |PeUc'U 1
PPUCCZUDDD PoGPD + 1P PD)IIE + |[(1 Po)
where the first equality holds since PcUc = Uc and PD = UT and the second equality follows from the Pythagorean theorem. Hence,
arg min |CCD rank-k Z
Gl} = arg min rank-k Z
-
Moreover,
1 - 2 - I
- [
where the first equality holds since UC and UDUD 1, and the second equality holds since Uc and UT are orthonormal. Hence,
arg min |CZD G = arg min || - UU.GUuu rank-k Z
rank-k 2
Next, we can find Z = GUTGUple by computing the SVD of GUD in form of Z, ZR where 7r, mRm' x and Zn E kx
13

---

Under review as a conference paper at ICLR 2021
process in time Oppm' + Runtime analysis. We can compute Uc and UD by the +rc))). Finally, the time to compute p' D'mrD), and GUD in time O ((min(rcp'pp+ (p)) (i.c., an SVD computation of Ug GÚ is O min(re. rD)), Since re & min(p, m'} and TD S min(y', mj. the total runtime to minimize Z in Eq. (A.7) is O(pm're + p/mr + pp' (rc + Tp)).
Lemma A.8. Ler 5 E (Poly(k/e)xdd R E [Rool((//exdd be CS matrices such that
min x
|AR USR Aslp-11p- S (1+ |A AAlE
(A.8)
Moreover, let S2 R and R2 R xd be CS matrices. Then, Algorithm 5 runs in (A)z((A) + (n + d) poly(k/e)) time and with constant probability gives a (1 + O(3 + 7))- approximate rank-k approximation of A.
Proof: The algorithm first computes C' = S2ART, D = SAR G = StARy which can be done in time (n))( As an example, we bound the time compute C = S2AR Note that since S₂ is a CS, S can be computed in (Anz((A) time and the number of non-zero entries in the resulting matrix is at most nnz(A). Hence, since is a CS as well, can be computed in time ( A) + (n)A)) = (()nz(A)). Then, it takes an extra poly(k/e) store C and G in matrix form. Next, as we showed in Lemma A.7, the time to compute Z in Algorithm 5 is 0 5 poly(k/c)). Finally, it takes (n + d) poly(k/c) time to compute Q = AR Z and P SRSA and return the solution in the form of P Qkxd- Hence, the total runtime is ((nzz()) A) + (n + d) poly(k/cc) + ks
k² time to z
334 pollyk//)).
The approximation guarantee follows from Eq. (A.8) and the fact that S2 and Rg are respectively affine 3-embedding matrices of AR and SA (see Lemma A.2).
Lemma A.9 (Avron et al. (2016); Lemma 25). Suppose that A E Rand and B Raxd' More- over, let S be an oblivious sparse affine e e-embedding matrix (i.e., a random CS) with (rank(A)²/E²) rows. Then with constant probubility. X arg rank X
||SAX- - -SB| satisfies
[AX - S (1 (1++) min ||A - -BI
rank-k X
In other words, in ((nz A) nnz (B) (d+ d') (rank(A))//" time, can reduce the problem to a smaller (mulli-response regression) problem with (rank(A) (A)2/e2 rows whose optimal solution a (1 (1+)))appraan) -)-pprooximate solution the original problem.
Proof of Theorem 4.3. Let So and Bo be CountSketch matrices of size poly(k/e) x n and poly(k/e) x d. Note that since rank(As) k and rank((So.A)T) S <oly(k/cc), So and Ro are respectively affine e-embedding matrices of (A,, A) and ((SoA) ). Then, by an application of Lemma A.6
min |IAROXSOA AII S (1+)()) | AlF rank-k X
(A.9)
Similarly, by the assumption that S,. and R,, finds a (1 + appproixiee solution of LRA over matrices A Aust, then
min |AR - A (1 0()) A rank-k X Next, let (Pr, Qr) and (Po, Qo) be respectively the rank-k approximations of A in factored form using (St, Rr.) and (So, Ro) (see Algorithm 5). Then, Eq. (A.10) together with Eq. (A.9) implies that
(A.10)
min(I - All:. IPooo - Al)²) = (1 + (min(e,, || - Al
(A.11)
Hence, it only remains to compute the minimum of P All2 and IPPo All efficiently and we proceed similarly to the proof of Theorem 4.1. We use two applications of Lemma A.3
I )2 S )x² Let Fz-- PtQr A. To Polo and I'M
with R.T
arg min R|I R|)T))))p) Hence,
MT [ < (1+ 0(3) M RI
c by Lemma A.3
$1 min[[] [Tol
b by Lemma A.3
< (1 + 01+0(3 min(s || A| All
by Eq. (A.
14

---

Under review as a conference paper at ICLR 2021
Runtime analysis. By Lemma A.8, Algorithm 5 computes QL and Po. Qo in ((nzz(A) + (n d) poly(2) /े. poly(
Next, it takes (A) (n + d) dk-k the amount of time required compute ,, SAR corresponding to Ar. Since S and R are sparse sketching matrices, SPr. Qr R and SAR can be computed in miz(SPr. + nnz(Qr. + nnz(A). Since 1r,, and Qt. R are respectively of size zr x and k- zz, in total it takes ( (A)+ k 3
& to compute A and Ao. As example, we bound 3o
to compute these three matrices. Then, we can compute S and [| in time ( g+
Hence, the total runtime of Algorithm 3 for LRA is ((mzz(A) (n+d) -pol)()))pll)))))))))))))))))))) # poly($)).
A.2 CORRESPONDING TO SECTION 4.3: k-MEANS
We restate notation and the main result below for ease of reference.
Notation. We define Am. the optimal rank-m approximation of A formed by truncated SVD:
Am Um VIT
Given a matrix U with orthogonal columns, let To AUUT, which is the projection of the rows of A onto col(U). Let Ct be the optimal k-means partition of To (A). Further, we let μCC.. denote the i-th cluster's center in the optimal k-means clustering on TT
We denote dist? (A, 14) the k-means loss given cluster centers (and their corresponding partition):
dist (A,,,)
Σ SE[R] JEC.
Σ
|
1
Likewise, cost(C) is the k-means loss given a partition:
cost(C)
Σ min Σ k
1|A -2ll2
Mi
JEC.
Definition A.10 (Projection-cost preserving sketch). is a projection-cost preserving sketch of A for any low rank projection matrix P and not dependent on P:
(1- - < Â- - P + + -
PA
Theorem (4.6: Sketch monotonicity property for k-means). Assume we have A E Rrxdd We also have random CountSkeich S E (ppolyy((k)))))) and define U (k/olpl)ykk/ee) with orthog- onal columms such that (U)) rowsp(SA). Then, any extension of S to S' (for example, concatenation with learned CountSketch S.,) yields a better approximate k-means approximation. Specifically, define W with orthogonal columns such that col(W)) row A). Let C* denote the optimal partition of Cv denote the optimal partition of no (A), and Cw denote the optimal partition of nw (A). Then
cost(Cw) S ((e)) S (100())
Proof:
cost(Cur) =
Σ
min
Σ
||Aj - 1ill
sE[k]
jECw.
|| - H.C:..ll
sE[k] jeCw.
Σ
Σ
|A, xw (A;))1 + (TT((A)) - NCw: 1
(A.12)
xE(R] JECW..
15

---

Under review as a conference paper at ICLR 2021
E
|Ag Fw (A;)ll
+||TT((A pCc.a ||?
(A.13)
[k(R] jECU..
E
Σ
Aí 1J--MOu.1 MOO.i 11
(A.14)
SEIH] JECU,
(1 Σ min
Σ
1|A, - Mill
(A.15)
Me
iEN
SECo,
(1 + c)cost(Cu)
- (1 + ((c))cost((")
(A.12): Pew E (W) so we can apply the Pythagorean Theorem.
(A.13): (New: Cw) is an optimal k-means clustering of the projected points Tw (4).
(A.14): Pey E colsp(U) c colsp(W),, so we can apply the Pythagorean Theorem.
(A.15): We apply Corollary A.14.
Remark A.11. Our result shows that the "sketch monotonicity" property holds for sketching ma- trices that provide strong coresets for k-means clustering. Besides strong coresets, an alternate approach to showing that the clustering objective is approximately preserved on sketched inputs is to show a weaker property: the clustering cost is preserved for all possible partitions the points into & groups Makarychev al. (2019). While the dimension reduction mappings satisfying strong coresets require poly(k/e) dimensions, Makarychev et al. (2019) shows that O(log k/e²) dimen- sions suffice to satisfy this "partition" guarantee. An interesting question for further research is if the sketch monotonicity guarantee also applies to the construction Makarychev et al. (2019).
Corollary A.12. Assume we have A E Raxd, e <> 0. Define min(O(poly(j/e)) d)). Let Am be (1 c)-approximation
Am of the form Å. AVVT where SA UEVT for CountSketch S E Rmxn, Let X E RRxx be a matrix whose columns are orthonormal, and let (dd be matrix with orthonormal columns that spans the orthogonal complement of (X)sp(X). Then
AXXX - A. X Sc. 11
Proof:
AXXX - -
VV = =1|A(T [IA(I-VVT))L) VXX 1/-
SS]|A(1 VV x|
(A.16)
S A/]| VV )12
(A.17)
S $5.0(5
114
A 11
(A.18)
min(n.d)
O(c)
Σ of
(A.19)
s O(c)
(A.20)
(A.16) Note that rank(X) = j. so rank (an VV (XX)XX) = j. We use this fact to bound the Frobenius norm by the operator norm.
(A.17) Using the fact that X X is projection.
(A.18) Using Lemma 18 from Cohen et (2015), where A, is the optimal rank-j approximation
A. We can apply this lemma because CountSketch is one of the eligible types of random projection matrices.
(A.19) Letting d; be the singular values of A.
(A.20) Emin(n,,) (n.d) of miny |Y|] for Y E Rdx(d-)) with orthonormal columns.
16

---

Under review as a conference paper at ICLR 2021
Theorem A.13. Assume we have A E R j Zt (0, 1]. Define (po(O))olyj/fe) d). Let Am be (1 + c)-aproximation Am of the form Am AVVT where SA = UVV for CountSketch E Rmmx. Then, for any non-empty set contained in j j-dimensional subspace, we have:
dist? 11) + +I|Am A dist d 1) S (A,
Proof: We follow the proof of Theorem 22 in Feldman et al. (2013), but substitute different analyses in place of Corollaries 16 and 20. The result of Feldman al. (2013) involves the best rank-m approximation of A, Ami we will show it for the approximate rank-m approximation, Am-
Define X E Rdxj with orthonormal columns such that colsp(X) = span(n)- Likewise, define Y E Rdx(d-1) with orthonormal columns such that colsp(Y) (p)n(u)+ By the Pythagorean
theorem:
dist (Ame ()) = [ + dist" (AmXX",, ,1)
and
dist (A, 11) = ||AY| + dist 2 (AXX" .1)..
(A.21)
Hence,
disi (lis)((Ams))) + | Am 1/2 dist (A,
A Am | (IIAY 17 dist .1))
|| 112 + dist (Am .1) +
|| 11? 1|A Am -11121 || 112 + |dist X .R) dist dist" (AXXT .1)
(A.22)
82
||A1|| + |dist (Am ('mXXXX XXT T 11) cist (AXX) m
(Am
(A.23)
dist (A, (1) + dist? (Am (AmXXT 14) dist? (AXX 11)
>> Used (A.21)
(A.24)
(A.22) Triangle inequality.
(A.23) Take e in Theorem 16 from Cohen et al. (2015) as 2/8. This theorem implies that Åm is projection-cost preserving sketch with the term A Am
F Specifically, it says AV is project-cost preserving sketch, which means Am = AVVT is too: has orthonormal columns so II(T - P) [|(T PAA--))VVTTT 1
By Corollary A.12,
2² AmXXT - AXXT < -88 -IAY
Since μ E (owsp(X), we have ||YI S dA, Using Corollary 21 from Feldman et al. (2013) while taking as e/4, A as Amx XXT and B AXXT yields
dAist"(An XX" ,1s (AXX ,1) < dist +)+ c /111 -
By A.21, dist? dist? (AXXT, 1) < dist (A, Finally, we combining the last two inequalities with (A.24):
(dist (Am, + -l|A-Amm|
Am 11m11/
dist?(A,
è (1 )) dist? (A, dist? (A, + dist dist? (A,
8
Se disdit (A, μ):
where in the last inequality we used the assumption S 1.
17

---

Under review as a conference paper at ICLR 2021
Corollary A.14, Assume we have A E Raxd and CountSketch S [((poly(k/))x Then, define U E (dd0O(ply(k)))) with orthogonal columns spanning row Also define A = au (T(A)) and py as the set of optimal cluster centers found A Now, assume e E (0, 1/3]. Then, py is (1 + e)approximattion to the optimal k--eans clustering of A. That is, defining Pe as the cluster centers which minimize the cost of partition Ct on A, have:
(A, S (1 + e) dist d (A,
Proof: By using in Theorem A.13 with j as k.
dist (A" mw) + ||A - A All dist (A, mb) S t "dist (A, (w))
which implies that
(1 - {) dist (A, pv) S dist? dist?(A" .uo) + +|AA- - Al ||A4-1411
(A.25)
Likewise, by Theorem A.13 on A and μi (and taking j as k).
dist (A" m() + A - A dist (A, mt))
S 6 dist? (A,
3
which implies that By (A.25) and (A.26) together, we have:
dist (A .b)) ||A A (1 () dist?(A,, pū)
(A.26)
(1 9) dist 2 (A, pc) s dist?
mb) + || A (A dist?(A'
S dist (A .bt) + || ^
$ (1 + 3 dist '(A (A, /
Now, 1te/3 I-e/3
< 1 e, so we have d mv) S (1 + e) dist 2 (A,
B SPECTRAL NORM GUARANTEE FOR ZIPFIAN MATRICES
In this section, we show that if the singular values of the input matrix A follow Zipfian distribution (i.e.. og x i for a constant a), we can find a (1 + E) rank-k & approximation of A with respect to the spectral norm. A key theorem in this section is the following.
Theorem B.1 (Cohen et al. (2015), Theorem 27). Given an input matrix A E Raxd there exists an algorithm that runs in (nnz) A) + (n++d) poly(k/ (c)) and returns a projection matrix P = QQ such that with constant probability the following holds:
||P
All s (1 + ( As + O ( Arl1
(B.1)
Next, we prove the main claim in this section. There are various ways to prove this, using, for example, a bound on the stable rank of A; we give the following short proof for completeness.
Theorem B.2. Given a matrix A [Rd whose singular values follow a Zipfian distribution (i.e., J; x i with a constant a 2 1/2, there exists an algorithm that computes rank-k matrix B (in factored form) such that || BIβ < (1 + <) A| Ax112
Proof: Note that since the singular values of A follow a Zipfian distribution with parameter a, for any value of k,
rank(A)
rank (A)
|| - A
Σ o, =
Σ
20
ba=
D
i 1++1
raak(A)
S C.
x 20 dr
18

---

Under review as a conference paper at ICLR 2021
(1+ 1 )
S k. 2a -1
+
O(k .
O(k ((K-I|A- ||A- Axl
(B.2)
By application of Theorem B.1, (A)z((A) + (n + d) poly(k/e)) such that with constant probability,
compute matrix B in factored form in time
| - 2 & (1+ +) |4 Au + 7 ||A-AAll
By Eq. (B.1)
(1 (())) I A
Eq. (B.2)
C GREEDY INITIALIZATION
In this section, analyze the performance of the greedy algorithm the two distributions men- tioned Theorem 4.5.
Preliminaries and Notation. Left-multiplying A by CountSketch S E I m is equivalent to hashing the rows of to m bins with coefficients in {-1,1}. The greedy algorithm proceeds through the rows of A (in order) and decides which bin to hash to, denoting this by adding an entry to S. We will denote the bins as d; and their summed contents as Wi.
C.I SPIKED COVARIANCE MODEL WITH SPARSE LEFT SINGULAR VECTORS.
To recap, every matrix A R from the distribution Asp(s, E) has < k "heavy" rows (Ar,..., AA,)) of norm & > 1. The indices of the heavy rows can be arbitrary, but must be the same for members the distribution and are unknown to the algorithm. The remaining rows (called "light" rows) have unit norm.
In other words: let (ri...., rs). For all rows A, E [n]:
L-
if R
A =
Ui
O.W.
where &; a uniformly random unit vector.
We also assume that Sr, S, E R and non-increasing row norm ordering for the greedy algorithm.
Proof sketch. First, we show that the greedy algorithm using a non-increasing row norm ordering will isolate heavy rows (i.e., each is alone in bin). Then, conclude by showing that this yields a 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.
is well known that a set of uniformly random vectors are e-almost orthogonal (i.e., the magnitudes of their pairwise inner products
most e).
Observation C.1. Let U1 ... Un E Rd be set of random unit vectors. Then with high probability (1r, uz)l < 2 log Vi Vi<is n.
We define E 22 logu
Observation C.2. Let un.... un be set of vectors such that for each pair of i < j S t,
[uzll S and gs. ... (11.1). Then,
| 2 Σ
2c
Σ isj
||ullla |15ll22 <
Σ u:u:
+ 2c Σ || |ullle i<g
(C.1)
Next, straightforward consequence of e-almost orthogonality is that can find OR-factorization of the matrix of such vectors where R (an upper diagonal matrix) has diagonal entries close to 1 and entries above the diagonal are close to
19

---

Under review as a conference paper at ICLR 2021
Lemma C.3. Let u1," u4 R be a set of unit vectors such that for any pair of i < j S t, Ku,, ug)l I e where E = = O)- There exists an orthonormal basis G1, ,et for the subspace
-i-1 spanned by "1,""" 24 such that for each i S t, u4 = -ド 1 where of. 2 1 1 2 for each j < i, al, S j2c. Proof: We follow the space spanned by ", up by first setting e = 21 and then processing U2 .* Ue one by one.
process to construct the orthonormal basis c Ct of the
The proof is by induction. We show that once the first j vectors "1,-.. , Uj are processed the statement of the lemma holds for these vectors. Note that the base case of the induction trivially holds 11 e1- Next, suppose that the induction hypothesis holds for the first l vectors U1, Ue-
Claim C.4. For each j S f, a2+1.j S..e.
Proof: The proof of the claim is itself by induction. Note that, for j = 1 and using the fact that Ku,, ue+1)l S c, the statement holds and 22t1,1 I e². Next, suppose that the statement holds for all j S i e, then by ((L+1,
((W++,, 1
10441,411 S u641| +
Σ lestrsl |as- |os-
S (e+ Σ
j²2))/ (/444114115
> by induction hypothesis on 04+1,5 for j $ i
j=1
(e+ Σ (²2)/(( Σ
> by induction hypothesis on a++1,8+1
j=1
/=1
(e++ Σ 5(?) (11 (1- Σ
22)/1
(1+2 2. Σ (0)
/=1
j=1
j=1
( + Σ (()) .1+2. Σ ((
j-1
j=1
S <E(() (e) (1 4e Eeee) +1)
5-1
j=1
se(i+1) >> = 00
Finally, since 1|u241 = 11
1, -1, 11,1--11 21 Ei
Corollary C.5. Suppose that = 2t There exists an orthonormal basis e1.-.. et for the
space spanned by the randomly picked vectors "1.... UU, of unit norm, so that for each Us Σj= aijej where oi,, - j2. and for each j i, oZ,1 < jz.
Proof: The proof follows from Lemma C.3 and the fact that the set of vectors 01, ".. Ut are 7-almost orthogonal (by Observation C.
The first main step is to show that the greedy algorithm (with non 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 set rank(SA)) k. the k--ank approximation cost is the simplified ex- pression LVV - AA] where UEVT = SA, rather than [[AV]& - AlI This is just the projection cost onto row(S)), Also, we observe that minimizing this projection cost is the same as maximizing the sum of squared projection coefficients:
min || - A--VVT min S
||A, ((Ar, v1 +... + (As, ox)uz)l2
#E|m]
min
(11.
Σ(A,, (A,, 13)3
*E|m]
jE(K]
20

---

Under review as a conference paper at ICLR 2021
max Σ ΣΣ jE(M)
(Ai, vz)²
In the following sections, will prove that our greedy algorithm makes certain choices by showing that these choices maximize the sum of squared projection coefficients.
Lemma C.6. For any matrix A or batch of matrices A, at the end of iteration k, the learned CountS- ketch matrix S maps each row to isolated bin. In particular, heavy rows mapped to isolated bins.
Proof: For any iteration I k, we consider the choice of assigning A, to an empty bin versus an occupied bin. Without loss of generality, let this occupied bin be b;-1, which already contains A;-
We consider the difference in cost for empty versus occupied. We will this cost comparison for Aj with j S - 2.j 2 i 1, and finally, E - 1,
First, we let ;ei....,ei) c; be an orthonormal basis for (AI...., ..). such that for each r S i, A, Ei= arjez where ar, > O. This exists by Lemma C.3. Let (ei....,es 2.-2,E (i- be an orthonormal
basis for {Ai,.... A142. A;-1 = Ad). Now, ē CyC--1 Cre, for some co, C1 because (A,-1 + A.) proj 2) (A,-1 + A,) E span(es-1,,) ei). We note that of + of = because we let e be
a unit vector. We can find cp, C1 to be:
1,i-- diii-1
co V (ai-1,1-1 + 12 + oi
C1
(ai-1,i-1 + arii-)) + 0
1. J.S i 2: The cost is zero for both cases because Aj E
(22))
2. j 2 i 1: We compare the rewards (sum of squared projection coefficients) and find that 04-2, 7 is no better than e.).
e) (Ao,A), G--1 + cA)
Cauchy-Schwarz incquality S (c (A + (Aj, e
-A)
3. j E {i 1i-1. 1,;): We compute the sum of squared projection coefficients of A41 and A, onto
100-1,1-(00--111-11 + 0),
(4- + 10, 12 a
(di.i-1(a ai,i-1) ai,(a),)))
((a--1.-1 + 01,i-1 (2
(C.2)
(1-1,i-11 1 + 0j
+ a + of 10, +
(C.3)
On the other hand, the sum of squared projection coefficients of A,-1 and A onto 6; Ue; is:
(4-1,1-1+ + (ap_ + ab, + + 11--1-11+4,,-11 + a²
(C.4)
Hence, the difference between the sum of squared projections of A and A, and e Ue is ((C.4) (C.3))
02,((a1-1,111 a1,i-1) )2 a2-1,6-1 + af,,-1 2a (ai-2,i-1 + di,i-1))
((1--1,-1 12 + +z,,)
20?,0? 1-1--11
>0
((a;-1,1-1 + 0i
Thus, we find that (e....., es is a strictly better basis than fei,..., ei.,.,,,,,;,,, e). This means the greedy algorithm will choose place A; in empty bin.
21

---

Under review as a conference paper at ICLR 2021
Next, we show that none of the rows left to be processed (all light rows) will be assigned to the same bin as a 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 decrease (before and after bin assignment of light row) in the sum of squared projection coefficients, lower-bounding it in the former case and upper-hounding it in the latter.
We introduce some results that will be used in Lemma C.10.
r E [1.....,2. n k] be light row not yet processed by the greedy algorithm. Claim C.7. Let Agtrrr
Let (e1....,ee e be the Gran-Schmidt basis for the current {wi..... Wg.. Let B = O( kn
upper bound the inner products of normalized Agtr, -..,,. Wk- Then, for any bin i, (ci, ARt1) g k2.
Proof: This is a straightforward application of Lemma C.3. From that, we have (ARtr, S 888
for i E [1....,k] ki,, which means (Ax+r: e,)2 S B².
Claim C.8. Let Aptr be a light row that has been processed by the greedy algorithm. Let fei,..., .ck be the Gram-Schmidt basis for the current (ui,....: w.. ;.1k...;kkk. J. If Agtr is assigned to bin 0k-1 (w..lo.g.). the squared projection coefficient of Agtr onto C;, + k - 1 is at most 4,3 k2, where 3 = O(n-1k-3) upper bounds the inner products of normalized Ak4r,W1:"'. W....., -.k.
Proof: Without loss of generality, it suffices to bound the squared projection of Astr onto the direction of wy that is orthogonal to the subspace spanned by w1.-. .Wk-1- Let c1.-.. Ce be an orthonormal basis of wi,... We guaranteed by Lemma C.3. Next, we expand the orthonormal basis to include ex+1 that we can write the normalized vector of A&tr Uk-r Ej=1 bjeg- By a similar approach to the proof of Lemma C.3, for each j S k- 2, bj S B2zz. Next, since Kwy, 0k+r)| B,
lbul $ Kekk ex)l
1
([ww, (k+r))+
Σ
lby (wkk ej)l)
j=1
1
18+ + Σ
1bb-11/1 $1 B² .² + (k - B)
V
jml
8 + E/= 8².
+ (k 1)3
V E/11 32.3'
B²(k 1)2
$ 2(kk- - 1)33
D> similar to the proof of Lemma C.3
|1-
1- 32.
<B k
Hence, the squared projection of Astr onto ek is at most 4.3. 11 We assumed ||A Claim C.9. We assume that the absolute values of the inner products of vectors in "1,'" Un are at most E < 1/(n2 E l| and the absolute values of the inner products of the normalized vectors of Wi,".. WWkk are at most B = 0 -). Suppose that bin b contains the row Aktr. Then, the squared projection of Autr onto the direction of w orthogonal to maann w.k
{w}) is at most
|| |
- OOn and is at least || O(n
||0||2
Proof: Without loss of generality, we assume that Aktr is mapped to by: w = Wk. First, we provide
an upper and a lower bound for (extrr wx)l where for each k we let w. wwlll denote the
normalized vector of w,. Recall that by definition uh+r
Autr.
TAutella*
A E Ed
El E||Ail|2
(Kk, Uetrr> <
|wlll2
A + n2 ||wll|2
-2 > by 7 Exee ||Arl12
|| + n
> >||wkll2 21
(C.5)
|wxel22
22

---

Under review as a conference paper at ICLR 2021
||
EA,eb, 1|Ail12 E
Kun; uktrll
1101e122
|| n -2
(C.6)
Iuxl12
Now, let {er.... .... ek) } be an orthonormal basis for the subspace spanned by (wI, ... , we} guaran- teed by Lemma C.3. Next, we expand the orthonormal basis to include ek++1 so that we can write Uh+r 21=1 byez. By similar approach to the proof of Lemma C.3, we can show that for each j S k - 1, 63 < 8² Moreover,
lbkl v
1
(T oktr 11
Σ j=1
1b3 (Txx eg)l)
(wks exحl
1-1
1
(1 lk+rr)
+
Σ
82
(9)
> by Lemma C.3
k- Σk=1
j= 32. .². j²
/=1
1
(n-2 -2 + ||At+ell2
+
1-1
Σ
B2 )²))
> by (C.5)
V Σ 32.
| lg
1=1
| l12
<B.k+ + /I--2k33 (n- +
similar to the proof of Lemma C.3
||
|| 112
S O(n + (1+ (m)
> by 3 = =Onn=
|
| + +O 2)
| < 1
12
|well2
and.
lbu|
2 ( exil
(((Tg Uk+r) Σ
bb;- (wxx c))l
k-1
2 (Th. Uk+rr
Σ 82 -²
> since (T ex) S 1
f=1
||
k-1 Σ
n2 -2
82 j2
D by (C.6)
||2kl22
1=1
| 1l|
(2
D by B = O(n 0n-53(n-) - 3)
|w 1|2
1 | 2 + +0(n and is at least ||
Hence, the squared projection of Aktr onto eg is at most Oln-2).
1w 1
Now, we show that at the end of the algorithm no light row will be assigned to the bins that contain heavy rows.
Lemma C.10. We assume that the absolute values of the inner products of vectors in U1," Un are at most ē < minn(n=2k- (n En will assign the light row Aktr A, Ew [ At each iteration k +r, the greedy algorithm bin that does not contain heavy row.
Proof: The proof is by induction. Lemma C.6 implies that no light row has been mapped to bin that contains a heavy row for the first & iterations. Next, we assume that this holds for the first k +r iterations and show that is also must hold for the (Is + r) iteration.
To this end, we compare the sum of squared projection coefficients when A avoids and collides with heavy row.
23

---

Under review as a conference paper at ICLR 2021
First, we upper bound 8 = max (wi, wy)l/([[ssll)) (12z the number of rows assigned to b; and by.
Let c2 and ej respectively denote
Kw,,, wyy1
G Cj C. -
max Cデ
b Observation C.2
|10l12 |
Va: - 2007 -2cc
162 /GE
D S n 25-
S n *
22 S n -5/3
5/3
1. If Ag+r is assigned to bin that contains c light rows and no heavy rows. In this case, the projection loss of the heavy rows AI, As onto row(S A) 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 wh denote the bin that contains ARtr- Since Sk-1 aoean(wi,... wx-1)) has not changed, we only need to bound the difference in cost between projecting onto the component of wy - Ak+r orthogonal to S and the component of wr orthog-
onal to Sx-1. respectively denoted as e and εg.
I. By Claim C.7, for the light rows that are not yet processed (i.e., Aj for j > k ++ r), the squared projection of each onto eg. is at most B² Hence, the total decrease in the squared projection is at most (n - r) k²
II. By Claim C.8, for the processed light rows that are not mapped to the last bin, the squared projection of each onto ex is at most 4,2 Hence, the total decrease in the squared projection cost is at most (r - 1) 2.32
III. For each row A; * Aktr that is mapped to the last bin, by Claim C.9 and the fact ||A:11 | = 1, the squared projection of A, onto e is most
1vk-Axirl + O(n=2) and the |4:11
squared projection of A, onto ε is at least |A
O( ?).
Moreover, the squared projection of Agtr onto e compared εµ increases by at least ||Au+ On - O(n | 1
O(n
Hence, the total squared projection the rows the bin by decreases by at least:
||A.112
|
Σ ||urk-
+ Onn -2n-22) 2)) - Σ
O(n-2))
Ar+1112
|w 112
A.EW
[u - -Ar+# + +(
1a 1 O(n + O(n
by Observation C.2
|lww - Ar+2112 2
|
<0(n
Hence, summing up the bounds in items I to III above, the total decrease in the sum of squared projection coefficients is at most On -1).
2. If A&tr is assigned to bin that contains heavy row. Without loss of generality, we can assume that A is mapped to by that contains the heavy row A.. In this case, the distance of heavy rows A Ag-1 onto the space spanned by the rows of SA is zero. Next, 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 (k- -)ddimensional space corresponding Wi,".. Wx-1 has not changed. Hence, we only need to bound the decrease in the projection distance of Aptr onto E. compared to ex (where ε ex are defined similarly as in the last part).
1. For the light rows other than Aptr, the squared projection of each onto ex is at most k²k. Hence, the total increase in the squared projection of light rows onto eg is at most (n k) 122 Ol1-1).
24

---

Under review as a conference paper at ICLR 2021
2. By Claim C.9, the sum of squared projections of Ag and Ag+r onto ej decreases by at least
[I 11 + || O(
|| 12
||A, + A 1/2
||
|
1/
|| 11
2 +|AAgl
+0(n
D by Observation C.2
||
2
|
k
12
2
n01
1| 14 & 112 (11,4,1
2
||
()
||A.112
O(n - O(n
||A,115
1
||
k
12 0(
1 & 11 [1I
1
|| A
12)
||A O(n-t) O(n
|| 11 +
+1 || 12
A| & 112 (1I 2
A 112)
-((n²t)
||A,012 + ||Ar+
| & 112 (1 (II 1/2 (|A.112) O(n
+ (1Ar++xl2 |1Ax112)
|| ( (n
1 1 6² 2
2 || A A
122
D
-
||A.112
Hence, in this case, the total decrease in the squared projection is at least
|| 122 O(n 1) = 1
|| || ) O(
|| (1
>> || =1
||all22
||
= -1-(1/V8 - -0(n
>> ||A.ll2 = vl
Thus, for sufficiently large value of L, the greedy algorithm will assign A&tr 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 C.11. The approximation loss of the best rank-k approximate solution in the rowspace S for A ~ (s, where A Raad for d 22n k* log n) and Sg is the CountSketch constructed by the greedy algorithm with non-increasing order is at most n 8.
Proof: First, we need to show that absolute values of the inner products of vectors in 01, . is most min(n 2k-2
2, (n ΣA. [|A:ll2) -113 so that we can apply Lemma C.10. To show this,
note that by Observation C.I, S 2 log < n n since d = Q(n*k K4 log n). The proof follows from Lemma C.6 and Lemma 10 Since all heavy rows are mapped to isolated bins, the projection loss of the light rows is at most n
Next, we bound the Frobenius norm error of the best mannk--approximaaiio solution constructed by the standard CountSketch with randomly chosen sparsity pattern.
Lemma C.12. Let = ok where 0.7 < < 1. The expected squared loss of the best rank-k approximate solution in the rowspace S,A E Raxd Asp(s, (s,{) E) where d 22 (2) and S, the sparsity pattern of CountSketch is chosen uniformly at random is al least n n zk de ax+
n -01))
We can interpret the randomized construction of the CountSketch as a "balls and bins" ex- periment. In particular, considering the heavy rows, compute the expected number of bins (i.e., rows in S, A) that contain heavy row. Note that the expected number of rows in S,A that do not contain any heavy row is ls (1 - 1) 2 k. _ Hence, the number of rows in S, A that contain heavy row of A is at most ke(1 - Thus, at least k(1 e - heavy rows 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 each such row is at least n -01).
Proof:
Claim C.13. Suppose that heavy rows Ar,.... Ar. are mapped to 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)e O(n
25

---

Under review as a conference paper at ICLR 2021
Proof: Let b denote the bin that contains the rows . Ar, and suppose that it has cf light rows as well. Note that by Claim C.8 and Claim C.9, the squared projection of each row Ar, onto the subspace spanned by the k bins is at most
|A w
+ O(nnn
6²
+ O(n
cl + d
20(c2 +ole+cvve+el
df n011) +n ((1)
by E S n 38
82
$ 2
(el + +n O(
1 - 5-+0(n
Hence, the total squared loss of these heavy rows is least of c. On O(n-1).
(=n0(11)) 1) 2 (c 1)E
Hence, the expected total squared loss of heavy rows is least:
e. (s - k(1 e )) 8 0(1)
28 k(a - - la 0(1)
D =s=a- (k- ((2-1) a.ka--k--
where 0.7 < a <
tk
2 2e
{ n -01)
D a 0.7
Ek O(
D assuming k > 4e
de
Next, we compute a lower bound on the expected squared loss of the light rows. Note that Claim C.8 and Claim C.9 imply that when light row collides with other rows, its contribution to the total squared loss (which the loss accounts for the amount it decreases from the squared projection the other rows in the bin well) is at least 1 - O(n Onn-)). Hence, the expected total squared loss of the light rows is at least:
sn-
2 (n (1 + a)) k) Onn
Hence, the expected squared loss of a CountSketch whose sparsity is picked at random is at least
tk O( 11 + n (1 a - O(n -1--)) nn + 4e (1 a1k -((=))
4e
Corollary C.14. Let s alk 1) where 0.7 < < 1 and let E (4e+1)n Let Sg be the ak 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 matrix ~ Asp where = & the 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 A returned by Sy by at least a constant factor.
Proof:
Ek:
Es.
min rank-k X (roows((S..)
||-
A1
2 n
a11 a: n -01))
Lemma C.12
4e
(4e 1)n ak
2 (1+ + (na)(n
= (1 + 1/0))
min rank-k ErowsppsS 4)
| - Al
Corollary C.11
26

---

Under review as a conference paper at ICLR 2021
C.2 ZIPFIAN ON SQUARED ROW NORMS.
Each matrix A E B ~ Azipt has rows which are uniformly random and orthogonal. Each A has 2++1 rows of squared norm n² /220 for i E (.....,(loo((n))) We also assume that each row has the same squared norm for all members of Asipf-
In this section, the s 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 a given index has the same norm for all members of Azipf. Also, assume that 2/2 and, for simplicity, minimum squared norm of a heavy row is n² /22h,, and the maximum squared norm of a light row is n²2 /22h,+2
Σin 2 for some h, Z+. That means the
The analysis of the greedy algorithm ordered by non-increasing row norms this family of matrices is similar to our analysis for the spiked covariance model. Here we 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 be the same.
As first step, by Lemma C.6, at the end of iteration k the first k rows are assigned 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 8 heavy rows.
Lemma C.15. At each iteration k the greedy algorithm picks the position the non-zero value in the (k + r)-th column of the CountSketch matrix S so that the light row Aktr is mapped to bin that does not contain any of top s heavy rows.
Proof: We prove the statement by induction. The base case r = 0 trivially holds as the first k rows are assigned to distinct bins. Next we assume that in none of the first k + iterations light row is assigned to bin that contains heavy row. Now, we consider the following cases:
1. If ARtr is assigned to a bin that only contains light rows. Without loss of generality we can assume that Astr is assigned by. Since the vectors are orthogonal, we only need to bound the difference in the projection of Ax+z and the light rows that are assigned to by onto the direction of wr before and after adding Astr to bg. In this case, the total squared loss corresponding to rows in by and A before and after adding A&+1 are respectively
Enyeb, ||A]112
1| )
before adding Au- to by: ||Ax+rIl
1|4,12
Σ A, Ebb
ArEb
|| 11
+ =AA,E, ||4
after adding A to by: ||
1|A,
| 12 Enjeb, ||A 2
AgEhe
Thus, the amount of increase in the squared loss is
En,eb ||Abll - || 1 + EA Abc ||4bll
|| Ea,ea 14 1|AA4r Ex,eo, ||A 1
Ed,ev ||A,
| 112 + EA, |A,112
(EA,Er || (1)
E,, Ebe || l12)
Σ Enyer, AAr15
Σn, Enyer cv A,l? ||Ax+r 12
= 1| 1122
Enyed || + || 12
EA,Ede || - || 12
< || 11.
Σ Eayeb || || l
(C.7)
2. If As+r is assigned to a bin that contains a heavy row. Without loss of generality and by the induction hypothesis, we assume that Aptr is assigned to bin 6 that only contains heavy row Ay- Since the rows are orthogonal, we only need to bound the difference in the projection of A and A, In this case, the total squared loss corresponding to A₁ and Agtr before and after adding Ar+1
27

---

Under review as a conference paper at ICLR 2021
to are respectively
before adding Ak+r to by: 1|
|| 12 + || 11
after adding Aktr to by! || + 111,,112
|| 12 + ||
Thus, the amount of increase in the squared loss is
|| || 1/ || 11 = || ||
||
|| Aktr|
(C.8)
||At+r
12 ||
|A 14 + ||
Then (C.8) is larger than (C.7) if || 2 E,,eb, | Next, we show that at every inductive iteration, there cxists 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 Am so that m > 2 2hm +2 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
22h,
g+
n
| = Σ
1A.2
S
2
22 +1)
A
n2
2n
22(h,+1) +1]
2h (k- 5)
n
n2
D s S k/2 and k > 2h+
22(h, +1)
22h
n
D ha h, +1
22h,+I
||
Hence, the greedy algorithm will map Aktr to a bin that only contains light rows.
Corollary C.16. The squared loss of the best rank-k approximate solution in the rowspace of Sy for E Baad Azsps where A E Raxd and Sg is the CountSketch constructed by the greedy n²
algorithm with non-increasing order, is < 2"
Proof: At the end of iteration k, the total squared loss is hh +imhh
21+1
2 After that, in each
iteration k + r, by (C.7), the squared loss increases by at most A Hence, the total squared loss the solution returned by S. is at most
2 2
An
1
2(
= 2n².
22r
2
hx +1
+11
4n
π1
2
2hg
Next, we bound the squared loss of the best rank-k- appro---ppr solution constructed by the standard CountSketch with a randomly chosen sparsity pattern.
Observation C.17. Let us assume that the orthogonal rows Ar....,, ,r. are mapped to same bin for each i S c, ||A > |A,, Then, the total squared loss of Ar, ... A,, after projecting - onto Ar, - ... A,. is at least ||A,,2 .... + ||A,,,1
Proof: Note that since Ar,,... Ar.. are orthogonal, for each i S c, the squared projection of Ar, 11// E5=1 ||,,, 1| Hence, the sum of squared projection coefficients
onto Ar, + ... + Ar. is ||
of Apys'.. Ar, onto A, + Ar
Σ c
Ar
4
< || Ar, 11
Arg
2
28

---

Under review as a conference paper at ICLR 2021
Hence, the total projection loss of Ar,.". Ar. onto A,, = .... + Ar. is at least
E Ar, 11 ||ral12 + A 1=1
In particular, Observation C.17 implies that whenever two rows are mapped into a same bin, the squared norm of the row with smaller norm normfuuly contributes to the total squared loss the solution.
Lemma C.18. For k 220 -2 the expected squared loss of the best rank-k approximate solution in the rowspace of Aor Anxd ~ Agipp where S, is the sparsity pattern of a CountSketch chosen uniformly at random, is at least 1.095n". 2-2-2
Proof: light of Observation C.17, we need to compute the expected number of collision between rows with "large" norm. We can interpret the randomized construction of the CountSketch "balls and bins" experiment.
For each 0 S j S hx, let A₁ denote the set of rows with squared norm 2-14--17 and let AAj
n²a
Ujeighe A. Note that for each j. |A5l = 2 and |A>i| = Emej+1 24x-i+1 hh--12 2(2ha-) 1). Moreover, note that & 2(2 1). Next, for row A, in Aj (0 S j < hg), compute the probability that at least one row in A> collide with A,-
Pr(at least one row in A>y collide with A,]
(1 - (1 (uassl)
2 (1 - LAzs!
~Ax -
(1 2b6t1-1)
2hr
2 e 2-2113)
since
2-
2h++1
Hence, by Observation C.17, the contribution of rows in Aj to the total squared loss at least
n²
n
- 11--22--7) (Ajl-
22(hx-1) =(1 e 2-2--
2ha-j
n²
=(1 ---
23-1
2ha-2
Thus, the contribution of rows with "large" squared norm, i.e. Aso. the total squared loss is least
h
n
20-1
22- 9)2 1.095
D-for hy
2ha -2
2ha-2
j=U
Corollary C.19. Let Sg be 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 n x d matrix A ~ Azipf, for sufficiently large value of k, the expected loss of the best rank-k k approximation of A returned by S, is worse than the approximation loss of the best rank-k approximation of A returned by Sq by at least constant factor.
Proof: The proof follows from Lemma C.18 and Corollary C.16.
Remark C.20. 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 chosen random. However, still other implementations of the greedy algorithm may result in better solution com- pared to the greedy with non-increasing order. To give an example, the following simple instance the greedy algorithm that check the rows of A in a random order (i.e., greedy with random order) achieves a rank-k & solution whose cost is a constant factor better than the solution returned by the greedy with non-increasing order.
'The numerical calculation computed by WolframAlpha.
29

---

Under review as a conference paper at ICLR 2021
+ Let A be matrix with four orthogonal rows u, where ||u|l2 1 and |||||| ||w||2 and suppose that the goal is to compute a rank-2 approximation A. Note that in the greedy with non-decreasing order, v and will be assigned to different bins and by simple calculation we show that the copies of и also will be assigned to different bins. Hence, the squared loss in the computed rank-2 solution is 2+(1+c)². However, the optimal solution will assign v and to one bin and the two copies of to the other bin which results in the squared loss of (1 + c)² which constant factor smaller than the solution returned by the greedy with non-increasing order for sufficiently small values of
(I+c)²
On the other hand, the greedy algorithm with random order, with a constant probability (3 $). the computed solution is the same as the optimal solution. Otherwise, the greedy algorithm a with random order returns the same solution as the greedy algorithm with 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 constant factor.
D EXPERIMENTS APPENDIX
D.1 BASELINES
We comment on two of our baselines:
Exact SVD: In the canonical, learning-free sketching setting (i.e., any matrix is equally probable), sketching using the top m singular vectors yields (1 + c)-approxination for both LRA and k-means Cohen al. (2015).
Column sampling: In the canonical, learning-free freee sketching setting, sketching via column sampling yields a (1 + <)-approximation for k-means Cohen et al. (2015).
D.2 EXPERIMENTAL PARAMETERS
For the tables in Section 5, we describe experimental parameters. First, we provide some general implementation details.
We implemented both the greedy (Algorithm 1) and stochastic gradient descent (Algorithm 2) algo- rithms in PyTorch. In the first case, Py PyTorch allowed us to harness GPUs to speed up computation on large matrices. We used several Nvidia GeForce GTX 1080 Ti machines. In the second case, PyTorch allowed us to effortlessly compute numerical gradients for each task's objective function. Specifically, PyTorch provides automatic differentiation. which is implemented by backpropagatioon through chained differentiable operators.
There also two points of note in the greedy algorithm implementation. First, noticed that for MRR and LRA, each iteration required computing the SVD for many rank-1 updates of the current S. Instead of computing the SVD from scratch for each of these variants, first computed the SVD of S and then used fast rank-1 SVD updates Brand (2006). This greatly improved the runtime of Algorithm 1. Second, we decided to set D., (the set of candidate row weights) 10 samples in [-2, 2] because we noticed most weights were in this range after running Algorithm 2.
D.3 RUNNING TIME
We examine the runtimes of our various sketching algorithms. In Table D.I, the times obtained for the LRA task with & = 30, m = 60 on the logo data set. However, similar trends should be expected for other combinations of task, task parameters, and data sets.
We define the inference runtime as the time to apply the sketching algorithm. The training runtime is the time to train a sketch on Atrain and only applies to learned sketches. Generally, the long training times are not problematic because training is only done once and can be completed offline. On the other hand, the inference runtime should be fast possible.
Note that inference was timed using matrix from Arest on an Nvidia Geforce GTX 1080 Ti GPU. The values averaged over 10 trials.
30

---

Under review as a conference paper at ICLR 2021
We observe that sparse sketches (such the ones used in GD only and Ours) have much lower inference runtimes than the dense sketches exact SVD.
Table D.1: Timing comparison of sketching algorithms for LRA (using Algorithm from Indyk al. (2019))
Time (sec)
random: inference
0.0114
exact SVD: inference
0.185
Icarned (random pattern): training
193 (3 min)
Icarned (random pattern): inference
0.0114
learned (greedy pattern): training
6300 (1 h 45 min)
learned (greedy pattern): inference
0.01
31


Reviews:
[{"review_id": "ap7nmPQA4Q0", "reviewer": "AnonReviewer1", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper introduces a general framework for learning and applying input-independent sparse sketches.  Then this paper instantiates this framework with three sketching applications: least-squares regression, low-rank approximation (LRA), and k-means clustering. Numerical experiments are given to demonstrate that the approach substantially decreases approximation error compared to classical and naively learned sketches. \n\nThe idea of learning sketches and the proposed greedy algorithms are simple. They are closely related to some of the existing related works on dictionary learning. The numerical results look promising. I would suggest to give more numerical results on different datasets to exemplfy the proposed approach, as that the main contribution of the paper is on algorithmic aspect. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "2IQtspaJDn", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: \n\nThe authors consider the problem of ``\"sketching\" – a popular compression technique in machine learning - used for reducing the size of the data enabling one to quickly compute an approximate solution using this compressed input. This paper introduces a general framework for learning and applying sparse sketching matrices. A two-stage procedure for learning sketches with the same sparsity pattern as CountSketch is proposed which involves first learning the sketch’s non-zero entries, and then optimizing their values. They then show how to apply the obtained sketch so that it has worst case approximation error guarantees. This procedure is applied to three applications, namely least squares regression, low rank approximation (LRA) and k-means clustering. Experimental results demonstrate a substantial reduction in the approximation error compared to other baseline approaches (including classically learned sketches). On the theoretical front, it is shown for regression and LRA that the proposed approach obtains improved error guarantees for fixed time complexity.  Additionally, it is shown for LRA (under certain input distributions) that including the first stage is strictly better than not including it. Finally, a more straightforward way of retaining worst case approximation guarantees for k-means is shown. \n\n----------------------------------------------------------\nReasons for score: \n\nOverall, I vote for accepting the paper. The procedure for learning and applying the sketching matrix seems novel and interesting, and the empirical results (Especially for LRA and least squares regression) on real datasets suggest that the proposed method does substantially improve over other baseline methods, including the classical CountSketch. The theoretical results for least squares regression and LRA also imply that their algorithm achieves a better accuracy than the classical CountSketch approach (given the same time complexity). I do have some questions mentioned in “Cons” and “Further remarks” below which I would like the authors to clarify in the rebuttal phase. \n\n------------------------------------------------------------------------------------\nPros:\n\n-\tPaper is in general written well barring some notational issues and explanations which were not clear to me (explained later in “Cons”)\n-\tThe two-stage idea of learning a sketching matrix with the same sparsity pattern as in CountSketch is interesting. Instantiating the framework to problems such as: least squares regression, LRA and k-means is done nicely, and the theoretical results appear to be strong.\n-\tThe empirical evaluation on real datasets makes a convincing case for the proposed algorithm and demonstrates that learning the CountSketch matrix can indeed be substantially better than using the classical CountSketch, on real data.\n\n----------------------------------------------------------------------------------\nCons:\n\n-\tThe sketching matrices considered in the paper are restricted to Count Sketches which have a particular sparsity pattern. So, the framework is not as general as is claimed in the abstract. \n-\tSome of the notation and explanations in Section 3 are quite unclear. For example, the loss function symbol is sometimes used as L(X), sometimes as L(S, I) etc. In Algorithm 1, I don’t understand what range(n) means, and also the summation over A_i \\in A_{train} seems strange. Algorithms 1,2 have a training set \\mathcal{A}_{train} as input but I don’t see how this appears for the examples in Section 4; is N =1 there? In stage 2, what does the operation S[p, range(n)] = v do? In Algorithm 3, it would be helpful to explain (with an example perhaps) what the symbol \\mathcal{M}(S,I) could look like. At the beginning of page 4, it is mentioned that learned sketches can encounter out of distribution inputs etc. But some of the guarantees in Section 4 (Theorems 4.1,4.2) apply for arbitrary inputs, isn’t it?\n-\tTheorems 4.1, 4.2 assume that the learned sketch matrices are affine \\beta embeddings and hence attain a (1+\\beta) approximation for the respective problems. But it is not mentioned whether this assumption can be satisfied.\n\n---------------------------------------------------------------------\nFurther remarks:\n\n-\tI don’t understand why in page 2 (“Our results” para) it is written that the sketches are input-independent, aren’t the sketches derived in the paper a function of the input data (for e.g. in Alg. 1)?  \n-\tFor theorems 4.1 and 4.2, it will be helpful to state the corresponding guarantee for the classical algorithm in the main text to enable a direct comparison.\n-\tIn the abstract and the introduction, it is mentioned that for LRA it strictly better to include the first optimization stage. But this is shown for a couple of specific input distributions, hence should be clearly specified as such.\n-\tIn Algorithms 1,2, sometimes \"\\mathcal{A}_{train}\" is used and sometimes \"A_{train}\". If sets are being denoted with calligraphic symbols, then mention this in the notations paragraph.\n\n----------------------------------------------------------\nPost rebuttal: I have read the authors response and will go with my original score for this paper.", "extra_scores": {}, "extra_text": {}}, {"review_id": "0PtJ41i63w", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors consider a specific sketching method, CountSketch, and three objective functions defined over the data design matrix: multiple-response regression (MMR), low-rank approximation (LRA), and k-means clustering. They compare the classical CountSketch with a random choices of the {-1,+1}-valued sketching matrix against: (1) Gradient descent optimization of the CountSketch weights. This was previously introduced for LRA and the authors extend it to MMR and k-means. (2) Greedy optimization of the positions of the CountSketch non-zero entries.\n\nIn order to avoid losing the worst-case guarantees of the original (non-optimized) sketch, a wrapper algorithm guarantees that the optimized sketch is considered only if it is not worse than the optimized sketch.\n\nThere are theoretical and empirical results. The first two theoretical results show that for MMR and LRA the wrapper algorithm can be implemented without having to solve the problem twice on the full dataset. Theorem 4.3 shows that greedy optimization produces a strictly better result for the LRA objective than CountSketch alone when data are sampled from a certain distribution. Theorem 4.4 finally shows that a different wrapping technique (used in a previous result for LRA), can be applied to the k-means objective.\n\nExperiments are performed on three real-world datasets comparing CountSketch, gradient optimization, and greedy+gradient optimization on the three objectives. The bottom line is that the greedy optimization step helps significantly.\n\nThe paper is not an easy read and the clarity can be significantly improved. There is not a lot of conceptual novelty compared to previous works: just the introduction of the greedy optimization step. There is technical work to extend gradient optimization and prove the theoretical results. However, the extent of the overall technical novelty remains unclear.\n\nIt is not explicitely written that the greedy optimization step is performed only on the training part of the data.\n\nIs greedy+gradient optimization fully deterministic?\n\nStandard deviations are not reported in the expertiments.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The work falls under the setting of learning-based sketching/compressive subsampling. It extends the work of Indyk et al 2019 (including sparsity pattern optimization and some theoretical enhancements).  The reviewers agree that while the conceptual novelty including the greedy optimization step is not too much, it is nonetheless interesting and is non-trivial. However, given the highly competitive submissions at ICLR, the current scores are not sufficient for acceptance. 
Decision: Reject