# Learning the Positions in CountSketch

Yi Li111Nanyang Technological University<EMAIL_ADDRESS>
Honghao Lin 222Carnegie Mellon University<EMAIL_ADDRESS>Simin Liu
333Carnegie Mellon University<EMAIL_ADDRESS>Ali Vakilian444Toyota
Technological Institute at Chicago<EMAIL_ADDRESS>David P.
Woodruff555Carnegie Mellon University<EMAIL_ADDRESS>
###### Abstract

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

##  1 Introduction

The work of [IVY19] investigated learning-based sketching algorithms for low-
rank approximation. A sketching algorithm is a method of constructing
approximate solutions for optimization problems via summarizing the data. In
particular, linear sketching algorithms compress data by multiplication with a
sparse “sketch matrix” and then use just the compressed data to find an
approximate solution. Generally, this technique results in much faster or more
space-efficient algorithms for a fixed approximation error. The pioneering
work of [IVY19] shows it is possible to learn sketch matrices for low-rank
approximation (LRA) with better average performance than classical sketches.

In this model, we assume inputs come from an unknown distribution and learn a
sketch matrix with strong expected performance over the distribution. This
distributional assumption is often realistic – there are many situations where
a sketching algorithm is applied to a large batch of related data. For
example, genomics researchers might sketch DNA from different individuals,
which is known to exhibit strong commonalities. The high-performance computing
industry also uses sketching, e.g., researchers at NVIDIA have created
standard implementations of sketching algorithms for CUDA, a widely used GPU
library. They investigated the (classical) sketched singular value
decomposition (SVD), but found that the solutions were not accurate enough
across a spectrum of inputs [CB19]. This is precisely the issue addressed by
the learned sketch paradigm where we optimize for “good” average performance
across a range of inputs.

While promising results have been shown using previous learned sketching
techniques, notable gaps remain. In particular, all previous methods work by
initializing the sketching matrix with a random sparse matrix, e.g., each
column of the sketching matrix has a single non-zero value chosen at a
uniformly random position. Then, the values of the non-zero entries are
updated by running gradient descent on a training data set, or via other
methods. However, the locations of the non-zero entries are held fixed
throughout the entire training process.

Clearly this is sub-optimal. Indeed, suppose the input matrix $A$ is an
$n\times d$ matrix with first $d$ rows equal to the $d\times d$ identity
matrix, and remaining rows equal to $0$. A random sketching matrix $S$ with a
single non-zero per column is known to require $m=\Omega(d^{2})$ rows in order
for $S\cdot A$ to preserve the rank of $A$ [NN14]; this follows by a birthday
paradox argument. On the other hand, it is clear that if $S$ is a $d\times n$
matrix with first $d$ rows equal to the identity matrix, then $\|S\cdot
Ax\|_{2}=\|Ax\|_{2}$ for all vectors $x$, and so $S$ preserves not only the
rank of $A$ but all important spectral properties. A random matrix would be
very unlikely to choose the non-zero entries in the first $d$ columns of $S$
so perfectly, whereas an algorithm trained to optimize the locations of the
non-zero entries would notice and correct for this. This is precisely the gap
in our understanding that we seek to fill.

#### Learned CountSketch Paradigm of [IVY19].

Throughout the paper, we assume our data $A\in\mathbbm{R}^{n\times d}$ is
sampled from an unknown distribution $\mathcal{D}$. Specifically, we have a
training set $\mathsf{Tr}=\left\\{A_{1},\ldots,A_{N}\right\\}\in\mathcal{D}$.
The generic form of our optimization problems is $\min_{X}f(A,X)$, where
$A\in\mathbbm{R}^{n\times d}$ is the input matrix. For a given optimization
problem and a set $\mathcal{S}$ of sketching matrices, define
$\mathsf{ALG}(\mathcal{S},A)$ to be the output of the classical sketching
algorithm resulting from using $\mathcal{S}$; this uses the sketching matrices
in $\mathcal{S}$ to map the given input $A$ and construct an approximate
solution $\hat{X}$. We remark that the number of sketches used by an algorithm
can vary and in its simplest case, $\mathcal{S}$ is a single sketch, but in
more complicated sketching approaches we may need to apply sketching more than
once—hence $\mathcal{S}$ may also denote a set of more than one sketching
matrix.

The learned sketch framework has two parts: (1) offline sketch learning and
(2) “online” sketching (i.e., applying the learned sketch and some sketching
algorithm to possibly unseen data). In offline sketch learning, the goal is to
construct a CountSketch matrix (abbreviated as $\mathsf{CS}$ matrix) with the
minimum expected error for the problem of interest. Formally, that is,

$\operatorname*{arg\,min}_{\mathsf{CS}~{}S}\operatorname*{{\mathbf{E}}}_{A\in\mathsf{Tr}}f(A,\mathsf{ALG}(S,A))-f(A,X^{*})=\operatorname*{arg\,min}_{\mathsf{CS}~{}S}\operatorname*{{\mathbf{E}}}_{A\in\mathsf{Tr}}f(A,\mathsf{ALG}(S,A)),\;$

where $X^{*}$ denotes the optimal solution. Moreover, the minimum is taken
over all possible constructions of $\mathsf{CS}$. We remark that when
$\mathsf{ALG}$ needs more than one $\mathsf{CS}$ to be learned (e.g., in the
sketching algorithm we consider for LRA), we optimize each $\mathsf{CS}$
independently using a surrogate loss function.

In the second part of the learned sketch paradigm, we take the sketch from
part one and use it within a sketching algorithm. This learned sketch and
sketching algorithm can be applied, again and again, to different inputs.
Finally, we augment the sketching algorithm to provide worst-case guarantees
when used with learned sketches. The goal is to have good performance on
$A\in\mathcal{D}$ while the worst-case performance on $A\not\in\mathcal{D}$
remains comparable to the guarantees of classical sketches. We remark that the
learned matrix $S$ is trained offline only once using the training data.
Hence, no additional computational cost is incurred when solving the
optimization problem on the test data.

#### Our Results.

In this work, in addition to learning the values of the non-zero entries, we
learn the locations of the non-zero entries. Namely, we propose three
algorithms that learn the locations of the non-zero entries in CountSketch.
Our first algorithm (Section 4) is based on a greedy search. The empirical
result shows that this approach can achieve a good performance. Further, we
show that the greedy algorithm is provably beneficial for LRA when inputs
follow a certain input distribution (Section F). However, one drawback of the
greedy algorithm is its much slower training time. We then fix this issue and
propose two specific approaches for optimizing the positions for the sketches
for low-rank approximation and second-order optimization, which run much
faster than all previous algorithms while achieving better performance.

For low-rank approximation, our approach is based on first sampling a small
set of rows based on their ridge leverage scores, assigning each of these
sampled rows to a unique hash bucket, and then placing each non-sampled
remaining row in the hash bucket containing the sampled row for which it is
most similar to, i.e., for which it has the largest dot product with. We also
show that the worst-case guarantee of this approach is strictly better than
that of the classical Count-Sketch (see Section 5).

For sketch-based second-order optimization where we focus on the case that
$n\gg d$, we observe that the actual property of the sketch matrix we need is
the subspace embedding property. We next optimize this property of the sketch
matrix. We provably show that the sketch matrix $S$ needs fewer rows, with
optimized positions of the non-zero entries, when the input matrix $A$ has a
small number of rows with a heavy leverage score. More precisely, while
CountSketch takes $O(d^{2}/(\delta\epsilon^{2}))$ rows with failure
probability $\delta$, in our construction, $S$ requires only
$O((d\mathop{\mathrm{polylog}}(1/\epsilon)+\log(1/\delta))/\epsilon^{2})$ rows
if $A$ has at most $d\mathop{\mathrm{polylog}}(1/\epsilon)/\epsilon^{2}$ rows
with leverage score at least $\epsilon/d$. This is a quadratic improvement in
$d$ and an exponential improvement in $\delta$. In practice, it is not
necessary to calculate the leverage scores. Instead, we show in our
experiments that the indices of the rows of heavy leverage score can be
learned and the induced $S$ is still accurate. We also consider a new learning
objective, that is, we directly optimize the subspace embedding property of
the sketching matrix instead of optimizing the error in the objective function
of the optimization problem in hand. This demonstrates a significant advantage
over non-learned sketches, and has a fast training time (Section 6).

We show strong empirical results for real-world datasets. For low-rank
approximation, our methods reduce the errors by $70\%$ than classical sketches
under the same sketch size, while we reduce the errors by $30\%$ than previous
learning-based sketches. For second-order optimization, we show that the
convergence rate can be reduced by 87% over the non-learned CountSketch for
the LASSO problem on a real-world dataset. We also evaluate our approaches in
the few-shot learning setting where we only have a limited amount of training
data [IWW21]. We show our approach reduces the error significantly even if we
only have one training matrix (Sections 7 and 8). This approach clearly runs
faster than all previous methods.

#### Additional Related Work.

In the last few years, there has been much work on leveraging machine learning
techniques to improve classical algorithms. We only mention a few examples
here which are based on learned sketches. One related body of work is data-
dependent dimensionality reduction, such as an approach for pair-wise/multi-
wise similarity preservation for indexing big data [WZSS17], learned sketching
for streaming problems [IVY19, AIV19, JLL+20, CGP20, EIN+21, IWW21], learned
algorithms for nearest neighbor search [DIRW20], and a method for learning
linear projections for general applications [HSYB15]. While we also learn
linear embeddings, our embeddings are optimized for the specific application
of low rank approximation. In fact, one of our central challenges is that the
theory and practice of learned sketches generally needs to be tailored to each
application. Our work builds off of [IVY19], which introduced gradient descent
optimization for LRA, but a major difference is that we also optimize the
locations of the non-zero entries.

##  2 Preliminaries

Notation. Denote the canonical basis vectors of $\mathbb{R}^{n}$ by
$e_{1},\dots,e_{n}$. Suppose that $A$ has singular value decomposition (SVD)
$A=U\Sigma V^{\top}$. Define $\left[A\right]_{k}=U_{k}\Sigma_{k}V^{\top}_{k}$
to be the optimal $\operatorname{rank-\mathnormal{k}}$ approximation to $A$,
computed by the truncated SVD. Also, define the Moore-Penrose pseudo-inverse
of $A$ to be $A^{\dagger}=V\Sigma^{-1}U^{\top}$, where $\Sigma^{-1}$ is
constructed by inverting the non-zero diagonal entries. Let $\mathsf{row}(A)$
and $\mathsf{col}(A)$ be the row space and the column space of $A$,
respectively.

CountSketch. We define $S_{C}\in\mathbbm{R}^{m\times n}$ as a classical
CountSketch (abbreviated as $\mathsf{CS}$). It is a sparse matrix with one
nonzero entry from $\\{\pm 1\\}$ per column. The position and value of this
nonzero entry are chosen uniformly at random. CountSketch matrices can be
succinctly represented by two vectors. We define
$p\in[m]^{n},v\in\mathbbm{R}^{n}$ as the positions and values of the nonzero
entries, respectively. Further, we let $\mathsf{CS}(p,v)$ be the CountSketch
constructed from vectors $p$ and $v$.

Below we define the objective function $f(\cdot,\cdot)$ and a classical
sketching algorithm $\mathsf{ALG}(\mathcal{S},A)$ for each individual problem.

Low-rank approximation (LRA). In LRA, we find a
$\operatorname{rank-\mathnormal{k}}$ approximation of our data that minimizes
the Frobenius norm of the approximation error. For $A\in\mathbb{R}^{n\times
d}$,
$\min_{\operatorname{rank-\mathnormal{k}}B}f_{\mathrm{LRA}}(A,B)=\min_{\operatorname{rank-\mathnormal{k}}X}\left\lVert
A-B\right\rVert_{F}^{2}$. Usually, instead of outputting the a whole
$B\in\mathbb{R}^{n\times d}$, the algorithm outputs two factors
$Y\in\mathbb{R}^{n\times k}$ and $X\in\mathbb{R}^{k\times d}$ such that $B=YX$
for efficiency.

The authors of [IVY19] considered Algorithm 1, which only compresses one side
of the input matrix $A$. However, in practice often both dimensions of the
matrix $A$ are large. Hence, in this work we consider Algorithm 2 that
compresses both sides of $A$.

Constrained regression. Given a vector $b\in\mathbb{R}^{n}$, a matrix
$A\in\mathbb{R}^{n\times d}$ ($n\gg d$) and a convex set $\mathcal{C}$, we
want to find $x$ to minimize the squared error

$\min_{x\in\mathcal{C}}f_{\mathrm{REG}}([A\ \
b],X)=\min_{x\in\mathcal{C}}\left\lVert Ax-b\right\rVert_{2}^{2}.$ (2.1)

Iterative Hessian Sketch. The Iterative Hessian Sketching (IHS) method [PW16]
solves the constrained least-squares problem by iteratively performing the
update

$x_{t+1}=\operatorname*{arg\,min}_{x\in\mathcal{C}}\left\\{\frac{1}{2}\left\lVert
S_{t+1}A(x-x_{t})\right\rVert_{2}^{2}-\langle
A^{\top}(b-Ax_{t}),x-x_{t}\rangle\right\\},$ (2.2)

where $S_{t+1}$ is a sketching matrix. It is not difficult to see that for the
unsketched version ($S_{t+1}$ is the identity matrix) of (2.2), the optimal
solution $x^{t+1}$ coincides with the optimal solution to the original
constrained regression problem (2.1). The IHS approximates the Hessian
$A^{\top}A$ by a sketched version $(S_{t+1}A)^{\top}(S_{t+1}A)$ to improve
runtime, as $S_{t+1}A$ typically has very few rows.

0: $\mathbf{A}\in\mathbb{R}^{n\times d}$, $S\in\mathbb{R}^{{m}\times n}$

1: $U,\Sigma,V^{\top}\leftarrow\text{\sc CompactSVD}(SA)$ $\quad\rhd$
{$r=\mathrm{rank}(SA),U\in\mathbb{R}^{m\times r},V\in\mathbb{R}^{d\times r}$}

2: Return: $[AV]_{k}V^{\top}$

Algorithm 1 Rank-$k$ approximation of $A$ using a sketch $S$ (see [CW09, Sec.
4.1.1])

0: $A\in\mathbb{R}^{n\times d}$, $S\in\mathbb{R}^{m_{S}\times n}$,
$R\in\mathbb{R}^{m_{R}\times d}$, $V\in\mathbb{R}^{m_{V}\times n}$,
$W\in\mathbb{R}^{m_{W}\times d}$

1: $U_{C}\begin{bmatrix}T_{C}&T_{C}^{\prime}\end{bmatrix}\leftarrow
VAR^{\top}$, $\begin{bmatrix}T_{D}^{\top}\\\
T_{D}^{\prime\top}\end{bmatrix}U^{\top}_{D}\leftarrow SAW^{\top}$ with
$U_{C},U_{D}$ orthogonal

2: $G\leftarrow VAW^{\top}$,
$Z_{L}^{\prime}Z_{R}^{\prime}\leftarrow[U^{\top}_{C}GU_{D}]_{k}$

3:
$Z_{L}\leftarrow\begin{bmatrix}Z_{L}^{\prime}(T_{D}^{-1})^{\top}&0\end{bmatrix},Z_{R}\leftarrow\begin{bmatrix}T_{C}^{-1}Z_{R}^{\prime}\\\
0\end{bmatrix}$

4: $Z\leftarrow Z_{L}Z_{R}$

5: return: $AR^{\top}ZSA$ in form $P_{n\times k},Q_{k\times d}$

Algorithm 2 $\mathsf{ALG}_{\mathrm{LRA}}$(Sketch-LowRank) [Sar06, CW17,
ACW17].

Learning-Based Algorithms in the Few-Shot Setting. Recently, [IWW21] studied
learning-based algorithms for LRA in the setting where we have access to
limited data or computing resources. Below we provide a brief explanation of
learning-based algorithms in the Few-Shot setting.

#### One-shot closed-form algorithm.

Given a sparsity pattern of a Count-Sketch matrix $S\in\mathbb{R}^{m\times
n}$, it partitions the rows of $A$ into $m$ blocks $A^{(1)},...,A^{(m)}$ as
follows: let $I_{i}=\\{j:S_{ij}=1\\}$. The block
$A^{(i)}\in\mathbb{R}^{|I_{i}|\times d}$ is the sub-matrix of $A$ that
contains the rows whose indices are in $I_{i}$. The goal here is for each
block $A^{(i)}$, to choose a (non-sparse) one-dimensional sketching vector
$s_{i}\in\mathbb{R}^{|I_{i}|}$. The first approach is to set $s_{i}$ to be the
top left singular vector of $A^{(i)}$, which is the algorithm 1Shot2Vec.
Another approach is to set $s_{i}$ to be a left singular vector of $A^{(i)}$
chosen randomly and proportional to its squared singular value. The main
advantage of the latter approach over the previous one is that it endows the
algorithm with provable guarantees on the LRA error. The 1Shot2Vec algorithm
combines both ways, obtaining the benefits of both approaches. The advantage
of these two algorithms is that they extract a sketching matrix by an analytic
computation, requiring neither GPU access nor auto-gradient functionality.

#### Few-shot SGD algorithm.

In this algorithm, the authors propose a new loss function for LRA, namely,

$\min_{\mathsf{CS}~{}S}\operatorname*{{\mathbf{E}}}_{A\in\mathsf{Tr}}\left\lVert
U_{k}^{\top}S^{\top}SU-I_{0}\right\rVert_{F}^{2}\;,$

where $A=U\Sigma V^{\top}$ is the SVD-decomposition of $A$ and
$U_{k}\in\mathbb{R}^{n\times k}$ denotes the submatrix of $U$ that contains
its first $k$ columns. $I_{0}\in\mathbb{R}^{k\times d}$ denotes the result of
augmenting the identity matrix of order $k$ with $d-k$ additional zero columns
on the right. This loss function is motivated by the analysis of prior LRA
algorithms that use random sketching matrices. It is faster to compute and
differentiate than the previous empirical loss in [IVY19]. In the experiments
the authors also show that this loss function can achieve a smaller error in a
shorter amount of time, using a small number of randomly sampled training
matrices, though the final error will be larger than that of the previous
algorithm in [IVY19] if we allow a longer training time and access to the
whole training set $\mathsf{Tr}$.

Leverage Scores and Ridge Leverage Scores. Given a matrix $A$, the leverage
score of the $i$-th row $a_{i}$ of $A$ is defined to be
$\tau_{i}:=a_{i}(A^{\top}A)^{\dagger}a_{i}^{\top}$, which is the squared
$\ell_{2}$-norm of the $i$-th row of $U$, where $A=U\Sigma V^{T}$ is the
singular value decomposition of $A$. Given a regularization parameter
$\lambda$, the ridge leverage score of the $i$-th row $a_{i}$ of $A$ is
defined to be $\tau_{i}:=a_{i}(A^{\top}A+\lambda I)^{\dagger}a_{i}^{\top}$.
Our learning-based algorithms employs the ridge leverage score sampling
technique proposed in [CMM17], which shows that sampling proportional to ridge
leverage scores gives a good solution to LRA.

##  3 Description of Our Approach

