# Adaptive Sketches for Robust Regression with Importance Sampling

Sepideh Mahabadi Microsoft Research, Redmond. E-mail<EMAIL_ADDRESS>David
P. Woodruff School of Computer Science, Carnegie Mellon University. E-mail:
<EMAIL_ADDRESS>Samson Zhou School of Computer Science, Carnegie Mellon
University. E-mail<EMAIL_ADDRESS>
###### Abstract

We introduce data structures for solving robust regression through stochastic
gradient descent (SGD) by sampling gradients with probability proportional to
their norm, i.e., importance sampling. Although SGD is widely used for large
scale machine learning, it is well-known for possibly experiencing slow
convergence rates due to the high variance from uniform sampling. On the other
hand, importance sampling can significantly decrease the variance but is
usually difficult to implement because computing the sampling probabilities
requires additional passes over the data, in which case standard gradient
descent (GD) could be used instead. In this paper, we introduce an algorithm
that approximately samples $T$ gradients of dimension $d$ from nearly the
optimal importance sampling distribution for a robust regression problem over
$n$ rows. Thus our algorithm effectively runs $T$ steps of SGD with importance
sampling while using sublinear space and just making a single pass over the
data. Our techniques also extend to performing importance sampling for second-
order optimization.

##  1 Introduction

Given a matrix ${\mathbf{A}}\in\mathbb{R}^{n\times d}$ with rows
${\mathbf{a}}_{1},\ldots,{\mathbf{a}}_{n}\in\mathbb{R}^{d}$ and a
measurement/label vector ${\mathbf{b}}\in\mathbb{R}^{n}$, we consider the
standard regression problem

$\underset{{\mathbf{x}}\in\mathbb{R}^{d}}{\min}{\mathcal{L}}({\mathbf{x}}):=\sum_{i=1}^{n}M(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i}),$

where $M:\mathbb{R}\to\mathbb{R}^{\geq 0}$ is a function, called a measure
function, that satisfies $M(x)=M(-x)$ and is non-decreasing in $|x|$. An _$M$
-estimator_ is a solution to this minimization problem and for appropriate
choices of $M$, can combine the low variance of $L_{2}$ regression with the
robustness of $L_{1}$ regression against outliers.

The Huber norm, for example, is defined using the measure function
$H(x)=\frac{x^{2}}{2\tau}$ for $|x|\leq\tau$ and $H(x)=|x|-\frac{\tau}{2}$ for
$|x|>\tau$, where $\tau$ is a threshold that governs the interpolation between
$L_{2}$ loss for small $|x|$ and $L_{1}$ loss for large $|x|$. Indeed, it can
often be more reasonable to have robust treatment of large residuals due to
outliers or errors and Gaussian treatment of small residuals [GS99]. Thus the
Huber norm is especially popular and “recommended for almost all situations”
[Zha97], because it is the “most robust” [Hub92] due to “the useful
computational and statistical properties implied by the convexity and
smoothness” [CW15] of its measure function, which is differentiable at all
points.

Since the measure function $H$ for the Huber norm and more generally, the
measure function $M$ for many common measure functions is convex, we can
consider the standard convex _finite-sum form_ optimization problem
$\underset{{\mathbf{x}}\in\mathbb{R}^{d}}{\min}F({\mathbf{x}}):=\frac{1}{n}\sum_{i=1}^{n}f_{i}({\mathbf{x}})$,
where $f_{1},\ldots,f_{n}:\mathbb{R}^{d}\to\mathbb{R}$ is a sequence of convex
functions that commonly represent loss functions. Whereas gradient descent
(GD) performs the update rule
${\mathbf{x}}_{t+1}={\mathbf{x}}_{t}-\eta_{t}\nabla F({\mathbf{x}}_{t})$ on
the iterate ${\mathbf{x}}_{t}$ at iterations $t=1,2,\ldots,T$, stochastic
gradient descent (SGD) [RM51, NY83, NJLS09] picks $i_{t}\in[n]$ in iteration
$t$ with probability $p_{i_{t}}$ and performs the update rule
${\mathbf{x}}_{t+1}={\mathbf{x}}_{t}-\frac{\eta_{t}}{np_{i_{t}}}\nabla
f_{i_{t}}({\mathbf{x}}_{t})$, where $\nabla f_{i_{t}}$ is the gradient (or a
subgradient) of $f_{i_{t}}$ and $\eta_{t}$ is some predetermined learning
rate. Effectively, training example $i_{t}$ is sampled with probability
$p_{i_{t}}$ and the model parameters are updated using the selected example.
The SGD update rule only requires the computation of a single gradient at each
iteration and provides an unbiased estimator to the full gradient, compared to
GD, which evaluates $n$ individual gradients in each iteration and is
prohibitively expensive for large $n$. However, since SGD is often performed
with uniform sampling, so that the probability $p_{i,t}$111In contrast to
$p_{i,t}$, the term $p_{i_{t}}$ denotes the probability associated with the
specific index $i_{t}$ chosen at time $t$. of choosing index $i\in[n]$ at
iteration $t$ is $p_{i,t}=\frac{1}{n}$ at all times, the variance introduced
by the randomness of sampling a specific vector function can be a bottleneck
for the convergence rate of this iterative process. Thus, the subject of
variance reduction beyond uniform sampling has been well-studied in recent
years [RSB12, JZ13, DBL14, RHS+15, ZZ15, DLH16, NSW16, SRJ17, JG18, KF18,
STC18, QRG+19].

A common technique to reduce variance is importance sampling, where the
probabilities $p_{i,t}$ are chosen so that vector functions with larger
gradients are more likely to be sampled. One such setting of importance
sampling is to set the probability of sampling a gradient with probability
proportional to its $L_{2}$ norm, so that

$p_{i,t}=\frac{\left\lVert\nabla
f_{i}({\mathbf{x}}_{t})\right\rVert_{2}}{\sum_{j\in[t]}\left\lVert\nabla
f_{j}({\mathbf{x}}_{t})\right\rVert_{2}}.$

Under these sampling probabilities, importance sampling gives variance

$\sigma_{opt,t}^{2}=\frac{1}{n^{2}}\left(\left(\sum_{i=1}^{n}\left\lVert\nabla
f_{i}({\mathbf{x}}_{t})\right\rVert_{2}\right)^{2}-n^{2}\cdot\left\lVert\nabla
F({\mathbf{x}}_{t})\right\rVert_{2}^{2}\right),$

where we define the variance of a random vector ${\mathbf{v}}$ to be
$\operatorname*{Var}({\mathbf{v}}):=\mathbb{E}\left[\left\lVert{\mathbf{v}}\right\rVert^{2}_{2}\right]-\left\lVert\mathbb{E}\left[{\mathbf{v}}\right]\right\rVert_{2}^{2}$,
and we define $\sigma_{opt,t}^{2}$ to be the variance of the random vector
${\mathbf{v}}$ produced at time $t$ by importance sampling.

By comparison, the probabilities for uniform sampling $p_{i,t}=\frac{1}{n}$
imply $\sigma_{t}^{2}=\operatorname*{Var}\left(\frac{1}{np_{i_{t},t}}\right)$
and thus the variance $\sigma_{uni,t}^{2}$ for uniform sampling satisfies

$\sigma_{uni,t}^{2}=\frac{1}{n^{2}}\left(n\sum_{i=1}^{n}\left\lVert\nabla
f_{i}({\mathbf{x}}_{t})\right\rVert_{2}^{2}-n^{2}\cdot\left\lVert\nabla
F({\mathbf{x}}_{t})\right\rVert_{2}^{2}\right).$

By the root mean square-arithmetic mean inequality, the variance of importance
sampling is always at most the variance of uniform sampling, and can be
significantly less. Hence $\sigma_{opt,t}^{2}\leq\sigma_{uni,t}^{2}$, so that
the variance at each step is reduced, possibly substantially, by performing
importance sampling instead of uniform sampling.

To see examples where uniform sampling an index performs significantly worse
than importance sampling, consider $\nabla
f_{i}({\mathbf{x}})=\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle\cdot{\mathbf{a}}_{i}$.
Then for ${\mathbf{A}}={\mathbf{a}}_{1}\circ\ldots\circ{\mathbf{a}}_{n}$:

###### Example 1.1

When the non-zero entries of the input ${\mathbf{A}}$ are concentrated in a
small number of vectors ${\mathbf{a}}_{i}$, uniform sampling will frequently
sample gradients that are small and make little progress, whereas importance
sampling will rarely do so. In an extreme case, the input ${\mathbf{A}}$ can
contain exactly one non-zero vector ${\mathbf{a}}_{i}$ and importance sampling
will always output the full gradient, whereas uniform sampling will only find
the non-zero row with probability $\frac{1}{n}$, so that
$\sigma_{uni,t}^{2}=n\cdot\sigma_{opt,t}^{2}$.

###### Example 1.2

It may be that all rows of ${\mathbf{A}}$ have large magnitude, but
${\mathbf{x}}$ is nearly orthogonal to most of the rows of ${\mathbf{A}}$, but
is well-aligned with row ${\mathbf{a}}_{r}$. Then
$\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle\cdot{\mathbf{a}}_{i}$ is small in
magnitude for most $i$, but
$\langle{\mathbf{a}}_{r},{\mathbf{x}}\rangle\cdot{\mathbf{a}}_{r}$ is large so
uniform sampling will often output small gradients while importance sampling
will output $\langle{\mathbf{a}}_{r},{\mathbf{x}}\rangle\cdot{\mathbf{a}}_{r}$
with high probability, so that it can be that
$\sigma_{uni,t}^{2}=\Omega(n)\cdot\sigma_{opt,t}^{2}$.

###### Example 1.3

More generally for a parameter $\nu\in[0,1]$, if a $\nu$-fraction of the $n$
gradients lengths are bounded by $\mathcal{O}\left(n\right)$ while the other
$1-\nu$ fraction of the $n$ gradient lengths are bounded by
$\operatorname*{poly}(d)\ll n$, then the variance for uniform sampling
satisfies $\sigma_{uni,t}^{2}=\mathcal{O}\left(\nu
n^{2}\right)+\operatorname*{poly}(d)$ while the variance for importance
sampling satisfies
$\mathcal{O}\left(\nu^{2}n^{2}\right)+\operatorname*{poly}(d)$.

In fact, it follows from the Cauchy-Schwarz inequality that the importance
sampling probability distribution is the _optimal_ distribution for variance
reduction.

However, computing the probability distribution for importance sampling
requires computing the gradients in each round, which creates a “chicken and
egg” problem because computing the gradients is too expensive in the first
place, or else it is feasible to just run gradient descent. Unfortunately,
computing the sampling probabilities in each iteration often requires
additional passes over the data, e.g., to compute the gradients in each step,
which is generally prohibitively expensive. This problem often prevents
importance sampling from being widely deployed.

In this paper, we overcome this problem by introducing efficient sketches for
a wide range of $M$-estimators that can enable importance sampling _without_
additional passes over the data. Using our sketches for various measure
functions, we give a time-efficient algorithm that _provably_ approximates the
optimal importance sampling distribution within a constant factor. Thus we can
surprisingly simulate $T$ steps of SGD with (nearly) the optimal sampling
distribution, while only using a single pass over the data, which avoids the
aforementioned problem.

###### Theorem 1.4

Given an input matrix ${\mathbf{A}}\in\mathbb{R}^{n\times d}$ whose rows
arrive sequentially in a data stream along with the corresponding labels of a
measurement vector ${\mathbf{b}}\in\mathbb{R}^{d}$, and a measure function $M$
whose derivative is a continuous union of piecewise constant or linear
functions, there exists an algorithm that performs $T$ steps of SGD with
variance within a constant factor of the optimal sampling distribution. The
algorithm uses $\tilde{\mathcal{O}}\left(nd^{2}+Td^{2}\right)$ pre-processing
time and $Td^{2}\,\operatorname*{polylog}(Tnd)$ words of space.

For $T$ iterations, both GD and optimal importance sampling SGD require $T$
passes over the data, while our algorithm only requires a single pass over the
data and uses sublinear space for $nd\gg Td^{2}$. We remark that although the
number $T$ of iterations for SGD may be large, a major advantage of GD and SGD
with importance sampling is a significantly smaller number of iterations than
SGD with uniform sampling, e.g., as in Example 1.1 and Example 1.2, so we
should expect $n\gg T$. In particular from known results about the convergence
of SGD, e.g., see Theorem 1.8, if the diameter of the search space and the
gradient lengths $\left\lVert\nabla f_{i}(x_{t})\right\rVert_{2}$ are both
bounded by $\operatorname*{poly}(d)$, then we should expect
$T\propto\operatorname*{poly}(d)$ even for uniform sampling. More generally,
if $\nu n$ of the gradients have lengths $\Theta(n)$, while the remaining
gradients have lengths $\operatorname*{poly}(d)\ll n$, then Example 1.3 and
Theorem 1.8 show that the number of steps necessary for convergence for
uniform sampling satisfies
$T\propto\mathcal{O}\left(\nu^{2}n^{4}\right)+\operatorname*{poly}(d)$, while
the number of steps necessary for convergence for importance sampling
satisfies $T\propto\mathcal{O}\left(\nu^{4}n^{4}\right)$. Thus for $\nu
n=\mathcal{O}\left(n^{C}\right)$ for $C<1$, i.e., a sublinear number of
gradients have lengths that exceed the input size, we have
$T=\mathcal{O}\left(n^{4C}\right)$ and hence for $C<\frac{1}{4}$, we have
roughly $T=o(n)$ steps are necessary for convergence for SGD with importance
sampling.

Finally, we show in the appendix that our techniques can also be generalized
to perform importance sampling for second-order optimization.

###  1.1 Our Techniques

In addition to our main conceptual contribution that optimal convergence rate
of importance sampling for SGD can surprisingly be achieved (up to constant
factors) without the “chicken and egg” problem of separately computing the
sampling probabilities, we present a number of technical contributions that
may be of independent interest. Our first observation is that if we were only
running a single step of importance sampling for SGD, then we just want a
subroutine that outputs a gradient
$G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-b_{i},{\mathbf{a}}_{i})$ with
probability proportional to its norm
$\|G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})\|_{2}$.

####  $G$-sampler.

In particular, we need an algorithm that reads a matrix
${\mathbf{A}}={\mathbf{a}}_{1}\circ\ldots\circ{\mathbf{a}}_{n}\in\mathbb{R}^{n\times
d}$ and a vector ${\mathbf{x}}\in\mathbb{R}^{d}$ given _after processing_ the
matrix ${\mathbf{A}}$, and outputs (a rough approximation to) a gradient
$G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-b_{i},{\mathbf{a}}_{i})$ with
probability roughly

$\frac{\|G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})\|_{2}}{\sum_{j=1}^{n}\|G(\langle{\mathbf{a}}_{j},{\mathbf{x}}\rangle-
b_{j},{\mathbf{a}}_{j})\|_{2}}.$

We call such an algorithm a _$G$ -sampler_ and introduce such a single-pass,
memory-efficient sampler with the following guarantees:

###### Theorem 1.5

