Year: 2020

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 naı̈vely 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.

heading: 1 INTRODUCTION
text: 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 to 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, a data compression 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 a 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 arbitrary inputs. With the ML perspective in mind, we examine if it is possible to construct sketches which also have low error in expectation over an input 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-specific. The sketch construction and approximation algorithm are tailored to individual applications, like robust regression or clustering (Sarlos, 2006; Clarkson & Woodruff, 2009; 2014; 2017; Cohen et al., 2015; Makarychev et al., 2019). Instead, we consider
three applications at once (regression, LRA, k-means) and propose generalizable methods, as well as extending previous application-specific work.
Our results. At a high level, our work’s aim is to make sketch learning more effective, general, and ultimately, practical. We propose a framework for constructing and using learned CountSketch. We chose CountSketch because it is a sparse, input-independent sketch (Charikar et al., 2002). Specifically, 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 can reuse the same CountSketch for different inputs. Our “learned” CountSketch will retain this characteristic sparsity pattern and input-independence1, 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 3 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 a 12, 20% boost for MRR, LRA.
Our sketch learning algorithm first places the sparse non-zero entries using a greedy strategy, and then learns their values using gradient descent. The resulting learned CountSketch is very different from the classical CountSketch: the non-zero entries no longer have random positions and ±1 values. As a result, the usual worst-case guarantees do not hold.
We sought a way to 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, k-means, 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 a 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 a diverse set of experiments to demonstrate the versatility and practicality of our approach. 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 counterparts. 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 a 31, 70% improvement over classical CountSketch for MRR, LRA.
Finally, we conducted ablation study of the components of our algorithm. In another novel theoretical result, we proved that including the time-consuming first optimization stage is 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 classical algorithms; we only mention a few examples here. One related body of work is data-dependent
1While learned CountSketch is data-dependent (it is optimized using sample input matrices), it is still considered input-independent because it is applied to unseen input matrices (test samples).
dimensionality reduction, such as an approach for pair-wise/multi-wise similarity preservation for indexing big data (Wang et al., 2017) and a method for learning linear projections for general applications (Hegde et al., 2015). We note that multiplying an input matrix on the left with a sparse 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 neighbor search problem (see Wang et al. for a 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 -embedding (Def. A.1).
In particular, we build off of the work of Indyk et al. (2019), which introduced gradient descent optimization 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 a more general method for worst-case guarantees and extend their method to k-means.

heading: 2 PRELIMINARIES
text: Our learned sketches have the sparsity pattern of the classical CountSketch. The construction of this sketch is described below. We also define affine -embeddings, which is a class of sketches that includes CountSketch. This -embedding property is desirable because it allows us to prove that certain sketching algorithms give (1 + )-approximations. Definition 2.1 (Classical CountSketch). The CountSketch (abbreviated as CS) matrix S has one non-zero entry in each column with a random location and random value in {±1}. Definition 2.2 (Affine Embedding). Given a pair of matrices A and B, a matrix S is an affine -embedding if for all X of the appropriate shape, ‖S(AX −B)‖2F = (1± ) ‖AX −B‖ 2 F .
Notation. We denote the singular value decomposition (SVD) of A by A = UΣV > with orthogonal U, V and diagonal Σ. Relatedly, the Moore-Penrose pseudo-inverse of A is A† = V Σ−1U>, where Σ−1 is constructed by inverting the non-zero diagonal entries.

heading: 3 FRAMEWORK
text: We describe a framework for learned CountSketch that can be adopted by many different applications, including least-squares regression (MRR), low-rank approximation (LRA), and k-means clustering. We will return to 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 a “fast comparison” step before the approximation algorithm, we can do this in optimal time.

heading: SKETCH OPTIMIZATION
text: Most applications are optimization problems. That is, they are defined by an objective function, L(·). For example, least-squares regression solves
min X ‖AX −B‖2F
given A ∈ Rn×d, B ∈ Rn×d′ . Of course, the optimal solution is a function of the inputs:
X∗ = arg min X ‖AX −B‖2F = A †B
However, in the sketching paradigm, we compute an approximately optimal solution as a function of the sketched (compressed) inputs. Taking CountSketch S ∈ Rm×n for m n, we have:
X̂∗ = (SA)†(SB)
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 to minimizing the
objective value of the approximate solution.
S∗ = arg min S∈CS E (A,B)∼D
[ LA,B((SA)†(SB))− LA,B(X∗) ] = arg min
S∈CS E (A,B)∼D
[ LA,B((SA)†(SB)) ] For ease of notation, we will define G(·) as a function which maps a sketch and inputs to an approximate solution. G(·) is defined by the application-specific approximation algorithm. For MRR, G(S, (A,B)) = (SA)†(SB). More generally, the sketch optimization objective is:
S∗ = arg min S∈CS E A∼D [LA(G(S,A))] (3.1)
If the application is regression, we let A be (A,B).
We will solve this constrained optimization in 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 CountSketch-sparse 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 to 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 placements. This is clearly intractable, so we consider a greedy alternative. In 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 ∈ Rm×n iteratively, placing one nonzero entry at a 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 of S, each of which contains a non-zero entry. Note that Sw,j = S + w(~ej ~ei
>) adds a single entry w to the i-th column, j-th row of the current, partially-constructed S.
Algorithm 1 GREEDY STAGE Require: Atrain = {A1, ..., AN} with Ai ∈ Rn×d; sketch dimension m
1: initialize S = Om×n 2: for i = 1 to n do 3: w∗, j∗ = arg min
w∈{±1},j∈[m] ∑ Ai∈Atrain LAi(G(Sw,j , Ai)) where Sw,j = S + w(~ej ~ei >)
4: S[j∗, i] = w∗ 5: end for
For some applications it can be inefficient to evaluate (3.1), since it requires computing the approximate solution. For MRR and LRA, the approximate solution has a closed form, but for k-means, it must be computed iteratively. This is prohibitively expensive, since we perform many evaluations. In this case, we recommend finding a surrogate L(·) with a 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 a vector of its nonzero entries, ~v ∈ Rn. We will denote H(~v) : Rn → Rm×n 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 E A∼D [LA(G(H(~v), A))] (3.1) with respect to ~v.
Algorithm 2 GRADIENT STAGE Require: Atrain = {A1, ..., AN} with Ai ∈ Rn×d; H(·) from Alg. 1; learning rate α
1: for i = 1 to niter do 2: S = Om×n 3: sample Abatch from Atrain
4: ~v ← ~v − α ( ∑ A∈Abatch ∂LA(G(H(~v)),A) ∂~v ) 5: end for

heading: LEARNED SKETCH WITH WORST-CASE GUARANTEES
text: We first run a fast comparison between our learned sketch and a classical one and then run the approximation algorithm with the “winner”. This allows us to compute an approximate solution (to the applications we consider here; i.e., MRR and LRA) that does not perform worse than classical CountSketch. In other words, our solution has the same worst-case guarantees as classical CountSketch. The benefit of this is that these guarantees hold for arbitrary inputs, so our method is protected from out-of-distribution inputs as well as in-distribution inputs which were not well-represented in the training data.
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 a rough estimate of the approximation error rather than the exact value is crucial here. It allows us 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 a classical sketch appended to the learned one) retains worst-case guarantees for LRA. We prove that this also works for k-means clustering (Theorem 4.6).
Algorithm 3 LEARNED-SKETCH-ALGORITHM Require: learned sketch SL; classical sketch SC ; trade-off parameter β; input data A
1: Define ML,MC such that LA(G(SL, A)) = ‖ML‖2F , LA(G(SC , A)) = ‖MC‖ 2 F 2: S ← CountSketch ∈ R 1 β2 ×n, R← CountSketch ∈ R 1 β2 ×d
3: ∆L ← ∥∥SMLR>∥∥2F , ∆C ← ∥∥SMCR>∥∥2F 4: if ∆L ≤ ∆C then 5: return G(SL, A) 6: end if 7: return G(SC , A)
This algorithm can be used for applications that minimize a Frobenius norm. In the MRR example,
LA(G(SL, A)) = ∥∥A [(SLA)†(SB)]−B∥∥2F
so ML = A [ (SLA) †(SB) ] − B. Note that β is a parameter which trades off the precision of the approximation error estimate and the runtime of Algorithm 3.

heading: 4 INSTANTIATIONS
text: For each of 3 problems (least-squares regression, low-rank approximation, k-means clustering), we define the problem’s objective, LA(·), and the approximation algorithm, G(·), which maps a sketch and inputs (S,A) to an approximate minimizer of LA(·).