We describe our contributions to the learning-based sketching paradigm which,
as mentioned, is to learn the locations of the non-zero values in the sketch
matrix. To learn a CountSketch for the given training data set, we locally
optimize the following in two stages:

$\min_{S}\operatorname*{{\mathbf{E}}}_{A\in\mathcal{D}}\left[f(A,\mathsf{ALG}(S,A))\right].$
(3.1)

(1) compute the positions of the non-zero entries, then (2) fix the positions
and optimize their values.

Stage 1: Optimizing Positions. In Section 4, we provide a greedy search
algorithm for this stage, as our starting point. In Section 5 and 6, we
provide our specific approaches for optimizing the positions for the sketches
for low-rank approximation and second-order optimization.

Stage 2: Optimizing Values. This stage is similar to the approach of [IVY19].
However, instead of the power method, we use an automatic differentiation
package, PyTorch [PGM+19], and we pass it our objective
$\textstyle\min_{v\in\mathbbm{R}^{n}}\operatorname*{{\mathbf{E}}}_{A\in\mathcal{D}}\left[f(A,\mathsf{ALG}(\mathsf{CS}(p,v),A))\right]$,
implemented as a chain of differentiable operations. It will automatically
compute the gradient using the chain rule. We also consider new approaches to
optimize the values for LRA (proposed in [IWW21], see Appendix A.3 for
details) and second-order optimization (proposed in Section 6).

Worst-Cases Guarantees. In Appendix D, we show that both of our approaches for
the above two problems can perform no worse than a classical sketching matrix
when $A$ does not follow the distribution $\mathcal{D}$. In particular, for
LRA, we show that the sketch monotonicity property holds for the time-optimal
sketching algorithm for low rank approximation. For second-order optimization,
we propose an algorithm which runs in input-sparsity time and can test for and
use the better of a random sketch and a learned sketch.

0: $f,\mathsf{ALG},\mathsf{Tr}=\\{A_{1},...,A_{N}\in\mathbb{R}^{n\times
d}\\}$; sketch dimension $m$

1: initialize $S_{L}=\mathbbm{O}^{m\times n}$

2: for $i=1$ to $n$ do

3:
$\displaystyle\overline{j}\\!\leftarrow\\!\operatorname*{arg\,min}_{j\in[m]}\\!\\!\sum_{A\in\mathsf{Tr}}\\!f(A,\mathsf{ALG}(S_{L}\pm
e_{j}e_{i}^{\top}\\!,A))$

4: $S_{L}\leftarrow S_{L}\pm(e_{\overline{j}}e_{i}^{\top})$

5: end for

6: return $p$ for $S_{L}=\mathsf{CS}(p,v)$

Algorithm 3 Position optimization: Greedy Search

##  4 Sketch Learning: Greedy Search

When $S$ is a CountSketch, computing $SA$ amounts to hashing the $n$ rows of
$A$ into the $m\ll n$ rows of $SA$. The optimization is a combinatorial
optimization problem with an empirical risk minimization (ERM) objective. The
naïve solution is to compute the objective value of the exponentially many
($m^{n}$) possible placements, but this is clearly intractable. Instead, we
iteratively construct a full placement in a greedy fashion. We start with $S$
as a zero matrix. Then, we iterate through the columns of $S$ in an order
determined by the algorithm, adding a nonzero entry to each. The best position
in each column is the one that minimizes Eq. (3.1) if an entry were to be
added there. For each column, we evaluate Eq. (3.1) $\mathcal{O}(m)$ times,
once for each prospective half-built sketch.

While this greedy strategy is simple to state, additional tactics are required
for each problem to make it more tractable. Usually the objective evaluation
(Algorithm 3, line 3) is too slow, so we must leverage our insight into their
sketching algorithms to pick a proxy objective. Note that we can reuse these
proxies for value optimization, since they may make gradient computation
faster too.

Proxy objective for LRA. For the two-sided sketching algorithm, we can assume
that the two factors $X,Y$ has the form $Y=AR^{\top}\tilde{Y}$ and
$X=\tilde{X}SA$, where $S$ and $R$ are both $\mathsf{CS}$ matrices, so we
optimize the positions in both $S$ and $R$. We cannot use
$f(A,\mathsf{ALG}(S,R,A))$ as our objective because then we would have to
consider combinations of placements between $S$ and $R$. To find a proxy, we
note that a prerequisite for good performance is for $\mathsf{row}(SA)$ and
$\mathsf{col}(AR^{\top})$ to both contain a good
$\operatorname{rank-\mathnormal{k}}$ approximation to $A$ (see proof of Lemma
C.5). Thus, we can decouple the optimization of $S$ and $R$. The proxy
objective for $S$ is $\left\lVert[AV]_{k}V^{\top}-A\right\rVert_{F}^{2}$ where
$SA=U\Sigma V^{\top}$. In this expression, $\hat{X}=[AV]_{k}V^{\top}$ is the
best $\operatorname{rank-\mathnormal{k}}$ approximation to $A$ in
$\mathsf{row}(SA)$. The proxy objective for $R$ is defined analogously.

In Appendix F, we show the greedy algorithm is provably beneficial for LRA
when inputs follow the spiked covariance or the Zipfian distribution. Despite
the good empirical performance we present in Section 7, one drawback is its
much slower training time. Also, for the iterative sketching method for
second-order optimization, it is non-trivial to find a proxy objective because
the input of the $i$-th iteration depends on the solution to the $(i-1)$-th
iteration, for which the greedy approach sometimes does not give a good
solution. In the next section, we will propose our specific approach for
optimizing the positions of the sketches for low-rank approximation and
second-order optimization, both of which achieve a very high accuracy and can
finish in a very short amount of time.

##  5 Sketch Learning: Low-Rank Approximation

0: $A\in\mathbb{R}^{n\times d}$: average of $\mathsf{Tr}$; sketch dim. $m$

1: initialize $S_{1},S_{2}=\mathbbm{O}^{m\times n}$

2: Sample a set $C=\\{C_{1}\cdots C_{m}\\}$ of rows using ridge leverage score
sampling (see Section 2).

3: for $i=1$ to $n$ do

4: $p_{i},v_{i}\leftarrow\operatorname*{arg\,max}\limits_{p\in[m],v\in\\{\pm
1\\}}\langle\frac{C_{p}}{\|C_{p}\|_{2}},v\frac{A_{i}}{\|A_{i}\|_{2}}\rangle$

5: $S_{1}[p_{i},i]\leftarrow v_{i}$

6: end for

7: for $i=1$ to $m$ do

8: $I_{i}\leftarrow\\{j\ |\ p_{j}=i\\}$

9: $A^{(i)}\leftarrow$ restriction of $A$ to rows in $I_{i}$

10: $u_{i}\leftarrow$ the top left singular vector of $A^{(i)}$

11: $S_{1}[i,I_{i}]\leftarrow u_{i}^{\top}$

12: end for

13: for $i=1$ to $m$ do

14: $q_{i}\leftarrow$ index such that $C_{i}$ is the $q_{i}$-th row of $A$

15: $S_{2}[i,q_{i}]\leftarrow 1$

16: end for

17: return $S_{1}$ or $[\begin{smallmatrix}S_{1}\\\ S_{2}\end{smallmatrix}]$

Algorithm 4 Position optimization: Inner Product

Now we present a conceptually new algorithm which runs much faster and
empirically achieves similar error bounds as the greedy search approach.
Moreover, we show that this algorithm has strictly better guarantees than the
classical Count-Sketch.

To achieve this, we need a more careful analysis. To provide some intuition,
if $\mathrm{rank}(SA)=k$ and $SA=U\Sigma V^{\top}$, then the rank-$k$
approximation cost is exactly $\left\lVert AVV^{\top}-A\right\rVert_{F}^{2}$,
the projection cost onto $\mathsf{col}(V)$. Minimizing it is equivalent to
maximizing the sum of squared projection coefficients:

$\operatorname*{arg\,min}_{S}\left\lVert
A-AVV^{\top}\right\rVert_{F}^{2}=\operatorname*{arg\,min}_{S}\sum_{i\in[n]}(\left\lVert
A_{i}\right\rVert_{2}^{2}-\sum_{j\in[k]}\langle
A_{i},v_{j}\rangle^{2})=\operatorname*{arg\,max}_{S}\sum_{i\in[n]}\sum_{j\in[k]}\langle
A_{i},v_{j}\rangle^{2}.$

As mentioned, computing $SA$ actually amounts to hashing the $n$ rows of $A$
to the $m$ rows of $SA$. Hence, intuitively, if we can put similar rows into
the same bucket, we may get a smaller error.

Our algorithm is given in Algorithm 4. Suppose that we want to form matrix $S$
with $m$ rows. At the beginning of the algorithm, we sample $m$ rows according
to the ridge leverage scores of $A$. By the property of the ridge leverage
score, the subspace spanned by this set of sampled rows contains an
approximately optimal solution to the low rank approximation problem. Hence,
we map these rows to separate “buckets” of $SA$. Then, we need to decide the
locations of the remaining rows (i.e., the non-sampled rows). Ideally, we want
similar rows to be mapped into the same bucket. To achieve this, we use the
$m$ sampled rows as reference points and assign each (non-sampled) row $A_{i}$
to the $p$-th bucket in $SA$ if the normalized row $A_{i}$ and $C_{p}$ have
the largest inner product (among all possible buckets).

Once the locations of the non-zero entries are fixed, the next step is to
determine the values of these entries. We follow the same idea proposed in
[IWW21]: for each block $A^{(i)}$, one natural approach is to choose the unit
vector $s_{i}\in\mathbb{R}^{|I_{i}|}$ that preserves as much of the Frobenius
norm of $A^{(i)}$ as possible, i.e., to maximize $\left\lVert
s_{i}^{\top}A^{(i)}\right\rVert_{2}^{2}$. Hence, we set $s_{i}$ to be the top
left singular vector of $A^{(i)}$. In our experiments, we observe that this
step reduces the error of downstream value optimizations performed by SGD.

To obtain a worst-case guarantee, we show that w.h.p., the row span of the
sampled rows $C_{i}$ is a good subspace. We set the matrix $S_{2}$ to be the
sampling matrix that samples $C_{i}$. The final output of our algorithm is the
vertical concatenation of $S_{1}$ and $S_{2}$. Here $S_{1}$ performs well
empirically, while $S_{2}$ has a worst-case guarantee for any input.

Combining Lemma E.2 and the sketch monotonicity for low rank approximation in
Section D, we get that $O(k\log k+k/\epsilon)$ rows is enough for a
$(1\pm\epsilon)$-approximation for the input matrix $A$ induced by
$\mathsf{Tr}$, which is better than the $\Omega(k^{2})$ rows required of a
non-learned Count-Sketch, even if its non-zero values have been further
improved by the previous learning-based algorithms in [IVY19, IWW21]. As a
result, under the assumption of the input data, we may expect that $S$ will
still be good for the test data. We defer the proof to Appendix E.1.

In Appendix A, we shall show that the assumptions we make in Theorem 5.1 are
reasonable. We also provide an empirical comparison between Algorithm 4 and
some of its variants, as well as some adaptive sketching methods on the
training sample. The evaluation result shows that only our algorithm has a
significant improvement for the test data, which suggests that both ridge
leverage score sampling and row bucketing are essential.

######  Theorem 5.1.

Let $S\in\mathbb{R}^{2m\times n}$ be given by concatenating the sketching
matrices $S_{1},S_{2}$ computed by Algorithm 4 with input $A$ induced by
$\mathsf{Tr}$ and let $B\in\mathbb{R}^{n\times d}$. Then with probability at
least $1-\delta$, we have
$\min_{\operatorname{rank-\mathnormal{k}}X:\mathsf{row}(X)\subseteq\mathsf{row}(SB)}\left\lVert
B-X\right\rVert_{F}^{2}\leq(1+\epsilon)\left\lVert
B-B_{k}\right\rVert_{F}^{2}$ if one of the following holds:.

  1. 1.

$m=O(\beta\cdot(k\log k+k/\epsilon))$, $\delta=0.1$, and
$\tau_{i}(B)\geq\frac{1}{\beta}\tau_{i}(A)$ for all $i\in[n]$.

  2. 2.

$m=O(k\log k+k/\epsilon)$, $\delta=0.1+1.1\beta$, and the total variation
distance $d_{\mathrm{tv}}(p,q)\leq\beta$, where $p,q$ are sampling
probabilities defined as $p_{i}=\frac{\tau_{i}(A)}{\sum_{i}\tau_{i}(A)}$ and
$q_{i}=\frac{\tau_{i}(B)}{\sum_{i}\tau_{i}(B)}$.

Time Complexity. As mentioned, an advantage of our second approach is that it
significantly reduces the training time. We now discuss the training times of
different algorithms. For the value-learning algorithms in [IVY19], each
iteration requires computing a differentiable SVD to perform gradient descent,
hence the runtime is at least $\Omega(n_{it}\cdot T)$, where $n_{it}$ is the
number of iterations (usually set $>500$) and $T$ is the time to compute an
SVD. For the greedy algorithm, there are $m$ choices for each column, hence
the runtime is at least $\Omega(mn\cdot T)$. For our second approach, the most
complicated step is to compute the ridge leverage scores of $A$ and then the
SVD of each submatrix. Hence, the total runtime is at most $O(T)$. We note
that the time complexities discussed here are all for training time. There is
no additional runtime cost for the test data.

##  6 Sketch Learning: Second-Order Optimization

In this section, we consider optimizing the sketch matrix in the context of
second-order methods. The key observation is that for many sketching-based
second-order methods, the crucial property of the sketching matrix is the so-
called subspace embedding property: for a matrix $A\in\mathbb{R}^{n\times d}$,
we say a matrix $S\in\mathbb{R}^{m\times n}$ is a $(1\pm\epsilon)$-subspace
embedding for the column space of $A$ if $(1-\epsilon)\left\lVert
Ax\right\rVert_{2}\leq\left\lVert
SAx\right\rVert_{2}\leq(1+\epsilon)\left\lVert Ax\right\rVert_{2}$ for all
$x\in\mathbb{R}^{d}$. For example, consider the iterative Hessian sketch,
which performs the update (2.2) to compute $\\{x_{t}\\}_{t}$. [PW16] showed
that if $S_{1},\dots,S_{t+1}$ are $(1+O(\rho))$-subspace embeddings of $A$,
then $\left\lVert A(x^{t}-x^{*})\right\rVert_{2}\leq\rho^{t}\left\lVert
Ax^{*}\right\rVert_{2}$. Thus, if $S_{i}$ is a good subspace embedding of $A$
and we will have a good convergence guarantee. Therefore, unlike [IVY19],
which treats the training objective in a black-box manner, we shall optimize
the subspace embedding property of the matrix $A$.

Optimizing positions. We consider the case that $A$ has a few rows of large
leverage score, as well as access to an oracle which reveals a _superset_ of
the indices of such rows. Formally, let $\tau_{i}(A)$ be the leverage score of
the $i$-th row of $A$ and $I^{\ast}=\left\\{i:\tau_{i}(A)\geq\nu\right\\}$ be
the set of rows with large leverage score. Suppose that a superset $I\supseteq
I^{\ast}$ is known to the algorithm. In the experiments we train an oracle to
predict such rows. We can maintain all rows in $I$ explicitly and apply a
Count-Sketch to the remaining rows, i.e., the rows in $[n]\setminus I$. Up to
permutation of the rows, we can write

$A=\begin{pmatrix}A_{I}\\\ A_{I^{c}}\end{pmatrix}\quad\text{and}\quad
S=\begin{pmatrix}I&0\\\ 0&S^{\prime}\end{pmatrix},$ (6.1)

where $S^{\prime}$ is a random Count-Sketch matrix of $m$ rows. Clearly $S$
has a single non-zero entry per column. We have the following theorem, whose
proof is postponed to Section E.2. Intuitively, the proof for Count-Sketch in
[CW17] handles rows of large leverage score and rows of small leverage score
separately. The rows of large leverage score are to be perfectly hashed while
the rows of small leverage score will concentrate in the sketch by the Hanson-
Wright inequality.

######  Theorem 6.1.

Let $\nu=\epsilon/d$. Suppose that
$m=O((d/\epsilon^{2})(\mathop{\mathrm{polylog}}(1/\epsilon)+\log(1/\delta)))$,
$\delta\in(0,1/m]$ and
$d=\Omega((1/\epsilon)\mathop{\mathrm{polylog}}(1/\epsilon)\log^{2}(1/\delta))$.
Then, there exists a distribution on $S$ of the form in (6.1) with $m+|I|$
rows such that $\Pr\big{\\{}\forall x\in\mathsf{col}(A),|\left\lVert
Sx\right\rVert_{2}^{2}-\left\lVert x\right\rVert_{2}^{2}|>\epsilon\left\lVert
x\right\rVert_{2}^{2}\big{\\}}\leq\delta$. In particular, when $\delta=1/m$,
the sketching matrix $S$ has
$O((d/\epsilon^{2})\mathop{\mathrm{polylog}}(d/\epsilon))$ rows.

Hence, if there happen to be at most
$d\mathop{\mathrm{polylog}}(1/\epsilon)/\epsilon^{2}$ rows of leverage score
at least $\epsilon/d$, the overall sketch length for embedding
$\operatorname{colsp}(A)$ can be reduced to
$O((d\mathop{\mathrm{polylog}}(1/\epsilon)+\log(1/\delta))/\epsilon^{2})$, a
quadratic improvement in $d$ and an exponential improvement in $\delta$ over
the original sketch length of $O(d^{2}/(\epsilon^{2}\delta))$ for Count-
Sketch. In the worst case there could be $O(d^{2}/\epsilon)$ such rows, though
empirically we do not observe this. In Section 8, we shall show it is possible
to learn the indices of the heavy rows for real-world data.

Optimizing values. When we fix the positions of the non-zero entries, we aim
to optimize the values by gradient descent. Rather than the previous black-box
way in [IVY19] that minimizes $\sum_{i}f(A,\mathsf{ALG}(S,A))$, we propose the
following objective loss function for the learning algorithm
$\mathcal{L}(S,\mathcal{A})=\sum_{A_{i}\in\mathcal{A}}\|(A_{i}R_{i})^{\top}A_{i}R_{i}-I\|_{F}$,
over all the training data, where $R_{i}$ comes from the QR decomposition of
$SA_{i}=Q_{i}R_{i}^{-1}$. The intuition for this loss function is given by the
lemma below, whose proof is deferred to Section E.3.

######  Lemma 6.2.

Suppose that $\epsilon\in(0,\frac{1}{2})$, $S\in\mathbb{R}^{m\times n}$,
$A\in\mathbb{R}^{n\times d}$ of full column rank, and $SA=QR$ is the QR-
decomposition of $SA$. If
$\|(AR^{-1})^{\top}AR^{-1}-I\|_{\mathrm{op}}\leq\epsilon$, then $S$ is a
$(1\pm\epsilon)$-subspace embedding of $\mathsf{col}(A)$.