Given an $(\alpha,\varepsilon)$-smooth gradient $G$, there exists an algorithm
Sampler that outputs a noisy vector ${\mathbf{v}}$ such that
$\|{\mathbf{v}}-{\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i})\|_{2}\leq\alpha\|{\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i})|\|_{2}$ and
$\mathbb{E}\left[{\mathbf{v}}\right]={\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i})$ is
$\left(1\pm\mathcal{O}\left(\varepsilon\right)\right)\frac{\|G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})\|_{2}}{\sum_{j\in[n]}\|G(\langle{\mathbf{a}}_{j},{\mathbf{x}}\rangle-
b_{j},{\mathbf{a}}_{j})\|_{2}}+\frac{1}{\operatorname*{poly}(n)}$. The
algorithm uses
$d^{2}\,\operatorname*{poly}\left(\log(nT),\frac{1}{\alpha}\right)$ update
time per arriving row and
$Td^{2}\,\operatorname*{poly}\left(\log(nT),\frac{1}{\alpha}\right)$ total
bits of space.

We say a gradient $G$ is $(\alpha,\varepsilon)$-smooth if a vector
${\mathbf{u}}$ that satisfies
$\|{\mathbf{u}}-{\mathbf{v}}\|_{2}\leq\alpha\|{\mathbf{v}}\|_{2}$ implies that
$(1-\varepsilon)\|G({\mathbf{v}})\|_{2}\leq\|G({\mathbf{u}})\|_{2}\leq(1+\varepsilon)\|G({\mathbf{v}})\|_{2}$.
In particular, the measure functions discussed in Section 1.2 have gradients
that are $(\mathcal{O}\left(\varepsilon\right),\varepsilon)$-smooth. For
example, the subgradient of the Huber estimator is
${\mathbf{a}}_{i}\cdot\operatorname*{sgn}(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i})$ for $|\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-b_{i}|>\tau$, which
may change sign when $\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-b_{i}$ is
close to zero, but its norm will remain the same. Moreover, the form
$G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-b_{i},{\mathbf{a}}_{i})$
necessitates that the gradient can be computed strictly from the two
quantities $\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-b_{i}$ and
${\mathbf{a}}_{i}$. Thus Theorem 1.5 implies that our algorithm can also
compute a noisy vector ${\mathbf{v}}^{\prime}$ such that
$\|{\mathbf{v}}^{\prime}-G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})\|_{2}\leq\varepsilon\|G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})\|_{2}$.

Observe that an instance of Sampler in Theorem 1.5 can be used to simulate a
single step of SGD with importance sampling and thus $T$ independent instances
of Sampler provide an oracle for $T$ steps of SGD with importance sampling.
However, this naïve implementation does not suffice for our purposes because
the overall runtime would be $\tilde{\mathcal{O}}\left(Tnd^{2}\right)$ so it
would be more efficient to just run $T$ iterations of GD. Nevertheless, our
$G$-sampler is a crucial subroutine towards our final algorithm and we briefly
describe it here.

An alternative definition of $G$-sampler is given in [JWZ22]. In their
setting, the goal is to sample a coordinate $i\in[n]$ of a frequency vector
$f$ with probability proportional to $G(f_{i})$, where $G$ in their notation
is a measure function rather than a gradient. However, because the $G$-sampler
of [JWZ22] is not a linear sketch, their approach cannot be easily generalized
to our setting where the sampling probability of each row ${\mathbf{a}}_{i}$
is a function of $\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle$, but the vector
${\mathbf{x}}$ arrives after the stream is already processed.

Furthermore, because the loss function $f$ may not be scale-invariant, then we
should also not expect its gradient to be scale invariant at any location
${\mathbf{x}}\in\mathbb{R}^{d}$, i.e., $\nabla f(C{\mathbf{x}})\neq
C^{p}\cdot\nabla f({\mathbf{x}})$ for any constants $p,C>0$. Hence, our
subroutine Sampler cannot use the standard $L_{p}$ sampler framework used in
[JST11, AKO11, JW18, MRWZ20], which generally rescales each row of
${\mathbf{a}}_{i}$ by the inverse of a uniform or exponential random variable.
A somewhat less common design for $L_{p}$ samplers is a level set and
subsampling approach [MW10, JLS+21], due to their suboptimal dependencies on
the accuracy parameter $\varepsilon$. Fortunately, because we require
$\varepsilon=\mathcal{O}\left(1\right)$ to achieve a constant factor
approximation, we can use the level set and subsampling paradigm as a starting
point for our algorithm. Because the algorithms of [MW10, JLS+21] only sample
entries of a vector implicitly defined from a data stream, our $G$-sampler
construction must (1) sample rows of a matrix implicitly defined from a data
stream and (2) permit updates to the sampling probabilities implicitly defined
through multiplication of each row ${\mathbf{a}}_{i}$ with a vector
${\mathbf{x}}$ that only arrives after the stream is processed.

####  $G$-sampler through level sets and subsampling.

To illustrate our method and simplify presentation here, we consider $L_{2}$
regression with gradient
${\mathbf{A}}_{i}{\mathbf{x}}:=\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle\cdot{\mathbf{a}}_{i}$,
by folding in the measurement vector ${\mathbf{b}}$ into a column of
${\mathbf{A}}$ – our full algorithm in Section 2 handles both sampling
distributions defined with respect to the norm of a general gradient $G$ in
the form of Theorem 1.5, as well as an independent measurement vector
${\mathbf{b}}$.

We first partition the rows of ${\mathbf{A}}$ into separate geometrically
growing classes based on their $L_{2}$ norms, so that for instance, class
$C_{k}$ contains the rows ${\mathbf{a}}_{i}$ of ${\mathbf{A}}$ such that
$2^{k}\leq\|{\mathbf{a}}_{i}\|_{2}<2^{k+1}$. We build a separate data
structure for each class $C_{k}$, which resembles the framework for $L_{p}$
norm estimation [IW05]. We would like to use the approximate contributions of
the level sets $\Gamma_{1},\ldots,\Gamma_{K}$, with
$K=\mathcal{O}\left(\frac{\log n}{\alpha}\right)$, toward the total mass
$F_{2}(S)=\sum_{i=1}^{n}\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}\right\rVert_{2}$,
where a level set $\Gamma_{j}$ is informally the set of rows
${\mathbf{a}}_{i}$ with
$\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}\right\rVert_{2}\in\left[\frac{F_{2}(S)}{(1+\alpha)^{j-1}},\frac{F_{2}(S)}{(1+\alpha)^{j}}\right]$
and the contribution of a level set $\Gamma_{j}$ is
$\sum_{i\in\Gamma_{j}}\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}\right\rVert_{2}$.
Then we could first sample a level set $\Gamma_{j}$ from a class $C_{k}$ and
then uniformly select a row ${\mathbf{a}}_{i}$ among those in $\Gamma_{j}$.
Indeed, we can run a generalized version of the $L_{2}$ heavy-hitter algorithm
CountSketch [CCF04] on the stream $S$ to identify the level set $\Gamma_{1}$,
since its rows will be heavy with respect to $F_{2}(S)$. However, the rows of
the level sets $\Gamma_{j}$ for large $j$ may not be detected by CountSketch.
Thus, we create $L=\mathcal{O}\left(\log n\right)$ substreams
$S_{1},\ldots,S_{L}$, so that substream $S_{\ell}$ samples each row of
${\mathbf{A}}$ with probability $2^{-\ell+1}$, and run an instance of
CountSketch on each substream $S_{\ell}$ to detect the rows of each level set
and thus estimate the contribution of each level set.

#### Sampling from level sets with small contribution.

However, there is still an issue – some level sets have contribution that is
too small to well-approximate with small variance. For example, if there is a
single row with contribution $\frac{F_{2}(S)}{(1+\alpha)^{j}}$, then it might
not survive the subsampling at a level $S_{\ell}$ that is used to detect it,
in which case it will never be sampled. Alternatively, if it is sampled, it
will be rescaled by a large amount, so that its level set will be sampled with
abnormally large probability. Instead of handling this large variance, we
instead add a number of dummy rows to each level set, to ensure that their
contributions are all “significant” and thus be well-approximated.

Now we have “good” approximations to the contributions of each level set
within a class, so we can first select a level set with probability
proportional to the approximate contributions of each level set and then
uniformly sample a row from the level set. Of course, we may uniformly sample
a dummy row, in which case we say the algorithm fails to acquire a sample. We
show that the contribution added by the dummy rows is a constant fraction, so
this only happens with a constant probability. Thus with
$\mathcal{O}\left(\log\frac{1}{\delta}\right)$ constant number of independent
samples, we can boost the probability of successfully acquiring a sample to
$1-\delta$ for any $\delta\in(0,1]$. We then set
$\delta=\frac{1}{\operatorname*{poly}(n,T,d)}$. An illustration of the entire
process can be viewed in Figure 1.

$G$-samplerClasses$\vdots$$C_{3}$$C_{2}$$C_{1}$Levels$\Gamma_{K}$$\vdots$$\Gamma_{2}$$\Gamma_{1}$RowsDummy
rows$G(\langle{\mathbf{a}}_{j},{\mathbf{x}}\rangle-
b_{j},{\mathbf{a}}_{j})$$\vdots$$G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})$ Fig. 1: $G$-sampler first samples a specific class,
then samples a specific level from the class, then samples a specific row from
the level. Sampled items are denoted by solid lines while rejected items are
denoted by dashed lines. The classes partition the rows ${\mathbf{a}}_{k}$ of
the input ${\mathbf{A}}$ matrix by $\|{\mathbf{a}}_{k}\|_{2}$, while the
levels further partition the rows ${\mathbf{a}}_{i}$ of each class based on
$G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-b_{i},{\mathbf{a}}_{i})$.

#### Unbiased samples.

Unfortunately, CountSketch using
$\mathcal{O}\left(\frac{1}{\alpha^{2}}\right)$ buckets only guarantees
additive $\alpha\,L_{2}(S)$ error to a particular row with constant
probability. To achieve the standard “for-all” guarantee across all $n$ rows,
an estimate for each row ${\mathbf{a}}_{i}$ is then output by taking the row
with the median length across $\mathcal{O}\left(\log n\right)$ independent
instances. However, the median row is no longer unbiased, which could
potentially affect the convergence guarantees of SGD. Instead, we use $d$
separate instances of CountSketch, so that each instance handles a separate
coordinate of the vector. Thus if the goal is to output a noisy estimate to
${\mathbf{a}}_{i}$, we have a separate CountSketch report each coordinate
$({\mathbf{a}}_{i})_{j}$, where $j\in[d]$. It can be shown that the median of
each estimated coordinate is an unbiased estimate to the true value
$({\mathbf{a}}_{i})_{j}$ of the coordinate because the probability mass
function is symmetric about the true value for each coordinate. Moreover, the
error to a single coordinate $({\mathbf{a}}_{i})_{j}$ may be large relative to
the value of the coordinate in the case that $({\mathbf{a}}_{i})_{j}$ is not
heavy with respect to $\\{({\mathbf{a}}_{i})_{j}\\}_{i\in[n]}$. However, we
show that the “overall” error to all coordinates of ${\mathbf{a}}_{i}$ is
small relative to $\|{\mathbf{a}}_{i}\|_{2}$, due to ${\mathbf{a}}_{i}$ being
a “heavy” row at the appropriate subsampling level.

#### Stochastic gradient descent with importance sampling.

The main problem with the proposed $G$-sampler is that it requires reading the
entire matrix ${\mathbf{A}}$ but it cannot be repeatedly used without
incurring dependency issues. In particular, if a sampler at the first
iteration of SGD outputs a gradient ${\mathbf{A}}_{i_{1}}{\mathbf{x}}_{1}$
that is used to construct ${\mathbf{x}}_{2}$, then ${\mathbf{x}}_{2}$ is not
independent of the sampler and thus the same sampler should not be used to
sample ${\mathbf{A}}_{i_{2}}{\mathbf{x}}_{2}$. This suggests that if we want
to perform $T$ steps of SGD with importance sampling, then we would require
$T$ separate data structures, which would require $Tnd$ time to construct for
dense matrices, but then we might as well just perform full gradient descent!

Instead in Section 3, we partition the matrix ${\mathbf{A}}$ among multiple
buckets and create a sampler for each bucket. Now as long as each bucket
should have been sampled a single time, then we will have a fresh sampler with
independent randomness for each time a new bucket is sampled. If we perform
$T$ steps of SGD with importance sampling, then roughly $T$ buckets should
suffice, but we cannot guarantee that each bucket is sampled a single time.
For example, if only a single ${\mathbf{A}}_{i}$ is non-zero, then whichever
bucket ${\mathbf{A}}_{i}$ is assigned to will be sampled every single time.

Now the challenge is identifying the submatrices
${\mathbf{A}}_{i}={\mathbf{a}}_{i}^{\top}{\mathbf{a}}_{i}$ that may be sampled
multiple times, since we do not know the values of the vectors
${\mathbf{x}}_{1},\ldots,{\mathbf{x}}_{T}$ a priori. Fortunately, we know that
$\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}_{t}\right\rVert_{2}$ can only be
large if ${\mathbf{a}}_{i}$ has high sensitivity, where we define the
sensitivity for a row ${\mathbf{a}}_{i}$ in ${\mathbf{A}}$ to be the quantity
$\max_{{\mathbf{x}}\in\mathbb{R}^{d}}\frac{\left\lVert{\mathbf{a}}_{i}^{\top}(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle)\right\rVert_{2}}{\sum_{j=1}^{n}\left\lVert{\mathbf{a}}_{j}^{\top}(\langle{\mathbf{a}}_{j},{\mathbf{x}}\rangle)\right\rVert_{2}}$.
Thus if a block is sampled multiple times, then one of its rows must have
large sensitivity.

Hence, we would like to identify the buckets that contain any row with
sensitivity at least $\frac{1}{T}$ and create $T$ independent samplers for
those buckets so that even if the same bucket is sampled all $T$ times, there
will be a fresh sampler available. Crucially, the process of building separate
buckets for the rows with the large sensitivities can be identified in just a
single pass over the data.

We remark that since each row has sensitivity
$\max_{{\mathbf{x}}\in\mathbb{R}^{d}}\frac{\left\lVert{\mathbf{a}}_{i}^{\top}(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle)\right\rVert_{2}}{\sum_{j=1}^{n}\left\lVert{\mathbf{a}}_{j}^{\top}(\langle{\mathbf{a}}_{j},{\mathbf{x}}\rangle)\right\rVert_{2}}$,
then it can be shown that the sum of the sensitivities is
$\mathcal{O}\left(d\log n\right)$ by partitioning the rows into
$\mathcal{O}\left(\log n\right)$ classes $C_{1},C_{2},\ldots$ of exponentially
increasing norm, so that ${\mathbf{a}}_{i}\in C_{\ell}$ if
$2^{\ell}\leq\left\lVert{\mathbf{a}}_{i}\right\rVert_{2}<2^{\ell+1}$. We then
note that the sensitivity of each row ${\mathbf{a}}_{i}\in C_{\ell}$ is upper
bounded by
$\max_{{\mathbf{x}}\in\mathbb{R}^{d}}\frac{|\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle|}{\sum_{{\mathbf{a}}_{j}\in
C_{\ell}}|\langle{\mathbf{a}}_{j},{\mathbf{x}}\rangle|}$. However, this latter
quantity is an $L_{1}$ sensitivity, whose sum is known to be bounded by
$\mathcal{O}\left(d\right)$, e.g., [CWW19]. Thus the sum of the sensitivities
in each class is at most $\mathcal{O}\left(d\right)$ and so for a matrix
${\mathbf{A}}$ whose entries are polynomially bounded by $n$, the sum of the
sensitivities is at most $\mathcal{O}\left(d\log n\right)$.