heading: 4.1 LEAST-SQUARES REGRESSION
text: We consider a generalized version of `2 regression called “multiple-response regression” (MRR). Given a matrix of observations (A ∈ Rn×d, with n d) and a matrix of corresponding values (B ∈ Rn×d′ , with n d′), the goal of MRR is to solve
min X L(A,B)(X) = min X ‖AX −B‖2F
Approximation algorithm. The algorithm is simply to sketch A,B with a sparse sketch S (e.g., CS) and compute the closed-form solution on SA, SB, which are small (Algorithm 4).
Algorithm 4 SKETCH-REGRESSION (Sarlos, 2006; Clarkson & Woodruff, 2017)

heading: Require: A ∈ Rn×d, B ∈ Rn×d′ , S ∈ Rm×n
text: 1: return: (SA)†(SB)
Sketch optimization. For both the greedy and gradient stages, we use the objective L(A,B)(G(S, (A,B))) = ‖A(SA)+(SB)−B‖ 2 F .
Learned sketch algorithm. For fixed accuracy, our learned sketch algorithm achieves stateof-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 worst-case guarantee). In the following theorem, we assume Alg. 3 uses a learned sketch which is an affine β-embedding and also a classical sketch of the same size which is an affine -embedding. If β < , the above statement is true. Theorem 4.1 (Learned sketching with guarantees for MRR). Given a learned sparse sketching matrix SL ∈ R d2 2 ×n which attains a (1 + γ)-approximation for MRR on A ∈ Rn×d, B ∈ Rn×d′ , Alg. 3 gives a (1 + O(β + min(γ, )))-approximation for MRR on A in O(nnz(A) + nnz(B) + d5d′ −4 + dβ−4) time where β is a trade-off parameter. Remark 4.2. By setting the trade-off parameter β−4 = O( −4d4d′), Alg. 3 has the same asymptotic runtime as the best (1 + )-approximation algorithm of MRR with classical sparse embedding matrices. Moreover, Alg. 3 for MRR achieves a strictly better approximation bound (1 + O(β + γ)) = (1 + o( )) when γ = o( ). On the other hand, in the worst case scenario when the learned sketch performs poorly (i.e., γ = Ω( )) the approximation guarantee of Alg. 3 remains (1 +O( )).

heading: 4.2 LOW-RANK APPROXIMATION
text: Given an input matrix A ∈ Rn×d and a desired rank k, the goal of LRA is to solve min
X,rank k LA(X) = min X,rank k ‖X −A‖2F
Approximation algorithm. We consider the time-optimal (up to 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-LOWRANK (Sarlos, 2006; Clarkson & Woodruff, 2017; Avron et al., 2016). Require: A ∈ Rn×d, S ∈ RmS×n, R ∈ RmR×d, S2 ∈ RmS2×n, R2 ∈ RmR2×d
1: UC [ TC T ′
C
] ← S2AR>, [ T>D T ′> D ] U>D ← SAR>2 with UC , UD orthogonal
2: G← S2AR>2 3: Z ′
LZ ′ R ← [U>CGUD]k 4: ZL = [ Z ′ LT −> D 0 ] , ZR = [ T−1C Z ′
R 0 ] 5: Z = ZLZR 6: return: AR>ZSA in form Pn×k, Qk×d
Sketch optimization. For the greedy stage, we optimize sketches S and R individually. However, we do not use LA(G((S,R, S2, R2), A)) = ‖X −A‖2F 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 a proxy objective. We observe that the proof that Alg. 5 is an -approximation uses the fact that the row space of SA and the column space ofAR both contain a good rank-k approximation toA. Thus, the proxy objectives for S and R are for k-rank approximation in the row/column space of SA/AR, respectively. For example, the proxy objective for S is LA(G(S,A)) =
∥∥[AV ]kV > −A∥∥2F where V is from the SVD SA = UΣV > and [·]k 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 (S,R, S2, R2) simultaneously using LA(G((S,R, S2, R2), A)) = ‖X −A‖2F with X from Algorithm 5, since it can easily be implemented 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 SL ∈ Rpoly( k )×n, RL ∈ Rpoly( k )×d which attain a (1 + γ)-approximation for LRA on A ∈ Rn×d, Alg. 3 gives a (1 + β + min(γ, ))-approximation for LRA on A in O(nnz(A) + (n + d) poly(k ) + k4
β4 · poly(k )) time where β is a trade-off parameter.
Remark 4.4. For k(n+d)−4, by setting the trade-off parameter β−4 = O(k(n+d)−4), Alg. 3 has the same asymptotic runtime as the best (1 + )-approximation algorithm of LRA with classical sparse embedding matrices. Moreover, Alg. 3 for LRA achieves a strictly better approximation bound (1 + O(β + γ)) = (1 + o( )) when γ = o( ). On the other hand, in the worst case scenario when the learned sketches perform poorly (i.e., γ = Ω( )) the approximation guarantee of Alg. 3 remains (1 +O( )).
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 is 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 SL denote a sparse sketch that Algorithm 1 has computed by iterating through indices in order of non-increasing row norms of A. Let SC denote a CountSketch matrix. Then, there is a fixed η > 0 such that minrank-kX∈rowsp(SLA) ‖X −A‖ 2 F ≤ (1− η) minrank-kX∈rowsp(SCA) ‖X −A‖ 2 F .
4.3 k-MEANS CLUSTERING
Let A ∈ Rn×d represent a set of n points, A1, . . . , An ∈ Rd. In k-means clustering, the goal is to find a partition of A1, ..., An into k clusters C = {C1, . . . , Ck} to
min C LA(C) = min C k∑ i=1 min µi∈Rd ∑ j∈Ci ‖Aj − µi‖2F
where µi denotes the center of cluster Ci.
Approximation algorithm. First, compress A into AV , where V is from SVD SA = UΣV > and S ∈ RO(k2/ 2)×n is a 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 withAV gives a (1+ )-approximation (Cohen et al., 2015) and Lloyd’s gives anO(log(k)) approximation, so we have a (1 + )O(log(k)) approximation overall.
Sketch optimization. k-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 a closed form solution, which means that evaluating the objective requires a (relatively) time-consuming iterative process. In the gradient stage, it is possible to differentiate through this iterative computation,
but propagating the gradient through such a nested expression is time-consuming. Thus, a proxy objective which is simple and quick to evaluate would be useful for both stages.
Cohen et al. showed that we get a good k-means solution if A is projected to an approximate top singular vector space. This suggests that a suitable proxy objective would be low-rank approximation: LA(G(S,A)) = ∥∥[AV ]kV > −A∥∥2F , 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 a 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 +O( ))-approximation.
Theorem 4.6 (Sketch monotonicity property for k-means). For a given classical CountSketch SC ∈ RO(poly(k/ ))×n, sketching with any extension of SC (i.e., by a learned sparse sketch SL) yields a better approximate solution for k-means than sketching with SC itself.

heading: 5 EVALUATION
text: 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 “naı̈vely” learned sketches for LRA, which are computed on just one sample from the training set.
• Exact SVD: sketch asAVm, where Vm contains the topm right singular vectors of random sample Ai ∈ Atrain
• Column sampling: sketch as AR, where R is computed from a randomly sampled Ai ∈ Atrain. Each column of R contains one entry. The location of this entry is sampled according to the squared column norms of Ai; the value of this entry is the inverse of the selected norm.
Data sets. We used high-dimensional data sets of three different types (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.
Evaluation metric. To evaluate the quality of a 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 tables denote
∆S = LA(G(S,Atest))− L∗A averaged over 10 trials.
Analysis of results. For least-squares regression and LRA, our method is the best. It is significantly 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 a 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 always 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 respectively. We should note that for k-means, the values generally do not differ much between methods.

heading: A APPENDIX

heading: CORRESPONDING TO SECTION 4.1: LEAST-SQUARES REGRESSION
text: Definition A.1 (Subspace Embedding). Given an n × d matrix A, a matrix S is an affine - embedding for (the column space of) A if for x ∈ Rd, ‖SAx‖22 = (1± ) ‖Ax‖ 2 2.
The following result is shown in Clarkson & Woodruff (2017) and sharpened with Nelson & Nguyên (2013); Meng & Mahoney (2013). Lemma A.2. Given matrices A,B with n rows, a sparse embedding matrix (i.e., random CS) with O(rank(A)2/ 2) rows is an affine -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 n × d matrix and let S ∈ RO( 1 2
)×n be a randomly chosen sparse embedding matrix (i.e., random CS). Then with constant probability, ‖SA‖2F = (1± ) ‖A‖ 2 F . Lemma A.4 (Sarlos (2006); Clarkson & Woodruff (2017)). Given classical CountSketch S ∈ Rm×n, SKETCH-REGRESSION(A,B, S) returns a (1 + )-approximation in time O(nnz(A) + nnz(B) + dd′m2 + min(d2m, dm2)). Proof: Since S is an affine -embedding matrix of A,B, then
‖SAX − SB‖2F = (1± ) ‖AX −B‖ 2 F
Next, by the normal equations, (SA)+(SB) is a minimizer of minX ‖SAX − SB‖2F and is a (1 + 3 )-approximate solution of minX ‖AX −B‖2F . To bound the runtime, note that since S is a CountSketch we can compute SA and SB in time O(nnz(A) + nnz(B)) and reduce the problem to an instance of multiple-response regression with m rows. Then, we can solve the reduced size problem in time O(dd′m2 + min(d2m, dm2)): O(min(dm2, d2m)) to compute (SA)+ and O(dd′m2) to compute (SA)+(SB).
Proof of Theorem 4.1. By Lemma A.2, a random CS SO with O(d 2
2 ) rows is an affine -embedding matrix of A,B with constant probability.
Next, let XL and XO be respectively the solutions returned by SKETCH-REGRESSION(A,B, SL) and SKETCH-REGRESSION(A,B, SO). By Lemma A.4, with constant probability,
‖AXO −B‖2F ≤ (1 + 3 ) minX ‖AX −B‖ 2 F .
Together with the assumption that the solution constructed from SL, which is denoted as XL, over Atrain is a γ-approximate solution,
min(‖AXL −B‖2F , ‖AXO −B‖ 2 F ) ≤ (1 + min(3 , γ)) minX ‖AX −B‖ 2 F . (A.1)
Hence, it only remains to compute the minimum of ‖AXL −B‖2F and ‖AXO −B‖ 2 F efficiently. Note that it takes Ω(n · d · d′) to compute these values exactly. However, for our purpose it suffices to compare (1 + β)-estimates of these values and return the minimum of the estimates. To achieve this, we use two applications of Lemma A.3 with R> ∈ RO( 1 β2 )×d′ , S ∈ RO( 1 β2
)×n. For any X ′ (in particular, both XO and XL), with constant probability,
‖S(AX ′ −B)R‖2F = ∥∥R>(AX ′ −B)>S>∥∥2 F = (1± β) ‖(AX ′ −B)‖2F (A.2)
Let ΓL and ΓO respectively denote AXL − B and AXO − B and let ΓM = arg min(‖SΓLR‖F , ‖SΓOR‖F ). By Eq. (A.2) and union bound over of XO and XL, with constant probability,
‖ΓM‖ ≤ (1 +O(β)) · ‖SΓMR‖2F B by Lemma A.3 ≤ (1 +O(β)) ·min(‖ΓO‖2F , ‖ΓL‖ 2 F ) B by Lemma A.3
≤ (1 +O(β + min( , β))) ‖AX −B‖2F B by Eq. (A.1)
Runtime Analysis. By Lemma A.4, XO an XL can be computed in O(nnz(A) + nnz(B) + d5d′ −4) time. Next, the time to compute ∆L and ∆O is
O(nnz(A) + nnz(XR) + nnz(B) + d
β4 ) = O(nnz(A) + nnz(B) + dd′ +
d
β4 ),
where X is eitherXO or XL and we use that fact that nnz(X) ≤ d ·d′ (i.e., the total number of cells in X).
Thus, the total runtime of Algorithm 3 is O(nnz(A) + nnz(B) + d5d′ −4 + dβ−4).
Theorem A.5 (Least Squares Regression). Suppose there exists a learned, sparse, subspace embedding matrix SL computed over Atrain with poly(d ) rows that attains a β-approximation over Atest. Then, there exists an algorithm that runs in time O(nnz(A) + poly(d/ )) that outputs a (1 + min(β, ))-approximation to the least squares regression problem.
Proof: Note that to solve minx∈Rd ‖Ax− b‖2, given a sparse embedding matrix S ∈ Rm×n for the columns of A together with the vector b, the problem can be solved within (1 + )-approximation in time nnz(A) + poly(d/ ) (e.g., see Theorem 2.14 Woodruff (2014)).
The proof outline is similar to the proof of Theorem 4.1. First, we compute solutions xL, xO ∈ Rd to the given instance using respectively a learned sketching matrix SL and a learning-free sketching algorithm. Then, we compare ‖AxL − b‖2 and ‖AxO − b‖2 and report the better. Note that since xL, xO are vectors, unlike the case of MRR, the naive comparison (i.e., without applying any sketching matrices) takes nnz(A). Hence, the algorithm runs in time nnz(A) + poly(d/ ) and returns a (1 + min( , β))-approximate 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/ 2 in the running time (e.g., see Section 2.5 of Woodruff (2014)).

heading: A.1 CORRESPONDING TO SECTION 4.2: LOW-RANK APPROXIMATION
text: Lemma A.6. Suppose that S ∈ RmS×n and R ∈ RmR×d are sparse affine -embedding matrices for (A>, A) and ((SA)>, A>). Then,
min rank-kX
∥∥AR>XSA−A∥∥2 F ≤ (1 + ) ‖Ak −A‖2F
Proof: Consider the following multiple-response regression problem:
min rank-kX
‖AkX −A‖2F . (A.3)
Note that since X = Ik is a feasible solution to Eq. (A.3), minrank-kX ‖AkX −A‖2F = ‖Ak −A‖2F . Let S ∈ RmS×n be a sketching matrix that satisfies the condition of Lemma A.9 for A := Ak and B := A. By the normal equations, the rank-k minimizer of ‖SAkX − SA‖2F is (SAk)
+SA. Hence, ∥∥Ak(SAk)+SA−A∥∥2F ≤ (1 + ) ‖Ak −A‖2F , (A.4) which in particular shows that a (1 + ) rank-k approximation of A exists in the row space of SA. In other words,
min rank-kX
‖XSA−A‖2F ≤ (1 + ) ‖Ak −A‖ 2 F . (A.5)
Next, let R ∈ RmR×d be a sketching matrix that satisfies the condition of Lemma A.9 for A := (SA)> and B := A>. Let Y denote the rank-k minimizer of ∥∥R(SA)>X> −RA>∥∥2 F
. Hence,∥∥(SA)>Y > −A>∥∥2 F ≤ (1 + ) min rank-kX ‖XSA−A‖2F B Lemma A.9
≤ (1 +O( )) ‖Ak −A‖2F B Eq. (A.5) (A.6)
Note that by the normal equations, again rowsp(Y >) ⊆ rowsp(RA>) and we can write Y = AR>Z where rank(Z) = k. Thus,
min rank-kX
∥∥AR>XSA−A∥∥2 F ≤ ∥∥AR>ZSA−A∥∥2 F
= ∥∥(SA)>Y > −A>∥∥2
F B Y = AR>Z
≤ (1 +O( )) ‖Ak −A‖2F B Eq. (A.6)
Lemma A.7 (Avron et al. (2016); Lemma 27). For C ∈ Rp×m′ , D ∈ Rm×p′ , G ∈ Rp×p′ , the following problem
min rank-k Z
‖CZD −G‖2F (A.7)
can be solved in O(pm′rC + p′mrD + pp′(rD + rC)) time, where rC = rank(C) ≤ min{m′, p} and rD = rank(D) ≤ min{m, p′}. Proof: Let UC and U>D be orthogonal bases for colsp(C) and rowsp(D), respectively, so that for each Z, CZD = UCZ ′U>D for some Z
′. Let PC and PD be the projection matrices onto the subspaces spanned by the rows of C> and D>, respectively: PC = UCU>C and PD = UDU > D . Then by the Pythagorean theorem,
‖CZD −G‖2F = ∥∥PCUCZ ′U>DPD −G∥∥2F
= ∥∥PCUCZ ′U>DPD − PCGPD∥∥2F + ‖PCG(I − PD)‖2F + ‖(I − PC)G‖2F ,
where the first equality holds since PCUC = UC and U>DPD = U > D and the second equality follows from the Pythagorean theorem. Hence,
arg min rank-k Z ‖CZD −G‖2F = arg min rank-k Z ∥∥PCUCZU>DPD − PCGPD∥∥2F . Moreover,∥∥PCUCZU>DPD − PCGPD∥∥2F = ∥∥UCZU>D − UCU>CGUDU>D∥∥2F = ∥∥Z − U>CGUD∥∥2F , where the first equality holds since U>C UC = I and U > DUD = I , and the second equality holds since UC and U>D are orthonormal. Hence,
arg min rank-k Z ‖CZD −G‖2F = arg min rank-k Z ∥∥Z − U>CGUD∥∥2F . Next, we can find Z = [U>CGUD]k by computing the SVD of U > CGUD in form of ZLZR where ZL ∈ Rm ′×k and ZR ∈ Rk×m.
Runtime analysis. We can computeUC andUD by the Gram-Schmidt process in timeO(pm′rC+ p′mrD), and U>CGUD in time O(min{rCp′(p+ rD), rDp(p′+ rC)}). Finally, the time to compute Z (i.e., an SVD computation of U>CGUD) is O(rCrD ·min{rC , rD}). Since rC ≤ min{p,m′} and rD ≤ min{p′,m}, the total runtime to minimize Z in Eq. (A.7) is O(pm′rC + p′mrD + pp′(rC + rD)).
Lemma A.8. Let S ∈ Rpoly(k/ )×d, R ∈ Rpoly(k/ )×d be CS matrices such that min
rank-kX
∥∥AR>XSA−A∥∥2 F ≤ (1 + γ) ‖Ak −A‖2F . (A.8)
Moreover, let S2 ∈ R k2 β2 ×n, and R2 ∈ R k2 β2 ×d be CS matrices. Then, Algorithm 5 runs in O(nnz(A) + (n + d) poly(k/ )) time and with constant probability gives a (1 + O(β + γ))approximate rank-k approximation of A. Proof: The algorithm first computes C = S2AR>, D = SAR>2 , G = S2AR > 2 which can be done in time O(nnz(A)). As an example, we bound the time to compute C = S2AR. Note that since S2 is a CS, S2A can be computed in O(nnz(A)) time and the number of non-zero entries in the resulting matrix is at most nnz(A). Hence, since R is a CS as well, C can be computed in time O(nnz(A) + nnz(S2A)) = O(nnz(A)). Then, it takes an extra poly(k/ ) · k 2
β2 time to store C,D and G in matrix form. Next, as we showed in Lemma A.7, the time to compute Z in Algorithm 5 isO( k 4
β4 ·poly(k/ )). Finally, it takes (n+d) poly(k/ ) time to computeQ = AR >ZL
and P = ZRSA and return the solution in the form of Pn×kQk×d. Hence, the total runtime is O(nnz(A) + (n+ d) poly(k/ ) + k 4
β4 · poly(k/ )).
The approximation guarantee follows from Eq. (A.8) and the fact that S2 and R2 are respectively affine β-embedding matrices of AR> and SA (see Lemma A.2).
Lemma A.9 (Avron et al. (2016); Lemma 25). Suppose that A ∈ Rn×d and B ∈ Rn×d′ . Moreover, let S be an oblivious sparse affine -embedding matrix (i.e., a random CS) with (rank(A)2/ 2) rows. Then with constant probability, X̃ = arg minrank-kX ‖SAX − SB‖
2 F , satisfies∥∥∥AX̃ −B∥∥∥2
F ≤ (1 + ) min rank-kX ‖AX −B‖2F .
In other words, in O(nnz(A) + nnz(B)) + (d+ d′)(rank(A)2/ 2) time, we can reduce the problem to a smaller (multi-response regression) problem with (rank(A)2/ 2) rows whose optimal solution is a (1 + )-approximate solution to the original problem. Proof of Theorem 4.3. Let SO and RO be CountSketch matrices of size poly(k/ ) × n and poly(k/ ) × d. Note that since rank(Ak) = k and rank((SOA)>) ≤ poly(k/ ), SO and RO are respectively affine -embedding matrices of (Ak, A) and ((SOA)>, A>). Then, by an application of Lemma A.6
min rank-kX ∥∥AR>OXSOA−A∥∥2F ≤ (1 +O( )) ‖Ak −A‖2F (A.9) Similarly, by the assumption that SL and RL finds a (1 + γ)-approximate solution of LRA over matrices A ∈ Atest, then
min rank-kX ∥∥AR>LXSLA−A∥∥2F ≤ (1 +O(γ)) ‖Ak −A‖2F (A.10) Next, let (PL, QL) and (PO, QO) be respectively the rank-k approximations of A in factored form using (SL, RL) and (SO, RO) (see Algorithm 5). Then, Eq. (A.10) together with Eq. (A.9) implies that
min(‖PLQL −A‖2F , ‖POQO −A‖ 2 F ) = (1 +O(min( , γ))) ‖Ak −A‖ 2 F (A.11)
Hence, it only remains to compute the minimum of ‖PLQL −A‖2F and ‖POQO −A‖ 2 F efficiently and we proceed similarly to the proof of Theorem 4.1. We use two applications of Lemma A.3 with R> ∈ RO( 1 β2 )×d , S ∈ RO( 1 β2
)×n. Let ΓL = PLQL − A, ΓO = POQO − A and ΓM = arg min(‖SΓLR‖F , ‖SΓOR‖F ). Hence,
‖ΓM‖2F ≤ (1 +O(β)) ‖SΓMR‖ 2 F B by Lemma A.3
≤ (1 +O(β)) ·min(‖ΓL‖2F , ‖ΓO‖ 2 F ) B by Lemma A.3 ≤ (1 +O(β + min( , γ))) ‖Ak −A‖2F B by Eq. (A.11)
Runtime analysis. By Lemma A.8, Algorithm 5 computes PL, QL and PO, QO in O(nnz(A) + (n+ d) poly(k ) + k4
β4 · poly( k )).
Next, it takes O(nnz(A) + (n + d) · k + kβ4 ) to compute ∆L and ∆O. As an example, we bound the amount of time required to compute SPLQLR − SAR corresponding to ∆L. Since S and R are sparse sketching matrices, SPL, QLR and SAR can be computed in nnz(SPL) + nnz(QLR) + nnz(A). Since SPL andQLR are respectively of size 1β2×k and k× 1 β2 , in total it takesO(nnz(A)+
k β2 ) to compute these three matrices. Then, we can compute SPLQLR and ‖SPLQLRSAR‖F in time O( kβ4 ).
Hence, the total runtime of Algorithm 3 for LRA isO(nnz(A)+(n+d) ·poly(k )+ k4
β4 ·poly( k )).
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 as the optimal rank-m approximation of A formed by truncated SVD: Am = UmΣmV > m .
Given a matrix U with orthogonal columns, let πU (A) = AUU>, which is the projection of the rows of A onto col(U). Let CU be the optimal k-means partition of πU (A). Further, we let µCU,i denote the i-th cluster’s center in the optimal k-means clustering on πU (A).
We denote dist2(A,µ) as the k-means loss given cluster centers (and their corresponding partition): dist2(A,µ) = ∑ i∈[k] ∑ j∈Ci ‖Aj − µi‖22
Likewise, cost(C) is the k-means loss given a partition: cost(C) = ∑ i∈[k] min µi ∑ j∈Ci ‖Aj − µi‖22
Definition A.10 (Projection-cost preserving sketch). Ã is a projection-cost preserving sketch of A if for any low rank projection matrix P and c not dependent on P :
(1− ) ‖A− PA‖2F ≤ ∥∥∥Ã− PÃ∥∥∥2
F + c ≤ (1 + ) ‖A− PA‖2F
Theorem (4.6: Sketch monotonicity property for k-means). Assume we have A ∈ Rn×d. We also have random CountSketch S ∈ RO(poly(k/ ))×n and define U ∈ Rd×O(poly(k/ )) with orthogonal columns such that colsp(U) = rowsp(SA). Then, any extension of S to S′ (for example, concatenation with a learned CountSketch SL) yields a better approximate k-means approximation. Specifically, define W with orthogonal columns such that col(W ) = row(S′A). Let C∗ denote the optimal partition of A, CU denote the optimal partition of πU (A), and CW denote the optimal partition of πW (A). Then
cost(CW ) ≤ (1 +O( ))cost(CU ) ≤ (1 +O( ))cost(C∗)
Proof:
cost(CW ) = ∑ i∈[k] min µi ∑ j∈CW,i ‖Aj − µi‖2
≤ ∑ i∈[k] ∑ j∈CW,i ∥∥Aj − µCW,i∥∥2 = ∑ i∈[k] ∑ j∈CW,i ‖Aj − πW (Aj)‖2 + ∥∥πW (Aj)− µCW,i∥∥2 (A.12)
≤ ∑ i∈[k] ∑ j∈CU,i ‖Aj − πW (Aj)‖2 + ∥∥πW (Aj)− µCU,i∥∥2 (A.13)
= ∑ i∈[k] ∑ j∈CU,i ∥∥Aj − µCU,i∥∥2 (A.14) ≤ (1 + )
∑ i∈k min µi ∑ j∈CU,i ‖Aj − µi‖2 (A.15)
= (1 + )cost(CU ) ≤ (1 +O( ))cost(C∗)
(A.12): µCW ∈ colsp(W ) so we can apply the Pythagorean Theorem. (A.13): (µCW , CW ) is an optimal k-means clustering of the projected points πW (A). (A.14): µCU ∈ colsp(U) ⊂ 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 matrices 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 of the points into k groups Makarychev et al. (2019). While the dimension reduction mappings satisfying strong coresets require poly(k/ ) dimensions, Makarychev et al. (2019) shows that O(log k/ 2) dimensions suffice to satisfy this “partition” guarantee. An interesting question for further research is if the sketch monotonicity guarantee also applies to the construction of Makarychev et al. (2019).
Corollary A.12. Assume we have A ∈ Rn×d, j ∈ Z+, > 0. Define m = min(O(poly(j/ )), d). Let Ãm be a (1 + )-approximation to Am of the form Ãm = AV V > where SA = UΣV > for CountSketch S ∈ Rm×n. Let X ∈ Rd×j be a matrix whose columns are orthonormal, and let Y ∈ Rd×(d−j) be a matrix with orthonormal columns that spans the orthogonal complement of colsp(X). Then ∥∥∥AXX> − ÃmXX>∥∥∥2
F ≤ · ‖AY ‖2F .
Proof: ∥∥∥AXX> − ÃmXX>∥∥∥2 F = ∥∥A(I − V V >)XX>∥∥2 F
≤ j ∥∥A(I − V V >)XX>∥∥2
2 (A.16) ≤ j ∥∥A(I − V V >)∥∥2
2 (A.17)
≤ j ·O( j ) ‖A−Aj‖2F (A.18)
= O( ) min(n,d)∑ i=j σ2i (A.19)
≤ O( ) ‖AY ‖2F (A.20)
(A.16) Note that rank(X) = j, so rank(A(I − V V >)XX>) = j. We use this fact to bound the Frobenius norm by the operator norm. (A.17) Using the fact that XX> is a projection. (A.18) Using Lemma 18 from Cohen et al. (2015), where Aj is the optimal rank-j approximation to A. We can apply this lemma because CountSketch is one of the eligible types of random projection matrices. (A.19) Letting σi be the singular values of A. (A.20) ∑min(n,d) i=j σ 2 i = minY ‖AY ‖ 2 F for Y ∈ Rd×(d−j) with orthonormal columns.
Theorem A.13. Assume we have A ∈ Rn×d, j ∈ Z+, ∈ (0, 1]. Define m = min(O(poly(j/ )), d). Let Ãm be a (1 + )-approximation to Am of the form Ãm = AV V > where SA = UΣV > for CountSketch S ∈ Rm×n. Then, for any non-empty set µ contained in a j-dimensional subspace, we have:∣∣∣∣dist2(Ãm, µ) + ∥∥∥Ãm −A∥∥∥2
F − dist2(A,µ) ∣∣∣∣ ≤ dist2(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 et al. (2013) involves the best rank-m approximation of A, Am; we will show it for the approximate rank-m approximation, Ãm.
Define X ∈ Rd×j with orthonormal columns such that colsp(X) = span(µ). Likewise, define Y ∈ Rd×(d−j) with orthonormal columns such that colsp(Y ) = span(µ)⊥. By the Pythagorean theorem:
dist2(Ãm, µ) = ‖ÃmY ‖2F + dist 2(ÃmXX T , µ)
and dist2(A,µ) = ‖AY ‖2F + dist
2(AXXT , µ). (A.21) Hence,∣∣∣∣(dist2(Ãm, µ) + ∥∥∥A− Ãm∥∥∥2F ) − dist2(A,µ) ∣∣∣∣ = ∣∣∣∣‖ÃmY ‖2F + dist2(ÃmXXT , µ) + ∥∥∥A− Ãm∥∥∥2 F − ( ‖AY ‖2F + dist 2(AXXT , µ) )∣∣∣∣
≤ ∣∣∣‖ÃmY ‖2F + ‖A− Ãm‖2F − ‖AY ‖2F ∣∣∣+ ∣∣∣dist2(ÃmXXT , µ)− dist2(AXXT , µ)∣∣∣ (A.22) ≤ε 2
8 · ‖AY ‖2F + ∣∣∣dist2(ÃmXXT , µ)− dist2(AXXT , µ)∣∣∣ (A.23) ≤ε 2
8 · dist2(A,µ) + ∣∣∣dist2(ÃmXXT , µ)− dist2(AXXT , µ)∣∣∣ B Used (A.21) (A.24) (A.22) Triangle inequality. (A.23) Take in Theorem 16 from Cohen et al. (2015) as 2/8. This theorem implies that Ãm is
a projection-cost preserving sketch with the c term as ∥∥∥A− Ãm∥∥∥2
F . Specifically, it says AV is a
project-cost preserving sketch, which means Ãm = AV V > is too: V has orthonormal columns so ‖(I − P )AV ‖2F = ∥∥(I − P )AV V >∥∥2 F .

heading: By Corollary A.12, ∥∥∥ÃmXXT −AXXT∥∥∥2
text: F ≤ 2 8 · ‖AY ‖2F .
Since µ ∈ rowsp(X), we have ‖AY ‖2F ≤ dist 2(A,µ). Using Corollary 21 from Feldman et al. (2013) while taking as /4, A as ÃmXXT , and B as AXXT yields
|dist2(ÃmXXT , µ)−dist2(AXXT , µ)| ≤ 4 ·dist2(AXXT , µ)+(1+4
)· ∥∥∥ÃmXXT −AXXT∥∥∥2
F
By A.21, dist2(AXXT , µ) ≤ dist2(A,µ). Finally, we combining the last two inequalities with (A.24):
∣∣∣∣(dist2(Ãm, µ) + ∥∥∥A− Ãm∥∥∥2 F ) − dist2(A,µ) ∣∣∣∣ ≤ ε 2
8 · dist2(A,µ) + 4 · dist2(A,µ) +
2
8 · (1 + 4 ) · dist2(A,µ)
≤ · dist2(A,µ), where in the last inequality we used the assumption ≤ 1.
Corollary A.14. Assume we have A ∈ Rn×d and CountSketch S ∈ RO(poly(k/ ))×n. Then, define U ∈ Rd×O(poly(k/ )) with orthogonal columns spanning row(SA). Also define AU = πU (A) and µU as the set of optimal cluster centers found on AU . Now, assume ∈ (0, 1/3]. Then, µU is a (1 + )-approximation to the optimal k-means clustering of A. That is, defining µ∗U as the cluster centers which minimize the cost of partition CU on A, we have:
dist2(A,µU ) ≤ (1 + ) dist2(A,µ∗U ) Proof: By using 3 in Theorem A.13 with j as k,∣∣∣dist2(AU , µU ) + ∥∥A−AU∥∥2F − dist2(A,µU )∣∣∣ ≤ 3 dist2(A,µU ) which implies that
(1− 3
) dist2(A,µU ) ≤ dist2(AU , µU ) + ∥∥A−AU∥∥2
F (A.25)
Likewise, by Theorem A.13 on AU and µ∗U (and taking j as k),∣∣∣dist2(AU , µ∗U ) + ∥∥A−AU∥∥2F − dist2(A,µ∗U )∣∣∣ ≤ 3 dist2(A,µ∗U ) which implies that
dist2(AU , µ∗U ) + ∥∥A−AU∥∥2 F ≤ (1 + 3 ) dist2(A,µ∗U ) (A.26)
By (A.25) and (A.26) together, we have:
(1− 3
) dist2(A,µU ) ≤ dist2(AU , µU ) + ∥∥A−AU∥∥2
F ≤ dist2(AU , µ∗U ) + ∥∥A−AU∥∥2
F
≤ (1 + 3 ) dist2(A,µ∗U )
Now, 1+ /31− /3 ≤ 1 + , so we have dist 2(A,µU ) ≤ (1 + ) dist2(A,µ∗U ).

heading: B SPECTRAL NORM GUARANTEE FOR ZIPFIAN MATRICES
text: In this section, we show that if the singular values of the input matrix A follow a Zipfian distribution (i.e., σi ∝ i−α for a constant α), we can find a (1 + ) 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 ∈ Rn×d, there exists an algorithm that runs inO(nnz(A)+(n+d) poly(k/ )) and returns a projection matrix P = QQ> such that with constant probability the following holds:
‖AP −A‖22 ≤ (1 + ) ‖A−Ak‖ 2 2 +O( k ) ‖A−Ak‖2F (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 ∈ Rn×d whose singular values follow a Zipfian distribution (i.e., σi ∝ i−α) with a constant α ≥ 1/2, there exists an algorithm that computes a rank-k matrix B (in factored form) such that ‖A−B‖22 ≤ (1 + ) ‖A−Ak‖ 2 2.
Proof: Note that since the singular values of A follow a Zipfian distribution with parameter α, for any value of k,
‖A−Ak‖2F = rank(A)∑ i=k+1 σ2i = C · rank(A)∑ i=k+1 i−2α B σi = √ C/iα
≤ C · ∫ rank(A) k x−2α dx
≤ k · 1 2α− 1 · (1 + 1 k )2α · C/(k + 1)2α
= O(k · σ2k+1) = O(k · ‖A−Ak‖22) (B.2)
By an application of Theorem B.1, we can compute a matrix B in a factored form in time O(nnz(A) + (n+ d) poly(k/ )) such that with constant probability,
‖B −A‖22 ≤ (1 + ) ‖A−Ak‖ 2 2 +O( k ) ‖A−Ak‖2F B By Eq. (B.1)
≤ (1 +O( )) ‖A−Ak‖22 B Eq. (B.2)

heading: C GREEDY INITIALIZATION
text: In this section, we analyze the performance of the greedy algorithm on the two distributions mentioned in Theorem 4.5.
Preliminaries and Notation. Left-multiplying A by CountSketch S ∈ Rm×n is equivalent to hashing the rows of A to m bins with coefficients in {−1, 1}. The greedy algorithm proceeds through the rows of A (in some order) and decides which bin to hash to, denoting this by adding an entry to S. We will denote the bins as bi and their summed contents as wi.

heading: C.1 SPIKED COVARIANCE MODEL WITH SPARSE LEFT SINGULAR VECTORS.
text: To recap, every matrix A ∈ Rn×d from the distribution Asp(s, `) has s < k “heavy” rows (Ar1 , · · · , Ars ) of norm ` > 1. The indices of the heavy rows can be arbitrary, but must be the same for all members of the distribution and are unknown to the algorithm. The remaining rows (called “light” rows) have unit norm.
In other words: letR = {r1, . . . , rs}. For all rows Ai, i ∈ [n]:
Ai = { ` · vi if i ∈ R vi o.w.
where vi is a uniformly random unit vector.
We also assume that Sr, Sg ∈ Rk×n 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 a bin). Then, we 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.
It is well known that a set of uniformly random vectors are -almost orthogonal (i.e., the magnitudes of their pairwise inner products are at most ). Observation C.1. Let v1, · · · , vn ∈ Rd be a set of random unit vectors. Then with high probability |〈vi, vj〉| ≤ 2 √ logn d ,∀ i < j ≤ n.
We define = 2 √
logn d .
Observation C.2. Let u1, · · · , ut be a set of vectors such that for each pair of i < j ≤ t, |〈 ui‖ui‖ , uj ‖uj‖ 〉| ≤ , and gi, · · · , gj ∈ {−1, 1}. Then,
t∑ i=1 ‖ui‖22 − 2 ∑ i<j≤t ‖ui‖2 ‖uj‖2 ≤ ∥∥∥∥∥ t∑ i=1 giui ∥∥∥∥∥ 2
2
≤ t∑ i=1 ‖ui‖22 + 2 ∑ i<j≤t ‖ui‖2 ‖uj‖2 (C.1)
Next, a straightforward consequence of -almost orthogonality is that we can find a QR-factorization of the matrix of such vectors where R (an upper diagonal matrix) has diagonal entries close to 1 and entries above the diagonal are close to 0.
Lemma C.3. Let u1, · · · , ut ∈ Rd be a set of unit vectors such that for any pair of i < j ≤ t, |〈ui, uj〉| ≤ where = O(t−2). There exists an orthonormal basis e1, · · · , et for the subspace spanned by u1, · · · , ut such that for each i ≤ t, ui = ∑i j=1 ai,jej where a 2 i,i ≥ 1 − ∑i−1 j=1 j 2 · 2 and for each j < i, a2i,j ≤ j2 2. Proof: We follow the Gram-Schmidt process to construct the orthonormal basis e1, · · · , et of the space spanned by u1, · · · , ut. by first setting e1 = u1 and then processing u2, · · · , ut, one by one. The proof is by induction. We show that once the first j vectors u1, · · · , uj are processed the statement of the lemma holds for these vectors. Note that the base case of the induction trivially holds as u1 = e1. Next, suppose that the induction hypothesis holds for the first ` vectors u1, · · · , u`. Claim C.4. For each j ≤ `, a2`+1,j ≤ j2 2. Proof: The proof of the claim is itself by induction. Note that, for j = 1 and using the fact that |〈u1, u`+1〉| ≤ , the statement holds and a2`+1,1 ≤ 2. Next, suppose that the statement holds for all j ≤ i < `, then by |〈ui+1, u`+1〉| ≤ ,
|a`+1,i+1| ≤ (|〈u`+1, ui+1|+ i∑
j=1
|a`+1,j | · |ai+1,j |)/|ai+1,i+1|
≤ ( + i∑
j=1
j2 2)/|ai+1,i+1| B by induction hypothesis on a`+1,j for j ≤ i
≤ ( + i∑
j=1
j2 2)/(1− i∑
j=1
j2 · 2)1/2 B by induction hypothesis on ai+1,i+1
≤ ( + i∑
j=1
j2 2) · (1− i∑
j=1
j2 · 2)1/2 · (1 + 2 · i∑
j=1
j2 2)
≤ ( + i∑
j=1
j2 2) · (1 + 2 · i∑
j=1
j2 2)
≤ (( i∑
j=1
j2 ) · (1 + 4 · i∑
j=1
j2 ) + 1)
≤ (i+ 1) B by = O(t−2)
Finally, since ‖u`+1‖22 = 1, a2`+1,`+1 ≥ 1− ∑` j=1 j 2 2.
Corollary C.5. Suppose that = O(t−2). There exists an orthonormal basis e1, · · · , et for the space spanned by the randomly picked vectors v1, · · · , vt, of unit norm, so that for each i, vi =∑i j=1 ai,jej where a 2 i,i ≥ 1− ∑i−1 j=1 j
2 · 2 and for each j < i, a2i,j ≤ j2 · 2. Proof: The proof follows from Lemma C.3 and the fact that the set of vectors v1, · · · , vt are -almost orthogonal (by Observation C.1).
The first main step is to show that the greedy algorithm (with non-increasing row norm ordering) will isolate rows into their own bins until all bins are filled. In particular, this means that the heavy rows (the first to be processed) will all be isolated.
We note that because we set rank(SA) = k, the k-rank approximation cost is the simplified expression ∥∥AV V > −A∥∥2 F , where UΣV > = SA, rather than ∥∥[AV ]kV > −A∥∥2F . This is just the projection cost onto row(SA). Also, we observe that minimizing this projection cost is the same as maximizing the sum of squared projection coefficients:
min S
∥∥A−AV V >∥∥2 F ∼ min
S ∑ i∈[n] ‖Ai − (〈Ai, v1〉v1 + . . .+ 〈Ai, vk〉vk)‖22
∼ min S ∑ i∈[n] (‖Ai‖22 − ∑ j∈[k] 〈Ai, vj〉2)
∼ max S ∑ i∈[n] ∑ j∈[k] 〈Ai, vj〉2
In the following sections, we will prove that our greedy algorithm makes certain choices by showing that these choices maximize the sum of squared projection coefficients. Lemma C.6. For any matrixA or batch of matricesA, at the end of iteration k, the learned CountSketch matrix S maps each row to an isolated bin. In particular, heavy rows are mapped to isolated bins.
Proof: For any iteration i ≤ k, we consider the choice of assigning Ai to an empty bin versus an occupied bin. Without loss of generality, let this occupied bin be bi−1, which already containsAi−1.
We consider the difference in cost for empty versus occupied. We will do this cost comparison for Aj with j ≤ i− 2, j ≥ i+ 1, and finally, j ∈ {i− 1, i}. First, we let {e1, . . . , ei} be an orthonormal basis for {A1, . . . , Ai} such that for each r ≤ i, Ar =∑r j=1 ar,jej where ar,r > 0. This exists by Lemma C.3. Let {e1, . . . , ei−2, e} be an orthonormal basis for {A1, . . . , Ai+2, Ai−1 ± Ai}. Now, e = c0ei−1 + c1ei for some c0, c1 because (Ai−1 ± Ai)− proj{e1,...,ei−2}(Ai−1 ±Ai) ∈ span(ei−1, ei). We note that c 2 0 + c 2 1 = 1 because we let e be a unit vector. We can find c0, c1 to be:
c0 = ai−1,i−1 + ai,i−1√
(ai−1,i−1 + ai,i−1)2 + a2i,i
, c1 = ai,i√
(ai−1,i−1 + ai,i−1)2 + a2i,i
1. j ≤ i− 2: The cost is zero for both cases because Aj ∈ span({e1, . . . , ei−2}).
2. j ≥ i + 1: We compare the rewards (sum of squared projection coefficients) and find that {e1, . . . , ei−2, e} is no better than {e1, . . . , ei}.
〈Aj , e〉2 = (c0〈Aj , ei−1〉+ c1〈Aj , ei〉)2
≤ (c21 + c20)(〈Aj , ei−1〉2 + 〈Aj , ei〉2) B Cauchy-Schwarz inequality = 〈Aj , ei−1〉2 + 〈Aj , ei〉2
3. j ∈ {i−1, i}: We compute the sum of squared projection coefficients of Ai−1 and Ai onto e:
( 1
(ai−1,i−1 + ai,i−1)2 + a2i,i ) · (a2i−1,i−1(ai−1,i−1 + ai,i−1)2
+ (ai,i−1(ai−1,i−1 + ai,i−1) + ai,iai,i) 2)
=( 1
(ai−1,i−1 + ai,i−1)2 + a2i,i ) · ((ai−1,i−1 + ai,i−1)2(a2i−1,i−1 (C.2)
+ a2i,i−1) + a 4 i,i + 2ai,i−1a 2 i,i(ai−1,i−1 + ai,i−1)) (C.3)
On the other hand, the sum of squared projection coefficients ofAi−1 andAi onto ei−1∪ei is:
( (ai−1,i−1 + ai,i−1) 2 + a2i,i (ai−1,i−1 + ai,i−1)2 + a2i,i ) · (a2i−1,i−1 + a2i,i−1 + a2i,i) (C.4)
Hence, the difference between the sum of squared projections of Ai−1 and Ai onto e and ei−1 ∪ ei is ((C.4) - (C.3))
a2i,i((ai−1,i−1 + ai,i−1) 2 + a2i−1,i−1 + a 2 i,i−1 − 2ai,i−1(ai−1,i−1 + ai,i−1))
((ai−1,i−1 + ai,i−1)2 + a2i,i)
= 2a2i,ia 2 i−1,i−1
((ai−1,i−1 + ai,i−1)2 + a2i,i) > 0
Thus, we find that {e1, . . . , ei} is a strictly better basis than {e1, . . . , ei−2, e}. This means the greedy algorithm will choose to place Ai in an empty bin.
Next, we show that none of the rows left to be processed (all light rows) will be assigned to the same bin as 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 a light row) in the sum of squared projection coefficients, lower-bounding it in the former case and upper-bounding it in the latter.
We introduce some results that will be used in Lemma C.10. Claim C.7. Let Ak+r, r ∈ [1, . . . , n− k] be a light row not yet processed by the greedy algorithm. Let {e1, . . . , ek} be the Gram-Schmidt basis for the current {w1, . . . , wk}. Let β = O(n−1k−3) upper bound the inner products of normalizedAk+r, w1, . . . , wk. Then, for any bin i, 〈ei, Ak+r〉2 ≤ β2 · k2. Proof: This is a straightforward application of Lemma C.3. From that, we have 〈Ak+r, ei〉2 ≤ i2β2, for i ∈ [1, . . . , k], which means 〈Ak+r, ei〉2 ≤ k2β2. Claim C.8. Let Ak+r be a light row that has been processed by the greedy algorithm. Let {e1, . . . , ek} be the Gram-Schmidt basis for the current {w1, . . . , wk}. If Ak+r is assigned to bin bk−1 (w.l.o.g.), the squared projection coefficient of Ak+r onto ei, i 6= k − 1 is at most 4β2 · k2, where β = O(n−1k−3) upper bounds the inner products of normalized Ak+r, w1, · · · , wk. Proof: Without loss of generality, it suffices to bound the squared projection of Ak+r onto the direction of wk that is orthogonal to the subspace spanned by w1, · · · , wk−1. Let e1, · · · , ek be an orthonormal basis of w1, · · · , wk guaranteed by Lemma C.3. Next, we expand the orthonormal basis to include ek+1 so that we can write the normalized vector of Ak+r as vk+r = ∑k+1 j=1 bjej . By a similar approach to the proof of Lemma C.3, for each j ≤ k − 2, bj ≤ β2j2. Next, since |〈wk, vk+r〉| ≤ β,
|bk| ≤ 1
|〈wk, ek〉| · (|〈wk, vk+r〉|+ k−1∑ j=1 |bj · 〈wk, ej〉|)
≤ 1√ 1− ∑k−1 j=1 β 2 · j2 · (β + k−2∑ j=1 β2 · j2 + (k − 1) · β) B |bk−1| ≤ 1
= β +
∑k−2 j=1 β
2 · j2√ 1− ∑k−1 j=1 β 2 · j2 + (k − 1)β
≤ 2(k − 1)β − β 2(k − 1)2√
1− ∑k−1 j=1 β 2 · j2 B similar to the proof of Lemma C.3
< 2β · k
Hence, the squared projection ofAk+r onto ek is at most 4β2 ·k2 ·‖Ak+r‖22. We assumed ‖Ak+r‖. Claim C.9. We assume that the absolute values of the inner products of vectors in v1, · · · , vn are at most < 1/(n2 ∑ Ai∈b ‖Ai‖2) and the absolute values of the inner products of the normalized vectors of w1, · · · , wk are at most β = O(n−3k− 3 2 ). Suppose that bin b contains the row Ak+r. Then, the squared projection of Ak+r onto the direction of w orthogonal to span({w1, · · · , wk} \ {w}) is at most ‖Ak+r‖ 4 2
‖w‖22 +O(n−2) and is at least ‖Ak+r‖ 4 2 ‖w‖22 −O(n−2).
Proof: Without loss of generality, we assume thatAk+r is mapped to bk; w = wk. First, we provide an upper and a lower bound for |〈vk+r, wk〉| where for each i ≤ k, we let wi = wi‖wi‖2 denote the normalized vector of wi. Recall that by definition vk+r =
Ak+r ‖Ak+r‖2 .
|〈wk, vk+r〉| ≤ ‖Ak+r‖2 +
∑ Ai∈bk ‖Ai‖2
‖wk‖2
≤ ‖Ak+r‖2 + n−2
‖wk‖2 B by < n−2∑ Ai∈bk ‖Ai‖2
≤ ‖Ak+r‖2 ‖wk‖2 + n−2 B ‖wk‖2 ≥ 1 (C.5)
|〈wk, vk+r〉| ≥ ‖Ak+r‖2 −
∑ Ai∈bk ‖Ai‖2 · ‖wk‖2
≥ ‖Ak+r‖2 ‖wk‖2 − n−2 (C.6)
Now, let {e1, · · · , ek} be an orthonormal basis for the subspace spanned by {w1, · · · , wk} guaranteed by Lemma C.3. Next, we expand the orthonormal basis to include ek+1 so that we can write vk+r = ∑k+1 j=1 bjej . By a similar approach to the proof of Lemma C.3, we can show that for each j ≤ k − 1, b2j ≤ β2j2. Moreover,
|bk| ≤ 1
|〈wk, ek〉| · (|〈wk, vk+r〉|+ k−1∑ j=1 |bj · 〈wk, ej〉|)
≤ 1√ 1− ∑k−1 j=1 β 2 · j2 · (|〈wk, vk+r〉|+ k−1∑ j=1 β2 · j2) B by Lemma C.3
≤ 1√ 1− ∑k−1 j=1 β 2 · j2 · (n−2 + ‖Ak+r‖2 ‖wk‖2 + k−1∑ j=1 β2 · j2) B by (C.5) < β · k + 1√ 1− β2k3 · (n−2 + ‖Ak+r‖2 ‖wk‖2 ) B similar to the proof of Lemma C.3
≤ O(n−2) + (1 +O(n−2)) ‖Ak+r‖2 ‖wk‖2
B by β = O(n−3k− 3 2 )
≤ ‖Ak+r‖2 ‖wk‖2 +O(n−2) B ‖Ak+r‖2 ‖wk‖2 ≤ 1
and,
|bk| ≥ 1
|〈wk, ek〉| · (|〈wk, vk+r〉| − k−1∑ j=1 |bj · 〈wk, ej〉|)
≥ |〈wk, vk+r〉| − k−1∑ j=1 β2 · j2 B since |〈wk, ek〉| ≤ 1
≥ ‖Ak+r‖2 ‖wk‖2 − n−2 − k−1∑ j=1 β2 · j2 B by (C.6)
≥ ‖Ak+r‖2 ‖wk‖2 −O(n−2) B by β = O(n−3k− 32 )
Hence, the squared projection ofAk+r onto ek is at most ‖Ak+r‖42 ‖wk‖22 +O(n−2) and is at least ‖Ak+r‖ 4 2 ‖wk‖22 − O(n−2).
Now, we show that at the end of the algorithm no light row will be assigned to the bins that contain heavy rows. Lemma C.10. We assume that the absolute values of the inner products of vectors in v1, · · · , vn are at most < min{n−2k− 53 , (n ∑ Ai∈w ‖Ai‖2)
−1}. At each iteration k+ r, the greedy algorithm will assign the light row Ak+r to a bin that does not contain a heavy row. Proof: The proof is by induction. Lemma C.6 implies that no light row has been mapped to a bin that contains a heavy row for the first k iterations. Next, we assume that this holds for the first k + r − 1 iterations and show that is also must hold for the (k + r)-th iteration. To this end, we compare the sum of squared projection coefficients when Ak+r avoids and collides with a heavy row.
First, we upper bound β = maxi 6=j≤k |〈wi, wj〉|/(‖wi‖2 ‖wj‖2). Let ci and cj respectively denote the number of rows assigned to bi and bj .
β = max i 6=j≤k |〈wi, wj〉| ‖wi‖2 ‖wj‖2 ≤ ci · cj · √ ci − 2 c2i · √ cj − 2 c2j
B Observation C.2
≤ 16 √cicj B ≤ n−2k−5/3
≤ n−1k− 53 B ≤ n−2k−5/3
1. If Ak+r is assigned to a bin that contains c light rows and no heavy rows. In this case, the projection loss of the heavy rows A1, · · · , As onto row(SA) remains zero. Thus, we only need to bound the change in the sum of squared projection coefficients of the light rows before and after iteration k + r.
Without loss of generality, let wk denote the bin that contains Ak+r. Since Sk−1 = span({w1, · · · , wk−1}) has not changed, we only need to bound the difference in cost between projecting onto the component of wk −Ak+r orthogonal to Sk−1 and the component of wk orthogonal to Sk−1, respectively denoted as ek and ek.
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 ek is at most β2k2. Hence, the total decrease in the squared projection is at most (n− k − r) · β2k2.
II. By Claim C.8, for the processed light rows that are not mapped to the last bin, the squared projection of each onto ek is at most 4β2k2. Hence, the total decrease in the squared projection cost is at most (r − 1) · 4β2k2.
III. For each row Ai 6= Ak+r that is mapped to the last bin, by Claim C.9 and the fact ‖Ai‖42 = ‖Ai‖22 = 1, the squared projection of Ai onto ek is at most
‖Ai‖22 ‖wk−Ak+r‖22 + O(n−2) and the
squared projection of Ai onto ek is at least ‖Ai‖22 ‖wk‖22 −O(n−2).
Moreover, the squared projection of Ak+r onto ek compared to ek increases by at least ( ‖Ak+r‖22 ‖wk‖22 −O(n−2))−O(n−2) = ‖Ak+r‖ 2 2 ‖wk‖22 −O(n−2).
Hence, the total squared projection of the rows in the bin bk decreases by at least:
( ∑
Ai∈wk/{Ar+k}
‖Ai‖22 ‖wk −Ar+k‖22 +O(n−2))− ( ∑ Ai∈wk ‖Ai‖22 ‖wk‖22 −O(n−2))
≤ ‖wk −Ar+k‖22 +O(n−1) ‖wk −Ar+k‖22 − ‖wk‖22 −O(n−1) ‖wk‖22 +O(n−1) B by Observation C.2 ≤O(n−1)
Hence, summing up the bounds in items I to III above, the total decrease in the sum of squared projection coefficients is at most O(n−1).
2. If Ak+r is assigned to a bin that contains a heavy row. Without loss of generality, we can assume thatAk+r is mapped to bk that contains the heavy rowAs. In this case, the distance of heavy rows A1, · · · , As−1 onto the space spanned by the rows of SA is zero. Next, we bound the amount of change in the squared distance of As and light rows onto the space spanned by the rows of SA.
Note that the (k − 1)-dimensional space corresponding to w1, · · · , wk−1 has not changed. Hence, we only need to bound the decrease in the projection distance of Ak+r onto ek compared to ek (where ek, ek are defined similarly as in the last part).
1. For the light rows other than Ak+r, the squared projection of each onto ek is at most β2k2. Hence, the total increase in the squared projection of light rows onto ek is at most (n− k) · β2k2 = O(n−1).
2. By Claim C.9, the sum of squared projections ofAs andAk+r onto ek decreases by at least
‖As‖22 − ( ‖As‖42 + ‖Ak+r‖ 4 2
‖As +Ar+k‖22 +O(n−1))
≥‖As‖22 − ( ‖As‖42 + ‖Ak+r‖ 4 2
‖As‖22 + ‖Ar+k‖ 2 2 − n−O(1)
+O(n−1)) B by Observation C.2
≥ ‖Ar+k‖22 (‖As‖ 2 2 − ‖Ak+r‖ 2 2)− ‖As‖ 2 2 ·O(n−1)
‖As‖22 + ‖Ar+k‖ 2 2 −O(n−1)
−O(n−1)
≥ ‖Ar+k‖22 (‖As‖ 2 2 − ‖Ak+r‖ 2 2)− ‖As‖ 2 2 ·O(n−1)
‖As‖22 + ‖Ar+k‖ 2 2
−O(n−1)
≥ ‖Ar+k‖22 (‖As‖ 2 2 − ‖Ak+r‖ 2 2)
‖As‖22 + ‖Ar+k‖ 2 2
−O(n−1)
≥ ‖Ar+k‖22 (1− (‖Ak+r‖ 2 2 / ‖As‖ 2 2))
1 + (‖Ar+k‖22 / ‖As‖ 2 2)
−O(n−1)
≥‖Ar+k‖22 (1− ‖Ak+r‖2 ‖As‖2 )−O(n−1) B 1− 2 1 + 2 ≥ 1−
Hence, in this case, the total decrease in the squared projection is at least
‖Ar+k‖22 (1− ‖Ak+r‖2 ‖As‖2 )−O(n−1) = 1− ‖Ak+r‖2 ‖As‖2 )−O(n−1) B ‖Ar+k‖2 = 1
= 1− (1/ √ `)−O(n−1) B ‖As‖2 = √ `
Thus, for a sufficiently large value of `, the greedy algorithm will assign Ak+r 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 SgA for A ∼ Asp(s, `) where A ∈ Rn×d for d = Ω(n4k4 log n) and Sg is the CountSketch constructed by the greedy algorithm with non-increasing order is at most n− s. Proof: First, we need to show that absolute values of the inner products of vectors in v1, · · · , vn is at most < min{n−2k−2, (n ∑ Ai∈w ‖Ai‖2) −1} so that we can apply Lemma C.10. To show this,
note that by Observation C.1, ≤ 2 √
logn d ≤ n −2k−2 since d = Ω(n4k4 log n). The proof follows from Lemma C.6 and Lemma C.10. Since all heavy rows are mapped to isolated bins, the projection loss of the light rows is at most n− s.
Next, we bound the Frobenius norm error of the best rank-k-approximation solution constructed by the standard CountSketch with a randomly chosen sparsity pattern. Lemma C.12. Let s = αk where 0.7 < α < 1. The expected squared loss of the best rank-k approximate solution in the rowspace SrA for A ∈ Rn×d ∼ Asp(s, `) where d = Ω(n6`2) and Sr is the sparsity pattern of CountSketch is chosen uniformly at random is at least n+ `k4e − (1 +α)k− n−O(1).
Proof: We can interpret the randomized construction of the CountSketch as a “balls and bins” experiment. In particular, considering the heavy rows, we compute the expected number of bins (i.e., rows in SrA) that contain a heavy row. Note that the expected number of rows in SrA that do not contain any heavy row is k · (1− 1k ) s ≥ k · e− s
k−1 . Hence, the number of rows in SrA that contain a heavy row of A is at most k(1 − e− s k−1 ). Thus, at least s − k(1 − e− s k−1 ) heavy rows are not mapped to an isolated bin (i.e., they collide with some other heavy rows). Then, it is straightforward to show that the squared loss of each such row is at least `− n−O(1). Claim C.13. Suppose that heavy rows Ar1 , · · · , Arc are mapped to the same bin via a CountSketch S. Then, the total squared distances of these rows from the subspace spanned by SA is at least (c− 1)`−O(n−1).
Proof: Let b denote the bin that contains the rowsAr1 , · · · , Arc and suppose that it has c′ light rows as well. Note that by Claim C.8 and Claim C.9, the squared projection of each row Ari onto the subspace spanned by the k bins is at most
‖Ahi‖ 4 2
‖w‖22 +O(n−1)
≤ ` 2
c`+ c′ − 2 (c2`+ cc′ √ `+ c′2) +O(n−1)
≤ ` 2
c`− n−O(1) + n−O(1) B by ≤ n−3`−1
≤ ` 2
c2`2 · (c`+O(n−1) +O(n−1)
≤ ` c +O(n−1)
Hence, the total squared loss of these c heavy rows is at least c` − c · ( `c + O(n −1)) ≥ (c − 1)` − O(n−1).
Hence, the expected total squared loss of heavy rows is at least:
` · (s− k(1− e− s k−1 ))− s · n−O(1)
≥` · k(α− 1 + e−α)− `α− n−O(1) B s = α · (k − 1) where 0.7 < α < 1
≥`k 2e − `− n−O(1) B α ≥ 0.7
≥`k 4e −O(n−1) B assuming k > 4e
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 a 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 of the other rows in the bin as well) is at least 1 − O(n−1). Hence, the expected total squared loss of the light rows is at least:
(n− s− k)(1−O(n−1)) ≥ (n− (1 + α) · k)−O(n−1)
Hence, the expected squared loss of a CountSketch whose sparsity is picked at random is at least
`k 4e −O(n−1) + n− (1 + α)k −O(n−1) ≥ n+ `k 4e − (1 + α)k −O(n−1)
Corollary C.14. Let s = α(k − 1) where 0.7 < α < 1 and let ` ≥ (4e+1)nαk . Let Sg be the CountSketch whose sparsity pattern is learned over a training set drawn from Asp via the greedy approach. Let Sr be a CountSketch whose sparsity pattern is picked uniformly at random. Then, for an n× d matrix A ∼ Asp where d = Ω(n6`2), the expected loss of the best rank-k approximation of A returned by Sr is worse than the approximation loss of the best rank-k approximation of A returned by Sg by at least a constant factor.
Proof:
ESr [ min rank-kX∈rowsp(SrA)
‖X −A‖2F ] ≥ n+ `k 4e − (1 + α)k − n−O(1) B Lemma C.12
≥ (1 + 1/α)(n− s) B ` ≥ (4e+ 1)n αk
= (1 + 1/α) min rank-kX∈rowsp(SgA)
‖X −A‖2F B Corollary C.11

heading: C.2 ZIPFIAN ON SQUARED ROW NORMS.
text: Each matrix A ∈ Rn×d ∼ Azipf has rows which are uniformly random and orthogonal. Each A has 2i+1 rows of squared norm n2/22i for i ∈ [1, . . . , O(log(n))]. We also assume that each row has the same squared norm for all members of Azipf . 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 1 − 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, we assume that s ≤ k/2 and, for simplicity, s = ∑hs i=1 2
i+1 for some hs ∈ Z+. That means the minimum squared norm of a heavy row is n2/22hs and the maximum squared norm of a light row is n2/22hs+2.
The analysis of the greedy algorithm ordered by non-increasing row norms on 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 a first step, by Lemma C.6, at the end of iteration k the first k rows are assigned to different bins. Then, via a similar inductive proof, we show that none of the light rows are mapped to a bin that contains one of the top s heavy rows.
Lemma C.15. At each iteration k+ r, the greedy algorithm picks the position of the non-zero value in the (k + r)-th column of the CountSketch matrix S so that the light row Ak+r is mapped to a 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+ r− 1 iterations a light row is assigned to a bin that contains a heavy row. Now, we consider the following cases:
1. If Ak+r is assigned to a bin that only contains light rows. Without loss of generality we can assume that Ak+r is assigned to bk. Since the vectors are orthogonal, we only need to bound the difference in the projection of Ak+r and the light rows that are assigned to bk onto the direction of wk before and after adding Ak+r to bk. In this case, the total squared loss corresponding to rows in bk and Ak+r before and after adding Ak+1 are respectively
before adding Ak+r to bk: ‖Ak+r‖22 + ∑ Aj∈bk ‖Aj‖22 − ( ∑ Aj∈bk ‖Aj‖ 4 2∑ Aj∈bk ‖Aj‖ 2 2 )
after adding Ak+r to bk: ‖Ak+r‖22 + ∑ Aj∈bk ‖Aj‖22 − ( ‖Ak+r‖42 +
∑ Aj∈bk ‖Aj‖ 4 2
‖Ak+r‖22 + ∑ Aj∈bk ‖Aj‖ 2 2 )
Thus, the amount of increase in the squared loss is
(
∑ Aj∈bk ‖Aj‖
4 2∑
Aj∈bk ‖Aj‖ 2 2
)− ( ‖Ak+r‖42 +
∑ Aj∈bk ‖Aj‖ 4 2
‖Ak+r‖22 + ∑ Aj∈bk ‖Aj‖ 2 2 ) = ‖Ak+r‖22 ·
∑ Aj∈bk ‖Aj‖ 4 2 − ‖Ak+r‖ 4 2 · ∑ Aj∈bk ‖Aj‖ 2 2
( ∑ Aj∈bk ‖Aj‖ 2 2)(‖Ak+r‖ 2 2 + ∑ Aj∈bk ‖Aj‖ 2 2)
= ‖Ak+r‖22 ·
∑ Aj∈bk
‖Aj‖42∑ Aj∈bk
‖Aj‖22 − ‖Ak+r‖22∑
Aj∈bk ‖Aj‖ 2 2 + ‖Ak+r‖ 2 2 ≤ ‖Ak+r‖22 · ∑ Aj∈bk ‖Aj‖ 2 2 − ‖Ak+r‖ 2 2∑
Aj∈bk ‖Aj‖ 2 2 + ‖Ak+r‖ 2 2
(C.7)
2. If Ak+r is assigned to a bin that contains a heavy row. Without loss of generality and by the induction hypothesis, we assume that Ak+r is assigned to a bin b that only contains a heavy row Aj . Since the rows are orthogonal, we only need to bound the difference in the projection of Ak+r and Aj In this case, the total squared loss corresponding to Aj and Ak+r before and after adding Ak+1
to b are respectively
before adding Ak+r to bk: ‖Ak+r‖22
after adding Ak+r to bk: ‖Ak+r‖22 + ‖Aj‖ 2 2 − (
‖Ak+r‖42 + ‖Aj‖ 4 2 ‖Ak+r‖22 + ‖Aj‖ 2 2 )
Thus, the amount of increase in the squared loss is
‖Aj‖22 − ( ‖Ak+r‖42 + ‖Aj‖ 4 2
‖Ak+r‖22 + ‖Aj‖ 2 2
) = ‖Ak+r‖22 · ‖Aj‖22 − ‖Ak+r‖ 2 2
‖Aj‖22 + ‖Ak+r‖ 2 2
(C.8)
Then (C.8) is larger than (C.7) if ‖Aj‖22 ≥ ∑ Ai∈bk ‖Ai‖ 2 2. Next, we show that at every inductive iteration, there exists a bin b which only contains light rows and whose squared norm is smaller than the squared norm of any heavy row. For each value m, define hm so that m = ∑hm i=1 2
i+1 = 2hm+2 − 2.
Recall that all heavy rows have squared norm at least n 2
22hs . There must be a bin b that only contains
light rows and has squared norm at most
‖w‖22 = ∑ Ai∈b ‖Ai‖22 ≤ n2 22(hs+1) + ∑hn i=hk+1 2i+1n2 22i k − s
≤ n 2
22(hs+1) +
2n2
2hk(k − s)
≤ n 2
22(hs+1) +
n2
22hk B s ≤ k/2 and k > 2hk+1
≤ n 2
22hs+1 B hk ≥ hs + 1 < ‖As‖22 Hence, the greedy algorithm will map Ak+r 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 SgA for A ∈ Rn×d ∼ Azipf where A ∈ Rn×d and Sg is the CountSketch constructed by the greedy algorithm with non-increasing order, is < n 2
2hk−2 . Proof: At the end of iteration k, the total squared loss is ∑hn i=hk+1 2i+1 · n 2
22i . After that, in each iteration k + r, by (C.7), the squared loss increases by at most ‖Ak+r‖22. Hence, the total squared loss in the solution returned by Sg is at most
2( hn∑ i=hk+1 2i+1n2 22i ) = 4n2 · hn∑ i=hk+1 1 2i
< 4n2
2hk =
n2
2hk−2
Next, we bound the squared loss of the best rank-k-approximate solution constructed by the standard CountSketch with a randomly chosen sparsity pattern. Observation C.17. Let us assume that the orthogonal rows Ar1 , · · · , Arc are mapped to same bin and for each i ≤ c, ‖Ar1‖ 2 2 ≥ ‖Ari‖ 2 2. Then, the total squared loss of Ar1 , · · · , Arc after projecting onto Ar1 ± · · · ±Arc is at least ‖Ar2‖ 2 2 + · · ·+ ‖Arc‖ 2 2. Proof: Note that since Ar1 , · · · , Arc are orthogonal, for each i ≤ c, the squared projection of Ari onto Ar1 ± · · · ±Arc is ‖Ari‖ 4 2 / ∑c j=1
∥∥Arj∥∥22. Hence, the sum of squared projection coefficients of Ar1 , · · · , Arc onto Ar1 ± · · · ±Arc is∑c
j=1 ∥∥Arj∥∥42∑c j=1 ∥∥Arj∥∥22 ≤ ‖Ar1‖22
Hence, the total projection loss of Ar1 , · · · , Arc onto Ar1 ± · · · ±Arc is at least c∑ j=1 ∥∥Arj∥∥22 − ‖Ar1‖22 = ‖Ar2‖22 + · · ·+ ‖Arc‖22 . 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 fully contributes to the total squared loss of the solution. Lemma C.18. For k > 210 − 2, the expected squared loss of the best rank-k approximate solution in the rowspace of SrA forAn×d ∼ Azipf , where Sr is the sparsity pattern of a CountSketch chosen uniformly at random, is at least 1.095n 2
2hk−2 .
Proof: In 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 as a “balls and bins” experiment.
For each 0 ≤ j ≤ hk, let Aj denote the set of rows with squared norm n 2 22(hk−j) and let A>j =⋃
j<i≤hk Ai. Note that for each j, |Aj | = 2 hk−j+1 and |A>j | = ∑hk i=j+1 2 hk−i+1 = ∑hk−j i=1 2 i = 2(2hk−j − 1). Moreover, note that k = 2(2hk+1 − 1). Next, for a row Ar in Aj (0 ≤ j < hk), we compute the probability that at least one row in A>j collide with Ar.
Pr[at least one row in A>j collide with Ar] = (1− (1− 1
k )|A>j |)
≥ (1− e− |A>j | k )
= (1− e− 2hk−j−1 2hk+1−1 )
≥ (1− e−2 −j−2 ) B since 2hk−j − 1 2hk+1 − 1 > 2−j−2
Hence, by Observation C.17, the contribution of rows in Aj to the total squared loss is at least
(1− e−2 −j−2 ) · |Aj | · n2
22(hk−j) =(1− e−2
−j−2 ) · n
2
2hk−j−1
=(1− e−2 −j−2 ) · n 2
2hk−2 · 2j−1
Thus, the contribution of rows with “large” squared norm, i.e., A>0, to the total squared loss is at least5
n2 2hk−2 · hk∑ j=0 2j−1 · (1− e−2 −j−2 ) ≥ 1.095 · n 2 2hk−2 Bfor hk > 8
Corollary C.19. Let Sg be a CountSketch whose sparsity pattern is learned over a training set drawn from Asp via the greedy approach. Let Sr be a CountSketch whose sparsity pattern is picked uniformly at random. Then, for an n × d matrix A ∼ Azipf , for a sufficiently large value of k, the expected loss of the best rank-k approximation ofA returned by Sr is worse than the approximation loss of the best rank-k approximation of A returned by Sg by at least a 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 a better rank-k solution compared to the CountSketch whose sparsity pattern is chosen at random. However, still other implementations of the greedy algorithm may result in a better solution compared to the greedy with non-increasing order. To give an example, in 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.
5The numerical calculation is computed by WolframAlpha.
Let A be a matrix with four orthogonal rows u, u, v, w where ‖u‖2 = 1 and ‖v‖2 = ‖w‖2 = 1 + and suppose that the goal is to compute a rank-2 approximation of A. Note that in the greedy with non-decreasing order, v and w will be assigned to different bins and by a simple calculation we can show that the copies of u also will be assigned to different bins. Hence, the squared loss in the computed rank-2 solution is 1 + (1+ ) 2
2+(1+ )2 . However, the optimal solution will assign v and w to one bin and the two copies of u to the other bin which results in the squared loss of (1 + )2 which is a constant factor smaller than the solution returned by the greedy with non-increasing order for sufficiently small values of .
On the other hand, in the greedy algorithm with random order, with a constant probability ( 13 + 1 8 ), 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 a non-increasing order. Hence, in expectation, the solution returned by the greedy with random order is better than the solution returned by the greedy algorithm with non-increasing order by a constant factor.

heading: D EXPERIMENTS - APPENDIX

heading: D.1 BASELINES
text: 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 a (1 + )-approximation for both LRA and k-means Cohen et al. (2015).
• Column sampling: In the canonical, learning-free sketching setting, sketching via column sampling yields a (1 + )-approximation for k-means Cohen et al. (2015).

heading: D.2 EXPERIMENTAL PARAMETERS
text: 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) algorithms in PyTorch. In the first case, 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 backpropagation through chained differentiable operators.
There are also two points of note in the greedy algorithm implementation. First, we 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, we 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 Dw (the set of candidate row weights) to 10 samples in [−2, 2] because we noticed most weights were in this range after running Algorithm 2.

heading: D.3 RUNNING TIME
text: We examine the runtimes of our various sketching algorithms. In Table D.1, the times are obtained for the LRA task with k = 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 onAtrain 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 as fast as possible.
Note that inference was timed using 1 matrix from Atest on an Nvidia Geforce GTX 1080 Ti GPU. The values were averaged over 10 trials.
We observe that sparse sketches (such as the ones used in GD only and Ours) have much lower inference runtimes than the dense sketches of exact SVD.

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

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.