Lemma 6.2 implies that if the loss function over
$\mathcal{A}_{\mathrm{train}}$ is small and the distribution of
$\mathcal{A}_{\mathrm{test}}$ is similar to $\mathcal{A}_{\mathrm{train}}$, it
is reasonable to expect that $S$ is a good subspace embedding of
$\mathcal{A}_{\mathrm{test}}$. Here we use the Frobenius norm rather than
operator norm in the loss function because it will make the optimization
problem easier to solve, and our empirical results also show that the
performance of the Frobenius norm is better than that of the operator norm.

##  7 Experiments: Low-Rank Approximation

In this section, we evaluate the empirical performance of our learning-based
approach for LRA on three datasets. For each, we fix the sketch size and
compare the approximation error $\left\lVert A-X\right\rVert_{F}-\left\lVert
A-A_{k}\right\rVert_{F}$ averaged over $10$ trials. In order to make position
optimization more efficient, in line 3 of Algorithm 3), instead of computing
many rank-$1$ SVD updates, we use formulas for fast rank-$1$ SVD updates
[Bra06]. For the greedy method, we used several Nvidia GeForce GTX 1080 Ti
machines. For the maximum inner product method, the experiments are conducted
on a laptop with a 1.90GHz CPU and 16GB RAM.

Datasets. We use the three datasets from [IVY19]: (1, 2) Friends, Logo
(image): frames from a short video of the TV show Friends and of a logo being
painted; (3) Hyper (image): hyperspectral images from natural scenes.
Additional details are in Table A.1.

Baselines. We compare our approach to the following baselines. Classical CS: a
random Count-Sketch. IVY19: a sparse sketch with learned values, and random
positions for the non-zero entries. Ours (greedy): a sparse sketch where both
the values and positions of the non-zero entries are learned. The positions
are learned by Algorithm 3. The values are learned similarly to [IVY19]. Ours
(inner product): a sparse sketch where both the values and the positions of
the non-zero entries are learned. The positions are learned by $S_{1}$ in
Algorithm 4. IVY19 and greedy algorithm use the full training set and our
Algorithm 4 takes the input as the average over the entire training matrix.

We also give a sensitivity analysis for our algorithm, where we compare our
algorithm with the following variants: Only row sampling (perform projection
by ridge leverage score sampling), $\ell_{2}$ sampling (Replace leverage score
sampling with $\ell_{2}$-norm row sampling and maintain the same downstream
step), and Randomly Grouping (Use ridge leverage score sampling but randomly
distribute the remaining rows). The result shows none of these variants
outperforms non-learned sketching. We defer the results of this part to
Appendix A.1.

Result Summary. Our empirical results are provided in Table LABEL:tab:LRA_2
for both Algorithm 2 and Algorithm 1, where the errors take an average over
$10$ trials. We use the average of all training matrices from $\mathsf{Tr}$,
as the input to the algorithm 4. We note that all the steps of our training
algorithms are done on the training data. Hence, no additional computational
cost is incurred for the sketching algorithm on the test data. Experimental
parameters (i.e., learning rate for gradient descent) can be found in Appendix
G. For both sketching algorithms, Ours are always the best of the four
sketches. It is significantly better than Classical CS, obtaining improvements
of around $70\%$. It also obtains a roughly $30\%$ improvement over IVY19.

Wall-Clock Times. The offline learning runtime is in Table
LABEL:tab:LRA_timing, which is the time to train a sketch on
$\mathcal{A}_{\mathrm{train}}$. We can see that although the greedy method
will take much longer (1h 45min), our second approach is much faster (5
seconds) than the previous algorithm in [IVY19] (3 min) and can still achieve
a similar error as the greedy algorithm. The reason is that Algorithm 4 only
needs to compute the ridge leverage scores on the training matrix once, which
is actually much cheaper than IVY19 which needs to compute a differentiable
SVD many times during gradient descent.

In Section A.3, we also study the performance of our approach in the few-shot
learning setting, which has been studied in [IWW21].

##  8 Experiments: Second-Order Optimization

In this section, we consider the IHS on the following instance of LASSO
regression:

$\textstyle x^{*}=\operatorname*{arg\,min}_{\left\lVert
x\right\rVert_{1}\leq\lambda}f(x)=\operatorname*{arg\,min}_{\left\lVert
x\right\rVert_{1}\leq\lambda}\frac{1}{2}\left\lVert Ax-b\right\rVert_{2}^{2},$
(8.1)

where $\lambda$ is a parameter. We also study the performance of the sketches
on the matrix estimation with a nuclear norm constraint problem, the fast
regression solver ([vdBPSW21]), as well as the use of sketches for first-order
methods. The results can be found in Appendix B. All of our experiments are
conducted on a laptop with a 1.90GHz CPU and 16GB RAM. The offline training is
done separately using a single GPU. The details of the implementation are
deferred to Appendix G.

Dataset. We use the
Electric666https://archive.ics.uci.edu/ml/datasets/ElectricityLoadDiagrams20112014
dataset of residential electric load measurements. Each row of the matrix
corresponds to a different residence. Matrix columns are consecutive
measurements at different times. Here $A^{i}\in\mathbb{R}^{370\times 9}$,
$b^{i}\in\mathbb{R}^{370\times 1}$, and $|(A,b)_{\mathrm{train}}|=320$,
$|(A,b)_{\mathrm{test}}|=80$. We set $\lambda=15$.

Experiment Setting. We compare the learned sketch against the classical Count-
Sketch777The framework of [IVY19] does not apply to the iterative sketching
methods in a straightforward manner, so here we only compare with the
classical CountSketch. For more details, please refer to Section B.. We choose
$m=6d,8d,10d$ and consider the error $f(x)-f(x^{*})$. For the heavy-row Count-
Sketch, we allocate 30% of the sketch space to the rows of the heavy row
candidates. For this dataset, each row represents a specific residence and
hence there is a strong pattern of the distribution of the heavy rows. We
select the heavy rows according to the number of times each row is heavy in
the training data. We give a detailed discussion about this in Appendix B.1.
We highlight that it is still possible to recognize the pattern of the rows
even if the row orders of the test data are permuted. We also consider
optimizing the non-zero values after identifying the heavy rows, using our new
approach in Section 6.

Results. We plot in Figures 8.1 the mean errors on a logarithmic scale. The
average offline training time is $3.67$s to find a superset of the heavy rows
over the training data and $66$s to optimize the values when $m=10d$, which
are both faster than the runtime of [IVY19] with the same parameters. Note
that the learned matrix $S$ is trained offline only once using the training
data. Hence, no additional computational cost is incurred when solving the
optimization problem on the test data.

We see all methods display linear convergence, that is, letting $e_{k}$ denote
the error in the $k$-th iteration, we have $e_{k}\approx\rho^{k}e_{1}$ for
some convergence rate $\rho$. A smaller convergence rate implies a faster
convergence. We calculate an estimated rate of convergence
$\rho=(e_{k}/e_{1})^{1/k}$ with $k=7$. We can see both sketches, especially
the sketch that optimizes both the positions and values, show significant
improvements. When the sketch size is small ($6d$), this sketch has a
convergence rate that is just 13.2% of that of the classical Count-Sketch, and
when the sketch size is large ($10d$), this sketch has a smaller convergence
rate that is just 12.1%.

Figure 8.1: Test error of LASSO in Electric dataset.

## Acknowledgements

Yi Li would like to thank for the partial support from the Singapore Ministry
of Education under Tier 1 grant RG75/21. Honghao Lin and David Woodruff were
supported in part by an Office of Naval Research (ONR) grant N00014-18-1-2562.
Ali Vakilian was supported by NSF award CCF-1934843.

## References

  * [AAB+19] Akshay Agrawal, Brandon Amos, Shane T. Barratt, Stephen P. Boyd, Steven Diamond, and J. Zico Kolter.  Differentiable convex optimization layers.  In Advances in Neural Information Processing Systems, pages 9558–9570, 2019.
  * [ACW17] Haim Avron, Kenneth L. Clarkson, and David P. Woodruff.  Sharper bounds for regularized data fitting.  In Approximation, Randomization, and Combinatorial Optimization. Algorithms and Techniques, (APPROX/RANDOM), pages 27:1–27:22, 2017.
  * [AIV19] Anders Aamand, Piotr Indyk, and Ali Vakilian.  (learned) frequency estimation algorithms under zipfian distribution.  arXiv preprint arXiv:1908.05198, 2019.
  * [BDN15] Jean Bourgain, Sjoerd Dirksen, and Jelani Nelson.  Toward a unified theory of sparse dimensionality reduction in Euclidean space.  Geometric and Functional Analysis, pages 1009–1088, 2015.
  * [Bra06] Matthew Brand.  Fast low-rank modifications of the thin singular value decomposition.  Linear Algebra and its Applications, 415.1, 2006.
  * [CB19] Lung-Sheng Chien and Samuel Rodriguez Bernabeu.  Fast singular value decomposition on gpu.  NVIDIA presentation at GPU Technology Conference, 2019.
  * [CD19] Graham Cormode and Charlie Dickens.  Iterative hessian sketch in input sparsity time.  In Proceedings of 33rd Conference on Neural Information Processing Systems (NeurIPS), Vancouver, Canada, 2019.
  * [CGP20] Edith Cohen, Ofir Geri, and Rasmus Pagh.  Composable sketches for functions of frequencies: Beyond the worst case.  In International Conference on Machine Learning, pages 2057–2067. PMLR, 2020.
  * [CMM17] Michael B. Cohen, Cameron Musco, and Christopher Musco.  Input sparsity time low-rank approximation via ridge leverage score sampling.  In Proceedings of the Twenty-Eighth Annual ACM-SIAM Symposium on Discrete Algorithms, (SODA), pages 1758–1777, 2017.
  * [CW09] 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), pages 205–214, 2009.
  * [CW17] Kenneth L Clarkson and David P Woodruff.  Low-rank approximation and regression in input sparsity time.  Journal of the ACM (JACM), 63(6):54, 2017.
  * [DIRW20] Yihe Dong, Piotr Indyk, Ilya Razenshteyn, and Tal Wagner.  Learning sublinear-time indexing for nearest neighbor search.  In International Conference on Learning Representations, 2020.
  * [EIN+21] Talya Eden, Piotr Indyk, Shyam Narayanan, Ronitt Rubinfeld, Sandeep Silwal, and Tal Wagner.  Learning-based support estimation in sublinear time.  In International Conference on Learning Representations, 2021.
  * [HSYB15] Chinmay Hegde, Aswin C. Sankaranarayanan, Wotao Yin, and Richard G. Baraniuk.  Numax: A convex approach for learning near-isometric linear embeddings.  In IEEE Transactions on Signal Processing, pages 6109–6121, 2015\.
  * [IVY19] Piotr Indyk, Ali Vakilian, and Yang Yuan.  Learning-based low-rank approximations.  In Advances in Neural Information Processing Systems, pages 7400–7410, 2019.
  * [IWW21] Piotr Indyk, Tal Wagner, and David Woodruff.  Few-shot data-driven algorithms for low rank approximation.  Advances in Neural Information Processing Systems, 34, 2021.
  * [JLL+20] Tanqiu Jiang, Yi Li, Honghao Lin, Yisong Ruan, and David P. Woodruff.  Learning-augmented data stream algorithms.  In International Conference on Learning Representations, 2020.
  * [MM13] Xiangrui Meng and Michael W Mahoney.  Low-distortion subspace embeddings in input-sparsity time and applications to robust linear regression.  In Proceedings of the forty-fifth annual ACM symposium on Theory of computing, pages 91–100, 2013.
  * [NN13] 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, pages 117–126, 2013.
  * [NN14] Jelani Nelson and Huy L. Nguyên.  Lower bounds for oblivious subspace embeddings.  In Automata, Languages, and Programming - 41st International Colloquium (ICALP), pages 883–894, 2014.
  * [PGM+19] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, and Trevor Killeen.  Pytorch: An imperative style, high-performance deep learning library.  2019\.
  * [PW16] Mert Pilanci and Martin J. Wainwright.  Iterative Hessian sketch: Fast and accurate solution approximation for constrained least-squares.  J. Mach. Learn. Res., 17:53:1–53:38, 2016.
  * [Sar06] Tamas Sarlos.  Improved approximation algorithms for large matrices via random projections.  In 47th Annual IEEE Symposium on Foundations of Computer Science (FOCS), pages 143–152, 2006.
  * [vdBPSW21] Jan van den Brand, Binghui Peng, Zhao Song, and Omri Weinstein.  Training (overparametrized) neural networks in near-linear time.  In James R. Lee, editor, 12th Innovations in Theoretical Computer Science Conference, ITCS, volume 185, pages 63:1–63:15, 2021.
  * [Ver12] Roman Vershynin.  Introduction to the non-asymptotic analysis of random matrices.  In Yonina C. Eldar and Gitta Kutyniok, editors, Compressed Sensing: Theory and Applications, page 210–268. Cambridge University Press, 2012.
  * [WZSS17] Jingdong Wang, Ting Zhang, Nicu Sebe, and Heng Tao ShenWang.  A survey on learning to hash.  In IEEE Transactions on Pattern Analysis and Machine Intelligence, pages 769–790, 2017.

##  Appendix A Additional Experiments: Low-Rank Approximation

The details (data dimension, $N_{\mathrm{train}}$, etc.) are presented in
Table A.1.

Name | Type | Dimension | $N_{\mathrm{train}}$ | $N_{\mathrm{test}}$
---|---|---|---|---
Friends | Image | $5760\times 1080$ | $400$ | $100$
Logo | Image | $5760\times 1080$ | $400$ | $100$
Hyper | Image | $1024\times 768$ | $400$ | $100$
Table A.1: Data set descriptions

###  A.1 Sensitivity Analysis of Algorithm 4

Sketch | Logo | Friends | Hyper
---|---|---|---
Ours(inner product) | 0.311 | 0.470 | 1.232
$\ell_{2}$ Sampling | 0.698 | 0.935 | 1.293
Only Ridge | 0.994 | 1.493 | 4.155
Randomly Grouping | 0.659 | 1.069 | 2.070
Table A.2: Sensitivity analysis for our approach (using Algorithm 1 from
[IVY19] with one sketch)

In this section we explore how sensitive the performance of our Algorithm 4 is
to the ridge leverage score sampling and maximum inner product grouping
process. We consider the following baselines:

  * •

$\ell_{2}$ norm sampling: we sample the rows according to their squared length
instead of doing ridge leverage score sampling.

  * •

Only ridge leverage score sampling: the subspace spanned by only the sampled
rows from ridge leverage score sampling.

  * •

Randomly grouping: we put the sampled rows into different buckets as before,
but randomly divide the non-sampled rows into buckets.

The results are shown in Table A.2. Here we set $k=30,m=60$ as an example. To
show the difference of the initialization method more clearly, we compare the
error using the one-sided sketching Algorithm 1 and do not further optimize
the non-zeros values. From the table we can see both that ridge leverage score
sampling and the downstream grouping process are necessary, otherwise the
error will be similar or even worse than that of the classical Count-Sketch.

###  A.2 Total Variation Distance

As we have shown in Theorem 5.1, if the total variation distance between the
row sampling probability distributions $p$ and $q$ is $O(1)$, we have a worst-
case guarantee of $O(k\log k+k/\epsilon)$, which is strictly better than the
$\Omega(k^{2})$ lower bound for the random CountSketch, even when its non-zero
values have been optimized. We now study the total variation distance between
the train and test matrix in our dataset. The result is shown in Figure A.1.
From the figure we can see that for all the three dataset, the total variation
distance is bounded by a constant, which suggests that the assumptions are
reasonable for real-world data.

Figure A.1: Total variation distance between train and test matrices. left:
Logo, middle: friend, right: Hyper.

###  A.3 Experiments: LRA in the few-shot setting

In the rest of this section, we study the performance of our second approach
in the few-shot learning setting. We first consider the case where we only
have one training matrix randomly sampled from $\mathsf{Tr}$. Here, we compare
our method with the 1Shot2Vec method proposed in [IWW21] in the same setting
($k=10,m=40$) as in their empirical evaluation. The result is shown in Table
A.4. Compared to 1Shot2Vec, our method reduces the error by around $50\%$, and
has an even slightly faster runtime.

[IWW21] also proposed a FewShotSGD algorithm which further improves the non-
zero values of the sketches after different initialization methods. We compare
the performance of this approach for different initialization methods: in all
initialization methods, we only use one training matrix and we use three
training matrices for the FewShotSGD step. The results are shown in Table A.4.
We report the minimum error of $50$ iterations of the FewShotSGD because we
aim to compare the computational efficiency for different methods. From the
table we see that our approach plus the FewShotSGD method can achieve a much
smaller error, with around a $50\%$ improvement upon [IWW21]. Moreover, even
without further optimization by FewShotSGD, our initialization method for
learning the non-zero locations in CountSketch obtains a smaller error than
other methods (even when they are optimized with 1ShotSGD or FewShotSGD
learning).

Algorithm | Dataset | Few-shot Error | Training Time
---|---|---|---
Classical CS | Logo | 0.331 | ✗
Friends | 0.524
Hyper | 1.082
1shot2Vec | Logo | 0.171 | 5.682
Friends | 0.306 | 5.680
Hyper | 0.795 | 1.054
Ours (inner product) | Logo | 0.065 | 4.515
Friends | 0.139 | 4.773
Hyper | 0.535 | 0.623

Table A.3: Test errors and training times for LRA in the one-shot setting
(using Alg. 1 with one sketch)

Sketch | Logo | Friends | Hyper
---|---|---|---
Ours (Initialization only) | 0.065 | 0.139 | 0.535
Ours + FewShotSGD | 0.048 | 0.125 | 0.443
1Shot1Vec only | 0.171 | 0.306 | 0.795
1Shot1Vec + FewShot SGD | 0.104 | 0.229 | 0.636
Classical CS | 0.331 | 0.524 | 1.082
Classical CS + FewShot SGD | 0.173 | 0.279 | 0.771

Table A.4: Test errors for LRA in the few-shot setting (using Alg. 1 from
[IVY19] with one sketch)

##  Appendix B Additional Experiments: Second-Order Optimization

As we mentioned in Section 8, despite the number of problems that learned
sketches have been applied to, they have not been applied to convex
optimization, or say, iterative sketching algorithms in general. To
demonstrate the difficulty, we consider the Iterative Hessian Sketch (IHS) as
an example. In that scheme, suppose that we have $k$ iterations of the
algorithm. Then we need $k$ independent sketching matrices (otherwise the
solution may diverge). A natural way is to follow the method in [IVY19], which
is to minimize the following quantity

$\min_{S_{1},...S_{k}}\operatorname*{{\mathbf{E}}}_{A\in\mathcal{D}}f(A,\mathsf{ALG}(S_{1},...,S_{k},A))\;,$

where the minimization is taken over $k$ Count-Sketch matrices
$S_{1},\dots,S_{k}$. In this case, however, calculating the gradient with
respect to $S_{1}$ would involve all iterations and in each iteration we need
to solve a constrained optimization problem. Hence, it would be difficult and
intractable to compute the gradients. An alternative way is to train $k$
sketching matrices sequentially, that is, learn the sketching matrix for the
$i$-th iteration using a local loss function for the $i$-th iteration, and
then using the learned matrix in the $i$-th iteration to generate the training
data for the $(i+1)$-st iteration. However, the empirical results suggest that
it works for the first iteration only, because in this case the training data
for the $(i+1)$-st iteration depends on the solution to the $i$-th iteration
and may become farther away from the test data in later iterations. The core
problem here is that the method proposed in [IVY19] treats the training
process in a black-box way, which is difficult to extend to iterative methods.