Unfortunately, since the sensitivities sum to $\mathcal{O}\left(d\log
n\right)$, there can be up to $Td$ rows with sensitivity at least
$\frac{1}{T}$, so creating $T$ independent samplers corresponding to each of
these rows would yield $\Omega(T^{2}d)$ samplers, which is a prohibitive
amount of space. Instead, we simply remove the rows with large sensitivities
from the buckets and store them explicitly. We then show this approach still
avoids any sampler from being used multiple times across the $T$ iterations
while also enabling the data structure to just use
$\tilde{\mathcal{O}}\left(Td\right)$ samplers. Now since we can explicitly
consider the rows with sensitivities roughly at least $\frac{1}{T}$, then we
can use $\Theta(T)$ buckets in total to ensure that the remaining non-zero
entries of ${\mathbf{A}}$ are partitioned evenly across buckets that will only
require $\Theta(\log(Td))$ independent samplers. Intuition for our algorithm
appears in Figure 2.

${\mathbf{A}}$LightRowsHeavyRows(Store rows
explicitly)$\mathcal{O}\left(\log(Td)\right)$$\mathcal{O}\left(T\right)$$G$-sampler$G$-sampler$G$-sampler$G$-sampler$\ldots$$\ldots$$\ldots$$\ldots$$G$-sampler$G$-sampler$G$-sampler$G$-samplerSGD
steps$1$$2$$3$$\vdots$$T-1$$T$$\vdots$ Fig. 2: Our SGD algorithm explicitly
stores the heavy rows of the input matrix ${\mathbf{A}}$, i.e., the rows with
sensitivity $\Omega\left(\frac{1}{T}\right)$. Our algorithm partitions the
remaining rows of the input matrix into $\mathcal{O}\left(T\right)$ buckets
and initiates $\mathcal{O}\left(\log(Td)\right)$ instances of $G$-samplers on
each bucket.

###  1.2 Applications

In this section, we discuss applications of our result to commonly used loss
functions, such as $L_{p}$ loss or various $M$-estimators, e.g., [CW15, CWW19,
TWZ+22, PSZ22].

####  $L_{1}$ and $L_{2}$ regression.

The $L_{p}$ regression loss function is defined using
$f_{i}({\mathbf{x}})=|{\mathbf{a}}_{i}^{\top}{\mathbf{x}}-b_{i}|^{p}$. The
case $p=2$ corresponds to the standard least squares regression problem, while
$p=1$ corresponds to least absolute deviation regression, which is more robust
to outliers than least squares, but also less stable and with possibly
multiple solutions. For $p=1$, the subgradient is
${\mathbf{a}}_{i}\cdot\operatorname*{sgn}(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i})$ while for $p=2$, the subgradient is
$2{\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-b_{i})$.

#### Huber estimator.

As previously discussed, Huber loss [Hub92] is commonly used, e.g., [Zha97,
CW15], to achieve Gaussian properties for small residuals [GS99] and robust
properties for large residuals due to outliers or errors. The Huber estimator
is also within a constant factor of other $M$-estimators that utilize the
advantage of the $L_{1}$ loss function to minimize the impact of large
errors/outliers and that of the $L_{2}$ loss function to be convex, such as
the $L_{1}$-$L_{2}$ estimator and the Fair estimator [BAG16]. Given a
threshold $\tau>0$, the Huber loss $H$ is defined by
$H(x)=\frac{x^{2}}{2\tau}$ for $|x|\leq\tau$ and $H(x)=|x|-\frac{\tau}{2}$ for
$|x|>\tau$. Thus the subgradient for $H$ is
$\frac{{\mathbf{a}}_{i}}{\tau}(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i})$ for $|\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-b_{i}|\leq\tau$ and
${\mathbf{a}}_{i}\cdot\operatorname*{sgn}(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i})$ for $|\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-b_{i}|>\tau$.

#### Ridge regression.

It is often desirable for a solution ${\mathbf{x}}$ to be sparse. The natural
approach to encourage sparse solutions is to add a regularization
$\lambda\|{\mathbf{x}}\|_{0}$ term to the loss function, for some parameter
$\lambda>0$. However, since $\|{\mathbf{x}}\|_{0}$ is not convex, ridge
regression is often used as a convex relaxation that encourage sparse
solutions. The ridge regression loss function satisfies
$f_{i}({\mathbf{x}})=({\mathbf{a}}_{i}^{\top}{\mathbf{x}}-b_{i})^{2}+\lambda\,\|{\mathbf{x}}\|_{2}^{2}$
for each $i\in[n]$, so that $\lambda$ regularizes the penalty term associated
with the squared magnitude of ${\mathbf{x}}$. Higher values of $\lambda$ push
the optimal solution towards zero, which leads to lower variance, as a
particular coordinate has a smaller effect on the prediction. The gradient for
the ridge regression loss function satisfies $\nabla
f_{i}({\mathbf{x}})=2{\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i})+2\lambda{\mathbf{x}}$.

#### Lasso.

Another approach that encourages sparsity is using the $L_{1}$ regularization
instead of the $L_{2}$ regularization. Least absolute shrinkage and selection
operator (Lasso) regression uses the loss function
$f_{i}({\mathbf{x}})=({\mathbf{a}}_{i}^{\top}{\mathbf{x}}-b_{i})^{2}+\lambda\,\|{\mathbf{x}}\|_{1}$.
Whereas the penalty term associated for ridge regression will drive down the
Euclidean norm of ${\mathbf{x}}$ for larger $\lambda$, solutions with large
$L_{1}$ norm are still possible if the mass of ${\mathbf{x}}$ is spread across
a large number of coordinates. By contrast, the penalty term associated for
Lasso drives down the total magnitude of the coordinates of ${\mathbf{x}}$.
Thus, in this sense, Lasso tends to drive coordinates to zero and encourages
sparsity, which does not usually happen for ridge regression. The subgradient
for the Lasso regression loss function satisfies $\nabla
f_{i}({\mathbf{x}})=2{\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},{\mathbf{x}}-b_{i})+2\lambda\operatorname*{sgn}({\mathbf{x}})$,
where we abuse notation by using $\operatorname*{sgn}({\mathbf{x}})$ to denote
the coordinate-wise sign of the entries of ${\mathbf{x}}$.

#### Group lasso.

[YL06] proposed Group Lasso as a generalization to Lasso. Suppose the weights
in ${\mathbf{x}}$ can be grouped into $m$ groups:
${\mathbf{x}}^{(1)},\ldots,{\mathbf{x}}^{(m)}$. We group the columns of
${\mathbf{A}}={\mathbf{a}}_{1}\circ\ldots\circ{\mathbf{a}}_{n}$ so that
${\mathbf{A}}^{(i)}$ is the set of columns that corresponds to the weights in
${\mathbf{x}}^{(i)}$. The Group Lasso function is defined as
$f_{i}({\mathbf{x}})=({\mathbf{a}}_{i}^{\top}{\mathbf{x}}-b_{i})^{2}+\lambda\sum_{j=1}^{m}\sqrt{G_{j}}\|{\mathbf{x}}^{(j)}\|_{2}$,
where $G_{j}$ represents the number of weights in ${\mathbf{x}}^{(j)}$. Note
that Group Lasso becomes Lasso for $m=n$.

###  1.3 Preliminaries

For an integer $n>0$, we use $[n]$ to denote the set $\\{1,2,\ldots,n\\}$. We
use boldfaced font for variables that represent either vectors of matrices and
plain font to denote variables that represent scalars. We use the notation
$\tilde{\mathcal{O}}\left(\cdot\right)$ to suppress $\operatorname*{polylog}$
factors, so that $f(T,n,d)=\tilde{\mathcal{O}}\left(g(T,n,d)\right)$ implies
that $f(T,n,d)\leq g(T,n,d)\operatorname*{polylog}(Tnd)$. Let
${\mathbf{A}}\in\mathbb{R}^{n\times d}$ and
${\mathbf{B}}\in\mathbb{R}^{m\times d}$. We use $\circ$ to denote vertical
concatenation, so that
${\mathbf{A}}\circ{\mathbf{B}}=\begin{bmatrix}{\mathbf{A}}\\\
{\mathbf{B}}\end{bmatrix}$, and $\otimes$ to denote outer product, so that the
$(i,j)$-th entry of the matrix
${\mathbf{u}}\otimes{\mathbf{v}}\in\mathbb{R}^{m\times n}$ for
${\mathbf{u}}\in\mathbb{R}^{m}$ and ${\mathbf{v}}\in\mathbb{R}^{n}$ is
$u_{i}v_{j}$. For a vector ${\mathbf{v}}\in\mathbb{R}^{n}$, we let
$\left\lVert{\mathbf{v}}\right\rVert^{p}_{p}=\sum_{i=1}^{n}v_{i}^{p}$ and
$\left\lVert{\mathbf{v}}\right\rVert_{\infty}=\max_{i}|v_{i}|$. For a matrix
${\mathbf{A}}\in\mathbb{R}^{n\times d}$, we denote the Frobenius norm of
${\mathbf{A}}$ by
$\left\lVert{\mathbf{A}}\right\rVert_{F}=\sqrt{\sum_{i=1}^{n}\sum_{j=1}^{d}A_{i,j}^{2}}$.
We also use
$\left\lVert{\mathbf{A}}\right\rVert_{p}=\left(\sum_{i=1}^{n}\sum_{j=1}^{d}|A_{i,j}|^{p}\right)^{\frac{1}{p}}$.
For a function $f$, we use $\nabla f$ to denote its gradient.

###### Definition 1.6

A function $f:\mathbb{R}^{d}\to\mathbb{R}$ is _convex_ if $f({\mathbf{x}})\geq
f({\mathbf{y}})+\langle\nabla
f({\mathbf{y}}),{\mathbf{x}}-{\mathbf{y}}\rangle$ for all
${\mathbf{x}},{\mathbf{y}}\in\mathbb{R}^{d}$.

###### Definition 1.7

A continuously differentiable function $f:\mathbb{R}^{d}\to\mathbb{R}$ is
_$\mu$ -smooth_ if

$\left\lVert\nabla f({\mathbf{x}})-\nabla
f({\mathbf{y}})\right\rVert_{2}\leq\mu\left\lVert{\mathbf{x}}-{\mathbf{y}}\right\rVert_{2},$

for all ${\mathbf{x}},{\mathbf{y}}\in\mathbb{R}^{d}$. Then it follows, e.g.,
by Lemma 3.4 in [Bub15], that for every
${\mathbf{x}},{\mathbf{y}}\in\mathbb{R}^{d}$,

$|f({\mathbf{y}})-f({\mathbf{x}})-\langle\nabla
f({\mathbf{x}}),{\mathbf{y}}-{\mathbf{x}}\rangle|\leq\frac{\mu}{2}\left\lVert{\mathbf{y}}-{\mathbf{x}}\right\rVert_{2}^{2}.$

Recall that SGD offers the following convergence guarantees for smooth convex
functions:

###### Theorem 1.8

[NJLS09, Mek17] Let $F$ be a $\mu$-smooth convex function and
${\mathbf{x}}_{\mathsf{opt}}=\operatorname*{argmin}F({\mathbf{x}})$. Let
$\sigma^{2}$ be an upper bound for the variance of the unbiased estimator
across all iterations and
$\overline{{\mathbf{x}}_{k}}=\frac{{\mathbf{x}}_{1}+\ldots+{\mathbf{x}}_{k}}{k}$.
Let each step-size $\eta_{t}$ be $\eta\leq\frac{1}{\mu}$. Then for SGD with
initial position ${\mathbf{x}}_{0}$, and any value of $k$,

$\mathbb{E}\left[F(\overline{{\mathbf{x}}_{k}})-F({\mathbf{x}}_{\mathsf{opt}})\right]\leq\frac{1}{2\eta
k}\left\lVert{\mathbf{x}}_{0}-{\mathbf{x}}_{\mathsf{opt}}\right\rVert_{2}^{2}+\frac{\eta\sigma^{2}}{2}.$

This means that
$k=\mathcal{O}\left(\frac{1}{\varepsilon^{2}}\left(\sigma^{2}+\mu\left\lVert{\mathbf{x}}_{0}-{\mathbf{x}}_{\mathsf{opt}}\right\rVert_{2}^{2}\right)^{2}\right)$
iterations suffice to obtain an $\varepsilon$-approximate optimal value by
setting $\eta=\frac{1}{\sqrt{k}}$.

##  2 $G$-Sampler Algorithm

In this section, we describe our $G$-sampler, which reads a matrix
${\mathbf{A}}={\mathbf{a}}_{1}\circ\ldots\circ{\mathbf{a}}_{n}\in\mathbb{R}^{n\times
d}$ and a vector ${\mathbf{x}}\in\mathbb{R}^{d}$ given after processing the
matrix ${\mathbf{A}}$, and outputs a gradient
$G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-b_{i},{\mathbf{a}}_{i})$ among
the $n$ gradients $\\{G(\langle{\mathbf{a}}_{1},{\mathbf{x}}\rangle-
b_{1},{\mathbf{a}}_{1}),\ldots,G(\langle{\mathbf{a}}_{n},{\mathbf{x}}\rangle-
b_{n},{\mathbf{a}}_{n})\\}$ with probability roughly
$\frac{\|G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})\|_{2}}{\sum_{j=1}^{n}\|G(\langle{\mathbf{a}}_{j},{\mathbf{x}}\rangle-
b_{j},{\mathbf{a}}_{j})\|_{2}}$. However, it is not possible to exactly return
$G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-b_{i},{\mathbf{a}}_{i})$ using
sublinear space; we instead return a vector ${\mathbf{v}}$ such that
$\mathbb{E}\left[{\mathbf{v}}\right]=G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})$ and
$\|{\mathbf{v}}-G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})\|\leq\varepsilon\|G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})\|_{2}$. To achieve our $G$-sampler, we first require a
generalization of the standard $L_{2}$-heavy hitters algorithm CountSketch
[CCF04], which we describe in Section 2.1. We then describe our $G$-sampler in
full in Section 2.2.

###  2.1 Heavy-Hitters

Before describing our generalization of CountSketch, we first require the
following $F_{G}$ estimation algorithm that generalizes both well-known
frequency moment estimation algorithm of [AMS99, TZ04] and symmetric norm
estimation algorithm of [BBC+17] by leveraging the linear sketches used in
those data structures to support “post-processing” with multiplication by any
vector ${\mathbf{x}}\in\mathbb{R}^{d}$.

###### Theorem 2.1

[BBC+17] Given a constant $\varepsilon>0$ and an $(\alpha,\varepsilon)$-smooth
gradient $G$, there exists a one-pass streaming algorithm Estimator that takes
updates to entries of a matrix ${\mathbf{A}}\in\mathbb{R}^{n\times d}$, as
well as vectors ${\mathbf{x}}\in\mathbb{R}^{d}$ and
${\mathbf{b}}\in\mathbb{R}^{d}$ that arrive after the stream, and outputs a
quantity $\hat{F}$ such that
$(1-\varepsilon)\sum_{i\in[n]}\|G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})\|_{2}\leq\hat{F}\leq(1+\varepsilon)\sum_{i\in[n]}\|G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})\|_{2}$. The algorithm uses
$\frac{d^{2}}{\alpha^{2}}\,\operatorname*{polylog}(nT)$ bits of space and
succeeds with probability at least $1-\frac{1}{\operatorname*{poly}(n,T)}$.

We now describe a straightforward generalization of the $L_{2}$-heavy hitter
algorithm CountSketch so that (1) it can find the “heavy rows” of a matrix
${\mathbf{A}}={\mathbf{a}}_{1}\circ\ldots\circ{\mathbf{a}}_{n}\in\mathbb{R}^{n\times
d}$ rather than the “heavy coordinates” of a vector and (2) it supports post-
processing multiplication by a vector ${\mathbf{x}}\in\mathbb{R}^{d}$ that
arrives only after ${\mathbf{A}}$ is processed. Let
${\mathbf{A}}_{i}={\mathbf{a}}_{i}\otimes{\mathbf{a}}_{i}\in\mathbb{R}^{d\times
d}$ for all $i\in[n]$. We define $\operatorname*{tail}(c)$ to be the $n-c$
rows that do not include the top $c$ values of
$\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}\right\rVert_{2}$. For a given
$\varepsilon>0$, we say a block ${\mathbf{A}}_{i}$ with $i\in[n]$ is _heavy_
if
$\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}\right\rVert_{2}\geq\varepsilon\sum_{i\in\operatorname*{tail}(2/\varepsilon^{2})}\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}\right\rVert_{2}$.

The standard CountSketch algorithm for finding the $L_{2}$-heavy hitters among
the coordinates of a vector ${\mathbf{v}}$ of dimension $n$ works by hashing
the universe $[n]$ across $\mathcal{O}\left(\frac{1}{\varepsilon^{2}}\right)$
buckets. Each coordinate $i\in[n]$ is also given a random sign $\sigma_{i}$
and so the algorithm maintains the
$\mathcal{O}\left(\frac{1}{\varepsilon^{2}}\right)$ signed sums
$\sum\sigma_{i}x_{i}$ across all the coordinates hashed to each bucket. Then
to estimate $x_{i}$, the algorithm simply outputs $\sigma_{i}C_{h(i)}$, where
$C_{h(i)}$ represents the counter corresponding to the bucket to which
coordinate $i$ hashes. It can be shown that
$\mathbb{E}\left[\sigma_{i}C_{h(i)}\right]=x_{i}$, where the expectation is
taken over the random signs $\sigma$ and the choices of the hash functions.
Similarly, the variance of the estimator can be bounded to show that with
constant probability, the estimator has additive error
$\mathcal{O}\left(\varepsilon\right)\,\|{\mathbf{x}}\|_{2}^{2}$ to $x_{i}$
with constant probability. Thus if
$x_{i}>\varepsilon\,\|{\mathbf{x}}\|_{2}^{2}$, the algorithm will be able to
identify coordinate $i$ as a heavy-hitter (in part by allowing some false
positives). We give the algorithm in full in Algorithm 1.

Algorithm 1 Output heavy vectors
$(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle){\mathbf{a}}_{i}$, where
${\mathbf{x}}$ can be a vector that arrives after ${\mathbf{A}}$ is processed

1:Matrix ${\mathbf{A}}\in\mathbb{R}^{n\times d}$, vector
${\mathbf{x}}\in\mathbb{R}^{d}$, accuracy parameter $\varepsilon>0$, failure
parameter $\delta\in(0,1]$.

2:Noisy vectors ${\mathbf{a}}_{i}^{\top}{\mathbf{a}}_{i}{\mathbf{x}}$ with
$\left\lVert{\mathbf{a}}_{i}^{\top}{\mathbf{a}}_{i}{\mathbf{x}}\right\rVert^{2}_{2}\geq\varepsilon^{2}\sum_{i\in\operatorname*{tail}(2/\varepsilon^{2})}\left\lVert{\mathbf{a}}_{i}^{\top}{\mathbf{a}}_{i}{\mathbf{x}}\right\rVert^{2}_{2}$.

3:$b\leftarrow\Omega\left(\frac{1}{\delta\varepsilon^{4}}\right)$

4:Let $\mathcal{T}$ contain $b$ buckets, each initialized to the all zeros
$\mathbb{R}^{d\times d}$ matrix.

5:Let $\sigma_{i}\in\\{-1,+1\\}$ be drawn from $4$-wise independent family for
$i\in[n]$.

6:Let $h:[n]\to[b]$ be $2$-wise independent

7:Process ${\mathbf{A}}$:

8:Let ${\mathbf{A}}={\mathbf{a}}_{1}\circ\ldots\circ{\mathbf{a}}_{n}$, where
each ${\mathbf{a}}_{i}\in\mathbb{R}^{d}$.

9:for each $j=1$ to $n$ do

10: ${\mathbf{A}}_{j}\leftarrow{\mathbf{a}}_{j}\otimes{\mathbf{a}}_{j}$

11: Add $\sigma_{j}{\mathbf{A}}_{j}$ to the matrix in bucket $h(j)$.

12:Let ${\mathbf{M}}_{j}$ be the matrix in bucket $j$ of $\mathcal{T}$ for
$i\in[r],j\in[b]$.

13:Process ${\mathbf{x}}$:

14:for $j\in[b]$ do

15: ${\mathbf{v}}_{j}\leftarrow{\mathbf{M}}_{j}{\mathbf{x}}$

16:On query $k\in[n]$, report $\sigma_{k}{\mathbf{v}}_{h(k)}$.

Thus Algorithm 1 can be used to give the following guarantee by taking the
median of the norms of $\mathcal{O}\left(\log(nT)\right)$ copies, as well as
the vector that realizes the median.

###### Lemma 2.2

[MRWZ20] There exists an algorithm that uses
$\mathcal{O}\left(\frac{d^{2}}{\varepsilon^{2}}\,\log^{2}n\right)$ space and
outputs a set $S$ of indices so that with probability
$1-\frac{1}{\operatorname*{poly}(n,T)}$, for all $i\in[n]$, $i\in S$ if
$\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}\right\rVert_{2}\geq\varepsilon\sum_{j\in\operatorname*{tail}(2/\varepsilon^{2})}\left\lVert{\mathbf{A}}_{j}{\mathbf{x}}\right\rVert_{2}$
and $i\notin S$ if
$\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}\right\rVert_{2}\leq\frac{\varepsilon}{2}\,\sum_{j\in\operatorname*{tail}(2/\varepsilon^{2})}\left\lVert{\mathbf{A}}_{j}{\mathbf{x}}\right\rVert_{2}$.
The algorithm uses
$\mathcal{O}\left(\frac{d^{2}}{\varepsilon^{2}}\log^{2}(nT)\right)$ space.

However, the vector that realizes the median of the norms may no longer be an
unbiased estimate to each heavy-hitter. Unfortunately, we shall require
unbiased estimates to each heavy-hitter, because we will use estimated heavy-
hitters as unbiased gradients as part of SGD with importance sampling. Thus we
give an additional algorithm so that for each $i\in S$ reported by Algorithm
1, the algorithm outputs an _unbiased_ estimate to the vector
$(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle){\mathbf{a}}_{i}$ with a “small”
error, in terms of the total mass
$\sum_{i\in\operatorname*{tail}(2/\varepsilon^{2})}\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}\right\rVert_{2}$
excluding the largest $\frac{2}{\varepsilon^{2}}$ rows.

To that end, we instead run $d$ separate instances of CountSketch to handle
the $d$ separate coordinates of each heavy-hitter
${\mathbf{A}}_{i}{\mathbf{x}}$. We show that the median of each estimated
coordinate is an unbiased estimate to the coordinate
$({\mathbf{A}}_{i}{\mathbf{x}})_{j}$, since the probability mass function is
symmetric about the true value for each coordinate. Furthermore, we show that
although the error to a single coordinate $({\mathbf{A}}_{i}{\mathbf{x}})_{j}$
may be large compared to $|({\mathbf{A}}_{i}{\mathbf{x}})_{j}|$, the error is
not large compared to
$\sum_{i\in\operatorname*{tail}(2/\varepsilon^{2})}\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}\right\rVert_{2}$.

###### Lemma 2.3

There exists an algorithm that uses
$\mathcal{O}\left(\frac{d^{2}}{\varepsilon^{2}}\log^{2}(nT)\right)$ space and
outputs a vector ${\mathbf{y}}_{i}$ for each index $i\in[n]$ so that
$|\left\lVert{\mathbf{y}}_{i}\right\rVert_{2}-\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}\right\rVert_{2}|\leq\varepsilon\sum_{i\in\operatorname*{tail}(2/\varepsilon^{2})}\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}\right\rVert_{2}$
and $\mathbb{E}\left[{\mathbf{y}}_{i}\right]={\mathbf{A}}_{i}{\mathbf{x}}$
with probability at least $1-\frac{1}{\operatorname*{poly}(n,T)}$.

Proof :  Given ${\mathbf{A}}\in\mathbb{R}^{n\times d}$, let
${\mathbf{A}}_{i}={\mathbf{a}}_{i}^{\top}{\mathbf{a}}_{i}$ for all $i\in[n]$.
For a fixed coordinate $k\in[d]$, we define a vector
${\mathbf{v}}^{(k)}\in\mathbb{R}^{n}$ so that for each $i\in[n]$, the $i$-th
coordinate of ${\mathbf{v}}^{(k)}$ is the $k$-th coordinate of
${\mathbf{A}}_{i}{\mathbf{x}}\in\mathbb{R}^{d}$.

Suppose we run a separate CountSketch instance on ${\mathbf{v}}^{(k)}$. For a
fixed index $i\in[n]$, let $h(i)$ be the bucket of $\mathcal{T}$ to which
$v^{(k)}_{i}$ hashes. For each $j\in[n]$, let $I_{j}$ be the indicator
variable for whether $v^{(k)}_{j}$ also hashes to bucket $h(i)$, so that
$I_{j}=1$ if $h(i)=h(j)$ and $I_{j}=0$ if $h(i)\neq h(j)$. Similarly for each
$j\in[n]$, let $s_{j}$ be a random sign assigned to $j$, so that the estimate
for $v^{(k)}_{i}$ by a single row of CountSketch is

$\sum_{j\in[n]}s_{i}s_{j}I_{j}v^{(k)}_{j}=v^{(k)}_{i}+\sum_{j:h(j)=h(i)}r_{j}v^{(k)}_{j},$

where $r_{j}=s_{i}s_{j}$ satisfies $r_{j}=1$ with probability $\frac{1}{2}$
and $r_{j}=-1$ with probability $\frac{1}{2}$. Thus if $y_{i}$ is the estimate
for $v^{(k)}_{i}$, then for any real number $u$, we have that

$\mathbf{Pr}\left[y_{i}=v^{(k)}_{i}+u\right]=\mathbf{Pr}\left[y_{i}=v^{(k)}_{i}-u\right],$

so that the probability mass function of $y_{i}$ is symmetric about
$v^{(k)}_{i}$. Thus given $\ell$ independent instances of CountSketch with
estimates $y_{i}^{(k,1)},\ldots,y_{i}^{(k,\ell)}$ for $v^{(k)}_{i}$ and any
real numbers $u^{(1)},\ldots,u^{(\ell)}$,

$\mathbf{Pr}\left[y_{i}^{(k,1)}=v^{(k)}_{i}+u^{(1)},\ldots,y_{i}^{(k,\ell)}=v^{(k)}_{i}+u^{(\ell)}\right]=\mathbf{Pr}\left[y_{i}^{(k,1)}=v^{(k)}_{i}-u^{(1)},\ldots,y_{i}^{(k,\ell)}=v^{(k)}_{i}-u^{(\ell)}\right].$

Therefore, the joint probability mass function is symmetric about
$(v^{(k)}_{i},\ldots,v^{(k)}_{i})$ and so the median across the $\ell$
instances of CountSketch is an unbiased estimator to $v^{(k)}_{i}$. Finally,
we have due to the properties of CountSketch that if each hash function $h$
maps to a universe of size $\mathcal{O}\left(\frac{1}{\varepsilon^{2}}\right)$
and $\ell=\mathcal{O}\left(\log(nT)\right)$, then with probability at least
$1-\frac{1}{\operatorname*{poly}(T,n)}$, the output estimate for $v^{(k)}_{i}$
has additive error at most
$\varepsilon\cdot\left(\sum_{j\in\operatorname*{tail}(2/\varepsilon^{2})}(v^{(k)}_{i})^{2}\right)^{1/2}$.

Thus using each of the estimated outputs across all $k\in[d]$, then for a
fixed $i\in[n]$, we can output a vector ${\mathbf{y}}_{i}$ such that
$\mathbb{E}\left[{\mathbf{y}}_{i}\right]={\mathbf{A}}_{i}{\mathbf{x}}$ and
with probability at least $1-\frac{1}{\operatorname*{poly}(T,n)}$,

$|\left\lVert{\mathbf{y}}_{i}\right\rVert_{2}-\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}\right\rVert_{2}|\leq\varepsilon\cdot\left(\sum_{i\in\operatorname*{tail}(2/\varepsilon^{2})}\left\lVert{\mathbf{A}}_{i}{\mathbf{x}}\right\rVert_{2}^{2}\right)^{1/2}.$

For a fixed $k\in[d]$, then our algorithm intends to hash the $k$-th
coordinate of ${\mathbf{A}}_{i}{\mathbf{x}}\in\mathbb{R}^{d}$. However, since
${\mathbf{x}}$ is only given after the data structure is already formed and in
particular, after ${\mathbf{A}}_{i}$ is given, then CountSketch must hash the
$k$-th row of ${\mathbf{A}}_{i}$ entirely, thus storing
$\mathcal{O}\left(\frac{d}{\varepsilon^{2}}\log^{2}(nT)\right)$ bits for each
coordinate $k\in[d]$. Hence across all $k\in[d]$, the algorithm uses the total
space $\mathcal{O}\left(\frac{d^{2}}{\varepsilon^{2}}\log^{2}(nT)\right)$.
$\Box$

However if say, we want to identify the heavy gradients
$(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-b_{i}){\mathbf{a}}_{i}$, then we
create separate data structures for the constant (in ${\mathbf{x}}$) term
$b_{i}{\mathbf{a}}_{i}$ and the linear term
$({\mathbf{a}}_{i}\otimes{\mathbf{a}}_{i}){\mathbf{x}}$, using the same
buckets, hash functions, and random signs. For the constant term data
structure, we hash the scaled rows $b_{i}{\mathbf{a}}_{i}$ into
$\mathcal{O}\left(\frac{1}{\varepsilon^{2}}\right)$ buckets, so that each
bucket contains a vector that represents the signed sum of the (scaled) rows
of ${\mathbf{A}}$ that hash to the bucket. For the linear term data structure,
we hash the outer products
${\mathbf{A}}_{i}:={\mathbf{a}}_{i}\otimes{\mathbf{a}}_{i}$ into
$\mathcal{O}\left(\frac{1}{\varepsilon^{2}}\right)$ buckets, so that each
bucket contains a vector that represents the signed sum of the matrices
${\mathbf{A}}_{i}$ that hash to the bucket. Once the vector ${\mathbf{x}}$
arrives after ${\mathbf{A}}$ is processed, then we can multiply each of the
matrices stored by each bucket by ${\mathbf{x}}$. Since the signed sum is a
linear sketch, this procedure is equivalent to originally taking the signed
sums of the vectors ${\mathbf{A}}_{i}{\mathbf{x}}$. Similarly, by linearity,
we can then take any linear combination of the two data structures to identify
the heavy gradients $(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i}){\mathbf{a}}_{i}$.