###  B.1 The Distribution of the Heavy Rows

In our experiments, we hypothesize that in real-world data there may be an
underlying pattern which can help us identify the heavy rows. In the Electric
dataset, each row of the matrix corresponds to a specific residence and the
heavy rows are concentrated on some specific rows.

To exemplify this, we study the heavy leverage score rows distribution over
the Electric dataset. For a row $i\in[370]$, let $f_{i}$ denote the number of
times that row $i$ is heavy out of $320$ training data points from the
Electric dataset, where we say row $i$ is heavy if $\ell_{i}\geq 5d/n$. Below
we list all $74$ pairs $(i,f_{i})$ with $f_{i}>0$.

(195,320), (278,320), (361,320), (207,317), (227,285), (240,284), (219,270),
(275,232), (156,214), (322,213), (193,196), (190,192), (160,191), (350,181),
(63,176), (42,168), (162,148), (356,129), (363,110), (362,105), (338,95),
(215,94), (234,93), (289,81), (97,80), (146,70), (102,67), (98,58), (48,57),
(349,53), (165,46), (101,41), (352,40), (293,34), (344,29), (268,21),
(206,20), (217,20), (327,20), (340,19), (230,18), (359,18), (297,14),
(357,14), (161,13), (245,10), (100,8), (85,6), (212,6), (313,6), (129,5),
(130,5), (366,5), (103,4), (204,4), (246,4), (306,4), (138,3), (199,3),
(222,3), (360,3), (87,2), (154,2), (209,2), (123,1), (189,1), (208,1),
(214,1), (221,1), (224,1), (228,1), (309,1), (337,1), (343,1)

Observe that the heavy rows are concentrated on a set of specific row indices.
There are only $30$ rows $i$ with $f_{i}\geq 50$. We view this as strong
evidence for our hypothesis.

#### Heavy Rows Distribution Under Permutation.

We note that even though the order of the rows has been changed, we can still
recognize the patterns of the rows. We continue to use the Electric dataset as
an example. To address the concern that a permutation may break the sketch, we
can measure the similarity between vectors, that is, after processing the
training data, we can instead test similarity on the rows of the test matrix
and use this to select the heavy rows, rather than an index which may simply
be permuted. To illustrate this method, we use the following example on the
Electric dataset, using locality sensitive hashing.

After processing the training data, we obtain a set $I$ of indices of heavy
rows. For each $i\in I$, we pick $q=3$ independent standard Gaussian vectors
$g_{1},g_{2},g_{3}\in\mathbb{R}^{d}$, and compute
$f(r_{i})=(g_{1}^{T}r_{i},g_{2}^{T}r_{i},g_{3}^{T}r_{i})\in\mathbb{R}^{3}$,
where $r_{i}$ takes an average of the $i$-th rows over all training sets. Let
$A$ be the test matrix. For each $i\in I$, let
$j_{i}=\operatorname{argmin}_{j}\|f(A_{j})-f(r_{i})\|_{2}$. We take the
$j_{i}$-th row to be a heavy row in our learned sketch. This method only needs
an additional $O(1)$ passes over the entries of $A$ and hence, the extra time
cost is negligible. To test the performance of the method, we randomly pick a
matrix from the test set and permute its rows. The result shows that when $k$
is small, we can roughly recover $70\%$ of the top-$k$ heavy rows, and we plot
below the regression error using the learned Count-Sketch matrix generated
this way, where we set $m=90$ and $k=0.3m=27$. We can see that the learned
method still obtains a significant improvement.

Figure B.1: Test error of LASSO on Electric dataset

###  B.2 Matrix Norm Estimation with a Nuclear Norm Constraint

In many applications, for the problem

$X^{*}:=\operatorname*{arg\,min}_{X\in\mathbb{R}^{d_{1}\times
d_{2}}}\left\lVert AX-B\right\rVert_{F}^{2}\;,$

it is reasonable to model the matrix $X^{*}$ as having low rank. Similar to
$\ell_{1}$-minimization for compressive sensing, a standard relaxation of the
rank constraint is to minimize the nuclear norm of $X$, defined as
$\left\lVert
X\right\rVert_{\ast}:=\sum_{j=1}^{\min\\{d_{1},d_{2}\\}}\sigma_{j}(X)$, where
$\sigma_{j}(X)$ is the $j$-th largest singular value of $X$.

Hence, the matrix estimation problem we consider here is

$X^{*}:=\operatorname*{arg\,min}_{X\in\mathbb{R}^{d_{1}\times
d_{2}}}\left\lVert AX-B\right\rVert_{F}^{2}\quad\text{such
that}\quad\left\lVert X\right\rVert_{\ast}\leq\rho,$

where $\rho>0$ is a user-defined radius as a regularization parameter.

We conduct Iterative Hessian Sketch (IHS) experiments on the following
dataset:

  * •

Tunnel888https://archive.ics.uci.edu/ml/datasets/Gas+sensor+array+exposed+to+turbulent+gas+mixtures:
The data set is a time series of gas concentrations measured by eight sensors
in a wind tunnel. Each $(A,B)$ corresponds to a different data collection
trial. $A^{i}\in\mathbb{R}^{13530\times 5},B^{i}\in\mathbb{R}^{13530\times
6}$, $|(A,B)|_{\mathrm{train}}=144$, $|(A,B)|_{\mathrm{test}}=36$. In our
nuclear norm constraint, we set $\rho=10$.

Experiment Setting: We choose $m=7d,10d$ for the Tunnel dataset. We consider
the error $\frac{1}{2}\left\lVert
AX-B\right\rVert_{2}^{2}-\frac{1}{2}\left\lVert
AX^{\ast}-B\right\rVert_{2}^{2}$. The leverage scores of this dataset are very
uniform. Hence, for this experiment we only consider optimizing the values of
the non-zero entries.

Results of Our Experiments: We plot on a logarithmic scale the mean errors of
the dataset in Figures B.2. We can see that when $m=7d$, the gradient-based
sketch, based on the first $6$ iterations, has a rate of convergence that is
48% of the random sketch, and when $m=10d$, the gradient-based sketch has a
rate of convergence that is 29% of the random sketch.

Figure B.2: Test error of matrix estimation with a nuclear norm constraint on
the Tunnel dataset

###  B.3 Fast Regression Solver

Figure B.3: Test error of the subroutine in fast regression on Electric
dataset. Figure B.4: Test error of fast regression on Electric dataset

Consider an unconstrained convex optimization problem $\min_{x}f(x)$, where
$f$ is smooth and strongly convex, and its Hessian $\nabla^{2}f$ is Lipschitz
continuous. This problem can be solved by Newton’s method, which iteratively
performs the update

$x_{t+1}=x_{t}-\operatorname*{arg\,min}_{z}\left\|(\nabla^{2}f(x_{t})^{1/2})^{\top}(\nabla^{2}f(x_{t})^{1/2})z-\nabla
f(x_{t})\right\|_{2},$ (B.1)

provided it is given a good initial point $x_{0}$. In each step, it requires
solving a regression problem of the form $\min_{z}\left\lVert
A^{\top}Az-y\right\rVert_{2}$, which, with access to $A$, can be solved with a
fast regression solver in [vdBPSW21]. The regression solver first computes a
preconditioner $R$ via a QR decomposition such that $SAR$ has orthonormal
columns, where $S$ is a sketching matrix, then solves
$\hat{z}=\operatorname*{arg\,min}_{z^{\prime}}\left\lVert(AR)^{\top}(AR)z^{\prime}-y\right\rVert_{2}$
by gradient descent and returns $R\hat{z}$ in the end. Here, the point of
sketching is that the QR decomposition of $SA$ can be computed much more
efficiently than the QR decomposition of $A$, since $S$ has only a small
number of rows.

In this section, We consider the unconstrained least squares problem
$\min_{x}f(x)$ with $f(x)=\frac{1}{2}\left\lVert Ax-b\right\rVert_{2}^{2}$
using the Electric dataset, using the above fast regression solver.

Training: Note that $\nabla^{2}f(x)=A^{\top}A$, independent of $x$. In the
$t$-th round of Newton’s method, by (B.1), we need to solve a regression
problem $\min_{z}\left\lVert A^{\top}Az-y\right\rVert_{2}^{2}$ with $y=\nabla
f(x_{t})$. Hence, we can use the same methods in the preceding subsection to
optimize the learned sketch $S_{i}$. For a general problem where
$\nabla^{2}f(x)$ depends on $x$, one can take $x_{t}$ to be the solution
obtained from the learned sketch $S_{t}$ to generate $A$ and $y$ for the
$(t+1)$-st round, train a learned sketch $S_{t+1}$, and repeat this process.

Experiment Setting: For the Electric dataset, we set $m=10d=90$. We observe
that the classical Count-Sketch matrix makes the solution diverge terribly in
this setting. To make a clearer comparison, we consider the following sketch
matrix:

  * •

Gaussian sketch: $S=\frac{1}{\sqrt{m}}G$, where $G\in\mathbb{R}^{m\times n}$
with i.i.d. $N(0,1)$ entries.

  * •

Sparse Johnson-Lindenstrauss Transform (SJLT): $S$ is the vertical
concatenation of $s$ independent Count-Sketch matrices, each of dimension
$m/s\times n$.

We note that the above sketching matrices require more time to compute $SA$
but need fewer rows to be a subspace embedding than the classical Count-Sketch
matrix.

For the step length $\eta$ in gradient descent, we set $\eta=1$ in all
iterations of the learned sketches. For classical random sketches, we set
$\eta$ in the following two ways: (a) $\eta=1$ in all iterations and (b)
$\eta=1$ in the first iteration and $\eta=0.2$ in all subsequent iterations.

Experimental Results: We examine the accuracy of the subproblem
$\min_{z}\left\lVert A^{\top}Az-y\right\rVert_{2}^{2}$ and define the error to
be $\left\lVert A^{\top}ARz_{t}-y\right\rVert_{2}/\left\lVert
y\right\rVert_{2}$. We consider the subproblems in the first three iterations
of the global Newton method. The results are plotted in Figure B.3. Note that
Count-Sketch causes a terrible divergence of the subroutine and is thus
omitted in the plots. Still, we observe that in setting (a) of $\eta$, the
other two classical sketches cause the subroutine to diverge. In setting (b)
of $\eta$, the other two classical sketches lead to convergence but their
error is significantly larger than that of the learned sketches, in each of
the first three calls to the subroutine. The error of the learned sketch is
less than $0.01$ in all iterations of all three subroutine calls, in both
settings (a) and (b) of $\eta$.

We also plot a figure on the convergence of the global Newton method. Here,
for each subroutine, we only run one iteration, and plot the error of the
original least squares problem. The result is shown in Figure B.4, which
clearly displays a significantly faster decay with learned sketches. The rate
of convergence using heavy-rows sketches is $80.6\%$ of that using Gaussian or
sparse JL sketches.

###  B.4 First-Order Optimization

In this section, we study the use of the sketch in first-order methods.
Particularly, let $QR^{-1}=SA$ be the QR-decomposition for $SA$, where $S$ is
a sketch matrix. We use $R$ as an (approximate) pre-conditioner and use
gradient descent to solve the problem $\min\left\lVert ARx-b\right\rVert_{2}$.
Here we use the Electric dataset where $A$ is $370\times 9$ and we set $S$ to
have $90$ rows. The result is shown in the following table, where the time
includes the time to compute $R$. We can see that if we use a learned sketch
matrix, the error converges very fast when we set the learning rate to be $1$
and $0.1$, while the classical Count-Sketch will lead to divergence.

Iteration | 1 | 10 | 100 | 500
---|---|---|---|---
Error (learned, $lr=1$) | 2.73 | 1.5e-7 |  |
Error (learned, $lr=0.1$) | 4056 | 605 | 4.04e-6 |
Error (learned, $lr=0.01$) | 4897 | 4085 | 667 | 0.217
Error (random, $lr=1$) | N.A | N.A | N.A | N.A
Error (random, $lr=0.1$)
Error (random, $lr=0.01$) | 4881 | 3790 | 685 | 1.52
Time | 0.00048 | 0.00068 | 0.0029 | 0.0013
Table B.1: Test Error for Gradient Descent

##  Appendix C Preliminaries: Theorems and Additional Algorithms

In this section, we provide the full description of the time-optimal sketching
algorithm for LRA in Algorithm 2. We also provide several definitions and
lemmas that are used in the proofs of our results for LRA.

######  Definition C.1 (Affine Embedding).

Given a pair of matrices $A$ and $B$, a matrix $S$ is an affine
$\epsilon$-embedding if for all $X$ of the appropriate shape,
$\left\|S(AX-B)\right\|^{2}_{F}=(1\pm\epsilon)\left\|AX-B\right\|^{2}_{F}$.

######  Lemma C.2 ([CW17]; Lemma 40).

Let $A$ be an $n\times d$ matrix and let
$S\in\mathbb{R}^{\mathcal{O}({1/\epsilon^{2}})\times n}$ be a CountSketch
matrix. Then with constant probability,
$\left\|SA\right\|^{2}_{F}=(1\pm\epsilon)\left\|A\right\|^{2}_{F}$.

The following result is shown in [CW17] and sharpened with [NN13, MM13].

######  Lemma C.3.

Given matrices $A,B$ with $n$ rows, a CountSketch with
$\mathcal{O}(\mathrm{rank}(A)^{2}/\epsilon^{2})$ rows is an affine
$\epsilon$-embedding matrix with constant probability. Moreover, the matrix
product $SA$ can be computed in $\mathcal{O}(\mathrm{nnz}(A))$ time, where
$\mathrm{nnz}(A)$ denotes the number of non-zero entries of matrix $A$.

######  Lemma C.4 ([Sar06, CW17]).

Suppose that $A\in\mathbb{R}^{n\times d}$ and $B\in\mathbb{R}^{n\times
d^{\prime}}$. Let $S\in\mathbb{R}^{m\times n}$ be a CountSketch with
$m=\frac{\mathrm{rank}(A)^{2}}{\epsilon^{2}}$. Let
$\tilde{X}=\operatorname*{arg\,min}_{\operatorname{rank-\mathnormal{k}}X}\left\|SAX-
SB\right\|^{2}_{F}$. Then,

  1. 1.

With constant probability,
$\left\|A\tilde{X}-B\right\|^{2}_{F}\leq(1+\epsilon)\min_{\operatorname{rank-\mathnormal{k}}X}\left\|AX-B\right\|^{2}_{F}$.
In other words, in
$\mathcal{O}(\mathrm{nnz}(A)+\mathrm{nnz}(B)+m(d+d^{\prime}))$ time, we can
reduce the problem to a smaller (multi-response regression) problem with $m$
rows whose optimal solution is a $(1+\epsilon)$-approximate solution to the
original instance.

  2. 2.

The $(1+\epsilon)$-approximate solution $\tilde{X}$ can be computed in time
$\mathcal{O}(\mathrm{nnz}(A)+\mathrm{nnz}(B)+mdd^{\prime}+\min(m^{2}d,md^{2}))$.

Now we turn our attention to the time-optimal sketching algorithm for LRA. The
next lemma is known, though we include it for completeness [ACW17]:

######  Lemma C.5.

Suppose that $S\in\mathbb{R}^{m_{S}\times n}$ and $R\in\mathbb{R}^{m_{R}\times
d}$ are sparse affine $\epsilon$-embedding matrices for $(A_{k},A)$ and
$((SA)^{\top},A^{\top})$, respectively. Then,

$\displaystyle\min_{\operatorname{rank-\mathnormal{k}}X}\left\|AR^{\top}XSA-A\right\|_{F}^{2}\leq(1+\epsilon)\left\|A_{k}-A\right\|_{F}^{2}$

###### Proof.

Consider the following multiple-response regression problem:

$\displaystyle\min_{\operatorname{rank-\mathnormal{k}}X}\left\|A_{k}X-A\right\|_{F}^{2}.$
(C.1)

Note that since $X=I_{k}$ is a feasible solution to Eq. (C.1),
$\min_{\operatorname{rank-\mathnormal{k}}X}\left\|A_{k}X-A\right\|_{F}^{2}=\left\|A_{k}-A\right\|_{F}^{2}$.
Let $S\in\mathbb{R}^{m_{S}\times n}$ be a sketching matrix that satisfies the
condition of Lemma C.4 (Item 1) for $A:=A_{k}$ and $B:=A$. By the normal
equations, the rank-$k$ minimizer of $\left\|SA_{k}X-SA\right\|_{F}^{2}$ is
$(SA_{k})^{+}SA$. Hence,

$\displaystyle\left\|A_{k}(SA_{k})^{+}SA-A\right\|_{F}^{2}$
$\displaystyle\leq(1+\epsilon)\left\|A_{k}-A\right\|_{F}^{2},$ (C.2)

which in particular shows that a $(1+\epsilon)$-approximate rank-$k$
approximation of $A$ exists in the row space of $SA$. In other words,

$\displaystyle\min_{\operatorname{rank-\mathnormal{k}}X}\left\|XSA-A\right\|_{F}^{2}\leq(1+\epsilon)\left\|A_{k}-A\right\|_{F}^{2}.$
(C.3)

Next, let $R\in\mathbb{R}^{m_{R}\times d}$ be a sketching matrix which
satisfies the condition of Lemma C.4 (Item 1) for $A:=(SA)^{\top}$ and
$B:=A^{\top}$. Let $Y$ denote the $\operatorname{rank-\mathnormal{k}}$
minimizer of $\left\|R(SA)^{\top}X^{\top}-RA^{\top}\right\|_{F}^{2}$. Hence,

$\displaystyle\left\|(SA)^{\top}Y^{\top}-A^{\top}\right\|_{F}^{2}$
$\displaystyle\leq(1+\epsilon)\min_{\operatorname{rank-\mathnormal{k}}X}\left\|XSA-A\right\|_{F}^{2}$
$\displaystyle\rhd\text{Lemma~{}\ref{lem:mult-reg}
(Item~{}\ref{itm:reduction})}$
$\displaystyle\leq(1+\mathcal{O}(\epsilon))\left\|A_{k}-A\right\|_{F}^{2}$
$\displaystyle\rhd\text{Eq.~{}\eqref{eq:sketch-s}}$ (C.4)

Note that by the normal equations, again
$\mathrm{rowsp}(Y^{\top})\subseteq\mathrm{rowsp}(RA^{\top})$ and we can write
$Y=AR^{\top}Z$ where $\mathrm{rank}(Z)=k$. Thus,

$\displaystyle\min_{\operatorname{rank-\mathnormal{k}}X}\left\|AR^{\top}XSA-A\right\|_{F}^{2}\leq\left\|AR^{\top}ZSA-A\right\|_{F}^{2}$
$\displaystyle=\left\|(SA)^{\top}Y^{\top}-A^{\top}\right\|_{F}^{2}$
$\displaystyle\rhd Y=AR^{\top}Z$
$\displaystyle\leq(1+\mathcal{O}(\epsilon))\left\|A_{k}-A\right\|_{F}^{2}$
$\displaystyle\rhd\text{Eq.~{}\eqref{eq:second-step}}$