###  2.2 $G$-Sampler Algorithm

In this section, we first describe our $G$-sampler algorithm, where we sample
a gradient $G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})$ with probability proportional to
$\|G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})\|_{2}$. Given an accuracy parameter $\varepsilon>0$,
let $\alpha$ be a constant, parametrized by $\varepsilon$, so that
$(1-\varepsilon)F_{G}({\mathbf{v}})\leq
F_{G}({\mathbf{u}})\leq(1+\varepsilon)F_{G}({\mathbf{v}})$, for any
${\mathbf{u}}$ with
$\|{\mathbf{u}}-{\mathbf{v}}\|_{2}\leq\alpha\|{\mathbf{v}}\|_{2}$. As our data
structure will be a linear sketch, we focus on the case where we fold the
measurement vector ${\mathbf{b}}$ into a column of ${\mathbf{A}}$, so that we
want to output a gradient
${\mathbf{A}}_{i}{\mathbf{x}}:=({\mathbf{a}}_{i}\otimes{\mathbf{a}}_{i}){\mathbf{x}}$.

Our algorithm first partitions the rows of ${\mathbf{A}}$ into classes, based
on their $L_{2}$ norm. For example, if all entries of ${\mathbf{A}}$ are
integers, then we define class
$C_{k}:=\\{{\mathbf{a}}_{i}\,:\,2^{k-1}\leq\|{\mathbf{a}}_{i}\|_{2}<2^{k}\\}$.
We create a separate data structure for each class. We will use the $F_{G}$
estimation algorithm on each class to first sample a particular class. It then
remains to sample a particular vector
$({\mathbf{a}}_{i}\otimes{\mathbf{a}}_{i}){\mathbf{x}}$ from a class.

Depending on the vector ${\mathbf{x}}$, the vectors
$({\mathbf{a}}_{i}\otimes{\mathbf{a}}_{i}){\mathbf{x}}$ in a certain class
$C_{k}$ can have drastically different $L_{2}$ norm. We define level set
$\Gamma_{j}$ as the vectors that satisfy
$(1+\varepsilon)^{j-1}\leq\|({\mathbf{a}}_{i}\otimes{\mathbf{a}}_{i}){\mathbf{x}}\|_{2}<(1+\varepsilon)^{j}$.
If we could estimate $|\Gamma_{j}|$, then we could estimate the contribution
of each level set $\Gamma_{j}$ toward the overall mass $\sum_{i\in
C_{k}}\|({\mathbf{a}}_{i}\otimes{\mathbf{a}}_{i}){\mathbf{x}}\|_{2}$, so we
can then sample a specific level set $\Gamma_{j}$ from the class $C_{k}$. To
that end, we create $L=\mathcal{O}\left(\log n\right)$ substreams,
$S_{1},\ldots,S_{L}$, so that we sample each row with probability
$\frac{1}{2^{\ell-1}}$ in substream $S_{\ell}$.

The point is that if the contribution of level set $\Gamma_{j}$ is
“significant”, then there exists a specific substream $S_{\ell}$ in which the
vectors of $\Gamma_{j}$ will be likely detected by the heavy-hitter algorithm
we introduced in Section 2.1, if they are sampled by $S_{\ell}$. We can then
use these vectors that are output by the heavy-hitter algorithm to estimate
the contribution of level set $\Gamma_{j}$. However, if the contribution of
level set $\Gamma_{j}$ is not significant, then there may not be any vectors
of $\Gamma_{j}$ that survive the sampling in substream $S_{\ell}$. Thus we add
a number of “dummy rows” to each level set to insist that all level sets are
significant, so that we can estimate their contributions.

We then sample a level set $\Gamma_{j}$ with probability proportional to its
contribution and uniformly select a (noisy) vector from the level set. If the
selected vector is one of the original rows of the matrix, then we output the
noisy vector. Otherwise, we say the sampler has failed. We show that the
sampler only fails with constant probability, so it suffices to run
$\mathcal{O}\left(\log\frac{1}{\delta}\right)$ independent instances to boost
the probability of success to any arbitrary $1-\delta$. The algorithm for
selecting a level set $\Gamma_{j}$ from a specific class $C_{k}$ appears in
Algorithm 2.

We first show that the dummy rows only contribute at constant multiple of the
mass $F_{G}(S)=\sum_{i=1}^{n}\|{\mathbf{A}}_{i}{\mathbf{x}}\|_{2}$, where we
assume for simplicity that all rows of ${\mathbf{A}}$ are in the same class.

###### Lemma 2.4

Let $S$ be the input data stream with subsamples $S_{1},\ldots,S_{L}$. Let
$\widetilde{S}$ be the input data stream with the additional dummy rows and
corresponding subsamples $\widetilde{S_{1}},\ldots,\widetilde{S_{L}}$. Then
$2F_{G}(S)\geq F_{G}(\widetilde{S})\geq F_{G}(S)$.

Proof :  Since $\widetilde{S}$ includes all the rows of $S$, then
$F_{G}(\widetilde{S})\geq F_{G}(S)$. Since each level $j\in[K]$ acquires
$\mathcal{O}\left(\frac{(1+\alpha)^{j}\alpha^{3}}{\log n}\right)$ dummy rows
that each contribute
$\mathcal{O}\left(\frac{\widehat{M}}{(1+\alpha)^{j}\alpha^{2}}\right)$ to
$F_{G}$ in $\widetilde{S}$, then each level of $F_{G}(S)$ contributes at most
$\mathcal{O}\left(\frac{\widehat{M}\cdot\alpha}{\log n}\right)$ more to
$F_{G}(\widetilde{S})$. Because $K=\mathcal{O}\left(\frac{\log
n}{\alpha}\right)$, then the total additional contribution by the dummy rows
is at most $\mathcal{O}\left(\widehat{M}\right)$. Since $\widehat{M}\leq
2F_{G}(S)$, then it follows that for sufficiently small constant in the
contribution of each dummy row, we have $F_{G}(\widetilde{S})-F_{G}(S)\leq
F_{G}(S)$ and thus, $F_{G}(\widetilde{S})\leq 2F_{G}(S)$. $\Box$

Algorithm 2 $G$-sampler for a single class of rows

1:Rows ${\mathbf{a}}_{1},\ldots,{\mathbf{a}}_{n}$ of a matrix
${\mathbf{A}}\in\mathbb{R}^{n\times d}$ with $2^{k}\leq\|a_{i}\|_{2}<2^{k+1}$
for all $i\in[n]$, function $G$, accuracy parameter $\alpha$ for sampling
parameter $\varepsilon$

2:Noisy row ${\mathbf{v}}$ with the correct sampling distribution induced by
$G$

3:$\gamma$ uniformly at random from $[1/2,1]$,
$K\leftarrow\mathcal{O}\left(\frac{\log n}{\alpha}\right)$,
$L\leftarrow\mathcal{O}\left(\log n\right)$

4:for $\ell\in[L]$ do $\triangleright$ Processing stage

5: Form a stream $S_{\ell}$ by sampling each row with probability
$2^{-\ell+1}$

6: Run ${\textsc{CountSketch}}^{(1)}_{\ell}$ with threshold
$\mathcal{O}\left(\frac{\alpha^{3}}{\log n}\right)$ and failure probability
$\frac{1}{\operatorname*{poly}(n,T)}$ by creating a table $A^{(1)}_{\ell}$
with entries ${\mathbf{a}}_{j}^{\top}{\mathbf{a}}_{j}$ in $S_{\ell}$ and a
table $B^{(1)}_{\ell}$ with entries ${\mathbf{a}}_{j}$ $\triangleright$
Identify heavy-hitters

7: Run ${\textsc{CountSketch}}^{(2)}_{\ell}$ with threshold
$\mathcal{O}\left(\frac{\alpha^{3}}{\log n}\right)$ and failure probability
$\frac{1}{\operatorname*{poly}(n,T)}$ by creating a table $A^{(2)}_{\ell}$
with entries ${\mathbf{a}}_{j}^{\top}{\mathbf{a}}_{j}$ in $S_{\ell}$ and a
table $B^{(2)}_{\ell}$ with entries ${\mathbf{a}}_{j}$ and separately
considering coordinates after post-processing $\triangleright$ Unbiased
estimates of heavy-hitters, see Lemma 2.3

8:for $\ell\in[L]$ do $\triangleright$ Post-processing

9: Set $C^{(i)}_{\ell}=A^{(i)}_{\ell}{\mathbf{x}}+B^{(i)}_{\ell}$ with post-
multiplication by ${\mathbf{x}}$ for $i\in\\{1,2\\}$

10: Query $\widehat{M}\in\left[M/2,2M\right]$, where
$M=\sum_{i=1}^{n}\|G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})\|_{2}$

11: for $j\in[K]$ do

12: if $j>\log_{(1+\alpha)}\frac{\log^{2}n}{\alpha^{3}}$ then

13: Add $\mathcal{O}\left(\frac{(1+\alpha)^{j}\alpha^{3}}{\log n}\right)$
dummy rows that each contribute
$\mathcal{O}\left(\frac{\widehat{M}}{(1+\alpha)^{j}\alpha^{2}}\right)$ to
$F_{G}$

14: Let $H^{(i)}_{\ell}$ be the heavy rows of $C^{(i)}_{\ell}$ for
$i\in\\{1,2\\}$ from ${\textsc{CountSketch}}^{(i)}_{\ell}$

15:for $j\in[K]$ do

16: $L_{j}\leftarrow\max\left(1,\log\frac{\alpha^{2}(1+\alpha)^{j}}{\log
n}\right)$

17: Let $X_{j}$ be the estimated heavy-hitters ${\mathbf{v}}$ from
$H^{(2)}_{j}$ that are reported by $H^{(1)}_{j}$ with $G({\mathbf{v}})$ in
$\left[\frac{8\gamma\widehat{M}}{(1+\alpha)^{j+1}},\frac{8\gamma\widehat{M}}{(1+\alpha)^{j}}\right)$

18: if $L_{j}=1$ then

19: $\widetilde{F_{G}}(\widetilde{S_{j}})\leftarrow\sum_{{\mathbf{v}}\in
X_{j}}\frac{8\gamma\widehat{M}}{(1+\alpha)^{j+1}}$

20: else if $L_{j}>1$ and $|X_{j}|>\frac{1}{\alpha^{2}}$ then

21: $\widetilde{F_{G}}(\widetilde{S_{j}})\leftarrow\sum_{{\mathbf{v}}\in
X_{j}}\frac{8\gamma\widehat{M}}{(1+\alpha)^{j+1}}\cdot 2^{L_{j}}$

22: else

23: $\widetilde{F_{G}}(\widetilde{S_{j}})\leftarrow 0$

24:Sample $j\in[K]$ with probability
$\frac{\widetilde{F_{G}}(\widetilde{S_{j}})}{\sum\widetilde{F_{G}}(\widetilde{S_{j}})}$

25:Sample ${\mathbf{v}}$ from $X_{j}$ with probability $\frac{1}{X_{j}}$

26:if ${\mathbf{v}}$ is a dummy row then

27: return $\bot$

28:else

29: return ${\mathbf{v}}$

We would now like to show that with high probability, each of the substreams
have exponentially smaller mass $F_{G}(S_{j})$. However, this may not be true.
Consider a single row ${\mathbf{a}}_{i}$ that contributes a constant fraction
of $F_{G}(S)$. Then even for $j=\log n$, the probability that
${\mathbf{a}}_{i}$ is sampled is roughly
$\frac{1}{n}\gg\frac{1}{\operatorname*{poly}(n)}$. Instead, we note that
CountSketch satisfies the stronger tail guarantee in Lemma 2.3. Hence for each
$j\in[K]$, we define $S_{j}^{\operatorname*{tail}(t)}$ to be the frequency
vector $S_{j}$ with its $t$ largest entries set to zero and we show an
exponentially decreasing upper bound on
$F_{G}(\widetilde{S_{j}^{\operatorname*{tail}(t)}})$.

###### Lemma 2.5

With high probability, we have that for all $j\in[K]$,
$F_{G}(\widetilde{S_{j}^{\operatorname*{tail}(t)}})\leq\frac{F_{G}(S)}{2^{j}}\,\log(nT)$
for $t=\mathcal{O}\left(\frac{\log n}{\alpha^{3}}\right)$.

Proof :  Observe that the number of rows that exceed
$\frac{\widehat{M}}{2^{j}}$ is at most $2^{j+1}$. Thus the expected number of
rows that exceed $\frac{\widehat{M}}{2^{j}}$ sampled by $S_{j}$ is at most
$\frac{1}{2}$. Hence by Chernoff bounds, the probability that the number of
rows that exceed $\frac{\widehat{M}}{2^{j}}$ sampled by $S_{j}$ is more than
$t=\mathcal{O}\left(\frac{\log n}{\alpha^{3}}\right)$ is
$\frac{1}{\operatorname*{poly}(nT)}$. $\Box$

We also show that the estimated contribution of each level set (after
incorporating the dummy rows) is a $(1+\alpha)$-approximation of the true
contribution.

###### Lemma 2.6

With high probability, we have that for all $j\in[K]$,
$(1-\varepsilon)F_{G}(\widetilde{S_{j}})\leq\widetilde{F_{G}}(\widetilde{S_{j}})\leq(1+\varepsilon)F_{G}(\widetilde{S_{j}})$.

Proof :  Suppose that for each $j\in[K]$, level $j$ consists of $N_{j}$ rows
and note that $N_{j}\geq\mathcal{O}\left(\frac{(1+\alpha)^{j}\alpha^{3}}{\log
n}\right)$ elements due to the dummy rows. Each element is sampled with some
probability $p_{L_{j}}$, where
$L_{j}=\max\left(1,\log\frac{\alpha^{2}(1+\alpha)^{j}}{\log n}\right)$ and
thus $p_{L_{j}}(1+\alpha)^{j}>1$ since $p_{L_{j}}=\frac{1}{2^{L_{j}}}$. Let
$\widehat{N_{j}}$ be the number of items sampled in $\widetilde{S_{L_{j}}}$.
We have $\mathbb{E}\left[2^{L_{j}}\cdot\widehat{N_{j}}\right]=N_{j}$ and the
second moment is at most $N_{j}\cdot 2^{L_{j}}\leq\frac{\alpha^{2}}{\log
n}(N_{j})^{2}$. Thus by Chernoff bounds with $\mathcal{O}\left(\log
n\right)$-wise independence, we have that with high probability,

$(1-\mathcal{O}\left(\alpha\right))N_{j}\leq
2^{L_{j}}\cdot\widehat{N_{j}}\leq(1+\mathcal{O}\left(\alpha\right))N_{j}.$

Each estimated row norm is a $(1+\alpha)$-approximation to the actual row norm
due to Lemma 2.5. Thus by Lemma 2.4, we have that
$F_{G}(\widetilde{S_{j}})\leq 2F_{G}(S_{j})$ so that each of the
$\widehat{N_{j}}$ rows will be detected by the threshold of CountSketch with
the tail guarantee, i.e., Lemma 2.3. Moreover, we assume that a noisy row with
$(1+\alpha)$-approximation to the row norm of the original vector suffices to
obtain a $(1+\varepsilon)$-approximation to the contribution of the row.
Therefore, the result then follows in an ideal scenario where
$G({\mathbf{v}})\in\left[\frac{\widehat{M}}{2^{j}},\frac{2\widehat{M}}{2^{j}}\right)$
if and only if the corresponding row ${\mathbf{a}}_{i}$ satisfies
$G({\mathbf{a}}_{i})\in\left[\frac{\widehat{M}}{2^{j}},\frac{2\widehat{M}}{2^{j}}\right)$.
Unfortunately, this may not be true because $G({\mathbf{a}}_{i})$ may lie near
the boundary of the interval
$\left[\frac{\widehat{M}}{2^{j}},\frac{2\widehat{M}}{2^{j}}\right)$ while the
estimate $G({\mathbf{v}})$ has a value that does not lie within the interval.
In this case, $G({\mathbf{v}})$ is used toward the estimation of some other
level set.

Hence, our algorithm randomizes the boundaries of the level sets
$\left[\frac{4\gamma\widehat{M}}{2^{j}},\frac{8\gamma\widehat{M}}{2^{j}}\right)$
by choosing $\gamma\in[1/2,1)$ uniformly at random. Since the threshold of
CountSketch is $\mathcal{O}\left(\frac{\alpha^{3}}{\log n}\right)$ then the
probability that each row ${\mathbf{a}}_{i}$ is misclassified over the choice
of $\gamma$ is at most $\mathcal{O}\left(\varepsilon\right)$. Moreover, if
${\mathbf{a}}_{i}$ is misclassified, then its contribution can only be
classified into level set $j-1$ or $j+1$, inducing an incorrect multiplicative
factor of at most two. Hence, the error due to the misclassification across
all rows is at most $\mathcal{O}\left(\varepsilon\right)$ fraction of
$F_{G}(S_{j})$ in expectation. By Markov’s inequality, this error is a most
$\varepsilon$-fraction of $F_{G}(S_{j})$ with probability at least $3/4$. Then
by taking the median across $\mathcal{O}\left(\log(nT)\right)$ independent
instances, we obtain high probability of success. $\Box$

Finally, we show that each row is sampled with the correct distribution and is
an unbiased estimate.

###### Lemma 2.7

Suppose that $2^{k}<\|{\mathbf{a}}_{i}\|_{2}\leq 2^{k+1}$ for all $i\in[n]$.
Then the probability that Algorithm 2 outputs a noisy vector ${\mathbf{v}}$
such that $\|{\mathbf{v}}-{\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i})\|_{2}\leq\alpha\|{\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i})|\|_{2}$ with
$\mathbb{E}\left[{\mathbf{v}}\right]={\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i})$ is
$p_{\mathbf{v}}=\left(1\pm\mathcal{O}\left(\varepsilon\right)\right)\frac{G({\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i}))}{F_{G}(\widetilde{S})}+\frac{1}{\operatorname*{poly}(nT)}$.

Proof :  Conditioned on the correctness of each of the estimates
$\widetilde{F_{G}}(\widetilde{S_{j}})$, which occurs with high probability by
Lemma 2.6, the probability that the algorithm selects $j\in[K]$ is
$\frac{\widetilde{F_{G}}(\widetilde{S_{j}})}{\sum_{j\in[K]}\widetilde{F_{G}}(\widetilde{S_{j}})}$.
Conditioned on the algorithm selecting $j\in[K]$, then either the algorithm
will choose a dummy row, or it will choose a row uniformly at random from the
rows ${\mathbf{v}}\in X_{j}$, where $X_{j}$ is the set of heavy-hitters
reported by $H_{j}$ with $L_{2}$ norm in
$\left[\frac{8\gamma\widehat{M}}{(1+\alpha)^{j+1}},\frac{8\gamma\widehat{M}}{(1+\alpha)^{j}}\right)$.
The latter event occurs with probability
$\frac{\widetilde{F_{G}}(S_{j})}{\widetilde{F_{G}}(\widetilde{S_{j}})}$. Due
to the tail guarantee of CountSketch in Lemma 2.3, we have that each heavy
hitter ${\mathbf{v}}\in X_{j}$ corresponds to a vector
${\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-b_{i})$ such that
$\|{\mathbf{v}}-{\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i})\|_{2}\leq\varepsilon\|{\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i})|\|_{2}$. Moreover, by Lemma 2.3, we have that
$\mathbb{E}\left[{\mathbf{v}}\right]={\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i})$. Hence the probability that vector ${\mathbf{v}}$ is selected is
$\frac{(1\pm\mathcal{O}\left(\alpha\right))G({\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i})}{\widetilde{F_{G}}(S_{j})}$. $\Box$

Putting things together, we have the guarantees of Theorem 1.5 for our
$G$-sampler.

See 1.5 Proof :  We define a class $C_{k}$ of rows as the subset of rows of
the input matrix ${\mathbf{A}}$ such that
$2^{k}\leq\|{\mathbf{a}}_{i}\|_{2}<2^{k+1}$. We first use the estimator
algorithm in Theorem 2.1 to sample a class $k$ of rows with probability
$\frac{\sum_{{\mathbf{a}}_{i}\in C_{k}}G(\langle{\mathbf{a}}_{i},x\rangle-
b_{i},{\mathbf{a}}_{i})}{\sum_{j\in[n]}G(\langle{\mathbf{a}}_{j},x\rangle-
b_{j},{\mathbf{a}}_{j})}$. Once a class $C_{k}$ is selected, then outputting a
row from $C_{k}$ under the correct distribution follows from Lemma 2.7. The
space complexity follows from storing a $d\times d$ matrix in each of the
$\mathcal{O}\left(\frac{\log^{2}(nT)}{\alpha^{3}}\right)$ buckets in
CountSketch for threshold
$\mathcal{O}\left(\frac{\alpha^{3}}{\log(nT)}\right)$ and high probability of
success. $\Box$

##  3 SGD Algorithm and Analysis

Before introducing our main SGD algorithm, we recall the following algorithm,
that essentially outputs noisy version of the rows with high “importance”.
Although Sampler outputs a (noisy) vector according to the desired probability
distribution, we also require an algorithm that automatically does this for
indices $i\in[n]$ that are likely to be sampled multiple times across the $T$
iterations. Equivalently, we require explicitly storing the rows with high
sensitivities.

###### Theorem 3.1

[BDM+20] Given a constant $\varepsilon>0$, there exists an algorithm Sens that
returns all indices $i\in[n]$ such that
$\sup_{x}\frac{|{\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i})|}{\sum_{j=1}^{n}|{\mathbf{a}}_{j}(\langle{\mathbf{a}}_{j},{\mathbf{x}}\rangle-
b_{j})|}\geq\frac{1}{200Td}$ for some ${\mathbf{x}}\in\mathbb{R}^{n}$, along
with the vector ${\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i})$. The algorithm requires a single pass over
${\mathbf{A}}={\mathbf{a}}_{1}\circ\ldots\circ{\mathbf{a}}_{n}$, uses
$\tilde{\mathcal{O}}\left(nd^{2}+Td^{2}\right)$ runtime and
$\tilde{\mathcal{O}}\left(Td^{2}\right)$ space, and succeeds with probability
$1-\frac{1}{\operatorname*{poly}(n)}$.

The quantity
$\sup_{x}\frac{\left\lVert{\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i})\right\rVert_{2}}{\sum_{j=1}^{n}\left\lVert{\mathbf{a}}_{j}(\langle{\mathbf{a}}_{j},{\mathbf{x}}\rangle-
b_{j})\right\rVert_{2}}$ can be considered the _sensitivity_ of row
${\mathbf{a}}_{i}$ and can be interpreted as a measure of “importance” of the
row ${\mathbf{a}}_{i}$ with respect to the other rows of ${\mathbf{A}}$.

We now proceed to describe our main SGD algorithm. For the finite-sum
optimization problem
$\underset{{\mathbf{x}}\in\mathbb{R}^{d}}{\min}F({\mathbf{x}}):=\frac{1}{n}\sum_{i=1}^{n}G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})$, where each $G$ is a piecewise function of a
polynomial with degree at most $1$, recall that we could simply use an
instance of Sampler as an oracle for SGD with importance sampling. However,
naïvely running $T$ SGD steps requires $T$ independent instances, which uses
$Tnd$ runtime by Theorem 1.5. Thus, as our main theoretical contribution, we
use a two level data structure by first implicitly partitioning the rows of
matrix ${\mathbf{A}}={\mathbf{a}}_{1}\circ\ldots\circ{\mathbf{a}}_{n}$ into
$\beta:=\Theta(T)$ buckets $B_{1},\ldots,B_{\beta}$ and creating an instance
of Estimator and Sampler for each bucket. The idea is that for a given query
${\mathbf{x}}_{t}$ in SGD iteration $t\in[T]$, we first query
${\mathbf{x}}_{t}$ to each of the Estimator data structures to estimate
$\sum_{i\in B_{j}}G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})$ for each $j\in[\beta]$. We then sample index
$j\in[\beta]$ among the buckets $B_{1},\ldots,B_{\beta}$ with probability
roughly $\frac{\sum_{i\in B_{j}}G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})}{\sum_{i=1}^{n}G(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i})}$. Once we have sampled index $j$, it would seem that
querying the instance Sampler corresponding to $B_{j}$ simulates SGD, since
Sampler now performs importance sampling on the rows in $B_{j}$, which gives
the correct overall probability distribution for each row $i\in[n]$. Moreover,
Sampler has runtime proportional to the sparsity of $B_{j}$, so the total
runtime across the $\beta$ instances of Sampler is
$\tilde{\mathcal{O}}\left(nd\right)$.

However, an issue arises when the same bucket $B_{j}$ is sampled multiple
times, as we only create a single instance of Sampler for each bucket. We
avoid this issue by explicitly accounting for the buckets that are likely to
be sampled multiple times. Namely, we show that if
$\frac{G(\langle{\mathbf{a}}_{i},{\mathbf{x}}_{t}\rangle-
b_{i},{\mathbf{a}}_{i})}{\sum_{j=1}^{n}G(\langle{\mathbf{a}}_{j},{\mathbf{x}}_{t}\rangle-
b_{j},{\mathbf{a}}_{j})}<\mathcal{O}\left(\frac{1}{T}\right)$ for all
$t\in[T]$ and $i\in[n]$, then by Bernstein’s inequality, the probability that
no bucket $B_{j}$ is sampled at least $2\log T$ times is at least
$1-\frac{1}{\operatorname*{poly}(T)}$. Thus we use Sens to separate all such
rows ${\mathbf{a}}_{i}$ whose sensitivities violate this property from their
respective buckets and explicitly track the SGD steps in which these rows are
sampled.

The natural approach would be to create $T$ samplers for each of the rows with
sensitivity at least $\Omega\left(\frac{1}{T}\right)$, ensuring that each of
these samplers has access to fresh randomness in each of the $T$ SGD steps.
However since the sensitivities sum to $\mathcal{O}\left(d\log n\right)$,
there can be up to $\mathcal{O}\left(Td\log n\right)$ rows with sensitivity at
least $\Omega\left(\frac{1}{T}\right)$, so creating $T$ samplers for each of
these rows could create up to $\Theta(T^{2}d\log n)$ samplers, which is
prohibitively expensive in $T$. Instead, we simply keep each row with
sensitivity at least $\Omega\left(\frac{1}{T}\right)$ explicitly, while not
including them in the bucket. Due to the monotonicity of sensitivities, the
sensitivity of each row may only decrease as the stream progresses. In the
case that a row had sensitivity at least $\Omega\left(\frac{1}{T}\right)$ at
some point, but then no longer exceeds the threshold at some later point, then
the row is given as input to the sampler corresponding to the bucket to which
the row hashes and then the explicit storage of the row is deleted. This
ensures we need only $\tilde{\mathcal{O}}\left(Td\right)$ samplers while still
avoiding any sampler from being used multiple times across the $T$ SGD steps.
We give the algorithm in full in Algorithm 3.

Algorithm 3 Approximate SGD with Importance Sampling

1:Matrix
${\mathbf{A}}={\mathbf{a}}_{1}\circ\ldots\circ{\mathbf{a}}_{n}\in\mathbb{R}^{n\times
d}$, parameter $T$ for number of SGD steps.

2:$T$ gradient directions.

3:Preprocessing Stage:

4:$\beta\leftarrow\Theta(T)$ with a sufficiently large constant in the
$\Theta$.

5:Let $h:[n]\to[\beta]$ be a uniformly random hash function.

6:Let ${\mathbf{B}}_{j}$ be the matrix formed by the rows ${\mathbf{a}}_{i}$
of ${\mathbf{A}}$ with $h(i)=j$, for each $j\in[\beta]$.

7:Create $\Theta(\log(Td))$ instances ${\textsc{Estimator}}_{j}$ and
${\textsc{Sampler}}_{j}$ for each ${\mathbf{B}}_{j}$ with $j\in[\beta]$ with
$\varepsilon=\frac{1}{2}$.

8:Run Sens to find a set $L_{0}$ of rows with sensitivity at least
$\Omega\left(\frac{1}{T}\right)$.

9:Gradient Descent Stage:

10:Randomly pick starting location ${\mathbf{x}}_{0}$

11:for $t=1$ to $T$ do

12: Let $q_{i}$ be the output of ${\textsc{Estimator}}_{j}$ on query
${\mathbf{x}}_{t-1}$ for each $i\in[\beta]$.

13: Sample $j\in[\beta]$ with probability
$p_{j}=\frac{q_{j}}{\sum_{i\in[\beta]}q_{i}}$.

14: if there exists $i\in L_{0}$ with $h(i)=j$ then

15: Use ${\textsc{Estimator}}_{j}$, $L_{0}$, and ${\textsc{Sampler}}_{j}$ to
sample gradient ${\mathbf{w}}_{t}=\widehat{\nabla
f_{i_{t}}({\mathbf{x}}_{t})}$

16: else

17: Use fresh ${\textsc{Sampler}}_{j}$ to sample gradient
${\mathbf{w}}_{t}=\widehat{\nabla f_{i_{t}}({\mathbf{x}}_{t})}$

18:
$\widehat{p_{i,t}}\leftarrow\frac{\left\lVert{\mathbf{w}}_{t}\right\rVert^{2}_{2}}{\sum_{j\in[\beta]}q_{j}}$

19:
${\mathbf{x}}_{t+1}\leftarrow{\mathbf{x}}_{t}-\frac{\eta_{t}}{n\widehat{p_{i,t}}}\cdot{\mathbf{w}}_{t}$

The key property achieved by Algorithm 3 in partitioning the rows and removing
the rows that are likely to be sampled multiple times is that each of the
Sampler instances are queried at most once.

###### Lemma 3.2

With probability at least $\frac{98}{100}$, each $t\in[T]$ uses a different
instance of ${\textsc{Sampler}}_{j}$.

Proof :  Let $C>0$ be a sufficiently large constant. For any $t\in[T]$ and
$i\in[n]$, $G({\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i}))\geq\frac{1}{CT}\sum_{j\in[n]}G({\mathbf{a}}_{j}(\langle{\mathbf{a}}_{j},x\rangle-
b_{j}))$ only if there exists a row in ${\mathbf{a}}_{i}\circ b_{i}$ whose
sensitivity is at least $\frac{1}{CT}$. However, we have explicitly stored all
rows ${\mathbf{a}}_{i}\circ b_{i}$ with sensitivity
$\Omega\left(\frac{1}{T}\right)$ and removed them from each $G$-sampler.