∎

######  Lemma C.6 ([ACW17]; Lemma 27).

For $C\in\mathbb{R}^{p\times m^{\prime}},D\in\mathbb{R}^{m\times
p^{\prime}},G\in\mathbb{R}^{p\times p^{\prime}}$, the following problem

$\displaystyle\min_{\operatorname{rank-\mathnormal{k}}Z}\left\|CZD-G\right\|_{F}^{2}$
(C.5)

can be solved in
$\mathcal{O}(pm^{\prime}r_{C}+p^{\prime}mr_{D}+pp^{\prime}(r_{D}+r_{C}))$
time, where $r_{C}=\mathrm{rank}(C)\leq\min\\{m^{\prime},p\\}$ and
$r_{D}=\mathrm{rank}(D)\leq\min\\{m,p^{\prime}\\}$.

######  Lemma C.7.

Let $S\in\mathbb{R}^{m_{S}\times d}$, $R\in\mathbb{R}^{m_{R}\times d}$ be
CountSketch (CS) matrices such that

$\displaystyle\min_{\operatorname{rank-\mathnormal{k}}X}\left\|AR^{\top}XSA-A\right\|^{2}_{F}\leq(1+\gamma)\left\|A_{k}-A\right\|_{F}^{2}.$
(C.6)

Let $V\in\mathbb{R}^{(m_{R}^{2}/\beta^{2})\times n}$, and
$W\in\mathbb{R}^{{m_{S}^{2}\over\beta^{2}}\times d}$ be CS matrices. Then,
Algorithm 2 gives a $(1+\mathcal{O}(\beta+\gamma))$-approximation in time
$\mathrm{nnz}(A)+\mathcal{O}(\frac{m_{S}^{4}}{\beta^{2}}+\frac{m_{R}^{4}}{\beta^{2}}+\frac{m_{S}^{2}m_{R}^{2}(m_{S}+m_{R})}{\beta^{4}}+k(nm_{S}+dm_{R}))$
with constant probability.

###### Proof.

The approximation guarantee follows from Eq. (C.6) and the fact that $V$ and
$W$ are affine $\beta$-embedding matrices of $AR^{\top}$ and $SA$,
respectively (see Lemma C.3).

The algorithm first computes $C=VAR^{\top},D=SAW^{\top},G=VAW^{\top}$ which
can be done in time $\mathcal{O}(\mathrm{nnz}(A))$. As an example, we bound
the time to compute $C=VAR$. Note that since $V$ is a CS, $VA$ can be computed
in $\mathcal{O}(\mathrm{nnz}(A))$ time and the number of non-zero entries in
the resulting matrix is at most $\mathrm{nnz}(A)$. Hence, since $R$ is a CS as
well, $C$ can be computed in time
$\mathcal{O}(\mathrm{nnz}(A)+\mathrm{nnz}(VA))=\mathcal{O}(\mathrm{nnz}(A))$.
Then, it takes an extra
$\mathcal{O}((m_{S}^{3}+m_{R}^{3}+m_{S}^{2}m_{R}^{2})/\beta^{2})$ time to
store $C,D$ and $G$ in matrix form. Next, as we showed in Lemma C.6, the time
to compute $Z$ in Algorithm 2 is
$\mathcal{O}(\frac{m_{S}^{4}}{\beta^{2}}+\frac{m_{R}^{4}}{\beta^{2}}+\frac{m_{S}^{2}m_{R}^{2}(m_{S}+m_{R})}{\beta^{4}})$.
Finally, it takes $\mathcal{O}(\mathrm{nnz}(A)+k(nm_{S}+dm_{R}))$ time to
compute $Q=AR^{\top}Z_{L}$ and $P=Z_{R}SA$ and to return the solution in the
form of $P_{n\times k}Q_{k\times d}$. Hence, the total runtime is

$\displaystyle\mathcal{O}(\mathrm{nnz}(A)+\frac{m_{S}^{4}}{\beta^{2}}+\frac{m_{R}^{4}}{\beta^{2}}+\frac{m_{S}^{2}m_{R}^{2}(m_{S}+m_{R})}{\beta^{4}}+k(nm_{S}+dm_{R}))$

∎

##  Appendix D Attaining Worst-Case Guarantees

###  D.1 Low-Rank Approximation

We shall provide the following two methods to achieve worst case guarantees:
MixedSketch—whose guarantee is via the sketch monotonicity property, and
approximate comparison method (a.k.a. ApproxCheck), which just approximately
evaluates the cost of two solutions and takes the better one. These methods
asymptotically achieve the same worst-case guarantee. However, for any input
matrix $A$ and any pair of sketches $S,T$, the performance of the MixedSketch
method on $(A,S,T)$ is never worse than the performance of its corresponding
ApproxCheck method on $(A,S,T)$, and can be much better.

######  Remark D.1.

Let $A=\mathrm{diag}(2,2,\sqrt{2},\sqrt{2})$, and suppose the goal is to find
a rank-$2$ approximation of $A$. Consider two sketches $S$ and $T$ such that
$SA$ and $TA$ capture $\operatorname{\mathsf{span}}(e_{1},e_{3})$ and
$\operatorname{\mathsf{span}}(e_{2},e_{4})$, respectively. Then for both $SA$
and $TA$, the best solution in the subspace of one of these two spaces is a
$(\frac{3}{2})$-approximation: $\left\lVert A-A_{2}\right\rVert_{F}^{2}=4$ and
$\left\lVert A-P_{SA}\right\rVert_{F}^{2}=\left\lVert
A-P_{TA}\right\rVert_{F}^{2}=6$ where $P_{SA}$ and $P_{TA}$ respectively
denote the best approximation of $A$ in the space spanned by $SA$ and $TA$.

However, if we find the best rank-$2$ approximation of $A$, $Z$, inside the
span of the union of $SA$ and $TA$, then $\left\lVert
A-Z\right\rVert_{F}^{2}=4$. Since ApproxCheck just chooses the better of $SA$
and $TA$ by evaluating their costs, it misses out on the opportunity to do as
well as MixedSketch.

Here, we show the sketch monotonicity property for LRA.

######  Theorem D.2.

Let $A\in\mathbbm{R}^{n\times d}$ be an input matrix, $V$ and $W$ be
$\eta$-affine embeddings, and $S_{1}\in\mathbbm{R}^{m_{S}\times
n},R_{1}\in\mathbbm{R}^{m_{R}\times n}$ be arbitrary matrices. Consider
arbitrary extensions to $S_{1},R_{1}$: $\overline{S},\overline{R}$ (e.g.,
$\overline{S}$ is a concatenation of $S_{1}$ with an arbitrary matrix with the
same number of columns). Then, $\left\lVert
A-\mathsf{ALG}_{\mathrm{LRA}}((\overline{S},\overline{R},V,W),A))\right\rVert_{F}^{2}\leq(1+\eta)^{2}\left\lVert
A-\mathsf{ALG}_{\mathrm{LRA}}((S_{1},R_{1},V,W),A)\right\rVert_{F}^{2}$

###### Proof.

We have $\left\lVert
A-\mathsf{ALG}_{\mathrm{LRA}}((\overline{S},\overline{R},V,W),A)\right\rVert_{F}^{2}\leq(1+\eta)\min_{\operatorname{rank-\mathnormal{k}}X}\left\lVert
A\overline{R}X\overline{S}A-A\right\rVert_{F}^{2}=(1+\eta)\min_{\operatorname{rank-\mathnormal{k}}X:X\in\mathsf{row}(\overline{S}A)\cap\mathsf{col}(A\overline{R})}\left\lVert
X-A\right\rVert_{F}^{2}$, which is in turn at most
$(1+\eta)\min_{\operatorname{rank-\mathnormal{k}}X:X\in\mathsf{row}(S_{1}A)\cap\mathsf{col}(AR_{1})}$
$\left\lVert
X-A\right\rVert_{F}^{2}=(1+\eta)\min_{\operatorname{rank-\mathnormal{k}}X}\left\lVert
AR_{1}XS_{1}A-A\right\rVert_{F}^{2}\leq(1+\eta)^{2}\left\lVert
A-\mathsf{ALG}_{\mathrm{LRA}}((S_{1},R_{1},V,W),A)\right\rVert_{F}^{2},$ where
we use the fact the $V,W$ are affine $\eta$-embeddings (Definition C.1), as
well as the fact that
$\left(\mathsf{col}(AR_{1})\cap\mathsf{row}(S_{1}A)\right)\subseteq\left(\mathsf{col}(A\overline{R})\cap\mathsf{row}(\overline{S}A)\right)$.
∎

#### ApproxCheck for LRA.

We give the pseudocode for the ApproxCheck method and prove that the runtime
of this method for LRA is of the same order as the classical time-optimal
sketching algorithm of LRA.

0: learned sketches $S_{L},R_{L},V_{L},W_{L}$; classical sketches
$S_{C},R_{C},V_{C},W_{C}$; $\beta$; $A\in\mathbbm{R}^{n\times d}$

1:
$P_{L},Q_{L}\leftarrow\mathsf{ALG}_{\mathrm{LRA}}(S_{L},R_{L},V_{L},W_{L},A)$,
$P_{C}Q_{C}\leftarrow\mathsf{ALG}_{\mathrm{LRA}}(S_{C},R_{C},V_{C},W_{C},A)$

2: Let $S\in\mathbb{R}^{\mathcal{O}({1/\beta^{2}})\times
n},R\in\mathbb{R}^{\mathcal{O}({1/\beta^{2}})\times d}$ be classical
CountSketch matrices

3:
$\Delta_{L}\leftarrow\left\|S\left(P_{L}Q_{L}-A\right)R^{\top}\right\|^{2}_{F}$,
$\Delta_{C}\leftarrow\left\lVert
S\left(P_{C}Q_{C}-A\right)R^{\top}\right\rVert^{2}_{F}$

4: if $\Delta_{L}\leq\Delta_{C}$ then

5: return $P_{L}Q_{L}$

6: end if

7: return $P_{C}Q_{C}$

Algorithm 5 LRA ApproxCheck

######  Theorem D.3.

Assume we have data $A\in\mathbb{R}^{n\times d}$, learned sketches
$S_{L}\in\mathbb{R}^{\mathop{\mathrm{poly}}({k\over\epsilon})\times
n},R_{L}\in\mathbb{R}^{\mathop{\mathrm{poly}}({k\over\epsilon})\times
d},V_{L}\in\mathbb{R}^{\mathop{\mathrm{poly}}({k\over\epsilon})\times
n},W_{L}\in\mathbb{R}^{\mathop{\mathrm{poly}}({k\over\epsilon})\times d}$
which attain a $(1+\mathcal{O}(\gamma))$-approximation, classical sketches of
the same size, $S_{C},R_{C},V_{C},W_{C}$, which attain a
$(1+\mathcal{O}(\epsilon))$-approximation, and a tradeoff parameter $\beta$.
Then, Algorithm 5 attains a $(1+\beta+\min(\gamma,\epsilon))$-approximation in
$\mathcal{O}(\mathrm{nnz}(A)+(n+d)\mathop{\mathrm{poly}}({k\over\epsilon})+{k^{4}\over\beta^{4}}\cdot\mathop{\mathrm{poly}}({k\over\epsilon}))$
time.

###### Proof.

Let $(P_{L},Q_{L})$, $(P_{C},Q_{C})$ be the approximate rank-$k$
approximations of $A$ in factored form using $(S_{L},R_{L})$ and
$(S_{O},R_{O})$. Then, clearly,

$\displaystyle\min(\left\|P_{L}Q_{L}-A\right\|_{F}^{2},\left\|P_{C}Q_{C}-A\right\|_{F}^{2})=(1+\mathcal{O}(\min(\epsilon,\gamma)))\left\|A_{k}-A\right\|_{F}^{2}$
(D.1)

Let $\Gamma_{L}=P_{L}Q_{L}-A$, $\Gamma_{C}=P_{C}Q_{C}-A$ and
$\Gamma_{M}=\operatorname*{arg\,min}(\left\|S\Gamma_{L}R\right\|_{F},\left\|S\Gamma_{C}R\right\|_{F})$.
Then,

$\displaystyle\left\|\Gamma_{M}\right\|_{F}^{2}$
$\displaystyle\leq(1+\mathcal{O}(\beta))\left\|S\Gamma_{M}R\right\|_{F}^{2}$
$\displaystyle\rhd\text{by Lemma~{}\ref{lem:fast-estimate}}$
$\displaystyle\leq(1+\mathcal{O}(\beta))\cdot\min(\left\|\Gamma_{L}\right\|^{2}_{F},\left\|\Gamma_{C}\right\|^{2}_{F})$
$\displaystyle\leq(1+\mathcal{O}(\beta+\min(\epsilon,\gamma)))\left\|A_{k}-A\right\|_{F}^{2}$
$\displaystyle\rhd\text{by Eq.~{}\eqref{eq:low-rank-best-sol}}$

Runtime analysis. By Lemma C.7, Algorithm 2 computes $P_{L},Q_{L}$ and
$P_{C},Q_{C}$ in time
$\mathcal{O}(\mathrm{nnz}(A)+\frac{k^{16}(\beta^{2}+\epsilon^{2})}{\epsilon^{24}\beta^{4}}+\frac{k^{3}}{\epsilon^{2}}(n+\frac{dk^{2}}{\epsilon^{4}}))$.
Next, once we have $P_{L},Q_{L}$ and $P_{C},Q_{C}$, it takes
$\mathcal{O}(\mathrm{nnz}(A)+{k\over\beta^{4}})$ time to compute $\Delta_{L}$
and $\Delta_{C}$.

$\displaystyle\mathcal{O}(\mathrm{nnz}(A)+\frac{k^{16}(\beta^{2}+\epsilon^{2})}{\epsilon^{24}\beta^{4}}+\frac{k^{3}}{\epsilon^{2}}(n+\frac{dk^{2}}{\epsilon^{4}})+\frac{k}{\beta^{4}})=\mathcal{O}(\mathrm{nnz}(A)+(n+d+{k^{4}\over\beta^{4}})\mathop{\mathrm{poly}}({k\over\epsilon})).$

∎

To interpret the above theorem, note that when $\epsilon\gg k(n+d)^{-4}$, we
can set $\beta^{-4}=\mathcal{O}(k(n+d)^{-4})$ so that Algorithm 5 has the same
asymptotic runtime as the best $(1+\epsilon)$-approximation algorithm for LRA
with the classical CountSketch. Moreover, Algorithm 5 is a
$(1+o(\epsilon))$-approximation when the learned sketch outperforms classical
sketches, $\gamma=o(\epsilon)$. On the other hand, when the learned sketches
perform poorly, $\gamma=\Omega(\epsilon)$, the worst-case guarantee of
Algorithm 5 remains $(1+\mathcal{O}(\epsilon))$.

###  D.2 Second-Order Optimization

For the sketches for second-order optimization, the monotonicity property does
not hold. Below we provide an input-sparsity algorithm which can test for and
use the better of a random sketch and a learned sketch. Our theorem is as
follows.

Algorithm 6 Solver for (D.2)

1: $S_{1}\leftarrow$ learned sketch, $S_{2}\leftarrow$ random sketch with
$\Theta(d^{2}/\epsilon^{2})$ rows

2: $(\hat{Z}_{i,1},\hat{Z}_{i,2})\leftarrow\textsc{Estimate}(S_{i},A)$,
$i=1,2$

3: $i^{\ast}\leftarrow\arg\min_{i=1,2}(\hat{Z}_{i,2}/\hat{Z}_{i,1})$

4: $\hat{x}\leftarrow$ solution of (D.2) with $S=S_{i^{\ast}}$

5: return $\hat{x}$

6: function Estimate($S,A$)

7: $T\leftarrow$ sparse $(1\pm\eta)$-subspace embedding matrix for
$d$-dimensional subspaces

8: $(Q,R)\leftarrow\textsc{QR}(TA)$

9: $\hat{Z}_{1}\leftarrow\sigma_{\min}(SAR^{-1})$

10: $\hat{Z}_{2}\leftarrow(1\pm\eta)$-approximation to
$\left\lVert(SAR^{-1})^{\top}(SAR^{-1})-I\right\rVert_{\mathrm{op}}$

11: return $(\hat{Z}_{1},\hat{Z}_{2})$

######  Theorem D.4.

Let $\epsilon\in(0,0.09)$ be a constant and $S_{1}$ a learned Count-Sketch
matrix. Suppose that $A$ is of full rank. There is an algorithm whose output
is a solution $\hat{x}$ which, with probability at least $0.98$, satisfies
that $\left\lVert A(\hat{x}-x^{\ast})\right\rVert_{2}\leq
O\big{(}\min\big{\\{}\frac{Z_{2}(S_{1})}{Z_{1}(S_{1})},\epsilon\big{\\}}\big{)}\left\lVert
Ax^{\ast}\right\rVert_{2}$, where
$x^{\ast}=\operatorname*{arg\,min}_{x\in\mathcal{C}}\left\lVert
Ax-b\right\rVert_{2}$ is the least-squares solution. Furthermore, the
algorithm runs in
$O(\mathrm{nnz}(A)\log(\frac{1}{\epsilon})+\mathop{\mathrm{poly}}(\frac{d}{\epsilon}))$
time.

Consider the minimization problem

$\min_{x\in\mathcal{C}}\left\\{\frac{1}{2}\left\lVert
SAx\right\rVert_{2}^{2}-\langle A^{\top}y,x\rangle\right\\},$ (D.2)

which is used as a subroutine for the IHS (cf. (2.2)). We note that in this
subroutine if we let $x\leftarrow x-x^{i-1},b\leftarrow
b-Ax^{i-1},\mathcal{C}\leftarrow\mathcal{C}-x^{i-1}$, we would get the
guarantee of the $i$-th iteration of the original IHS. To analyze the
performance of the learned sketch, we define the following quantities
(corresponding exactly to the unconstrained case in [PW16])

$\displaystyle
Z_{1}(S)=\inf_{v\in\operatorname{colsp}(A)\cap\mathbb{S}^{n-1}}\left\lVert
Sv\right\rVert_{2}^{2},$ $\displaystyle
Z_{2}(S)=\sup_{u,v\in\operatorname{colsp}(A)\cap\mathbb{S}^{n-1}}\left\langle
u,(S^{\top}S-I_{n})v\right\rangle.$

When $S$ is a $(1+\epsilon)$-subspace embedding of $\operatorname{colsp}(A)$,
we have $Z_{1}(S)\geq 1-\epsilon$ and $Z_{2}(S)\leq 2\epsilon$.

For a general sketching matrix $S$, the following is the approximation
guarantee of $\hat{Z}_{1}$ and $\hat{Z}_{2}$, which are estimates of
$Z_{1}(S)$ and $Z_{2}(S)$, respectively. The main idea is that $AR^{-1}$ is
well-conditioned, where $R$ is as calculated in Algorithm 6.

######  Lemma D.5.