Thus, for all $j\in[\beta]$ so that $h(i)\neq j$ for any index $i\in[n]$ such
that $G({\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i}))\leq\frac{1}{CT}\sum_{k\in[n]}G({\mathbf{a}}_{k}(\langle{\mathbf{a}}_{k},x\rangle-
b_{k}))$, we have

$\sum_{i:h(i)=j}G({\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i}))\leq\frac{\log(Td)}{200T}\sum_{k\in[n]}G({\mathbf{a}}_{k}(\langle{\mathbf{a}}_{k},x\rangle-
b_{k})),$

with probability at least $1-\frac{1}{\operatorname*{poly}(Td)}$ by
Bernstein’s inequality and a union bound over $j\in[\beta]$, where
$\beta=\Theta(T)$ is sufficiently large. Intuitively, by excluding the hash
indices containing “heavy” matrices, the remaining hash indices contain only a
small fraction of the mass with high probability.

We analyze the probability that any bucket containing rows with sensitivity
less than $\mathcal{O}\left(\frac{1}{T}\right)$ are sampled more than
$\Omega(T\log(Td))$ times, since we create $\mathcal{O}\left(T\log(Td)\right)$
separate $G$-samplers for each of these buckets. By a coupling argument and
Chernoff bounds, the probability that any $j\in[\beta]$ with
$\sum_{i:h(i)=j}G({\mathbf{a}}_{i}(\langle{\mathbf{a}}_{i},x\rangle-
b_{i}))\leq\frac{\log(Td)}{200T}\sum_{k\in[n]}G({\mathbf{a}}_{k}(\langle{\mathbf{a}}_{k},x\rangle-
b_{k}))$ is sampled more than $200\log(Td)$ times is at most
$\frac{1}{\operatorname*{poly}(Td)}$ for any $t\in[T]$, provided there is no
row with $h(i)=j$ whose sensitivity is at least $\frac{1}{CT}$. Thus, the
probability that some bucket $j\in[\beta]$ is sampled more than $200\log(Td)$
times across $T$ steps is at most $\frac{1}{\operatorname*{poly}(Td)}$.

In summary, we would like to maintain $T$ separate instances of $G$-samplers
for the heavy matrices and $\Theta(\log(Td))$ separate instances of
$G$-samplers for each hash index that does not contain a heavy matrix, but
this creates a $\Omega(T^{2})$ space dependency. Instead, we explicitly store
the heavy rows with sensitivity $\Omega\left(\frac{1}{T}\right)$, removing
them from the heavy matrices, and manually perform the sampling, rather than
rely on the $G$-sampler subroutine. There can be at most
$\mathcal{O}\left(Td\log n\right)$ such rows, resulting in
$\mathcal{O}\left(Td^{2}\log n\right)$ overall space for storing these rows
explicitly. Since the resulting matrices are light by definition, we can
maintain $\Theta(\log(Td))$ separate instances of $G$-samplers for each of the
$\Theta(T)$ buckets, which results in $\tilde{\mathcal{O}}\left(Td^{2}\right)$
space overall. With probability at least $\frac{98}{100}$, any hash index not
containing a heavy matrix is sampled only once, so each time $t\in[T]$ has
access to a fresh $G$-sampler. $\Box$

Theorem 1.4 then follows from Lemma 3.2 and the sampling distribution
guaranteed by each subroutine in Lemma 2.7. In particular, Lemma 3.2 crucially
guarantees that each step $t\in[T]$ of SGD will receive a vector with fresh
independent randomness. Moreover, we have that each (noisy) vector has small
variance and is an unbiased estimate of a subgradient sampled from nearly the
optimal importance sampling probability distribution.

See 1.4 Proof :  Consider Algorithm 3. By Lemma 3.2, each time $t\in[T]$ uses
a fresh instance of ${\textsc{Sampler}}_{j}$, so that independent randomness
is used. A possible concern is that each instance ${\textsc{Estimator}}_{j}$
is not using fresh randomness, but we observe that the Estimator procedures
are only used in sampling a bucket $j\in[\beta]$; otherwise the sampling uses
fresh randomness whereas the sampling is built into each instance of
${\textsc{Sampler}}_{j}$. By Theorem 1.5, each index $i$ is sampled with
probability within a factor $2$ of the importance sampling probability
distribution. By Theorem 2.1, we have that $\widehat{p_{i,t}}$ is within a
factor $4$ of the probability $p_{i,t}$ induced by optimal importance sampling
SGD. Note that
${\mathbf{w}}_{t}=\widehat{G(\langle{\mathbf{a}}_{i},{\mathbf{x}}_{t}\rangle-
b_{i},{\mathbf{a}}_{i})}$ is an unbiased estimator of
$G(\langle{\mathbf{a}}_{i},{\mathbf{x}}_{t}\rangle-b_{i},{\mathbf{a}}_{i})$
and $G({\mathbf{w}}_{t})$ is a $2$-approximation to $G({\mathbf{x}}_{t})$ by
Theorem 1.5. Hence, the variance at each time $t\in[T]$ of Algorithm 3 is
within a constant factor of the variance
$\sigma^{2}=\left(\sum_{i=1}^{n}G(\langle{\mathbf{a}}_{i},{\mathbf{x}}_{t}\rangle-
b_{i},{\mathbf{a}}_{i})\right)^{2}-\sum_{i=1}^{n}G(\langle{\mathbf{a}}_{i},{\mathbf{x}}_{t}\rangle-
b_{i},{\mathbf{a}}_{i})^{2}$ of optimal importance sampling SGD.

By Theorem 1.5, Theorem 2.1, and Theorem 3.1, the preprocessing time is
$d^{2}\,\operatorname*{polylog}(nT)$ for
$\varepsilon=\mathcal{O}\left(1\right)$ and $\beta=\Theta(T)$, but
partitioning the non-zero entries of ${\mathbf{A}}$ across the $\beta$ buckets
and the space used by the algorithm is
$\tilde{\mathcal{O}}\left(Td^{2}\right)$. Once the gradient descent stage of
Algorithm 3 begins, it takes $Td^{2}\,\operatorname*{polylog}(n)$ time in each
step $t\in[T]$ to query the $\beta=\Theta(T)$ instances of Sampler and
Estimator, for total time $Td^{2}\,\operatorname*{polylog}(n)$. $\Box$

#### Derandomization of the algorithm.

To derandomize our algorithm, we first recall the following formulation of
Nisan’s pseudorandom generator.

######  Theorem 3.3 (Nisan’s PRG)

[Nis92] Let $\mathcal{A}$ be an algorithm that uses $S=\Omega(\log n)$ space
and $R$ random bits. Then there exists a pseudorandom generator for
$\mathcal{A}$ that succeeds with high probability and runs in
$\mathcal{O}\left(S\log R\right)$ bits.

The goal of Nisan’s PRG is to fool a small space tester by generating a number
of pseudorandom bits in a read-once tape in place of a number of truly random
bits. In the row-arrival model, the updates to each row ${\mathbf{a}}_{i}$ of
${\mathbf{A}}\in\mathbb{R}^{n\times d}$ arrive sequentially, so it suffices to
use a read-once input tape. Thus a tester that is only allowed to $S$ space
cannot distinguish between the output of our algorithm using true randomness
and pseudorandom bits generated by Nisan’s PRG. Since our algorithm uses
$S=Td^{2}\,\operatorname*{polylog}(Tnd)$ bits of space and
$R=\operatorname*{poly}(n,T,d)$ bits of randomness, then it can be randomized
by Nisan’s PRG while using $Td^{2}\,\operatorname*{polylog}(Tnd)$ total space.

## Acknowledgements

David P. Woodruff and Samson Zhou were supported by a Simons Investigator
Award and by the National Science Foundation under Grant No. CCF-1815840.

## References

  * [AKO11] Alexandr Andoni, Robert Krauthgamer, and Krzysztof Onak.  Streaming algorithms via precision sampling.  In IEEE 52nd Annual Symposium on Foundations of Computer Science, FOCS, pages 363–372, 2011.
  * [AMS99] Noga Alon, Yossi Matias, and Mario Szegedy.  The space complexity of approximating the frequency moments.  J. Comput. Syst. Sci., 58(1):137–147, 1999.
  * [BAG16] Michael Bosse, Gabriel Agamennoni, and Igor Gilitschenski.  Robust estimation and applications in robotics.  Found. Trends Robotics, 4(4):225–269, 2016.
  * [BBC+17] Jaroslaw Blasiok, Vladimir Braverman, Stephen R. Chestnut, Robert Krauthgamer, and Lin F. Yang.  Streaming symmetric norms via measure concentration.  In Proceedings of the 49th Annual ACM SIGACT Symposium on Theory of Computing, STOC, pages 716–729, 2017.
  * [BDM+20] Vladimir Braverman, Petros Drineas, Cameron Musco, Christopher Musco, Jalaj Upadhyay, David P. Woodruff, and Samson Zhou.  Near optimal linear algebra in the online and sliding window models.  In 61st IEEE Annual Symposium on Foundations of Computer Science, FOCS, pages 517–528, 2020.
  * [Bub15] Sébastien Bubeck.  Convex optimization: Algorithms and complexity.  Foundations and Trends in Machine Learning, 8(3-4):231–357, 2015\.
  * [CCF04] Moses Charikar, Kevin C. Chen, and Martin Farach-Colton.  Finding frequent items in data streams.  Theor. Comput. Sci., 312(1):3–15, 2004.
  * [CW15] Kenneth L. Clarkson and David P. Woodruff.  Sketching for _M_ -estimators: A unified approach to robust regression.  In Proceedings of the Twenty-Sixth Annual ACM-SIAM Symposium on Discrete Algorithms, SODA, pages 921–939, 2015.
  * [CWW19] Kenneth L. Clarkson, Ruosong Wang, and David P. Woodruff.  Dimensionality reduction for tukey regression.  In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML, pages 1262–1271, 2019.
  * [DBL14] Aaron Defazio, Francis R. Bach, and Simon Lacoste-Julien.  SAGA: A fast incremental gradient method with support for non-strongly convex composite objectives.  In Advances in Neural Information Processing Systems 27: Annual Conference on Neural Information Processing Systems, pages 1646–1654, 2014.
  * [DLH16] Hadi Daneshmand, Aurélien Lucchi, and Thomas Hofmann.  Starting small - learning with adaptive sample sizes.  In Proceedings of the 33nd International Conference on Machine Learning, ICML, pages 1463–1471, 2016.
  * [GS99] Antoine Guitton and William W Symes.  Robust and stable velocity analysis using the huber function.  In SEG Technical Program Expanded Abstracts 1999, pages 1166–1169. Society of Exploration Geophysicists, 1999.
  * [Hub92] Peter J Huber.  Robust estimation of a location parameter.  In Breakthroughs in statistics, pages 492–518. Springer, 1992.
  * [IW05] Piotr Indyk and David P. Woodruff.  Optimal approximations of the frequency moments of data streams.  In Proceedings of the 37th Annual ACM Symposium on Theory of Computing, pages 202–208, 2005.
  * [JG18] Tyler B. Johnson and Carlos Guestrin.  Training deep models faster with robust, approximate importance sampling.  In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems, NeurIPS, pages 7276–7286, 2018.
  * [JLS+21] Yifei Jiang, Yi Li, Yiming Sun, Jiaxin Wang, and David P. Woodruff.  Single pass entrywise-transformed low rank approximation.  In Proceedings of the 38th International Conference on Machine Learning, ICML, pages 4982–4991, 2021.
  * [JST11] Hossein Jowhari, Mert Saglam, and Gábor Tardos.  Tight bounds for lp samplers, finding duplicates in streams, and related problems.  In Proceedings of the 30th ACM SIGMOD-SIGACT-SIGART Symposium on Principles of Database Systems, PODS, pages 49–58, 2011.
  * [JW18] Rajesh Jayaram and David P. Woodruff.  Perfect lp sampling in a data stream.  In 59th IEEE Annual Symposium on Foundations of Computer Science, FOCS, pages 544–555, 2018.
  * [JWZ22] Rajesh Jayaram, David P. Woodruff, and Samson Zhou.  Truly perfect samplers for data streams and sliding windows.  In PODS: International Conference on Management of Data, pages 29–40, 2022.
  * [JZ13] Rie Johnson and Tong Zhang.  Accelerating stochastic gradient descent using predictive variance reduction.  In Advances in Neural Information Processing Systems 26, Proceedings., pages 315–323, 2013.
  * [KF18] Angelos Katharopoulos and François Fleuret.  Not all samples are created equal: Deep learning with importance sampling.  In Proceedings of the 35th International Conference on Machine Learning, ICML, pages 2530–2539, 2018.
  * [Mek17] Raghu Meka.  Cs289ml: Algorithmic machine learning notes, 2017.
  * [MRWZ20] Sepideh Mahabadi, Ilya P. Razenshteyn, David P. Woodruff, and Samson Zhou.  Non-adaptive adaptive sampling on turnstile streams.  In Proccedings of the 52nd Annual ACM SIGACT Symposium on Theory of Computing, STOC, pages 1251–1264, 2020.
  * [MW10] Morteza Monemizadeh and David P. Woodruff.  1-pass relative-error l${}_{\mbox{p}}$-sampling with applications.  In Proceedings of the Twenty-First Annual ACM-SIAM Symposium on Discrete Algorithms, SODA, pages 1143–1160, 2010.
  * [Nis92] Noam Nisan.  Pseudorandom generators for space-bounded computation.  Combinatorica, 12(4):449–461, 1992.
  * [NJLS09] Arkadi Nemirovski, Anatoli B. Juditsky, Guanghui Lan, and Alexander Shapiro.  Robust stochastic approximation approach to stochastic programming.  SIAM Journal on Optimization, 19(4):1574–1609, 2009.
  * [NSW16] Deanna Needell, Nathan Srebro, and Rachel Ward.  Stochastic gradient descent, weighted sampling, and the randomized kaczmarz algorithm.  Math. Program., 155(1-2):549–573, 2016.
  * [NY83] Arkadi Semenovich Nemirovsky and David Borisovich Yudin.  Problem complexity and method efficiency in optimization, 1983.
  * [PSZ22] Eric Price, Sandeep Silwal, and Samson Zhou.  Hardness and algorithms for robust and sparse optimization.  In International Conference on Machine Learning, ICML, pages 17926–17944, 2022.
  * [QRG+19] Xun Qian, Peter Richtárik, Robert M. Gower, Alibek Sailanbayev, Nicolas Loizou, and Egor Shulgin.  SGD with arbitrary sampling: General analysis and improved rates.  In Proceedings of the 36th International Conference on Machine Learning, ICML, volume 97, pages 5200–5209, 2019.
  * [RHS+15] Sashank J. Reddi, Ahmed Hefny, Suvrit Sra, Barnabás Póczos, and Alexander J. Smola.  On variance reduction in stochastic gradient descent and its asynchronous variants.  In Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems, pages 2647–2655, 2015.
  * [RM51] Herbert Robbins and Sutton Monro.  A stochastic approximation method.  The annals of mathematical statistics, pages 400–407, 1951.
  * [RSB12] Nicolas Le Roux, Mark Schmidt, and Francis R. Bach.  A stochastic gradient method with an exponential convergence rate for finite training sets.  In Advances in Neural Information Processing Systems 25: 26th Annual Conference on Neural Information Processing Systems., pages 2672–2680, 2012.
  * [SRJ17] Sebastian U. Stich, Anant Raj, and Martin Jaggi.  Safe adaptive importance sampling.  In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems, pages 4381–4391, 2017.
  * [STC18] Farnood Salehi, Patrick Thiran, and L. Elisa Celis.  Coordinate descent with bandit sampling.  In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems, NeurIPS, pages 9267–9277, 2018.
  * [TWZ+22] Murad Tukan, Xuan Wu, Samson Zhou, Vladimir Braverman, and Dan Feldman.  New coresets for projective clustering and applications.  In International Conference on Artificial Intelligence and Statistics, AISTATS, 2022.
  * [TZ04] Mikkel Thorup and Yin Zhang.  Tabulation based 4-universal hashing with applications to second moment estimation.  In Proceedings of the Fifteenth Annual ACM-SIAM Symposium on Discrete Algorithms, SODA, pages 615–624. SIAM, 2004.
  * [YL06] Ming Yuan and Yi Lin.  Model selection and estimation in regression with grouped variables.  Journal of the Royal Statistical Society: Series B (Statistical Methodology), 68(1):49–67, 2006.
  * [Zha97] Zhengyou Zhang.  Parameter estimation techniques: a tutorial with application to conic fitting.  Image Vis. Comput., 15(1):59–76, 1997.
  * [ZZ15] Peilin Zhao and Tong Zhang.  Stochastic optimization with importance sampling for regularized loss minimization.  In Proceedings of the 32nd International Conference on Machine Learning ICML, pages 1–9, 2015.

##  Appendix A Second-Order Optimization

In this section, we repurpose our data stucture that performs importance
sampling for SGD to instead perform importance sampling for second-order
optimization. Given a second-order optimization algorithm that requires a
sampled Hessian ${\mathbf{H}}_{t}$, possibly along with additional inputs such
as the current iterate ${\mathbf{x}}_{t}$ and the gradient ${\mathbf{g}}_{t}$
of $F$, we model the update rule by an oracle
${\mathbb{O}}({\mathbf{H}}_{t})$, suppressing other inputs to the oracle in
the notation. For example, the oracle ${\mathbb{O}}$ corresponding to the
canonical second-order algorithm Newton’s method can be formulated as

${\mathbf{x}}_{t+1}={\mathbb{O}}({\mathbf{x}}_{t}):={\mathbf{x}}_{t}-[{\mathbf{H}}_{t}]^{-1}{\mathbf{g}}_{t}.$

We can incorporate our sampling oracle into the update rule of any second-
order optimization algorithm; thus we can focus our attention to the running
time of sampling a Hessian with nearly the optimal probability distribution.

Thus we use variations of our $G$-sampler for Hessians. In particular, we can
store ${\mathbf{a}}_{i}^{\otimes 3}$ rather than
${\mathbf{a}}_{i}^{\top}{\mathbf{a}}_{i}={\mathbf{a}}_{i}^{\otimes 2}$ for
each entry of the $b$ buckets of GenCountSketch in Algorithm 1, where the
notation ${\mathbf{v}}^{\otimes p}$ for a vector
${\mathbf{v}}\in\mathbb{R}^{d}$ represents the $p$-fold tensor of
${\mathbf{v}}$. We first require the following generalization of Theorem 2.1.

###### Corollary A.1

[BBC+17] Given a constant $\varepsilon>0$ and a Hessian $H$ of a loss function
that is the piecewise union of degree $p$ polynomials, there exists a one-pass
streaming algorithm HEstimator that takes updates to entries of a matrix
${\mathbf{A}}\in\mathbb{R}^{n\times d}$, as well as vectors
${\mathbf{x}}\in\mathbb{R}^{d}$ and ${\mathbf{b}}\in\mathbb{R}^{d}$ that
arrive after the stream, and outputs a quantity $\hat{F}$ such that
$(1-\varepsilon)\sum_{i\in[n]}\|H(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i}))\|_{2}\leq\hat{F}\leq(1+\varepsilon)\sum_{i\in[n]}\|H(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i}))\|_{2}$. The algorithm uses
$\frac{d^{p}}{\alpha^{2}}\,\operatorname*{polylog}(nT)$ bits of space and
succeeds with probability at least $1-\frac{1}{\operatorname*{poly}(n,T)}$.

We remark that HEstimator with parameter $p$ can be interpreted similarly as
Estimator with parameter $p$. Namely, HEstimator maintains a signed sum of
tensors ${\mathbf{a}}_{i}^{\otimes p}$, which can be reshaped into a matrix
with dimension $\mathbb{R}^{d\times d^{p-1}}$.

Then we have the following analog of Lemma 2.3 to find the “heavy” matrices,
corresponding to the “important” Hessians.

###### Lemma A.2

Given an integer $p\geq 2$ and an integer $q\in(0,p]$, there exists an
algorithm that uses
$\mathcal{O}\left(\frac{d^{p}}{\varepsilon^{2}}\log^{2}(nT)\right)$ space and
outputs a vector ${\mathbf{y}}_{i}$ for each index $i\in[n]$ so that
$|\left\lVert{\mathbf{y}}_{i}\right\rVert_{2}-\left\lVert{\mathbf{a}}_{i}^{\otimes
p}{\mathbf{x}}^{\otimes
q}\right\rVert_{2}|\leq\varepsilon\sum_{i\in\operatorname*{tail}(2/\varepsilon^{2})}\left\lVert{\mathbf{a}}_{i}^{\otimes
p}{\mathbf{x}}^{\otimes q}\right\rVert_{2}$ and
$\mathbb{E}\left[{\mathbf{y}}_{i}\right]={\mathbf{a}}_{i}^{\otimes
p}{\mathbf{x}}^{\otimes q}$ with probability at least
$1-\frac{1}{\operatorname*{poly}(n,T)}$.

We now have the following analog to Theorem 1.5.

###### Corollary A.3

Given a constant $\varepsilon>0$ and a Hessian $H$ of a loss function that is
the piecewise union of degree $p$ polynomials, there exists an algorithm
HSampler that outputs a noisy matrix ${\mathbf{V}}$ such that
$\|{\mathbf{V}}-H(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i}))\|_{F}\leq\alpha\|H(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i}))\|_{F}$ and
$\mathbb{E}\left[{\mathbf{V}}\right]=H(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i}))$ is
$\left(1\pm\mathcal{O}\left(\varepsilon\right)\right)\frac{\|H(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i}))\|_{F}}{\sum_{i\in[n]}\|H(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle-
b_{i},{\mathbf{a}}_{i}))\|_{F}}+\frac{1}{\operatorname*{poly}(n)}$. The
algorithm uses
$d^{p}\,\operatorname*{poly}\left(\log(nT),\frac{1}{\varepsilon}\right)$
update time per arriving row and
$Td^{p}\,\operatorname*{poly}\left(\log(nT),\frac{1}{\varepsilon}\right)$
total bits of space.

Once the input is processed, the input ${\mathbf{x}}\in\mathbb{R}^{d}$ arrives
and induces multiplication by a matrix with dimension
$\mathbb{R}^{d^{p-1}\times d}$, which represents the tensor
${\mathbf{x}}^{\otimes p}$, to obtain a sampled matrix with dimension
$\mathbb{R}^{d\times d}$ for second-order optimization. By comparison,
Estimator induces multiplication by a vector with dimension
$\mathbb{R}^{d^{p-1}}$, which represents the tensor
${\mathbf{x}}^{\otimes(p-1)}$, to obtain a sampled vector with dimension
$\mathbb{R}^{d}$ for first-order optimization. Similarly, the following two
statements are simple corollaries of Theorem 1.5 and Theorem 3.1.

###### Corollary A.4

[BDM+20] For a fixed $\varepsilon>0$ and polynomial $f$ of degree $p$, there
exists an algorithm HSens that returns all indices $i\in[n]$ such that
$\frac{(1\pm\varepsilon)\cdot\left\lVert
f(\langle{\mathbf{a}}_{i},{\mathbf{x}}-b_{i}\rangle)\cdot{\mathbf{a}}_{i}^{\top}{\mathbf{a}}_{i}\right\rVert_{F}}{\sum_{j=1}^{n}\left\lVert
f(\langle{\mathbf{a}}_{j},{\mathbf{x}}-b_{j}\rangle)\cdot{\mathbf{a}}_{i}^{\top}{\mathbf{a}}_{i}\right\rVert_{F}}\geq\frac{1}{200Td}$
for some ${\mathbf{x}}\in\mathbb{R}^{n}$, along with a matrix
${\mathbf{U}}_{i}:=f(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle)\cdot{\mathbf{a}}_{i}^{\top}{\mathbf{a}}_{i}+{\mathbf{V}}_{i}$,
where
$\left\lVert{\mathbf{V}}_{i}\right\rVert_{F}\leq\varepsilon\cdot\left\lVert
f(\langle{\mathbf{a}}_{i},{\mathbf{x}}\rangle)\cdot{\mathbf{a}}_{i}^{\top}{\mathbf{a}}_{i}\right\rVert_{F}$.
The algorithm requires a single pass over
${\mathbf{A}}={\mathbf{a}}_{1}\circ\ldots\circ{\mathbf{a}}_{n}$, uses
$\tilde{\mathcal{O}}\left(\operatorname*{nnz}({\mathbf{A}})+T\operatorname*{poly}(d)\right)$
runtime and $\tilde{\mathcal{O}}\left(Td^{p}\right)$ space, and succeeds with
probability $1-\frac{1}{\operatorname*{poly}(n)}$.

Thus the procedures HEstimator, HSampler, and HSens generalize Estimator,
Sampler, and Sens in a straightforward manner by storing
${\mathbf{a}}_{i}^{\otimes p}$ where previously
${\mathbf{a}}_{i}^{\top}{\mathbf{a}}_{i}$ was stored for the SGD subroutines.
Note that ${\mathbf{a}}_{i}^{\otimes p}$ can be reshaped into a $d\times
d^{p-1}$ matrix so that subsequently, multiplication is done by
${\mathbf{x}}^{\otimes p}$, which can be reshaped as a $d^{p-1}\times d$
matrix.

Algorithm 4 Second-Order Optimization with Importance Sampling

1:Matrix
${\mathbf{A}}={\mathbf{a}}_{1}\circ\ldots\circ{\mathbf{a}}_{n}\in\mathbb{R}^{n\times
d}$, parameter $T$ for number of sampled Hessians, oracle ${\mathbb{O}}$ that
performs the update rule.

2:$T$ approximate Hessians.

3:Preprocessing Stage:

4:$\beta\leftarrow\Theta(T)$ with a sufficiently large constant in the
$\Theta$.

5:Let $h:[n]\to[\beta]$ be a uniformly random hash function.

6:Let ${\mathbf{B}}_{j}$ be the matrix formed by the rows ${\mathbf{a}}_{i}$
of ${\mathbf{A}}$ with $h(i)=j$, for each $j\in[\beta]$.

7:Create an instance ${\textsc{HEstimator}}_{j}$ and ${\textsc{HSampler}}_{j}$
for each ${\mathbf{B}}_{j}$ with $j\in[\beta]$ with $\varepsilon=\frac{1}{2}$.

8:Run HSens to find a set $L_{0}$ of row indices and corresponding (noisy)
outer products.

9:Second-Order Optimization Stage:

10:Randomly pick starting location ${\mathbf{x}}_{0}$

11:for $t=1$ to $T$ do

12: Let $q_{i}$ be the output of ${\textsc{HEstimator}}_{j}$ on query
${\mathbf{x}}_{t-1}$ for each $i\in[\beta]$.

13: Sample $j\in[\beta]$ with probability
$p_{j}=\frac{q_{j}}{\sum_{i\in[\beta]}q_{i}}$.

14: if there exists $i\in L_{0}$ with $h(i)=j$ then

15: Use ${\textsc{HEstimator}}_{j}$, HSens, and ${\textsc{HSampler}}_{j}$ to
sample Hessian ${\mathbf{H}}_{t}$.

16: else

17: Use ${\textsc{HSampler}}_{j}$ to sample Hessian
${\mathbf{H}}_{t}=\widehat{\nabla f_{i_{t}}({\mathbf{x}}_{t})}$.

18:
$\widehat{p_{i,t}}\leftarrow\frac{\left\lVert{\mathbf{H}}_{t}\right\rVert_{F}^{2}}{\sum_{j\in[\beta]}q_{j}}$

19:
${\mathbf{x}}_{t+1}\leftarrow{\mathbb{O}}\left(\frac{1}{n\widehat{p_{i,t}}}{\mathbf{H}}_{t}\right)$

As before, observe that we could simply run an instance of HSampler to sample
a Hessian through importance sampling, but sampling $T$ Hessians requires $T$
independent instances, significantly increasing the total runtime. We thus use
the same two level data structure that partitions the rows of matrix
${\mathbf{A}}={\mathbf{a}}_{1}\circ\ldots\circ{\mathbf{a}}_{n}$ into
$\beta:=\Theta(T)$ buckets $B_{1},\ldots,B_{\beta}$. We then create an
instance of HEstimator and HSampler for each bucket. For an iterate
${\mathbf{x}}_{t}$, we sample $j\in[\beta]$ among the buckets
$B_{1},\ldots,B_{\beta}$ with probability roughly $\frac{\sum_{i\in
B_{j}}\left\lVert
f(\langle{\mathbf{a}}_{i},{\mathbf{x}}_{t}\rangle)\cdot{\mathbf{a}}_{i}^{\top}{\mathbf{a}}_{i}\right\rVert_{F}}{\sum_{i=1}^{n}\left\lVert
f(\langle{\mathbf{a}}_{i},{\mathbf{x}}_{t}\rangle)\cdot{\mathbf{a}}_{i}^{\top}{\mathbf{a}}_{i}\right\rVert_{F}}$
using HEstimator and then querying ${\textsc{HSampler}}_{j}$ at
${\mathbf{x}}_{t}$ to sample a Hessian among the indices partitioned into
bucket $B_{j}$. As before, this argument fails when the same bucket $B_{j}$ is
sampled multiple times, due to dependencies in randomness, but this issue can
be avoided by using HSens to decrease the probability that each bucket is
sampled. We give the algorithm in full in Algorithm 4, which provides the
following guarantees:

###### Theorem A.5

Given an input matrix ${\mathbf{A}}\in\mathbb{R}^{n\times d}$ whose rows
arrive sequentially in a data stream along with the corresponding labels of a
measurement vector ${\mathbf{b}}\in\mathbb{R}^{d}$ and a loss function $M$
that is a polynomial of degree at most $p$, there exists an algorithm that
outputs $T$ sequential Hessian samples with variance within a constant factor
of the optimal sampling distribution. n. The algorithm uses
$\tilde{\mathcal{O}}\left(nd^{p}+Td^{p}\right)$ pre-processing time and
$Td^{p}\,\operatorname*{polylog}(Tnd)$ words of space.

We remark that Algorithm 4 can be generalized to handle oracles ${\mathbb{O}}$
corresponding to second-order methods that require batches of subsampled
Hessians in each iteration. For example, if we want to run $T$ iterations of a
second-order method that requires $s$ subsampled Hessians in each batch, we
can simply modify Algorithm 4 to sample $s$ Hessians in each iteration as
input to ${\mathbb{O}}$ and thus $Ts$ Hessians in total.