Suppose that $\eta\in(0,\frac{1}{3})$ is a small constant, $A$ is of full rank
and $S$ has $\mathop{\mathrm{poly}}(d/\eta)$ rows. The function
$\textsc{Estimate}(S,A)$ returns in
$O((\mathrm{nnz}(A)\log\frac{1}{\eta}+\mathop{\mathrm{poly}}(\frac{d}{\eta}))$
time $\hat{Z}_{1},\hat{Z}_{2}$ which with probability at least $0.99$ satisfy
that $\frac{Z_{1}(S)}{1+\eta}\leq\hat{Z}_{1}\leq\frac{Z_{1}(S)}{1-\eta}$ and
$\frac{Z_{2}(S)}{(1+\eta)^{2}}-3\eta\leq\hat{Z}_{2}\leq\frac{Z_{2}(S)}{(1-\eta)^{2}}+3\eta$.

###### Proof.

Suppose that $AR^{-1}=UW$, where $U\in\mathbb{R}^{n\times d}$ has orthonormal
columns, which form an orthonormal basis of the column space of $A$. Since $T$
is a subspace embedding of the column space of $A$ with probability $0.99$, it
holds for all $x\in\mathbb{R}^{d}$ that

$\frac{1}{1+\eta}\left\lVert TAR^{-1}x\right\rVert_{2}\leq\left\lVert
AR^{-1}x\right\rVert_{2}\leq\frac{1}{1-\eta}\left\lVert
TAR^{-1}x\right\rVert_{2}.$

Since

$\left\lVert TAR^{-1}x\right\rVert_{2}=\left\lVert
Qx\right\rVert_{2}=\left\lVert x\right\rVert_{2}$

and

$\left\lVert Wx\right\rVert_{2}=\left\lVert UWx\right\rVert_{2}=\left\lVert
AR^{-1}x\right\rVert_{2}$ (D.3)

we have that

$\frac{1}{1+\eta}\left\lVert x\right\rVert_{2}\leq\left\lVert
Wx\right\rVert_{2}\leq\frac{1}{1-\eta}\left\lVert x\right\rVert_{2},\quad
x\in\mathbb{R}^{d}.$ (D.4)

It is easy to see that

$Z_{1}(S)=\min_{x\in\mathbb{S}^{d-1}}\left\lVert
SUx\right\rVert_{2}=\min_{y\neq 0}\frac{\left\lVert
SUWy\right\rVert_{2}}{\left\lVert Wy\right\rVert_{2}},$

and thus,

$\min_{y\neq 0}(1-\eta)\frac{\left\lVert SUWy\right\rVert_{2}}{\left\lVert
y\right\rVert_{2}}\leq Z_{1}(S)\leq\min_{y\neq 0}(1+\eta)\frac{\left\lVert
SUWy\right\rVert_{2}}{\left\lVert y\right\rVert_{2}}.$

Recall that $SUW=SAR^{-1}$. We see that

$(1-\eta)\sigma_{\min}(SAR^{-1})\leq
Z_{1}(S)\leq(1+\eta)\sigma_{\min}(SAR^{-1}).$

By definition,

$Z_{2}(S)=\left\lVert U^{T}(S^{\top}S-I_{n})U\right\rVert_{\mathrm{op}}.$

It follows from (D.4) that

$(1-\eta)^{2}\left\lVert
W^{T}U^{T}(S^{T}S-I_{n})UW\right\rVert_{\mathrm{op}}\leq
Z_{2}(S)\leq(1+\eta)^{2}\left\lVert
W^{T}U^{T}(S^{T}S-I_{n})UW\right\rVert_{\mathrm{op}}.$

and from (D.4), (D.3) and Lemma 5.36 of [Ver12] that

$\left\lVert(AR^{-1})^{\top}(AR^{-1})-I\right\rVert_{\mathrm{op}}\leq 3\eta.$

Since

$\left\lVert
W^{T}U^{T}(S^{T}S-I_{n})UW\right\rVert_{\mathrm{op}}=\left\lVert(AR^{-1})^{\top}(S^{T}S-I_{n})AR^{-1}\right\rVert_{\mathrm{op}}$

and

$\displaystyle\quad\
\left\lVert(AR^{-1})^{\top}S^{T}SAR^{-1}-I\right\rVert_{\mathrm{op}}-\left\lVert(AR^{-1})^{\top}(AR^{-1})-I\right\rVert_{\mathrm{op}}$
$\displaystyle\leq\left\lVert(AR^{-1})^{\top}(S^{T}S-I_{n})AR^{-1}\right\rVert_{\mathrm{op}}$
$\displaystyle\leq\left\lVert(AR^{-1})^{\top}S^{T}SAR^{-1}-I\right\rVert_{\mathrm{op}}+\left\lVert(AR^{-1})^{\top}(AR^{-1})-I\right\rVert_{\mathrm{op}},$

it follows that

$\displaystyle\quad\
(1-\eta)^{2}\left\lVert(SAR^{-1})^{\top}SAR^{-1}-I\right\rVert_{\mathrm{op}}-3(1-\eta)^{2}\eta$
$\displaystyle\leq Z_{2}(S)$
$\displaystyle\leq(1+\eta)^{2}\left\lVert(SAR^{-1})^{\top}SAR^{-1}-I\right\rVert_{\mathrm{op}}+3(1+\eta)^{2}\eta.$

We have so far proved the correctness of the approximation and we next analyze
the runtime below.

Since $S$ and $T$ are sparse, computing $SA$ and $TA$ takes
$O(\mathrm{nnz}(A))$ time. The QR decomposition of $TA$, which is a matrix of
size $\mathop{\mathrm{poly}}(d/\eta)\times d$, can be computed in
$\mathop{\mathrm{poly}}(d/\eta)$ time. The matrix $SAR^{-1}$ can be computed
in $\mathop{\mathrm{poly}}(d)$ time. Since it has size
$\mathop{\mathrm{poly}}(d/\eta)\times d$, its smallest singular value can be
computed in $\mathop{\mathrm{poly}}(d/\eta)$ time. To approximate $Z_{2}(S)$,
we can use the power method to estimate
$\left\lVert(SAR^{-1})^{T}SAR^{-1}-I\right\rVert_{op}$ up to a
$(1\pm\eta)$-factor in
$O((\mathrm{nnz}(A)+\mathop{\mathrm{poly}}(d/\eta))\log(1/\eta))$ time. ∎

Now we are ready to prove Theorem D.4.

###### Proof of Theorem D.4.

In Lemma D.5, we have with probability at least $0.99$ that

$\frac{\hat{Z}_{2}}{\hat{Z}_{1}}\geq\frac{\frac{1}{(1+\epsilon)^{2}}Z_{2}(S)-3\epsilon}{\frac{1}{1-\epsilon}Z_{1}(S)}\geq\frac{1-\epsilon}{(1+\epsilon)^{2}}\frac{Z_{2}(S)}{Z_{1}(S)}-\frac{3\epsilon(1-\epsilon)}{Z_{1}(S)}.$

and similarly,

$\frac{\hat{Z}_{2}}{\hat{Z}_{1}}\leq\frac{\frac{1}{(1-\epsilon)^{2}}Z_{2}(S)+3\epsilon}{\frac{1}{1+\epsilon}Z_{1}(S)}\leq\frac{1+\epsilon}{(1-\epsilon)^{2}}\frac{Z_{2}(S)}{Z_{1}(S)}+\frac{3\epsilon(1+\epsilon)}{Z_{1}(S)}.$

Note that since $S_{2}$ is an $\epsilon$-subspace embedding with probability
at least $0.99$, we have that $Z_{1}(S_{2})\geq 1-\epsilon$ and
$Z_{2}(S_{2})/Z_{1}(S_{2})\leq 2.2\epsilon$. Consider $Z_{1}(S_{1})$.

First, we consider the case where $Z_{1}(S_{1})<1/2$. Observe that
$Z_{2}(S)\geq 1-Z_{1}(S)$. We have in this case
$\widehat{Z}_{1,2}/\widehat{Z}_{1,1}>1/5\geq 2.2\epsilon\geq
Z_{2}(S_{2})/Z_{1}(S_{2})$. In this case our algorithm will choose $S_{2}$
correctly.

Next, assume that $Z_{1}(S_{1})\geq 1/2$. Now we have with probability at
least $0.98$ that

$(1-3\epsilon)\frac{Z_{2}(S_{i})}{Z_{1}(S_{i})}-3\epsilon\leq\frac{\widehat{Z}_{i,2}}{\widehat{Z}_{i,1}}\leq(1+4\epsilon)\frac{Z_{2}(S_{i})}{Z_{1}(S_{i})}+4\epsilon,\quad
i=1,2.$

Therefore, when $Z_{2}(S_{1})/Z_{1}(S_{1})\leq c_{1}Z_{2}(S_{2})/Z_{1}(S_{2})$
for some small absolute constant $c_{1}>0$, we will have
$\widehat{Z}_{1,2}/\widehat{Z}_{1,1}<\widehat{Z}_{2,2}/\widehat{Z}_{2,1}$, and
our algorithm will choose $S_{1}$ correctly. If $Z_{2}(S_{1})/Z_{1}(S_{1})\geq
C_{1}\epsilon$ for some absolute constant $C_{1}>0$, our algorithm will choose
$S_{2}$ correctly. In the remaining case, both ratios
$Z_{2}(S_{2})/Z_{1}(S_{2})$ and $Z_{2}(S_{1})/Z_{1}(S_{1})$ are at most
$\max\\{C_{2},3\\}\epsilon$, and the guarantee of the theorem holds
automatically.

The correctness of our claim then follows from Proposition 1 of [PW16],
together with the fact that $S_{2}$ is a random subspace embedding. The
runtime follows from Lemma D.5 and Theorem 2.2 of [CD19]. ∎

##  Appendix E Sketch Learning: Omitted Proofs

###  E.1 Proof of Theorem 5.1

We need the following lemmas for the ridge leverage score sampling in [CMM17].

######  Lemma E.1 ([CMM17, Lemma 4]).

Let $\lambda=\left\lVert A-A_{k}\right\rVert_{F}^{2}/k$. Then we have
$\sum_{i}\tau_{i}(A)\leq 2k$.

######  Lemma E.2 ([CMM17, Theorem 7]).

Let $\lambda=\left\lVert A-A_{k}\right\rVert_{F}^{2}/k$ and
$\tilde{\tau}_{i}\geq\tau_{i}(A)$ be an overestimate to the $i$-th ridge
leverage score of $A$. Let $p_{i}=\tilde{\tau}_{i}/\sum_{i}\tilde{\tau}_{i}$.
If $C$ is a matrix that is constructed by sampling $t=O((\log
k+\frac{\log(1/\delta)}{\epsilon})\cdot\sum_{i}\tilde{\tau}_{i})$ rows of $A$,
each set to $a_{i}$ with probability $p_{i}$, then with probability at least
$1-\delta$ we have

$\min_{\operatorname{rank-\mathnormal{k}}X:\mathsf{row}(X)\subseteq\mathsf{row}(C)}\left\lVert
A-X\right\rVert_{F}^{2}\leq(1+\epsilon)\left\lVert
A-A_{k}\right\rVert_{F}^{2}.$

Recall that the sketch monotonicity for low-rank approximation says that
concatenating two sketching matrices $S_{1}$ and $S_{2}$ will not increase the
error compared to the single sketch matrix $S_{1}$ or $S_{2}$, Now matter how
$S_{1}$ and $S_{2}$ are constructed. (see Section D.1 and Section 4 in
[IVY19])

###### Proof.

We first consider the first condition. From the condition that
$\tau_{i}(B)\geq\frac{1}{\beta}\tau_{i}(A)$ we know that if we sample
$m=O(\beta\cdot(k\log k+k/\epsilon))$ rows according to $\tau_{i}(A)$. The
actual probability that the $i$-th row of $B$ gets sampled is

$1-(1-\tau_{i}(A))^{m}=O(m\cdot\tau_{i}(A))=O\left((k\log
k+k/\epsilon)\cdot\tau_{i}(B)\right)\;.$

From $\sum_{i}\tau_{i}(B)\leq 2k$ and Lemma E.2 (recall the sketch
monotonicity property for LRA), we have that with probability at least $9/10$,
$S_{2}$ is a matrix such that

$\min_{\operatorname{rank-\mathnormal{k}}X:\mathsf{row}(X)\subseteq\mathsf{row}(S_{2}B)}\left\lVert
B-X\right\rVert_{F}^{2}\leq(1+\epsilon)\left\lVert
B-B_{k}\right\rVert_{F}^{2}.$

Hence, since $S=[\begin{smallmatrix}S_{1}\\\ S_{2}\end{smallmatrix}]$, from
the the sketch monotonicity property for LRA we have that

$\min_{\operatorname{rank-\mathnormal{k}}X:\mathsf{row}(X)\subseteq\mathsf{row}(SB)}\left\lVert
B-X\right\rVert_{F}^{2}\leq(1+\epsilon)\left\lVert
B-B_{k}\right\rVert_{F}^{2}\;.$

Now we consider the second condition. From the coupling lemma, there exists a
joint distribution $\mu=(Z,W)$ such that (i) the marginal distributions of $Z$
and $W$ are the sampling probability distributions $p$ and $q$, respectively,
where $p_{i}=\frac{\tau_{i}(A)}{\sum_{i}\tau_{i}(A)}$ and
$q_{i}=\frac{\tau_{i}(B)}{\sum_{i}\tau_{i}(B)}$, and (ii)
$\Pr_{(Z,W)\sim\mu}\left[Z\neq W\right]=d_{\mathrm{TV}}(p,q)\;,$

Suppose that $\\{Z_{i}\\}_{i\leq m}$ and $\\{W_{i}\\}_{i\leq m}$ are a
sequence of $m=O(k\log k+k/\epsilon)$ samples from $[n]$ according to joint
distribution $\mu$. Let $S$ be the set of index $i$ such that $Z_{i}\neq
W_{i}$. From the assumption that $d_{\mathrm{tv}}(p,q)\leq\beta$, we get that

$\Pr\left[Z_{i}\neq W_{i}\right]=d_{\mathrm{tv}}(p,q)\leq\beta\;,$

and

$\mathbb{E}[|S|]=\sum_{i}\Pr\left[Z_{i}\neq W_{i}\right]\leq\beta m.$

From Markov’s inequality we get that with probability at least $1-1.1\beta$,
$|S|\leq 1/(1.1\beta)\cdot\beta m=\frac{10}{11}m$. Let $T$ be the set of index
$i$ such that $Z_{i}=W_{i}$. We have that with probability at least
$1-1.1\beta$, $|T|\geq m-\frac{10}{11}m=\Omega(k\log k+k/\epsilon)$. Because
that $\\{W_{i}\\}_{i\in T}$ is i.i.d samples according to $q$ and the actual
sample we take is $\\{Z_{i}\\}_{i\in T}$ from $p$. From Lemma E.2 we get that
with probability at least $9/10$, the row space of $B_{T}$ satisfies

$\min_{\operatorname{rank-\mathnormal{k}}X:\mathsf{row}(X)\subseteq\mathsf{row}(B_{T})}\left\lVert
B-X\right\rVert_{F}^{2}\leq(1+\epsilon)\left\lVert
B-B_{k}\right\rVert_{F}^{2}\;.$

Similarly, from the the sketch monotonicity property we have that with
probability at least $0.9-1.1\beta$

$\min_{\operatorname{rank-\mathnormal{k}}X:\mathsf{row}(X)\subseteq\mathsf{row}(SB)}\left\lVert
B-X\right\rVert_{F}^{2}\leq(1+\epsilon)\left\lVert
B-B_{k}\right\rVert_{F}^{2}\;.\qed$

###  E.2 Proof of Theorem 6.1

First we prove the following lemma.

######  Lemma E.3.

Let $\delta\in(0,1/m]$. It holds with probability at least $1-\delta$ that

$\sup_{x\in\operatorname{colsp}(A)}\left|\left\lVert
Sx\right\rVert_{2}^{2}-\left\lVert
x\right\rVert_{2}^{2}\right|\leq\epsilon\left\lVert x\right\rVert_{2}^{2},$

provided that

$\displaystyle m\gtrsim\epsilon^{-2}((d+\log
m)\min\\{\log^{2}(d/\epsilon),\log^{2}m\\}+d\log(1/\delta)),$ $\displaystyle
1\gtrsim\epsilon^{-2}\nu((\log
m)\min\\{\log^{2}(d/\epsilon),\log^{2}m\\}+\log(1/\delta))\log(1/\delta).$

###### Proof.

We shall adapt the proof of Theorem 5 in [BDN15] to our setting. Let $T$
denote the unit sphere in $\operatorname{colsp}(A)$ and set the sparsity
parameter $s=1$. Observe that $\left\lVert Sx\right\rVert_{2}^{2}=\left\lVert
x_{I}\right\rVert_{2}^{2}+\left\lVert Sx_{I^{c}}\right\rVert_{2}^{2}$, and so
it suffices to show that

$\Pr\left\\{\left|\left\lVert
S^{\prime}x_{I^{c}}\right\rVert_{2}^{2}-\left\lVert
x_{I^{c}}\right\rVert_{2}^{2}\right|>\epsilon\right\\}\leq\delta$

for $x\in T$. We make the following definition, as in (2.6) of [BDN15]:

$A_{\delta,x}:=\sum_{i=1}^{m}\sum_{j\in I^{c}}\delta_{ij}x_{j}e_{i}\otimes
e_{j},$

and thus, $S^{\prime}x_{I^{c}}=A_{\delta,x}\sigma$. Also by
$\operatorname*{{\mathbf{E}}}\left\lVert
S^{\prime}x_{I^{c}}\right\rVert_{2}^{2}=\left\lVert
x_{I^{c}}\right\rVert_{2}^{2}$, one has

$\sup_{x\in T}\left|\left\lVert
S^{\prime}x_{I^{c}}\right\rVert_{2}^{2}-\left\lVert
x_{I^{c}}\right\rVert_{2}^{2}\right|=\sup_{x\in T}\left|\left\lVert
A_{\delta,x}\sigma\right\rVert_{2}^{2}-\operatorname*{{\mathbf{E}}}\left\lVert
A_{\delta,x}\sigma\right\rVert_{2}^{2}\right|.$ (E.1)

Now, in (2.7) of [BDN15] we instead define a semi-norm

$\left\lVert x\right\rVert_{\delta}=\max_{1\leq i\leq m}\left(\sum_{j\in
I^{c}}\delta_{ij}x_{j}^{2}\right)^{1/2}.$

Then (2.8) continues to hold, and (2.9) as well as (2.10) continue to hold if
the supremum in the left-hand side is replaced with the left-hand side of
(E.1). At the beginning of Theorem 5, we define $U^{(i)}$ to be $U$, but each
row $j\in I^{c}$ is multiplied by $\delta_{ij}$ and each row $j\in I$ is
zeroed out. Then we have in the first step of (4.5) that

$\sum_{j\in I^{c}}\delta_{ij}\left|\sum_{k=1}^{d}g_{k}\langle
f_{k},e_{j}\rangle\right|^{2}\leq\left\lVert U^{(i)}g\right\rVert_{2}^{2},$

instead of equality. One can verify that the rest of (4.5) goes through. It
remains true that
$\left\lVert\cdot\right\rVert_{\delta}\leq(1/\sqrt{s})\left\lVert\cdot\right\rVert_{2}$,
and thus (4.6) holds. One can verify that the rest of the proof of Theorem 5
in [BDN15] continues to hold if we replace $\sum_{j=1}^{n}$ with $\sum_{j\in
I^{c}}$ and $\max_{1\leq j\leq n}$ with $\max_{j\in I^{c}}$, noting that

$\operatorname*{{\mathbf{E}}}\sum_{j\in I^{c}}\delta_{ij}\left\lVert
P_{E}e_{j}\right\rVert_{2}^{2}=\frac{s}{m}\sum_{j\in I^{c}}\langle
P_{E}e_{j},e_{j}\rangle\leq\frac{s}{m}d$

and

$\operatorname*{{\mathbf{E}}}(U^{(i)})^{\ast}U^{(i)}=\sum_{j\in
I^{c}}(\operatorname*{{\mathbf{E}}}\delta_{ij})u_{j}u_{j}^{\ast}\preceq\frac{1}{m}.$

Thus, the symmetrization inequalities on

$\left\lVert\sum_{j\in I^{c}}\delta_{ij}\left\lVert
P_{E}e_{j}\right\rVert_{2}^{2}\right\rVert_{L_{\delta}^{p}}\quad\text{and}\quad\left\lVert\sum_{j\in
I^{c}}\delta_{ij}u_{j}u_{j}^{\ast}\right\rVert_{L_{\delta}^{p}}$

continue to hold. The result then follows, observing that $\max_{j\in
I^{c}}\left\lVert P_{E}e_{j}\right\rVert^{2}\leq\nu$. ∎

The subspace embedding guarantee now follows as a corollary.

######  Theorem 6.1.

Let $\nu=\epsilon/d$. Suppose that
$m=\Omega((d/\epsilon^{2})(\mathop{\mathrm{polylog}}(1/\epsilon)+\log(1/\delta)))$,
$\delta\in(0,1/m)$ and
$d=\Omega((1/\epsilon)\mathop{\mathrm{polylog}}(1/\epsilon)\log^{2}(1/\delta))$.
Then, there exists a distribution on $S$ with $m+|I|$ rows such that

$\Pr\left\\{\forall x\in\operatorname{colsp}(A),\left|\left\lVert
Sx\right\rVert_{2}^{2}-\left\lVert
x\right\rVert_{2}^{2}\right|>\epsilon\left\lVert
x\right\rVert_{2}^{2}\right\\}\leq\delta.$

###### Proof.

One can verify that the two conditions in Lemma E.3 are satisfied if

$\displaystyle
m\gtrsim\frac{d}{\epsilon^{2}}\left(\mathop{\mathrm{polylog}}(\frac{d}{\epsilon})+\log\frac{1}{\delta}\right),$
$\displaystyle
d\gtrsim\frac{1}{\epsilon}\left(\log\frac{1}{\delta}\right)\left(\mathop{\mathrm{polylog}}(\frac{d}{\epsilon})+\log\frac{1}{\delta}\right).$

The last condition is satisfied if

$d\gtrsim\frac{1}{\epsilon}\left(\log^{2}\frac{1}{\delta}\right)\mathop{\mathrm{polylog}}\left(\frac{1}{\epsilon}\right).\qed$

###  E.3 Proof of Lemma 6.2

###### Proof.

On the one hand, since $Q=SAR$ is an orthogonal matrix, we have

$\left\lVert x\right\rVert_{2}=\left\lVert Qx\right\rVert_{2}=\left\lVert
SARx\right\rVert_{2}.$ (E.2)

On the other hand, the assumption implies that

$\left\lVert(ARx)^{T}(ARx)-x^{T}x\right\rVert_{2}\leq\epsilon\left\lVert
x\right\rVert_{2}^{2},$

that is,

$(1-\epsilon)\left\lVert x\right\rVert_{2}^{2}\leq\left\lVert
ARx\right\rVert_{2}^{2}\leq(1+\epsilon)\left\lVert x\right\rVert_{2}^{2}.$
(E.3)

Combining both (E.2) and (E.3) leads to

$\sqrt{1-\epsilon}\left\lVert SARx\right\rVert_{2}\leq\left\lVert
ARx\right\rVert_{2}\leq\sqrt{1+\epsilon}\left\lVert SARx\right\rVert_{2},\\\
\forall x\in\mathbb{R}^{d}.$

Equivalently, it can be written as

$\frac{1}{\sqrt{1+\epsilon}}\left\lVert SAy\right\rVert_{2}\leq\left\lVert
Ay\right\rVert_{2}\leq\frac{1}{\sqrt{1-\epsilon}}\left\lVert
SAy\right\rVert_{2},\quad\forall y\in\mathbb{R}^{d}.$

The claimed result follows from the fact that $1/\sqrt{1+\epsilon}\geq
1-\epsilon$ and $1/\sqrt{1-\epsilon}\leq 1+\epsilon$ whenever
$\epsilon\in(0,\frac{\sqrt{5}-1}{2}]$. ∎

##  Appendix F Location Optimization in CountSketch: Greedy Search

While the position optimization idea is simple, one particularly interesting
aspect is that it is provably better than a random placement in some scenarios
(Theorem. F.1). Specifically, it is provably beneficial for LRA when inputs
follow the spiked covariance model or Zipfian distributions, which are common
for real data.

Spiked covariance model. Every matrix $A\in\mathbb{R}^{n\times d}$ from the
distribution $\mathcal{A}_{sp}(s,\ell)$ has $s<k$ “heavy” rows
$A_{r_{1}},\cdots,A_{r_{s}}$ of norm $\ell>1$. The indices of the heavy rows
can be arbitrary, but must be the same for all members of
$\mathcal{A}_{sp}(s,\ell)$ and are unknown to the algorithm. The remaining
(“light”) rows have unit norm. In other words, let
$\mathcal{R}=\\{r_{1},\ldots,r_{s}\\}$. For all rows $A_{i},i\in[n]$,
$A_{i}=\ell\cdot v_{i}$ if $i\in\mathcal{R}$ and $A_{i}=v_{i}$ otherwise,
where $v_{i}$ is a uniformly random unit vector.

Zipfian on squared row norms. Every $A\in\mathbb{R}^{n\times
d}\sim\mathcal{A}_{zipf}$ has rows which are uniformly random and orthogonal.
Each $A$ has $2^{i+1}$ rows of squared norm $n^{2}/2^{2i}$ for
$i\in[1,\ldots,\mathcal{O}(\log(n))]$. We also assume that each row has the
same squared norm for all members of $\mathcal{A}_{zipf}$.

######  Theorem F.1.

Consider a matrix $A$ from either the spiked covariance model or a Zipfian
distribution. Let $S_{L}$ denote a CountSketch constructed by Algorithm 3 that
optimizes the positions of the non-zero values with respect to $A$. Let
$S_{C}$ denote a CountSketch matrix. Then there is a fixed $\eta>0$ such that,
$\min_{\operatorname{rank-\mathnormal{k}}X\in\mathrm{rowsp}(S_{L}A)}\left\|X-A\right\|_{F}^{2}\leq(1-\eta)\min_{\operatorname{rank-\mathnormal{k}}X\in\mathrm{rowsp}(S_{C}A)}\left\|X-A\right\|_{F}^{2}$

######  Remark F.2.

Note that the above theorem implicitly provides an upper bound on the
generalization error of the greedy placement method on the two distributions
that we considered in this paper. More precisely, for each of these two
distributions, if $\Pi$ is learned via our greedy approach over a set of
sampled training matrices, the solution returned by the sketching algorithm
using $\Pi$ over any (test) matrix $A$ sampled from the distribution has error
at most
$(1-\eta)\min_{\operatorname{rank-\mathnormal{k}}X\in\mathrm{rowsp}(S_{C}A)}\left\|X-A\right\|_{F}^{2}$.

A key structural property of the matrices from these two distributions that is
crucial in our analysis is the $\epsilon$-almost orthogonality of their rows
(i.e., (normalized) pairwise inner products are at most $\epsilon$). Hence, we
can find a QR-factorization of the matrix of such vectors where the upper
diagonal matrix $R$ has diagonal entries close to $1$ and entries above the
diagonal are close to $0$.

To state our result, we first provide an interpretation of the location
optimization task as a selection of hash function for the rows of $A$. Note
that left-multiplying $A$ by CountSketch $S\in\mathbb{R}^{m\times n}$ is
equivalent to hashing the rows of $A$ to $m$ bins with coefficients in $\\{\pm
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$. The
intuition is that our greedy approach separates heavy-norm rows (which are
important “directions” in the row space) into different bins.

Proof Sketch of Theorem F.1 The first step is to observe that in the greedy
algorithm, when rows are examined according to a non-decreasing order of
squared norms, the algorithm will isolate rows into their singleton bins until
all bins are filled. In particular, this means that the heavy norm rows will
all be isolated—e.g., for the spiked covariance model, Lemma F.8 presents the
formal statement.

Next, we show that none of the rows left to be processed (all light rows) will
be assigned to the same bin 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. This is the main place we use the properties of the aforementioned
distributions and the fact that each heavy row is already mapped to a
singleton bin. Overall, we show that at the end of the algorithm no light row
will be assigned to the bins that contain heavy rows—the formal statement and
proof for the spiked covariance model is in Lemma F.12.

Finally, we can interpret the randomized construction of CountSketch as a
“balls and bins” experiment. In particular, considering the heavy rows, we
compute the expected number of bins (i.e., rows in $S_{C}A$) that contain a
heavy row. Note that the expected number of rows in $S_{C}A$ that do not
contain any heavy row is $k\cdot(1-{1\over k})^{s}\geq k\cdot e^{-{s\over
k-1}}$. Hence, the number of rows in $S_{C}A$ that contain a heavy row of $A$
is at most $k(1-e^{-{s\over k-1}})$. Thus, at least $s-k(1-e^{-{s\over 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 the solution corresponding to $S_{C}$ is larger than the squared loss of
the solution corresponding to $S_{L}$, the CountSketch constructed by
Algorithm 3—please see Lemma F.14 for the formal statement of its proof.

#### Preliminaries and notation.

Left-multiplying $A$ by a CountSketch $S\in\mathbb{R}^{m\times n}$ is
equivalent to hashing the rows of $A$ to $m$ bins with coefficients in $\\{\pm
1\\}$. The greedy algorithm proceeds through the rows of $A$ (in some order)
and decides which bin to hash to, which we can think of as adding an entry to
$S$. We will denote the bins by $b_{i}$ and their summed contents by $w_{i}$.

###  F.1 Spiked covariance model with sparse left singular vectors.

To recap, every matrix $A\in\mathbb{R}^{n\times d}$ from the distribution
$\mathcal{A}_{sp}(s,\ell)$ has $s<k$ “heavy” rows
($A_{r_{1}},\cdots,A_{r_{s}}$) of norm $\ell>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: let $\mathcal{R}=\\{r_{1},\ldots,r_{s}\\}$. For all rows
$A_{i},i\in[n]$:

$\displaystyle A_{i}=\left\\{\begin{array}[]{ll}\ell\cdot v_{i}&\mbox{if
}i\in\mathcal{R}\\\ v_{i}&\mbox{o.w.}\end{array}\right.$

where $v_{i}$ is a uniformly random unit vector.

We also assume that $S_{r},S_{g}\in\mathbb{R}^{k\times n}$ and that the greedy
algorithm proceeds in a non-increasing row norm order.

#### 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 is $\epsilon$-almost
orthogonal (i.e., the magnitudes of their pairwise inner products are at most
$\epsilon$).

######  Observation F.3.

Let $v_{1},\cdots,v_{n}\in\mathbb{R}^{d}$ be a set of random unit vectors.
Then with probability $1-1/\mathop{\mathrm{poly}}(n)$, we have $|\langle
v_{i},v_{j}\rangle|\leq 2\sqrt{\log n\over d},\forall\;i<j\leq n$.

We define $\overline{\epsilon}=2\sqrt{\log n\over d}$.

######  Observation F.4.

Let $u_{1},\cdots,u_{t}$ be a set of vectors such that for each pair of
$i<j\leq t$, $|\langle\frac{u_{i}}{\left\lVert
u_{i}\right\rVert},\frac{u_{j}}{\left\lVert
u_{j}\right\rVert}\rangle|\leq\epsilon$, and
$g_{i},\cdots,g_{j}\in\\{-1,1\\}$. Then,

$\displaystyle\sum_{i=1}^{t}\left\|u_{i}\right\|^{2}_{2}-2\epsilon\sum_{i<j\leq
t}\left\|u_{i}\right\|_{2}\left\|u_{j}\right\|_{2}\leq\left\|\sum_{i=1}^{t}g_{i}u_{i}\right\|_{2}^{2}\leq\sum_{i=1}^{t}\left\|u_{i}\right\|^{2}_{2}+2\epsilon\sum_{i<j\leq
t}\left\|u_{i}\right\|_{2}\left\|u_{j}\right\|_{2}$ (F.1)

Next, a straightforward consequence of $\epsilon$-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 F.5.

Let $u_{1},\cdots,u_{t}\in\mathbb{R}^{d}$ be a set of unit vectors such that
for any pair of $i<j\leq t$, $|\langle u_{i},u_{j}\rangle|\leq\epsilon$ where
$\epsilon=O(t^{-2})$. There exists an orthonormal basis $e_{1},\cdots,e_{t}$
for the subspace spanned by $u_{1},\cdots,u_{t}$ such that for each $i\leq t$,
$u_{i}=\sum_{j=1}^{i}a_{i,j}e_{j}$ where $a_{i,i}^{2}\geq
1-\sum_{j=1}^{i-1}j^{2}\cdot\epsilon^{2}$ and for each $j<i$, $a_{i,j}^{2}\leq
j^{2}\epsilon^{2}$.

###### Proof.

We follow the Gram-Schmidt process to construct the orthonormal basis
$e_{1},\cdots,e_{t}$ of the space spanned by $u_{1},\cdots,u_{t}$, by first
setting $e_{1}=u_{1}$ and then processing $u_{2},\cdots,u_{t}$, one-by-one.

The proof is by induction. We show that once the first $j$ vectors
$u_{1},\cdots,u_{j}$ are processed, the statement of the lemma holds for these
vectors. Note that the base case of the induction trivially holds as
$u_{1}=e_{1}$. Next, suppose that the induction hypothesis holds for the first
$\ell$ vectors $u_{1},\cdots,u_{\ell}$.

######  Claim F.6.

For each $j\leq\ell$, $a_{\ell+1,j}^{2}\leq j^{2}\epsilon^{2}$.

###### Proof.

The proof of the claim is itself by induction. Note that, for $j=1$ and using
the fact that $|\langle u_{1},u_{\ell+1}\rangle|\leq\epsilon$, the statement
holds and $a_{\ell+1,1}^{2}\leq\epsilon^{2}$. Next, suppose that the statement
holds for all $j\leq i<\ell$. Then using that $|\langle
u_{i+1},u_{\ell+1}\rangle|\leq\epsilon$,

$\displaystyle|a_{\ell+1,i+1}|$ $\displaystyle\leq(|\langle
u_{\ell+1},u_{i+1}|+\sum_{j=1}^{i}|a_{\ell+1,j}|\cdot|a_{i+1,j}|)/|a_{i+1,i+1}|$
$\displaystyle\leq(\epsilon+\sum_{j=1}^{i}j^{2}\epsilon^{2})/|a_{i+1,i+1}|\quad\rhd\text{by
the induction hypothesis on $a_{\ell+1,j}$ for $j\leq i$}$
$\displaystyle\leq(\epsilon+\sum_{j=1}^{i}j^{2}\epsilon^{2})/({1-\sum_{j=1}^{i}j^{2}\cdot\epsilon^{2}})^{1/2}\quad\rhd\text{by
the induction hypothesis on $a_{i+1,i+1}$}$
$\displaystyle\leq(\epsilon+\sum_{j=1}^{i}j^{2}\epsilon^{2})\cdot({1-\sum_{j=1}^{i}j^{2}\cdot\epsilon^{2}})^{1/2}\cdot(1+2\cdot\sum_{j=1}^{i}j^{2}\epsilon^{2})$
$\displaystyle\leq(\epsilon+\sum_{j=1}^{i}j^{2}\epsilon^{2})\cdot(1+2\cdot\sum_{j=1}^{i}j^{2}\epsilon^{2})$
$\displaystyle\leq\epsilon((\sum_{j=1}^{i}j^{2}\epsilon)\cdot(1+4\epsilon\cdot\sum_{j=1}^{i}j^{2}\epsilon)+1)$
$\displaystyle\leq\epsilon(i+1)\quad\rhd\text{by $\epsilon=O(t^{-2})$}$

∎

Finally, since $\left\|u_{\ell+1}\right\|^{2}_{2}=1$,
$a_{\ell+1,\ell+1}^{2}\geq 1-\sum_{j=1}^{\ell}j^{2}\epsilon^{2}$. ∎

######  Corollary F.7.

Suppose that $\overline{\epsilon}=O(t^{-2})$. There exists an orthonormal
basis $e_{1},\cdots,e_{t}$ for the space spanned by the randomly picked
vectors $v_{1},\cdots,v_{t}$, of unit norm, so that for each $i$,
$v_{i}=\sum_{j=1}^{i}a_{i,j}e_{j}$ where $a_{i,i}^{2}\geq
1-\sum_{j=1}^{i-1}j^{2}\cdot\overline{\epsilon}^{2}$ and for each $j<i$,
$a_{i,j}^{2}\leq j^{2}\cdot\overline{\epsilon}^{2}$.

###### Proof.

The proof follows from Lemma F.5 and the fact that the set of vectors
$v_{1},\cdots,v_{t}$ is $\overline{\epsilon}$-almost orthogonal (by
Observation F.3). ∎

The first main step is to show that the greedy algorithm (with non-increasing
row norm ordering) will isolate rows into their own bins until all bins are
filled. In particular, this means that the heavy rows (the first to be
processed) will all be isolated.

We note that because we set $\mathrm{rank}(SA)=k$, the $k$-rank approximation
cost is the simplified expression $\left\lVert
AVV^{\top}-A\right\rVert_{F}^{2}$, where $U\Sigma V^{\top}=SA$, rather than
$\left\lVert[AV]_{k}V^{\top}-A\right\rVert_{F}^{2}$. This is just the
projection cost onto $\mathsf{row}(SA)$. Also, we observe that minimizing this
projection cost is the same as maximizing the sum of squared projection
coefficients:

$\displaystyle\operatorname*{arg\,min}_{S}\left\lVert
A-AVV^{\top}\right\rVert_{F}^{2}$
$\displaystyle=\operatorname*{arg\,min}_{S}\sum_{i\in[n]}\left\lVert
A_{i}-(\langle A_{i},v_{1}\rangle v_{1}+\ldots+\langle A_{i},v_{k}\rangle
v_{k})\right\rVert_{2}^{2}$
$\displaystyle=\operatorname*{arg\,min}_{S}\sum_{i\in[n]}(\left\lVert
A_{i}\right\rVert_{2}^{2}-\sum_{j\in[k]}\langle A_{i},v_{j}\rangle^{2})$
$\displaystyle=\operatorname*{arg\,max}_{S}\sum_{i\in[n]}\sum_{j\in[k]}\langle
A_{i},v_{j}\rangle^{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 F.8.

For any matrix $A$ or batch of matrices $\mathcal{A}$, 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\leq k$, we consider the choice of assigning $A_{i}$ to an
empty bin versus an occupied bin. Without loss of generality, let this
occupied bin be $b_{i-1}$, which already contains $A_{i-1}$.

We consider the difference in cost for empty versus occupied. We will do this
cost comparison for $A_{j}$ with $j\leq i-2$, $j\geq i+1$, and finally,
$j\in\\{i-1,i\\}$.

First, we let $\\{e_{1},\ldots,e_{i}\\}$ be an orthonormal basis for
$\\{A_{1},\ldots,A_{i}\\}$ such that for each $r\leq i$,
$A_{r}=\sum_{j=1}^{r}a_{r,j}e_{j}$ where $a_{r,r}>0$. This exists by Lemma
F.5. Let $\\{e_{1},\ldots,e_{i-2},\overline{e}\\}$ be an orthonormal basis for
$\\{A_{1},\ldots,A_{i+2},A_{i-1}\pm A_{i}\\}$. Now,
$\overline{e}=c_{0}e_{i-1}+c_{1}e_{i}$ for some $c_{0},c_{1}$ because
$(A_{i-1}\pm
A_{i})-\operatorname{proj}_{\\{e_{1},\ldots,e_{i-2}\\}}(A_{i-1}\pm
A_{i})\in\operatorname{\mathsf{span}}(e_{i-1},e_{i})$. We note that
$c_{0}^{2}+c_{1}^{2}=1$ because we let $\overline{e}$ be a unit vector. We can
find $c_{0},c_{1}$ to be:

$\displaystyle
c_{0}={a_{i-1,i-1}+a_{i,i-1}\over\sqrt{(a_{i-1,i-1}+a_{i,i-1})^{2}+a_{i,i}^{2}}},\quad
c_{1}={a_{i,i}\over\sqrt{(a_{i-1,i-1}+a_{i,i-1})^{2}+a_{i,i}^{2}}}$

  1. 1.

$j\leq i-2$: The cost is zero for both cases because
$A_{j}\in\operatorname{\mathsf{span}}(\\{e_{1},\ldots,e_{i-2}\\})$.

  2. 2.

$j\geq i+1$: We compare the rewards (sum of squared projection coefficients)
and find that $\\{e_{1},\ldots,e_{i-2},\overline{e}\\}$ is no better than
$\\{e_{1},\ldots,e_{i}\\}$.

$\displaystyle\langle A_{j},\overline{e}\rangle^{2}$
$\displaystyle=(c_{0}\langle A_{j},e_{i-1}\rangle+c_{1}\langle
A_{j},e_{i}\rangle)^{2}$ $\displaystyle\leq(c_{1}^{2}+c_{0}^{2})(\langle
A_{j},e_{i-1}\rangle^{2}+\langle A_{j},e_{i}\rangle^{2})$
$\displaystyle\rhd\text{Cauchy-Schwarz inequality}$ $\displaystyle=\langle
A_{j},e_{i-1}\rangle^{2}+\langle A_{j},e_{i}\rangle^{2}$

  3. 3.

$j\in\\{i-1,i\\}$: We compute the sum of squared projection coefficients of
$A_{i-1}$ and $A_{i}$ onto $\overline{e}$:

$\displaystyle({1\over(a_{i-1,i-1}+a_{i,i-1})^{2}+a_{i,i}^{2}})\cdot(a_{i-1,i-1}^{2}(a_{i-1,i-1}+a_{i,i-1})^{2}+(a_{i,i-1}(a_{i-1,i-1}+a_{i,i-1})+a_{i,i}a_{i,i})^{2})$
$\displaystyle=({1\over(a_{i-1,i-1}+a_{i,i-1})^{2}+a_{i,i}^{2}})\cdot((a_{i-1,i-1}+a_{i,i-1})^{2}(a_{i-1,i-1}^{2}+a_{i,i-1}^{2})$
$\displaystyle\quad+a_{i,i}^{4}+2a_{i,i-1}a_{i,i}^{2}(a_{i-1,i-1}+a_{i,i-1}))$
(F.2)

On the other hand, the sum of squared projection coefficients of $A_{i-1}$ and
$A_{i}$ onto $e_{i-1}\cup e_{i}$ is:

$\displaystyle({(a_{i-1,i-1}+a_{i,i-1})^{2}+a_{i,i}^{2}\over(a_{i-1,i-1}+a_{i,i-1})^{2}+a_{i,i}^{2}})\cdot(a_{i-1,i-1}^{2}+a_{i,i-1}^{2}+a_{i,i}^{2})$
(F.3)

Hence, the difference between the sum of squared projections of $A_{i-1}$ and
$A_{i}$ onto $\overline{e}$ and $e_{i-1}\cup e_{i}$ is ((F.3) - (F.2))

$\displaystyle\quad{a_{i,i}^{2}((a_{i-1,i-1}+a_{i,i-1})^{2}+a_{i-1,i-1}^{2}+a_{i,i-1}^{2}-2a_{i,i-1}(a_{i-1,i-1}+a_{i,i-1}))\over(a_{i-1,i-1}+a_{i,i-1})^{2}+a_{i,i}^{2}}$
$\displaystyle={2a_{i,i}^{2}a_{i-1,i-1}^{2}\over(a_{i-1,i-1}+a_{i,i-1})^{2}+a_{i,i}^{2}}>0$

Thus, we find that $\\{e_{1},\ldots,e_{i}\\}$ is a strictly better basis than
$\\{e_{1},\ldots,e_{i-2},\overline{e}\\}$. This means the greedy algorithm
will choose to place $A_{i}$ 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 sum of squared projection coefficients, lower-bounding it
in the former case and upper-bounding it in the latter.

We introduce some results that will be used in Lemma F.12.

######  Claim F.9.

Let $A_{k+r},r\in[1,\ldots,n-k]$ be a light row not yet processed by the
greedy algorithm. Let $\\{e_{1},\ldots,e_{k}\\}$ be the Gram-Schmidt basis for
the current $\\{w_{1},\ldots,w_{k}\\}$. Let $\beta=\mathcal{O}(n^{-1}k^{-3})$
upper bound the inner products of the normalized $A_{k+r},w_{1},\ldots,w_{k}$.
Then, for any bin $i$, $\langle e_{i},A_{k+r}\rangle^{2}\leq\beta^{2}\cdot
k^{2}$.

###### Proof.

This is a straightforward application of Lemma F.5. From that, we have
$\langle A_{k+r},e_{i}\rangle^{2}\leq i^{2}\beta^{2}$, for $i\in[1,\ldots,k]$,
which means $\langle A_{k+r},e_{i}\rangle^{2}\leq k^{2}\beta^{2}$. ∎

######  Claim F.10.

Let $A_{k+r}$ be a light row that has been processed by the greedy algorithm.
Let $\\{e_{1},\ldots,e_{k}\\}$ be the Gram-Schmidt basis for the current
$\\{w_{1},\ldots,w_{k}\\}$. If $A_{k+r}$ is assigned to bin $b_{k-1}$
(w.l.o.g.), the squared projection coefficient of $A_{k+r}$ onto $e_{i},i\neq
k-1$ is at most $4\beta^{2}\cdot k^{2}$, where
$\beta=\mathcal{O}(n^{-1}k^{-3})$ upper bounds the inner products of
normalized $A_{k+r},w_{1},\cdots,w_{k}$.

###### Proof.

Without loss of generality, it suffices to bound the squared projection of
$A_{k+r}$ onto the direction of $w_{k}$ that is orthogonal to the subspace
spanned by $w_{1},\cdots,w_{k-1}$. Let $e_{1},\cdots,e_{k}$ be an orthonormal
basis of $w_{1},\cdots,w_{k}$ guaranteed by Lemma F.5. Next, we expand the
orthonormal basis to include $e_{k+1}$ so that we can write the normalized
vector of $A_{k+r}$ as $v_{k+r}=\sum_{j=1}^{k+1}b_{j}e_{j}$. By a similar
approach to the proof of Lemma F.5, for each $j\leq k-2$,
$b_{j}\leq\beta^{2}j^{2}$. Next, since $|\langle
w_{k},v_{k+r}\rangle|\leq\beta$,

$\displaystyle|b_{k}|$ $\displaystyle\leq{1\over|\langle
w_{k},e_{k}\rangle|}\cdot(|\langle
w_{k},v_{k+r}\rangle|+\sum_{j=1}^{k-1}|b_{j}\cdot\langle w_{k},e_{j}\rangle|)$
$\displaystyle\leq{1\over\sqrt{1-\sum_{j=1}^{k-1}\beta^{2}\cdot
j^{2}}}\cdot(\beta+\sum_{j=1}^{k-2}\beta^{2}\cdot j^{2}+(k-1)\cdot\beta)$
$\displaystyle\rhd|b_{k-1}|\leq 1$
$\displaystyle={\beta+\sum_{j=1}^{k-2}\beta^{2}\cdot
j^{2}\over\sqrt{1-\sum_{j=1}^{k-1}\beta^{2}\cdot j^{2}}}+(k-1)\beta$
$\displaystyle\leq
2(k-1)\beta-{\beta^{2}(k-1)^{2}\over\sqrt{1-\sum_{j=1}^{k-1}\beta^{2}\cdot
j^{2}}}$ $\displaystyle\rhd\text{similar to the proof of
Lemma~{}\ref{lem:almost-orth}}$ $\displaystyle<2\beta\cdot k$

Hence, the squared projection of $A_{k+r}$ onto $e_{k}$ is at most
$4\beta^{2}\cdot k^{2}\cdot\left\|A_{k+r}\right\|^{2}_{2}$. We assumed
$\left\lVert A_{k+r}\right\rVert=1$; hence, the squared projection of
$A_{k+r}$ onto $e_{k}$ is at most $4\beta^{2}\cdot k^{2}$. ∎

######  Claim F.11.

We assume that the absolute values of the inner products of vectors in
$v_{1},\cdots,v_{n}$ are at most $\overline{\epsilon}<1/(n^{2}\sum_{A_{i}\in
b}\left\|A_{i}\right\|_{2})$ and the absolute values of the inner products of
the normalized vectors of $w_{1},\cdots,w_{k}$ are at most
$\beta=\mathcal{O}(n^{-3}k^{-{3\over 2}})$. Suppose that bin $b$ contains the
row $A_{k+r}$. Then, the squared projection of $A_{k+r}$ onto the direction of
$w$ orthogonal to
$\operatorname{\mathsf{span}}(\\{w_{1},\cdots,w_{k}\\}\setminus\\{w\\})$ is at
most
${\left\|A_{k+r}\right\|^{4}_{2}\over\left\|w\right\|^{2}_{2}}+\mathcal{O}(n^{-2})$
and is at least
${\left\|A_{k+r}\right\|^{4}_{2}\over\left\|w\right\|^{2}_{2}}-\mathcal{O}(n^{-2})$.

###### Proof.

Without loss of generality, we assume that $A_{k+r}$ is mapped to $b_{k}$;
$w=w_{k}$. First, we provide an upper and a lower bound for $|\langle
v_{k+r},\overline{w}_{k}\rangle|$ where for each $i\leq k$, we let
$\overline{w}_{i}={w_{i}\over\left\|w_{i}\right\|_{2}}$ denote the normalized
vector of $w_{i}$. Recall that by definition
$v_{k+r}={A_{k+r}\over\left\|A_{k+r}\right\|_{2}}$.

$\displaystyle|\langle\overline{w}_{k},v_{k+r}\rangle|$
$\displaystyle\leq{\left\|A_{k+r}\right\|_{2}+\sum_{A_{i}\in
b_{k}}\overline{\epsilon}\left\|A_{i}\right\|_{2}\over\left\|w_{k}\right\|_{2}}$
$\displaystyle\leq{\left\|A_{k+r}\right\|_{2}+n^{-2}\over\left\|w_{k}\right\|_{2}}$
$\displaystyle\rhd\text{by $\overline{\epsilon}<{n^{-2}\over\sum_{A_{i}\in
b_{k}}\left\|A_{i}\right\|_{2}}$}$
$\displaystyle\leq{\left\|A_{k+r}\right\|_{2}\over\left\|w_{k}\right\|_{2}}+n^{-2}$
$\displaystyle\rhd\text{$\left\|w_{k}\right\|_{2}\geq 1$}$ (F.4)
$\displaystyle|\langle\overline{w}_{k},v_{k+r}\rangle|$
$\displaystyle\geq{\left\|A_{k+r}\right\|_{2}-\sum_{A_{i}\in
b_{k}}\left\|A_{i}\right\|_{2}\cdot\overline{\epsilon}\over\left\|w_{k}\right\|_{2}}$
$\displaystyle\geq{\left\|A_{k+r}\right\|_{2}\over\left\|w_{k}\right\|_{2}}-n^{-2}$
(F.5)

Now, let $\\{e_{1},\cdots,e_{k}\\}$ be an orthonormal basis for the subspace
spanned by $\\{w_{1},\cdots,w_{k}\\}$ guaranteed by Lemma F.5. Next, we expand
the orthonormal basis to include $e_{k+1}$ so that we can write
$v_{k+r}=\sum_{j=1}^{k+1}b_{j}e_{j}$. By a similar approach to the proof of
Lemma F.5, we can show that for each $j\leq k-1$,
$b_{j}^{2}\leq\beta^{2}j^{2}$. Moreover,

$\displaystyle|b_{k}|$
$\displaystyle\leq{1\over|\langle\overline{w}_{k},e_{k}\rangle|}\cdot(|\langle\overline{w}_{k},v_{k+r}\rangle|+\sum_{j=1}^{k-1}|b_{j}\cdot\langle\overline{w}_{k},e_{j}\rangle|)$
$\displaystyle\leq{1\over\sqrt{1-\sum_{j=1}^{k-1}\beta^{2}\cdot
j^{2}}}\cdot(|\langle\overline{w}_{k},v_{k+r}\rangle|+\sum_{j=1}^{k-1}\beta^{2}\cdot
j^{2})$ $\displaystyle\rhd\text{by Lemma~{}\ref{lem:almost-orth}}$
$\displaystyle\leq{1\over\sqrt{1-\sum_{j=1}^{k-1}\beta^{2}\cdot
j^{2}}}\cdot(n^{-2}+{\left\|A_{k+r}\right\|_{2}\over\left\|w_{k}\right\|_{2}}+\sum_{j=1}^{k-1}\beta^{2}\cdot
j^{2})$ $\displaystyle\rhd\text{by~{}\eqref{eq:up}}$ $\displaystyle<\beta\cdot
k+{1\over\sqrt{1-\beta^{2}k^{3}}}\cdot(n^{-2}+{\left\|A_{k+r}\right\|_{2}\over\left\|w_{k}\right\|_{2}})$
$\displaystyle\rhd\text{similar to the proof of Lemma~{}\ref{lem:almost-
orth}}$
$\displaystyle\leq\mathcal{O}(n^{-2})+(1+\mathcal{O}(n^{-2})){\left\|A_{k+r}\right\|_{2}\over\left\|w_{k}\right\|_{2}}$
$\displaystyle\rhd\text{by $\beta=\mathcal{O}(n^{-3}k^{-{3\over 2}})$}$
$\displaystyle\leq{\left\|A_{k+r}\right\|_{2}\over\left\|w_{k}\right\|_{2}}+\mathcal{O}(n^{-2})$
$\displaystyle\rhd{\left\|A_{k+r}\right\|_{2}\over\left\|w_{k}\right\|_{2}}\leq
1$

and,

$\displaystyle|b_{k}|$
$\displaystyle\geq{1\over|\langle\overline{w}_{k},e_{k}\rangle|}\cdot(|\langle\overline{w}_{k},v_{k+r}\rangle|-\sum_{j=1}^{k-1}|b_{j}\cdot\langle\overline{w}_{k},e_{j}\rangle|)$
$\displaystyle\geq|\langle\overline{w}_{k},v_{k+r}\rangle|-\sum_{j=1}^{k-1}\beta^{2}\cdot
j^{2}$ $\displaystyle\rhd\text{since
$|\langle\overline{w}_{k},e_{k}\rangle|\leq 1$}$
$\displaystyle\geq{\left\|A_{k+r}\right\|_{2}\over\left\|w_{k}\right\|_{2}}-n^{-2}-\sum_{j=1}^{k-1}\beta^{2}\cdot
j^{2}$ $\displaystyle\rhd\text{by~{}\eqref{eq:low}}$
$\displaystyle\geq{\left\|A_{k+r}\right\|_{2}\over\left\|w_{k}\right\|_{2}}-\mathcal{O}(n^{-2})$
$\displaystyle\rhd\text{by $\beta=\mathcal{O}(n^{-3}k^{-{3\over 2}})$}$

Hence, the squared projection of $A_{k+r}$ onto $e_{k}$ is at most
${\left\|A_{k+r}\right\|^{4}_{2}\over\left\|w_{k}\right\|^{2}_{2}}+\mathcal{O}(n^{-2})$
and is at least
${\left\|A_{k+r}\right\|^{4}_{2}\over\left\|w_{k}\right\|^{2}_{2}}-\mathcal{O}(n^{-2})$.
∎

Now, we show that at the end of the algorithm no light row will be assigned to
the bins that contain heavy rows.

######  Lemma F.12.

We assume that the absolute values of the inner products of vectors in
$v_{1},\cdots,v_{n}$ are at most $\overline{\epsilon}<\min\\{n^{-2}k^{-{5\over
3}},(n\sum_{A_{i}\in w}\left\|A_{i}\right\|_{2})^{-1}\\}$. At iteration $k+r$,
the greedy algorithm will assign the light row $A_{k+r}$ to a bin that does
not contain a heavy row.

###### Proof.

The proof is by induction. Lemma F.8 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
$A_{k+r}$ avoids and collides with a heavy row.

First, we upper bound $\beta=\max_{i\neq j\leq k}|\langle
w_{i},w_{j}\rangle|/(\left\|w_{i}\right\|_{2}\left\|w_{j}\right\|_{2})$. Let
$c_{i}$ and $c_{j}$ respectively denote the number of rows assigned to $b_{i}$
and $b_{j}$.

$\displaystyle\beta=\max_{i\neq j\leq k}{|\langle
w_{i},w_{j}\rangle|\over\left\|w_{i}\right\|_{2}\left\|w_{j}\right\|_{2}}$
$\displaystyle\leq{c_{i}\cdot
c_{j}\cdot\overline{\epsilon}\over\sqrt{c_{i}-2\overline{\epsilon}c_{i}^{2}}\cdot\sqrt{c_{j}-2\overline{\epsilon}c_{j}^{2}}}$
$\displaystyle\rhd~{}\text{Observation \ref{obser:norm}}$
$\displaystyle\leq{16\overline{\epsilon}\sqrt{c_{i}c_{j}}}$
$\displaystyle\rhd\overline{\epsilon}\leq n^{-2}k^{-5/3}$ $\displaystyle\leq
n^{-1}k^{-{5\over 3}}$ $\displaystyle\rhd\overline{\epsilon}\leq
n^{-2}k^{-5/3}$

#### 1\. If $A_{k+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 $A_{1},\cdots,A_{s}$ onto
$\mathsf{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