Adaptive Single-Pass Stochastic Gradient Descent in Input Sparsity Time
Authors: Sepideh Mahabadi
David Woodruff
Samson Zhou
Venue:  (2021)
Keywords: stochastic gradient descent
streaming algorithm
stochastic optimization

Abstract: We study sampling algorithms for variance reduction methods for stochastic optimization. Although stochastic gradient descent (SGD) is widely used for large scale machine learning, it sometimes experiences slow convergence rates due to the high variance from uniform sampling. In this paper, we introduce an algorithm that approximately samples a gradient from the optimal distribution for a common finite-sum form with $n$ terms, while just making a single pass over the data, using input sparsity time, and $\tO{Td}$ space. Our algorithm can be implemented in big data models such as the streaming and distributed models. Moreover, we show that our algorithm can be generalized to approximately sample Hessians and thus provides variance reduction for second-order methods as well. We demonstrate the efficiency of our algorithm on large-scale datasets. 

Under review as a conference paper at ICLR 2021
ADAPTIVE SINGLE-PASS STOCHASTIC GRADIENT DE- SCENT IN INPUT SPARSITY TIME
Anonymous authors Paper under double-blind review
ABSTRACT
We study sampling algorithms for variance reduction methods for stochastic op- timization. Although stochastic gradient descent (SGD) is widely used for large scale machine learning, it sometimes experiences slow convergence rates due the high variance from uniform sampling. In this paper, we introduce an algorithm that approximately samples gradient from the optimal distribution for a common finite-sum form with terms, while just making single pass over the data, using input sparsity time, and O (Td) space. Our algorithm can be implemented in big data models such as the streaming and distributed models. Morcover, we show that our algorithm can be generalized to approximately sample Hessians and thus provides variance reduction for second-order methods as well. We demonstrate the efficiency of our algorithm on sange-ccale datasets.
INTRODUCTION
There has recently been tremendous progress in variance reduction methods for stochastic gradient
descent (SGD) methods for the standard convex finite-sum form optimization problem min F(x) ==
xER
EE= fx).. where fi. IRd + IR is set of convex functions that commonly represent
loss functions. Whereas gradient descent (GD) performs the update rule X X,
1). Fx(,) on
the iterative solution X, at iterations = 1, 2, SGD (Robbins & Monro 1951 Nemirovsky &
Yudin 1983 Nemirovski et al. 2009) picks E [n] iteration with probability p;, and performs xpp VJe, (x+), where Vfe, is the gradient (or subgradient) of Si, and Ne is some predetermined learning rate. Effectively, training example is is sampled with probability Pi, and the model parameters are updated using the selected example.
the update rule X
The SGD update rule only requires the computation of single gradient each iteration and provides an unbiased estimator to the full gradient, compared to GD, which evaluates n gradients at each iteration and is prohibitively expensive for large n.. However, since SGD is often performed with uniform sampling so that the probability Pi, of choosing index E [r]] at iteration is pi,, 1 at all times, the variance introduced by the randomness of sampling specific vector function can bottleneck for the convergence rate of the iterative process. Thus the subject of variance reduction beyond uniform sampling has been well-studied in recent years (Roux et al.) [2012) Johnson & Zhang
2013 Defazio et al. 2014 Reddi et al. [015 Zhao & Zhang Daneshaad et al. 2016 leedell
et al.
2016 (tich et al. 2017 Johnson & Guestrin 2018
Katharopoulos & Fleuret 2018 Salehi
al.
2018
Qian al."
2019).
common technique to reduce variance is importance sampling, where the probabilities pi, are chosen so that vector functions with larger gradients are more likely to be sampled. Thus for Var(v) = E
||]]|] -||E [v]]22 for random vector v, then Pi, = for uniform sampling implies
07 Var
1
= Vfrn
1
[vx(())) -22-1IVF((x))
n
ipPia.
n²

---

Under review as a conference paper at ICLR 2021
whereas importance sampling with Pi,
Σラ= [x(())] gives
V [x((x)]
2 ( I1 - - .1II 1=11
of = Var
n2?
which is most n n Σ (Vx)) (1 by the Root-Mean Square-Arithmetic Mean Inequality, and can be significantly less. Hence the variance at each step is reduced, possibly substantially, e.g.. [Example 1.3 and [Examll 1.4] by performing importance sampling instead of uniform sampling. In fact, it follows from the Cauchy-Schward inequality that the above importance sampling probability distribution is the optimal distribution for variance reduction. However, comput- ing the probability distribution for importance sampling requires computing the gradients in each round, which is too expensive in the first place.
Second-Order Methods. Although first-order methods such as SGD are widely used, they do sometimes have issues such as sensitivity to the choice of hyperparameters, stagnation at high training errors, and difficulty in escaping saddle points. By considering second-order information such as curvature, second-order optimization methods are known to be robust to several of these issues, such as ill-conditioning. For example, Newton's method can achieve locally super-linear convergence rate under certain conditions, independent of the problem. Although naïve second-order methods are generally too slow compared common first-order methods, stochastic Newton-type methods such as Gauss- Newssn have shown to be scalable in the scientific computing community (Roosta-Khorasal Khorasani al.
2014 Roosta- Khorasani & Mahoney 2016a b. Xu al. 2019 2020).
Our Contributions. We give a time efficient algorithm that provably approximates the optimal importance sampling using a small space data structure. Remarkably, our data structure can be implemented in big data models such as the streaming model, which just takes a single pass over the data, and the distributed model, which requires just a single round of communication between parties holding each loss function. For VF 1 Σ f(()), where each f. = f((a, x)) a, for some polynomial and vector a; d. let nnz((A A) be the number of nonzero entries of A :- ai o... ca. Thus for T iterations, where < T < n, GD has runtime Ô (T' mnz(()) while our algorithm has runtime poly(d, log n) + Ô (unz(A)), where we use (-) to suppress polylogarithmic terms.
Theorem 1.1 Let VF = Σ (xx), where each V = f((a; x)) a; for some polynomial f a, For
and vector a, Bdd and let anz(A) be the number of nonzero entries of A a1 & < n, there exists an algorithum that performs steps of SGD and at each step samples a gradient within constant factor of the optimal probability distribution. The algorithm requires single pass over A and uses O (nnz(A (A)) pre-processing time and Ô (Td) space.
Theorem I.I can be used to immediately obtain improved convergence guarantees for a class of functions whose convergence rate depends on the variance of, such as p-smooth functions or strongly convex functions. Recall that SGD offers the following convergence guarantees for smooth functions:
Theorem 1.2 (Nemirovski al. 2009 Meka 2017) Let F be a p-smooth convex function and Xopt = argmin Fxx). Let 0² be an upper bound for the variance of the unbiased estimator across all iterations and X2 -xt Let each step-size ne be r S }.- Then for SGD with initial position
Xg,
1
11²
E [X F [Xxopt)]
Ixo Xopt 122
so that 0
+ 02 + 1 |xo xxpell)) xxppe)))) iterutions suffices to obtain an e-approximate optimal
value by setting 1 V
In the convergence guarantees of Theorem 1.2, variance Tpp from optimal importance sampling, which can be significantly better than the
we obtain a constant factor approximation to the
We use the notation a b denote the vertical concatenation
a b
2

---

Under review as a conference paper at ICLR 2021
variance Juni from uniform sampling in standard SGD. We first show straightforward examples where uniform sampling an index performs significantly worse than importance sampling. For example, if fx() = xa a, then for A a1 o... a,
Example 1.3
When the nonzero entries of the input A are concentrated in small number of vectors a, uniform sampling will frequently sample gradients that are small and make little progress, whereas importance sampling will rarely do so. an extreme case, the A can contain exactly one nonzero vector a, and importance sampling will always oulput the full gradient whereas uniform sampling will only find the nonzero row with probability :.
Example 1.4 It may be that all rows of A have large magnitude, but x is nearly orthogonal to most of the rows of and heavily the direction of now a,- Then (a,; x) a, is small in magniuuee for most i, but (ar,x) ar is large uniform sampling will often ouiput small gradients while importance sampling will output (ar, x) aa with high probability.
Thus Example 1.3] shows that naïve SGD with uniform sampling can suffer up to multiplicative n factor loss in the convergence rate of Theorem 1.2 compared to that of SGD with importance sampling whereas EExampee 1.4 shows a possible additive n factor loss.
Unlike a number of previous variance reduction methods, we do not require distributional assump-
tions (Bouchard et al. 2015: Frostig et al., 2015; Gopal 2016; Jothimurugesan et al. (0018) or offline
access to the data (Roux et al. [ Johnson & Zhang 2013; Defazio et al. [0144] Reddi et al. 2015
Zhao & Zhang
et | Needell et al. 2016 Stich al. [201 |217 1 &
Guestrin 2018 Katharopoulos & Fleuret 2018 Salehi et 2018] Qian et al. (019). On the other
hand, for applications such neural nets which the parameters in the loss function can change,
can use second-ordd -oneeerrrerrreeee approximation for number of iterations, then reread the data to build a new
second-order approximation when necessary.
We complement our main theoretical result with empirical evaluations comparing our algorithm to SGD with uniform sampling for logistic regression on the a9a Adult dataset collected by UCI and retrieved from LibSVM (Chang & Lin [011]. Our evaluations demonstrate that for various step-sizes, our algorithm has significantly better performance than uniform sampling across both the number of SGD iterations and surprisingly, wall-clock time.
We then show that our same framework can also be reworked to approximate importance sampling for the Hessian, thereby performing variance reduction for second-order optimization methods. (Xu et al. 2016) reduce the bottleneck of many second-order optimization methods to the task of sampling
[[((()))))) aTa] a
rows of A- a1
a, so that a row a; is sampled with probability Σn x)) "
for some fixed function f that the Hessian H has the form H == V Σ (f x))a, a, (Xu et al., 2016) show that this finite-sum form arises frequently in machine learning problems such as logistic regression with least squares loss.
Theorem 1.5 Let V (f(x), where each V = f(aa, a, a, for some polynomial ff and vector a E Rdd and let unz(A) be the number of nonzero entries of A == ai o.... o an- For < T & n, there exists an algorithm that subsamples T Hessians within a constant factor of the optimal probability distribution. The algorithm requires a single pass over A and uses Ô (nnz(A) pre-processing time and Ô (Td) space.
2 SGD ALGORITHM
We first introduce a number of algorithms that will be used in our final SGD algorithm, along with their guarantees. We defer all formal proofs to the appendix.
Le polynomial inner product sketch. approximation to | [f ((a,,)) aill, for any query R we call such an algorithm an L2 polynomial inner product sketch and give such an algorithm with the following guarantee:
For fixed polynomial f. we first require a constant-facee factor
Theorem 2.1 For fixed > and polynomial f. there exists a data structure ESTIMATOR that outputs a (1 + approxximation E (I) x)) a ai112 1122.oo for any query x E Rdd The data structure
3

---

Under review as a conference paper at ICLR 2021
requires single pass over A a1 an (possibly through turnstile update D can be built in O (nnz( A) + 2 time and O ( space, uses query time poly (d, n)g, and succeeds with probability poly(n)
The Le polynomial inner product sketch ESTIMATOR is a generalization of AMS variants Alon et al. 1999); Mahabadi et al. (220)) and is simple to implement. For intuition, observe that for and the identity function J. the matrix A E Rand reduces to a vector of length n so that estimating Σi= [I,((a x)) a;l1 is just estimating the squared norm of vector in sublinear space, For degree p polynomial ʃ. ESTIMATOR generates random sign matrices So, Si,..., Sp with ò (+) rows and maintains S, S,A. To estimate Ei=1 ||a ((a, x a for an integer E [0. pl and scalar Og on a given query x, ESTIMATOR creates the q-fold tensor Y = y for each row y of S and the (q - 1)-fold tensor X xo(g-1), Note that X and Y can be refolded into dimensions F-t" and Rdxd"
that YX R and ||a, YXI is an unbiased estimator We give this algorithm in full Algorithm Thus, taking the average over O (+²) instances of the sums of the tensor products for rows y across the sketches A ..... SpA gives (1 + a)aprooximaaio to E [ x)) aaa] l12 with constant probability. The success probability of success can then be boosted to poly(n) by taking the median of O (log n) such outputs.
of E || ((a,x)) x)) . 2-122
Algorithm Basic algorithm ESTIMATOR that outputs (1 + c) -approximation to
E Li (a, x) a 2il where x is post-processing vector
Input: Matrix A = a1 .... E Rexd post-processing vector x E Rd,, integer 2 0, constant parameter > 0.
Output: + -)aprooxinittion E (a, x))²
I: - O(log n.) with a sufficiently large constant. "op..aill. aill².
2: < 2 (.2) with a sufficiently large constant.
3: Let T be an x b table of buckets, where each bucket stores an Rd vector, initialized to the zeros vector.
4: Let E {-11 1, +1} be 4-wise independent for i E 'n].
5: Let h; : [n] + [b] be 4-wise independent for E [r).
6: Let Uig be the all zeros vector for each i E [r]. j [6].
7: for each j = to do
8: for each i 1 to do
9:
Add s,a, to the vector in bucket hz( of row i.
10: Let ViJ be the vector in row i, bucket j of T for i E [r], E [b].
11: Process x:
12: for i E [r].j
do S(P-1)
13: Uij
14: return median 1 EEEN |
L polynomial inner product sampler. Given a matrix A = a1 o... o a, E RRu and a fixed function data structure that takes query Ed and outputs an index i [n] with probability roughly
If((a, a
E |f((a, x)) a
called an L2 polynomial inner product sampler. We give such a data structure in Section A.
Theorem 2.2 For fixed > 0 and polynomial f. there exists data structure SAMPLER that takes any query dgd and outputs index E [n]] with probability Σn ((a.))) poly(nj' along
with vector u := f((a,, a, + where E [v] = 0 and | 1l2 < e ((, x)) a aillz- The data
Turnstile updates are defined as sequential updates the entries of A.
4

---

Under review as a conference paper at ICLR 2021
structure requires single pass over A .... a, (possibly through turnstile updates), can be built in O (miz(A) + 2 time and O d
(2) space, uses query time poly (d, log n),, and succeeds
with probability poly(x))
We remark that T independent instances of SAMPLER provide an oracle for T steps of SGD with importance sampling. but the overall rumtime would be T nnz (A) it would be just efficient to run T' iterations of GD. The subroutine SAMPLER is significantly more challenging to describe and analyze,
defer its discussion |Section A.T, though it can be as combination of ESTIMATOR and a generalized CountSketch Charikar et al. (2004); Nelson & Nguyen (2013); Mahabadi al. 20200) variant and is nevertheless relatively straightforward to implement.
Leverage score sampler. Although SAMPLER outputs (noisy) vector according to the desired probability distribution, we also require an algorithm that automatically does this for indices E [n] that are likely to be sampled multiple times across the T iterations. Equivalently, we require explicitly storing the rows with high leverage scores, but we defer the formal discussion and algorithmic presentation to Section A.2 For our purposes, the following suffices:
Theorem 2.3 There exists an algorithm LEVERAGE that returns all indices i E [in] such that (l
Σ: (x). 2 1 2 for some x E R", along with a vector u; := f((a; x)) a; + v, where
IVill2 S (( x)) aill2. The algorithm requires a single pass over A a1 o... a, (possibly
through turnstile updates), uses Ô (nnz(A) + ) runtime (where w denotes the exponent of square matrix multiplication) and Ô ( space, and succeeds with probability 1 poly(n]
2.1 SGD ALGORITHM AND ANALYSIS
For the finite-sur sum optimization problem min F(x) ==
1 n. Em= f((x), where each V = =(((a,,))
XER
a, recall that we could simply 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 T nnz(A) A) runtime by Theorem 2.2 Thus we use a two level data structure by first implicitly partition the rows of matrix A ai .... a, into == ((Td) buckets BI...., B3 and creating an instance of ESTIMATOR and SAMPLER for each bucket. The idea is that for a given query X in SGD iteration t E [T]. we first query x; to each of the ESTIMATOR data structures to estimate EiEB, (If((a,,)) a 2 for each j E [@]. We then sample index j E [3] among the buckets
By,... Bg with probability roughly Eien, ir, mm.l1 E._,II[((a::> (()(a:xx)))
Once we have sampled index j. it would seem that querying the instance SAMPLER corresponding to B, simulates SGD, since SAMPLER now performs importance sampling the rows in B,, which gives the correct overall probability distribution for each row E [n]. Moreover, SAMPLER has runtime proportional to the sparsity of B,, so the total runtime across the 3 instances of SAMPLER is Ô (Anz(A)).
However, an issue arises when the same bucket By is sampled multiple times, as we only create 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 Ei,li xx))
(.(((a.x.))...
1 2007d for all E [T] and E [n.], then by Bernstein's 's inequality, the probability that no bucket Bj .99
is sampled multiple times is at least T0D Thus we use LEVER to separate all such rows a, that violate this property from their respective buckets and explicitly track the SGD steps in which these rows are sampled. We give the algorithm in full in Algorithm
The key property achieved by Algorithm in partitioning the rows and removing the rows that likely to be sampled multiple times is that each of the SAMPLER instances are queried at most once.
98 Lemma 2.4 With probubility at least T0o each & E [T] uses a different instance of SAMPLER..
Consider Algorithm 2 By Lemma 2.4, each time [T] uses fresh instance SAMPLER,, so that independent randomness is used. A possible concern that each instance ESTIMATOR; using fresh randomness, but observe that ESTIMATOR procedures
Proof of Theorem 1.1]
5

---

Under review as a conference paper at ICLR 2021
Algorithm 2 Approximate SGD with Importance Sampling
Input: Matrix A ai o... a, E Rexd parameter T for number SGD steps.
Output: 7 gradient directions.
1: Preprocessing Stage:
2: B f (TTd) with a sufficiently large constant.
3: Let h: in] - [3] be uniformly random hash function.
4: Let B, be the matrix formed by the
a, of A with h(t) for each [@].
5: Create instance ESTIMATOR; and SAMPLER, for each B, with j [3] with
6: Run LEVERAGE to find a set Lo of row indices and corresponding (noisy) vectors.
7: Gradient Descent Stage:
8: Randomly pick starting location X
9: for 1 to T do
10: Let q be the output of ESTIMATOR; on query x,-1 for each i E [3].
Sample j E [3] with probability Pj Eieto 9:
if there exists E Lo with h(i) = then
13:
Use ESTIMATORJ. LEVERAGE, and SAMPLER, sample gradient w. Vfi, (xe)
14:
else
Use SAMPLER, to sample gradient w, VJa (x,)
[|w.ll2
16:
Pi.t
17:
X f- Xt
Wt
np.,
only used in sampling a bucket E [B]] as an L2 polynomial inner product sketch; otherwise the sampling uses fresh randomness whereas the sampling is built into each instance of SAMPLERJ. By (Theorem 2.2 each index sampled with probability within a factor of the importance sampling probability distribution. By Theorem 2.1. we have that Pi, is within factor 4 of the
probability Pi.t induced by optimal importance sampling SGD. Note that W, = (x), is an unbiased estimator of V (xx:) and || |w, is a a a-appoxxiaation
fx(,)) by Theorem 2.2 Hence, the variance at each time t E [T] of Algoorihmm 2 is within a constant factor of the variance o²
( (xxx))))) [V )11 of optimal importance sampling SGD,
Theorem 2.2 and Theorem 2.3 the preprocessing time is Ô (Inz(A)) + T poly(d, log n.) due to the choices c (1) and / = (Tdd), but partitioning the nonzero entries of A across the /3 buckets. Similarly, the space used by the algorithm is (Td). Once the gradient descent stage of Algorritm beggis. takes poly(d) time in each step E [T] query the 3 = (TTd) instances of SAMPLER and ESTIMATOR, for total time T poly(d, log n).
By Theorem 2.1
3 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 sampled Hessian H,, possibly along with additional inputs such as the current iterate x, and the gradient g: of F, we model the update rule by an oracle O(H,), ). suppressing other inputs to the oracle in the notation. For example, the oracle O corresponding to the canonical recond-order algorithm Newton's 's method can be formulated as
X = 0(x.) = x; [H.] 'g.
By black-boxing the update rule of any second-order optimization algorithm into the oracle, we can focus our attention to the running time of sampling Hessian with nearly the optimal probability distribution. Thus we prove generalizations of the T-2 polynomial inner product sketch, the La polynomial inner product sampler, and the leverage score sampler for Hessians.
Theorem 3.1 For a fixed e > and polynomial f, there exists a data structure HESTIMATOR that outputs (1 <)-approximatton to Σn [|](((,,))) a. as for any query RRd The data

---

Under review as a conference paper at ICLR 2021
structure requires single pass over A a1 .... a, (possibly through turnstile updates), can be built in O (nnz(A) + 2 time and O d space, uses query time poly (d, Log n),, and succeeds with probability poly(x))
Theorem 3.2 For a fixed > 0 and polynomial f. there exists data structure HS that takes any query x E Rd and outputs an index E [n]] with probability (1 ((mmm))s(mmm))))))))))))))))
Σ (((a, a.l
polyinj' along with a matrix U- := [((a;, (()a a a; + V, where E [V] = and [ 6
[f((at,)) a, aill The data structure requires a single pass over A a G .... an (possibly aTaalls through turnstile updates), can be built in Ô (unz(A) + z) time and Ô space, uses query time poly (d. { ,,log), and succeeds with probability poly(n)'
Theorem 3.3 There exists an algorithm HLEVERAGE that returns all indices i E [n] such that (146) ((m,)))) a. a.
Eis x)) a///a Z a some RR"", along with a matrix U, := f((a,, x))x a, V,, where [Vel S e x)) a. a pr The algorithm uses requires a single pass over A = a1 .... an (possibly through turnstile updates), uses Ô (nnz(A) + =) runtime (where denotes the exponent of square matrix multiplication) and Ô ($) space, and succeeds with probability
1
poly(n)
We remark that HSAMPLER and LEVERAGE are generalizations of ESTIMATOR and SAMPLER that simply return an outer product of noisy vector rather than the noisy vector itself.
As before, observe that could simply run an instance of HS AMPLER 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 A- aj o a, into 3 == @ d) buckets B B,. We then create an instance of HESTIMATOR and HSAMPLER for each bucket. For an iterate xt: we sample j E [$] among the
buckets Bi,..., BB with probability roughly Exen, [(())))))))) [|]l? using HESTIMATOR and then Σ. [ [((aa))))) x.)) aa/a)) a.
||
querying HSAMPLER, at x, to sample Hessian among the indices partitioned into bucket By. As before, this argument fails when the same bucket B, is sampled multiple times, due to dependencies in randomness, but this issue be avoided by using HLEVERAGE to decrease the probability that each bucket sampled. We give the algorithm in full in Algorithm 3)
We remark that Algorithm 3 can be generalized to handle oracles O corresponding second-order methods that require batches of subsampled Hessians in each iteration. For example, if we want to run T iterations of second-order method that requires 8 subsampled Hessians in each batch, can simply modify Allgrrihhm 3 to sample Hessians in each iteration as input to O and thus T's Hessians in total.
EMPIRICAL EVALUATIONS
Our primary contribution is the theoretical design of nearly input sparsity time algorithm that approximates optimal importance sampling SGD. In this section we implement a scaled-down version of our algorithm and compare its performance on large-scale real world datasets SGD with uniform sampling on logistic regression. We also consider both linear regression and support-vector machines (SVMs) in the supplementary material. Because most rows have roughly uniformly small leverage scores in real-world data, we assume that no bucket contains a row with a significantly large leverage score and thus the implementation of our importance sampling algorithm does not create multiple samplers for any buckets. By similar reasoning, our implementation uniformly samples number of indices and estimates ET [If((a,, ()) ail12 by rescaling. Observe that although these simplifications to our algorithm decreases the wall-clock running time and the total space used by our algorithm, they only decrease the quality of our solution for each SGD iteration. We also consider two bybrid SGD sampling algorithms; the first takes the better gradient obtained at each iteration from both uniform sampling and importance sampling while the second performs 25 iterations of
7

---

Under review as a conference paper at ICLR 2021
Algorithm 3 Second-Order Optimization with Importance Sampling
Input: Matrix A ai o... o a, E RRax parameter T for number sampled Hessians, oracle !
that performs the update rule.
Output: T approximate Hessians.
I: Preprocessing Stage:
2: B + (Tdd with a sufficiently large constant.
3: Let h i]] [$] be uniformly random hash function.
4: Let Bj be the matrix formed by the rows a, of A with h(t) j. for each [3].
5: Create an instance HESTIMATOR, and HS A for each B, with j E [3] with = 3-
6: Run HLEVERAGE find set Lo of row indices and corresponding (noisy) outer products. Optimization Stage:
7: Second-Orde
8: Randomly pick starting location Xo
9: for 1 to 1" do
Let qi be the output of HESTIMATOR, on query X for each E [8].
11: Sample j E [6] with probability Pj Eieto
12: if there exists E Lo with h(t) then
Use HESTIMATOR,, HLEVERAGE, and HSAMPLER, to sample Hessian H,,
14:
else
15:
Use HSAMPLER, to sample Hessian H. = (x,).
16:
Pit
E
%
17:
X++1
O
L H 1.P..
7.5
Atrege Obyecive Xiher
Aureng Ofgicrive Nahee
Aeoeage Otpective
30
100
130
300
150
30
100
133
300
250
30
100
150
350
SGD Demazoo
SCED Benazion
SCID Beratim
(a) Step size 0.1
(b) Step size 0.01
(c) Step size 0.001
2.4
2.4
24
TrantTuve
Awerge Tree
Average Treil Time
1.2
1.2
1.2
0.8
30
103
130
230
250
0
30
103
1.0
230
250
30
102
130
330
250
SCD lecratio
SGD Beration
SCDD Beratim:
(d) Step size 0.1
(e) Step size 0.01
(i) Step size 0.001
Fig. 1: Comparison of objective values and runtimes for importance sampling (in blue squares), uniform sampling (in red triangles), hybrid sampling that chooses the better gradient at each step (in purple circles), and hybrid sampling that performs 25 steps of importance sampling followed by uniform sampling (in teal X's) over various step-sizes for logistic regression on a9a Adult dataset from UCI, across 250 iterations, averaged over 10 repetitions.
importance sampling before using uniform sampling for the remaining iterations. Surprisingly, SGD importance sampling implementation not only significantly improves upon SGD with uniform sampling. but are also competitive with the two hybrid algorithms. We do not consider other SGD variants due to either their distributional assumptions or lack of known flexibility to big data models. The experiments performed in Python 3.6.9
Intel Core 17-8700K 3.70 GHz CPU with
8

---

Under review as a conference paper at ICLR 2021
10
Average Objective Value
6
0.7
1.4
2.1
2.8
3.5
Average Total Time
-10"
Fig. 2: Comparison of objective values and wall-clock time for importance sampling (in blue squares), uniform sampling (in red triangles), and hybrid sampling that chooses the better gradient at each step (in purple circles) over step-size 0.1 for logistic regression on a9a Adult dataset from UCI, averaged over 3 repetitions across approximately 15 minute total computation time.
12 cores and 64GB DDR4 memory, using a Nvidia Geforce GTX 1080 I IIB GPU. Our code is publicly available at httts: com SGD
www.o///m/aaaaaaaaaa.aaaaaaaaaae adthD ietpappiveeimmopoee www.ieeeeiimttaaeee code
Logistic Regression. We performed logistic regression on the a9a Adult data set collected by UCI and retrieved from LibSVM (Chang & Lin 2011). The features correspond responses from the 1994 Census database and the prediction task to determine whether a person makes over 50K USD year. We trained using data batch of 32581 points and 123 features and tested the performance on a separate batch of 16281 data points. For each evaluation, we generated 10 random initial positions shared for importance sampling and uniform sampling. We then ran 250 iterations SGD for each of the four algorithms, creating only 250 buckets for the importance sampling algorithm and computed the average performance on each iteration across these 10 separate instances. The relative average performance of all algorithms was relatively robust to the step-size. Although uniform sampling used significantly less time overall, our importance sampling SGD algorithm actually had better performance when considering either number of iterations wall-clock time across all tested step-sizes. For example, uniform sampling had average objective value 20680 at iteration 250 using 0.0307 seconds with step-size 0.1, but importance sampling had average objective value 12917 at iteration using 0.025 seconds. We give our results for logistic regression in (ngiiuure We repeat our experiments in Figure to explicitly compare the objective value of each algorithm with respect wall-clock time, rather than SGD iterations. Thus our results in Figure 2 lempirically demonstrate the ad advantages of our algorithm across the most natural metrics. For additional experiments, see Section B
5 CONCLUSION AND FUTURE WORK
We have given variance reduction methods for both first-order and second-order stochastic optimiza- tion. Our algorithms require a single pass over the data, which may even arrive implicitly the form of turnstile updates, and use input sparsity time and Ô (Td) space. Our algorithms are also amenable to big data models such the streaming and distributed models and supported by empirical evaluations on large-scale datasets. We believe there are many interesting future directions to explore. For example, can generalize our techniques to show provable guarantees for other SGD variants and accelerated methods? A very large-scale empirical study of these methods would also be quite interesting.
REFERENCES
Noga Alon, Yossi Matias, and Mario Szegedy. The space complexity of approximating the frequency moments. Comput. Syst. Sci., 58(1):137-147, 1999. 12

---

Under review as a conference paper at ICLR 2021
Alexandr Andoni, Robert Krauthgamer, and Krzysztof Onak. Streaming algorithms via precision sampling. In IEEE 52nd Anmual Symposium on Foundations of Computer Science, FOCS, pp- 363-372, 2011.
13
Guillaume Bouchard, Théo Trouillon, Julien Perez, and Adrien Gaidon. Online learning to sample. CoRR, abs/ 50d./9996.9.... 2015. 3
Chih-Chung Chang and Chih-Jen Lin. LIBSVM: A library for support vector machines. ACM Transactions on Intelligent Systems and Technology, 2:27:1-27:27, 2011 3
Moses Charikar, Kevin C. Chen, and Martin Farach-Colton Finding frequent items in data streams. Theor, Comput. Sci., 312(1):3-15, 2004. 5
Kenneth L. Clarkson and David P. Woodruff. Low rank approximation and regression in input sparsity time. In Symposium on Theory of Computing Conference, STOC, pp. 81-90, 2013. [6
Hadi Daneshmand, Aurélien Lucchi, and Thomas Hofmann. Starting small learning with adaptive sample sizes. In Proceedings of the 33nd International Conference Machine Learning, ICML, pp. 1463--147,, 2016. 3
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,
pp. 1646-1654, 2014. 13
Petros Drineas, Malik Magdon-Ismail, Michael W. Mahoney, and David P. Woodruff. Fast approxi- mation of matrix coherence and statistical leverage. J. Mach. Learn. Res., 13:3475-3506, 2012.
16 17
Roy Frostig, Rong Ge, Sham M. Kakade, and Aaron Sidford. Competing with the empirical risk minimizer in a single pass. In Proceedings of The 28th Conference on Learning Theory, COLT, pp. 728-763, 2015.
Siddharth Gopal. Adaptive sampling for SGD by exploiting side information. In Proceedings of the 33nd International Conference on Machine Learning, ICMI., pp. 364-372, 2016.
Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. In Advances Neural Information Processing Systems 26, Proceedings. pp. 315-323, 2013.
Tyler Johnson and Carlos Guestrin. Training deep models faster with robust, approximate impor- tance sampling. In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems, NeurlPS, pp. 7276-7286, 2018
Ellango Jothimurugesan, Ashraf Tahmasbi, Phillip B. Gibbons, and Srikanta Tirthapura. Variance- reduced stochastic gradient descent on streaming data. Advances Neural Information Pro- cessing Systems 31: Annual Conference on Neural Information Processing Systems, NeurIPS, pp. 9928 9938-9937-10 2018.
Hossein Jowhari, Mert Saglam, and Gábor Tardos. Tight bounds for Ip samplers, finding duplicates in streams, and related problems. In Proceedings of the 30th ACM CCGMOD-SSGAATT-SIAARR Symposium on Principles of Database Systems, PODS, pp. 49-58, 2011. 13
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, pp. 2530-2539, 2018
Sepideh Mahabadi, Ilya P. Razenshteyn, David Woodruff, and Samson Zhou. Non-adaptive adaptive sampling on turnstile streams. In Proccedings of the 52nd Anmual ACM SIGACT Symposium on Theory of Computing, STOC, pp. 1251-1264, 2020.
17
Raghu Meka. Cs289ml: Algorithmic machine learning notes, 2017. URL httphhttt///ettt hat::///rrgummaka
gitbbb. Cttubb.////ae
fdf
10

---

Under review as a conference paper at ICLR 2021
Deanna Needell, Nathan Srebro, and Rachel Ward. Stochastic gradient descent, weighted sampling, and the randomized kaczmarz algorithm. Math. Program., 155(1-2):549-577 -5757, 2016. 3
Jelani Nelson and Huy L. Nguyen. OSNAP:
faster numerical linear algebra algorithms via sparser subspace embeddings. In 54th Anmual IEEE Symposium on Foundations of Computer Science, FOCS, pp. 117-126. IEEE Computer Society, 2013.
16 17
Arkadi Nemirovski, Anatoli B. Juditsky, Guanghui Lan, and Alexander Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal Optimization, 19(4))1574- 1609, 2009.
Arkadi Semenovich Nemirovsky and David Borisovich Yudin. Problem complexity and method efficiency in optimization, 1983.
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 Machine Learning, ICML, volume 97, pp. 5200 -200-5209, 2019.
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, pp. 2647-2655, 2015.
3
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pp. 400 400-407, 1951.
Farbod Roosta- Khorasani and Michael W. Mahoney. Sub-sampled newton methods I: globally convergent algorithms. CoRR, abs/1 .04737, 201 2
Farbod KoastaaKhooasaai and Michael W. Mahoney. Sub-sampled newton methods II: local conver- gence rates. CoRR, abs/1601 004338, 2016b. [
Farbod KhossaaaK Kees Van Den Doel, and Uri Ascher. Stochastic algorithms for inverse problems involving pdes and many measurements. SIAM Journal on Scientific Computing, 36(5):
S3-S22,,4 2014.
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 Neural Information Processing Systems.. pp. 2672-2680, 2012.
Farnood Salchi, Patrick Thiran, and L. Elisa Celis. Coordinate descent with bandit sampling. In Advances in Neural Information Processing Systems 31: Annual Conference Neural Information Processing Systems, NeurlPS, pp. 9267-9277, 2018.
Sebastian U. Stich, Anant Raj, and Martin Jaggi. Safe adaptive importance sampling. In Advances in Neural Information Processing Systems 30: Annual Conference Neural Information Processing Systems, pp. 4381-4391, 2017. 3
Peng Xu, Jiyan Yang, Farbod Goosta-Khorasani Christopher Ré, and Michael W. Mahoney. Sub- sampled newton methods with non-uniform sampling. In Advances in Neural Information Process- ing Systems 29: Annual Conference on Neural Information Processing Systems, pp. 3000-3008, 2016. 3
Peng Xu, Fred Roosta, and Michael W Mahoney. Newton-type methods for non-convex optimization under inexact hessian information. Mathematical Programming, pp. 1-36, 2019. 2
Peng Xu, Fred Roosta, and Michael W. Mahoney. Second-order optimization for non-convex machine learning: an empirical study. In Proceedings of the 2020 SIAM International Conference on Data Mining, SDM 2020, Cincinnati, Ohio, USA, May 7-9, 2020, pp. 199-207. SIAM, 2020. 2
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. pp. 1-9, 2015.
11

---

natural generalizations Under review as a conference paper at ICLR 2021
A DISCUSSION, FULL ALGORITHMS, AND PROOFS
For the sake of presentation, we consider the case where = 2: higher dimensions follow from the same approach, using tensor representation instead of matrix representation. Instead of viewing the input matrix A aj an E Ruxd as a number of rows, we instead view the matrix A A₁ .... An Rndxd, where each matrix A; a, X a; is the outer product of the row a, with itself.
A.1 L2 POLYNOMIAL INNER PRODUCT SAMPLER
For ease of discussion, we describe this section a data structure that allows sampling an index i E [n] with probability approximately ||Ax||ı,2.4 ||A,x|]1,2.a
in linear time and sublinear space, where for
a matrix A E Rndxd, we use ||A]] 1,2,d to denote Di=1 ||A, where each A, E Rddd and A Aoo...oAr The generalization to L2 polynomial inner product sampler follows immediately, Notably, our data structure can be built simply given access to A, and will still sample from the correct distribution when X is given as a post-processing vector. We first describe in Section 1.1 some necessary subroutines that sampler requires. These subroutines of the well-known frequency moment estimation algorithm of Alon et al. (1999) and heavy hitter detection algorithm of Charikar et al. 2004). We then give the L1,2,d sampler in full in Section A.I.2
A.1.1 FREQUENCY MOMENT AND HEAVY HITTER GENERALIZATIONS
We first recall a generalization to the frequency moment estimation algorithm by Alon et al. (999) that also supports post processinn multiplication by any vector x E RR
Lemma A.1 MMahabadi et al. 2020) Given constant >> O, there exists a one-pass streaming algorithm AMS that takes updates to entries of matrix A E Ruxd, as well as query access to post- processing vectors E Rdd and E Rd that arrive after the stream, and outputs a quantity F such that (1 || vl12 < f < (1 +) Ax vl1z- The algorithm uses O (log" + log }))
bits of spuce and succeeds with probability at least
Algorithm 4 Basic algorithm COUNTSKETCH that outputs heavy submatrices of |A Ax where x is a processog Vector
Input: Matrix A E Rndxd, post-proee processoocegg vector Rd constant parameter >> 0. Output: Slight perturbations of the vector A,x for which | 2 e e|||x|||1,2
I: + -Oolog n) with a sufficiently large constant.
2: - with a sufficiently large constant.
3: Let T be an r x table of buckets, where each bucket stores an gdxd matrix, initialized to the
zeros matrix.
4: Let s; E f-1, +1} be 4-wise independent for i E in].
5: Let hę [n] > [b] be 4-wise independent for E [r].
6: Process A:
7: Let
A o... ..... a An, where each A; Rdd
8: for each j = 1 to do
9: for each i
to r do Add A, to the matrix in bucket h2()) of row i.
10:
11: Let M,, be the matrix in row i, bucket j of T for e [r].j e [b].
12: Process
13: for i E [r].j E [b] do
14: M.j < M,,x
15: On query & E [n). report mediangetr) M M,
Let Ai.... An E d and A A₁ A, Banxd Let d be post-processing vector that revealed only after has been completely processed. For given e > 0, we say a block A, with i [ni is heavy if ||xx]] Ax 1,2 We show in Algorithm 4 an algorithm that
12

---

Under review as a conference paper at ICLR 2021
processes A into sublinear space data structure and identifies the heavy blocks of A once is given. Moreover, for each heavy block A, the algorithm outputs a vector y that is a good approximation to X, The algorithm is natural generalization of the CountSketch heavy-hitter algorithm introduced by Charikar et al. (2004)
Vioi((b) denote with the blocks of d rows of v with the largest For vector Rudx1, lg norm set to zeros.
Lemma A.2 There exists an algorithm that uses O d (log² log tt) space that outputs
vector Yi for each index i [n] so that lly:ll2 ||Axx|12 I (Ax) (+) S
(Ax) (3) 1,2,d with probability at least 1 - Moreover if Y = Y1 o Yn then
:3) 1 S Ax Ax- 112 Ax-V11, I =2((Ax))aut (au))) with probability least 5, where
mast (Ax)
(Ax)
Y= Y Y (as(())) denotes the top Eor blocks of Y by 82 norm.
Fix an index i E [n]. Consider the estimate of || in row of the CountSketch table T. Then (i) is the bucket of T in row which A,X hashes. Let Ei be the event that the 2 blocks of size of Ax with the largest & norm are not hashed ha (i). Observe that for b Ω ( with sufficiently large constant, & occurs with probability least 12 by a union bound.
Proof:
Let v be the sum the vectors representing the blocks that are hashed to bucket (i) excluding A,x, so that is the noise for the estimate of A;x in row a. Conditioned &,, we can bound the expected
squared norm of the noise in bucket ha (1) for sufficiently large by E ||v
2 2 S (Ax) mail( $ )
Hence we have Var(IIv;ll2) s (Ax) boill (3) inequality and conditioning on E
Thus from Jensen's 's inequality, Chebyshev's
Pr |Ivella 2 (Ax) (3)
2
2
12 3
(Ax) tail( (3)113 - (Ax)
The first claim then follows from the observation that
and tail (3) 1,2,
noting that we can boost the probability of success to poly(n) by repeating for each of the Olog n) rows and taking the median.
Finally, observe that ((A)) taii( (3) < |x - Y 2 since Y has at most & nonzero blocks, while
(Ax) hail (3) has all zeros in the 2 blocks of Ax with the largest lg norm. Since Ax Y alters at most 3 rows of Ax, each by at most
(Ax) 3) then
2/2²
Ax Y
<
Σ
E (Ax) tast (Ax)
1)1 (aa))))))))))
=2
noil (ax)))))))
A.1.2 SAMPLING ALGORITHM
Our approach is similar to b sampling techniques in Andoni et al. 2011 Jowhari et al. 2011), who consider sampling indices in vectors, and almost verbatim to Mahabadi et al. al. (2020). who consider sampling rows in matrices given post-processing multiplication by vector.
The high level idea is to note that if f E [0, 1] chosen uniformly at random, then
Al
||A.ll
Pr
2 1,,,
||A|| 1,2,d
Thus if B, A: and there exists exactly one index such that ||B:]] 2,. 2 ||A|],, then the task would reduce to outputting Bj that maximizes |B over all j E [m). In fact, we can show that
13

---

Under review as a conference paper at ICLR 2021
B, is O (e)-heavy hitter of B with respect to the L,,2,d norm. Hence, generalization of COUNTSKETCH to identify the heavy hitters of B, approximate the maximum index i, and check whether [B | is at least (an estimate of) |
Unfortunately, this argument might fail due to several reasons. Firstly, there might exists zero or
multiple indices such that 11,, 2 || 1,2,0 Then the probability distribution that an index satisfies ||B:ll,2.a 2 ||Al1,2,dd
and that | ||B,ll,,,ad > [|B]ll,2, for all other E [v) is not the same as the desired distribution. Fortunately, we show that this only happens with small probability, slightly perturbing the probability of returning each [n].
Another possibility is that the error in COUNTSKETCH is large enough to misidentify whether [ 2 || Using statistical test, this case usually be identified and the algorithm will be prevented from outputting sample in this case. Crucially, the probability that the algorithm is aborted by the statistical test is roughly independent of which index achieves the maximum. As
result, the probability of returning each E [n]] is within (1 + <) factor of |A, 1A.lla.24 when the algorithm does not abort.
We show that the probability that algorithm succeeds is (c)) so then running O (log !) instances of the algorithm suffices output some index from the desired distribution with constant probability, or abort otherwise. Because the underlying data structure is a linear sketch, then it is also robust to post-processing multiplication by any vector refers to the scaling factors i as independent random variables, our analysis shows that they only need be O wiie independent and thus we can generate the scaling factors small space in the streaming model. We give the L,,2,d sampler in Algorithm
Rd. Finally, note that although our presentation
Algorithm 5 1,22aa Sampler
Input: Matrix A E Rdndd with A = A₁ o..... An. where each A; E IRdxd, vector x E IRdx1 that arrives after processing A, constant parameter e 0.
Output: Noisy A, of Ax sampled roughly proportional to ||A,x||
Pre-processing Stage:
2: b - Q ( + O(logn) with sufficiently large constants
3: For i E [n], generate independent scaling factors t, E [0, 1] uniformly at random.
4: Let B be the matrix consisting matrices B, AAi
5: Let ESTIMATOR and AMS track the L1,2,a norm of Ax and Frobenius norm of Bx, respectively,
6: Let COUNTSKETCH be b table, where each entry is matrix in Rdxd
7: for each submatrix A, do
8:
Process A:
9: Update COUNTSKETCH with B; A,
10: Update linear sketch ESTIMATOR with A;. Update linear sketch AMS with B, 1
8, A.A.
Post- Post-process x in AMS, COUNTSKETCH, and ESTIMATOR. 13: Sample a submatrix:
PPPocess x:
14: Use ESTIMATOR to compute F with ||Ax|]] 2, < F S <2||Ax|| 11,20
Extract the & (noisy) blocks of rows of Bx with the largest estimated & norms by COUNTSKETTCH.
16: Let M [gndx1 the -bolok sparse matrix consisting of these top (noisy) block.
17: Use AMS to compute S with |Bxx MIl, S < 2 Bx M
18: Let r; be the (noisy) block of rows in COUNTSKETCH with the largest norm.
19: if S > F /1ogo or ||r:lla < 1 then
20:
Return FAIL.
21:
else
22:
Return r f,r,
We first show that the probability that Algorithm 5 returns FAIL is independent of which index i E [n]] achieves argmax;e(n) T, A
14

---

Under review as a conference paper at ICLR 2021
Lemma A.3 Letie E [n] and fix value of [o, 1] uniformly random. Then conditioned on the value of tz.
Pr > F log
= (E)) - poly(n)
Proof: We first observe that we upper bound S by 4 (Bx)
(ax) tail 3)
and lower bound F by
AA |Ax 1,2,
then it suffices to show that the probability of 4 (Bx) rail ))12
vlog?
||Ax||
is small. Thus define E as the event that:
(1) [Ax||1.2 S F s 221|Ax(122.
(2) Bx M S 2 2 |Bx M
(3) Bx taill
- [Bx M| 2 [xx) tail
Note that by Theorem 2.1 Lemma A.2 and Lemma A. Er holds with high probability.
Let U |A For each block Ayx, we define yj to be the indicator variable for whether the scaled block B,x heavy, so that yj = 1 if B, > U and y, = otherwise. We also define 2j E [0, 1] as a scaled random variable for whether Byx is light and how much squared mass it contributes, =j - t7Z |
Ej My be the total number of heavy blocks besides B,x and - j 2, be the total scaled squared mass of the small rows. Let h E IRnd be the vector that contains the heavy blocks so that coordinates (j 1)d + 1 through ja of h correspond to B,x x if yΙ and they are all zeros otherwise. Hence, h contains at most Y nonzero blocks and thus
Let (1-))))
most (Y + 1)d nonzero entries. Moreover, U2 |Bx hl2 and (Bx) (3)1| S SUZ
unless 2 3.
Thus if we define & be the event that Y > 2 and &g to be the event that 7> 160 log ! || 1,2,
then - A EEs implies 4 (Bx) (3)1 /log | so it suffices to bound the proba- ||Ax|1,,,,a
bility of the events Eg and Eg by O (c). Intuitively, if the number of heavy rows is small (~E2) and the total contribution of the small rows is small (-Ex), then the tail estimator small, so the probability of failure due to the tail estimator is small.
To analyze E2, note that y = if and only if 1/1 | > U, so [MJ]
A and thus U
E [Y] < 1 since Y = EjØ; My and U = ||Ax|1 -, |Azl|2 We also have Var(Y) < 1 so that Pr [E2] O (c) for sufficiently small by Chebyshev's 's inequality.
To analyze E,, recall that zj -zz B 2 (1 yj). Thus zj 0 only y or equivalently, ||,x Therefore.
|B,x||] s U. Since B,x t, x then 2j > 0 only if ||Ax|],,2.
1 1
|
E [ <
zj zjd
||A, de, < | 1,2,d
75
U2
||A,0|1//1 Ax 1...
| 1.2
Since Z Σ j Zj. then E [2]] < and similarly Var(Z) < Hence by Bernstein's 's inequality,
Pr 2 T6
log = = O (e), so then Pr (&]] = O (e). Thus Pr E V & Ea] = O + poly(n)
as desired.
We now show that Algorithm 5 outputs noisy approximation to A;x, where E [r)] is drawn from approximately the correct distribution, i.e., the probability of failure does not correlate with the index that achieves the maximum value.
Lemma A.4 For a fixed value of F, the probability that Algorithm 5 outputs (noisy) submatrix A,x (1 + O (e))
|Axx|
F
poly(x)'
15

---

Under review as a conference paper at ICLR 2021
Proof: Let be the event that 4 Al so that Pr Pr[8] =||A,x|la F
1 F Let Ei be the event that COUNTSKETCH, AMS, or ESTIMATOR fails so that Pr [&1] = poly(n) by Lemma A.2 Lemma A.1
Let Eg be the event that S > F Vlog so that Pr [&2] = O (e) by Lemma A.3
and Theorem 2.1
Let Eg be the event that multiple rows B,x exceeding the threshold are observed in the CountSketch data structure and & be the event that |[Bx|]|, exceeds the threshold but is not reported due to noise in the CountSketch data structure. Observe that &g and & are essentially two sides of the same coin, where error is incurred due to the inaccuracies of CountSketch.
To analyze E3, note that row j * i can be reported as exceeding the threshold if |B,xII Z _ - Fv Voloz log 2. which occurs with probability at most O ( (=llAxxll2)
) By union bound over all
rows j E [n]] with j # i, then Pr [Eg] = O (c).
To analyze Es, we first condition on - and -E,, so that Bxx M SS F log + Then by
Lemma A.2 the estimate B,x for B,x output by the sampler satisfies
|B B 2
S c
[xx)
hail (3) sc M 5 of log
Hence, E can only occur for
F < [ S F + eF log 1
which occurs with probability at most O (e²).
To put things together, & occurs with probability ellA,llaa in which case the L1,2,a sampler should output A,x. However, this may not happen due to any of the events &,, Ez, Es. or &g. Since
Pr & V Ea 8]
O (e) and Pr [E4] = O (e²). then we have Pr [E [8] (e). Moreover,
poly(n) so that index i is sampled with probability (1 (c)) ||Axll,, Finally by
Pr [
F
Lemma A.2 [,, |
eF log 1 and [ 2
2 f.. Hence
B is e)
2
approximation to /,,l|| and therefore, to B.x is a (1 + e) approximation to ||,x
2
Thus we have the following full guarantees for our 2 sampler.
Theorem A.5 Given > 0, there exists an algorithm that takes matrix A E Ranxdd which can be written as A = A1 o .... a An where each A E Rdxdd After A is processed, the algorithm is given a query vector
Rd and outputs (noisy) vector A.x with probability
A,x O (+)) [Ax]
polytn] The algorithm uses log 1 {-nnz( Apol (d, 2. logni) time, O (d (poly (L. log + log 4)) bits of space, and succeeds with probability as least 5.
Proof: By Lemma A.4 and Theorem 2.1 then |AP|| 1,2 S S 2 1,2 with high proba-
bility and so each vector A, is sampled with probability (1++) ||A,x12 + poly(n) conditioned [A
on the sampler succeeding. The probability that the sampler succeeds is ( so the sampler can be repeated O (: log n)) times to obtaín probability of success at least
poly(n) Since each instance of AMS, ESTIMATOR, and COUNTSKETCH use log 1 mnz(A) + poly (d. A log n) time and O (d (poly log m) log!)) bits of space, then the total time and space complexity follow.
By adding sketches corresponding to different polynomial degrees, Theorem A.5 implies Theorem 2.2
and Theorem 3.2]
A.2 LEVERAGE SCORE SAMPLER
Our starting point is the input sparsity time algorithm of (Nelson & N Nguyen 2013) for approximating the leverage scores, which is in turn modification of Drineas al. 2012 Clarkson & Woodruff
16

---

Under review as a conference paper at ICLR 2021
2013) Given imput matrix A. (Nelson & Nguyen 2013) randomly samples sparse matrix II with O rows and O (t) signs per column, setting the remaining entries to be zero. (Nelson & Nguyen 2013 maintains A A and poot-procsssing. computes R- so that II, AR has orthonor- mal co c.llums. Previous work of (Drineas et al.] 2012) had shown that the squared row norms of AR are + +1 p-approximations the leverage scores of A. Hence for JL matrix II2 that gives (1 + -approoxmmmtso to the row norms of AR we can compute A (R-) and output the row norms of ARII2 as the approximate leverage scores for each row. Due to the sparsity of II1 and II2. the total runtime Ô ( mnz(A) (A)) Computing R R-1 takes additional Ö (=#) runtime.
Now since the squared it suffices to take the rows of AR -1 with large squared norms. To that effect, we randomly sample CountSketch matrix T and maintain TA. Once R is computed, we can post-processing right multiply to obtain TAR similar to Algorithm 5| It follows that any row of TAR that is at least 2002' heavy (with respect to squared Frobenius norm) has leverage score at least 1000d" Thus we can obtain these rows by querying the CountSketch data structure while using space O (Td). Due to the sparsity of the CountSketch matrix, the total runtime is Ô (nnz(A) + 22 Finally, (Mahabadi eet al. 2020) show that the error guarantee on each reported heavy row required by [Theremm 2.3 By reporting the outer products of each of the heavy rows rather than the heavy rows, we obtain Theorem 3.3
norms of AR-1 are (1 -approximations the leverage scores of A,
d"
A.3 APPROXIMATE SGD WITH IMPORTANCE SAMPLING
For any E (T) and i E [m]. ||A, Z 1oora || only if there exists a row in A; whose leverage score is at least TonTa since there are d rows in A,, Algorithm calculates 2-approximation to each leverage score and maintains T separate instances of the 2,,2,d samplers for any matrix containing these indices E [n]. we maintain T separate instances of the 2,,2,d samplers for A, by explicitly maintaining the heavy row.
Proof of Lemma 2.4;
with approximate leverage score at least 1007 Thus for
Otherwise, for all j [8] so that h(i) # for any index i [n]] such that |A, ToUra |xx 1 we have
Σ |Axx1 S 100T | (:(())
with probability at least TOO with sufficiently high constant. Intuitively, by excluding the hash indices containing "heavy" matrices, the remaining hash indices contain only a small fraction of the mass with high probability. Then the probability that any j E [8] with Σ= |xxell2 S 10T || is sampled more than once most 100T for any E [I'] provided there is no row in any A; with h(()) = j whose l2 leverage score is least 100ra- Thus, the probability that some bucket j [3] is sampled twice across T
99 by Bernstein's 's inequality and union bound over j [3] for 3 3 O(Td)
steps is at most (100T)2 Σ 1 100
In summary, we maintain T separate instances of L1,2,a samplers for the heavy matrices and one L 1,, sampler for each hash index that does not contain heavy matrix. With probability least 98 100 any hash index not containing heavy matrix is sampled only once, so each time E [T] has access to a fresh L,, sampler.
B EMPIRICAL EVALUATIONS
We again emphasize that our primary contribution is the theoretical design of a nearly input sparsity time streaming algorithm that simulates the optimal importance sampling distribution for variance reduction in stochastic gradient descent without computing the full gradient. Thus our theory is optimized to minimize the number of SGD iterations without asymptotic wall-clock time penalties; we do not attempt to further optimize wall-clock runtimes. Nevertheless, in this section implement a scaled-down version of our algorithm and compare its performance across multiple iterations on large-scale real world data sets to SGD with uniform sampling on both linear regression and
17

---

Under review as a conference paper at ICLR 2021
Oliguire Vaho
Obsectiv
0
0
SGD
(a) Step size
x 10 13
(b) Step size v
5 x 10
(c) Step size 1
x 10
Fig. Comparison of importance sampling (in blue squares) and uniform sampling (in red triangles) over various step-sizes for linear regression on CIFAR-10.
support-vector machines (SVMs). Because most rows have roughly uniformly small leverage scores in real-world data, we assume that no bucket contains row with a significantly large leverage score our importance sampling algorithm does not create multiple samplers implementation uniformly samples number of indices i
and thus the implementation for any buckets. By similar reasoning, E, || 2.2 by scaling up ||A, Observe that although these simplifications to our algorithm decreases the wall-clock running time and the total space used by our algorithm, they only decrease the quality of our solution for each SGD iteration. Nevertheless, our implementations significantly improve upon SGD with uniform sampling. The experiments in this section were performed on a Dell Inspiron 15-7579 device with an Intel Core 15-7500U dual core processor, clocked at 2.70 GHz and 2.90 GHz, in contrast to the logistic regression experiments that were performed on GPU.
and estimates ||Ax|| 2,
Linear Regression. We performed linear regression on the CIFAR-10 dataset to compare the performance of our importance sampling algorithm to the uniform sampling SGD algorithm. We trained using data batch of 100000 points and 3072 features and tested the performance separate batch of data points. We aggregated the objective values across 10 separate instances. Fach instance generated random starting location as initial position for both importance sampling and uniform sampling. We then ran 40 iterations of SGD for each algorithm and observed the objective value on the test data for each of these iterations. Finally, we computed the average performance on each iteration across these 10 separate instances. As we ran our algorithm for 40 iterations, we created 1600 buckets that partitioned the data values for the importance sampling algorithm.
The sampled gradients were generally large in magnitude for both importance sampling and uniform sampling and thus we required small step-size. For step-sizes 77 x 10 13, 77 x 10 10-12 and x 10-12 12, the objective value of the solution output by our importance sampling algorithm quickly and significantly improved over the objective value of the solution output by uniform sampling. Our algorithm performance is much more sensitive to the choice of larger step-sizes, as choices of step-sizes larger than 5 x 10 generally caused the importance sampling algorithm to diverge, while the uniform sampling algorithm still slowly converged. We give results in Figure 3
Support-Vectuon Machines. We also compared the performance of our importance sampling al- gorithm to the uniform sampling SGD algorithm using support-vector machines (SVM) on the a9a Adult data set collected by UCI and retrieved from LibSVM (Chang & Lin 2011). The features correspond to responses from the 1994 Census database and the prediction task to determine whether a person makes over 50K USD a year. We trained using data batch of 32581 points and 123 features and tested the performance on a separate batch of 16281 data points. We assume the data is not linearly separable and thus use the hinge loss function so that we aim to minimize 1
D. max(0. (0, 3: (w X; b)) + > ||w [|2 where X is the data matrix, y/ is the corresponding label, and w is the desired maximum-manne -anginn hyperplane. For each evaluation, we generated 10 random initial positions shared for both importance sampling and uniform sampling. We then ran 75 iterations of SGD for each algorithm, creating 125 buckets for the importance sampling algorithm and computed the average performance each iteration across these 5 separate instances.
18

---

Under review as a conference paper at ICLR 2021
The sampled gradients were generally smaller than those from linear regression CIFAR- 10 and thus we were able to choose significantly larger step-sizes. Nevertheless, our algorithm performance was sensitive to both the step-size and the regularization parameter. For step-sizes 1j 0.25, 17 0.5 and regularization parameters À = 0, = 0.001 and À 0.0001, the objective value of the solution output by importance sampling algorithm quickly and significantly improved over the objective value of the solution output by uniform sampling. We give our results in Figure 5 Our algorithm performance degraded with larger values of A, as well step-sizes larger than 7)
We also compared step-size = and regularization parameters 0, A - 0.001 and À - 0.0001 with a hybrid sampling scheme that selects the better gradient between importance sampling and uniform sampling at each step, as well as a hybrid sampling scheme that uses few steps of importance sampling, followed by uniform sampling the remaining steps. Our experiments show that the hybrid sampling algorithms perform better at the beginning and thus our importance sampling algorithm may be used in conjunction with existing techniques in offine settings accelerate SGD. Surprisingly, the hybrid sampling algorithms do not necessarily remain better than our importance sampling algorithm thus indicating that even uniform sampling were run for a significantly larger number of iterations, its performance may not exceed our importance sampling algorithm. We give our results in Figure
Finally, we compare wall-clock times of each of the aforementioned sampling schemes with step-size 7 1 and regularization 0 across 100 iterations. Our results in FFigure 4 show that as expected, uniform sampling has the fastest running time. However, each iteration of importance sampling takes about 15 iterations of uniform sampling, which empirically shows that even using wall-clock times for comparison, rather than total number of SGD iterations, the performance of our importance sampling algorithm still surpasses that of uniform sampling. Moreover, the runtime experiments reveal the main bottleneck of our experiments: each of the 100 iterations took approximately 70 seconds on average after including the evaluation of the objective on each gradient step.
0.18
0.18
Ronine
75-1002
25-10-3
Fig. 4: Runtimes comparison for SVM a9a Adult dataset from LibSVM/UCI with step size 1.0 and regularization 0, averaged over 100 iterations: importance sampling (in blue squares), uniform sampling (in red triangles), hybrid sampling that chooses the better gradient at each step (in purple circles), and hybrid sampling that performs 25 steps of importance sampling followed by uniform sampling (in teal X's). By comparison, the average total runtime over 100 iterations was 72.5504 seconds, including the computation of the scores.
19

---

Under review as a conference paper at ICLR 2021
10
10
125 207
1.50
1.25
Oigativ Wite
Objestiv
0.5
0.
0.5
8:50
0.20
0.25
(a) Step size 0.25, Regularization (b) Step size 0.25, Regularization (c) Step size 0.25, Regularization 0.001
0.0001
10
10P
20F
Objsdive Wale
0.76
0.75
0.75
Ofieottv
0.0
0.3
0.5
0.25
SCDD
(d) Step size 0.5, Regularization (e) Step size 0.5, Regularization (f) Step size 0.5, Regularization 0.001
0.0001
Fig. 5: Comparison of importance sampling (in blue squares) and uniform sampling (in red triangles) over various step-sizes for SVM on a9a Adult dataset from LibSVM/UCI, averaging across ten iterations, averaged over 10 iterations.
1.50
1.25
1.35
Obgatie Valee
Obgedie
0.25
0.25
0.25
20
20
SGD Zenatios
SGD Bertioss
(a) Step size 1.0, Regularization (b) Step size 1.0, Regularization (c) Step size 1.0, Regularization 0.001
0.0001
Fig. 6: Comparison of importance sampling (in blue squares), uniform sampling (in red triangles), hybrid sampling that chooses the better gradient at each step (in purple circles), and hybrid sampling that performs 25 steps importance sampling followed by uniform sampling (in teal X's) over various step-sizes for SVM on a9a Adult dataset from LibSVM/UCI, averaged over 10 iterations
20


Reviews:
[{"review_id": "mVhau9M9HTD", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The present paper shows how to perform variance reduction for SGD (as well as stochastic second-order methods) in a streaming setting. Their algorithms generally run in roughly O(nnz + poly(d)) time, where nnz is the number of non-zero entries in the data and poly(d) is an unspecified polynomial in the input dimension (I believe it is often d^2, but I hope the authors can clarify; see questions below).\n\nStreaming variance reduction is an important practical problem that could potentially speed up model training substantially. As far as I can tell, the theoretical algorithm is correct and the analysis, while a bit terse, is overall well-presented (especially given space constraints). My main concerns are around whether the algorithm really delivers in practice.\n\nI have two concerns along this line. The first is that the theory requires the gradient of the loss to be of the form f(<a,x>) * a, where f is a polynomial. I have trouble thinking of cases where this is the case (even in logistic regression and SVMs, f is not a polynomial; and for neural nets the loss doesn't even have this form). It would help if the authors could clarify this point, since I didn't see it explained in the experiments on logistic regression.\n\nMy second concern is over whether the algorithm is really \"nearly linear-time\" in a practically meaningful sense, given that there is a poly(d) component to the running time, and the exponent in d is not specified. Looking through the actual algorithm and analysis, I believe it is O(d^p) where p is the degree of f, which if we take a second-order approximation to f means at least O(d^2). In many applications this would be infeasible (many models have d > n, so d^2 is worse than n*d). Figure 1 suggests the running time is indeed much slower compared to vanilla SGD. I'd like it if the authors could clarify the actual exponent in the running time, and also present a version of Figure 1 plotting accuracy against wall clock time, which seems like the more relevant comparison between the two algorithms.\n\nI currently give the paper a weak reject, both because I am not currently convinced of the practical relevance of the algorithm, and because I feel the paper somewhat brushes these practical issues under the rug, which could be confusing to readers. However, I would increase my score if the author response to my questions is able to sway me on the practical import question.", "extra_scores": {}, "extra_text": {}}, {"review_id": "gf0xxD0r-8v", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThis paper develops an efficient streaming algorithm to approximate the optimal importance sampling weights for variance reduction in finite-sum SGD. The optimal weights are proportional to each sample's gradient norm; this work uses AMS-like moment estimation to sketch gradient norms which take the form of a bounded-degree polynomial, in time linear in the input sparsity and polynomial in the dimension d, iteration count T, and the log of the number of the samples n. A second-order analogue is derived for approximating optimal importance weights for sampling the Hessian. Some experiments are shown with more simplistic importance weight estimators (not the proposed algorithm), to demonstrate the advantage over uniform sampling.\n\nPros:\n- To my knowledge, this is a novel way to bridge the tools of moment sketching/leverage score sampling with the problems and subproblems in large-scale stochastic optimization.\n- The technical contribution is overall solid and interesting, despite some caveats pointed out below. The decoupling of T from nnz(A) is striking.\n\nCons:\n- This algorithm is somewhat far from being useful (as acknowledged in the first paragraph of Appendix B), despite the last sentence in the abstract. The unconventional regime of d << T << n seems necessary to experience speedups; d-by-d matrix inversions are needed the compute the leverage scores. The \"Estimator\" routine has to maintain a d^(deg f)-sized tensor. This doesn't count against the paper, but it should thus be evaluated as a theoretical work.\n- My primary concern, and the reason I believe this work could use a round of significant revisions before publication, is with clarity. The relationship with prior work could stand to be clearer; the results are presented in a rather unconventional way (see comments below).\n\nDetailed comments:\n- A source of confusion: at the beginning of Section 2, under \"L_2 polynomial inner product sketch\", it's claimed that we require a constant-factor approximation to the norm of the population gradient (sum inside the norm). Then, it seems that all theorems in the main paper are concerned with a sum on the outside, of the norms squared. Then, in the appendix, where the claims are proven, this issue goes away, as the _{1,2,d} norm is a summation of the non-squared norms. A clarification would be appreciated, if there was some reason to switch between norms and squared norms in the main paper, or if some subset of these are typos. Given the above, the proof of Theorem 1.1 in the main paper is hard to follow, in how it uses Theorem 2.2 (top of page 6).\n- \"The generalization to a L_2 polynomial inner product sampler follows immediately.\" Though it might be straightforward, it's probably best to include the full analysis in the appendix, and show the dependences on the degree of the polynomial. (If I understand correctly, these are straightforward analogues in the same way that the tensor procedure in Alg. 1 is a straightforward analogue of scalar moment sketching.)\n- There is no explicit convergence theorem accompanying Algorithm 2; only that it performs T steps of gradient descent. However, this is a minor point, as it should suffice to push the variance bound through the analysis found in the cited work [Zhao & Zhang '15]. Including an explicit convergence rate would be helpful, since the paper culminates in an optimization algorithm.\n- Specific discussions on the cited works on variance reduction, with comparisons of assumptions and convergence rates, would be helpful for placing this work in the context of those results.\n- The paper derives much of its technical content from [Mahabadi et al. '20]; the exposition would benefit from more specific pointers to the theorems and lemmas from that work, as well as where this paper goes beyond those results, as opposed to direct applications.\n- There is also no explicit theorem accompanying the second-order algorithm. The black-box oracle model with Hessian estimators is problematic, since their inverses may not be usable, so approximation guarantees don't give convergence rates for free. See [1,2] for some work on getting stochastic second-order optimization methods to converge.\n\n[1] Erdogdu & Montanari, \"Convergence Rates of Sub-Sampled Newton Methods\".\n[2] Agarwal et al., \"Second-Order Stochastic Optimization for Machine Learning in Linear Time\".\n\n\n*** post-response ***\n\nThanks for the response and extensive modifications. I have increased my score to \"weak accept\" to reflect the improved clarity. I think the paper remains on the borderline, for the following reasons:\n\n- I still doubt the practicality of this algorithm in any realistic settings. The experiments demonstrate that one can gain on the variance term from importance sampling in SGD, but not the practicality of the moment-sketching one-pass estimator of multiple gradient queries on practical problem sizes. Thus I disagree with the \"vastly superior\" characterization; it is well-known that the gradient estimator this paper seeks to efficiently approximate is a good choice; a convincing proof-of-concept is that it's useful to use sketching to efficiently approximate it (whether it's worth the computational overhead + approximation error). The wall-clock time experiment is independently interesting, perhaps showing that you don't even need the sketching ideas from this paper to benefit from importance-sampled stochastic gradients. That said, my evaluation discounts the experimental part and evaluates this as a purely theoretical work.\n- The theoretical work is an application of [Mahabadi et al. '20], without discernible major theoretical innovations. Though, in my opinion, the application to SGD makes it creative and relevant enough for publication, despite practicality concerns.\n- Some clarity issues remain, rendering the paper hard to digest: squared-norms in main paper vs. norms in appendix, hidden dependences on p, \"follows immediately\". After thinking about it, the \"implicitly computed\" tensor contractions in the comment to R4 makes sense, but this should probably be pointed out clearly in the paper, if I correctly understand that it removes all d^p factors from the analysis.", "extra_scores": {}, "extra_text": {}}, {"review_id": "S5faLhsgvQx", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "#### Problem statement\n\nConsider the scenario where we wish to compute the gradient of \n$$ F(x) = 1/n\\sum_i f_i(x)$$\nwhich is given by \n$$ \\nabla F(x) = 1/n\\sum_i \\nabla f_i(x). $$\n\nStochastic gradient descent (SGD) is a computationally efficient alternative to full -blown gradient descent (GD), which gives an unbiased estimator of the gradient, we pick $i$ uniformly at random and compute $\\nabla f_i$. The estimator might have huge variance, hence variance reduction approaches are studied, which sample each $f_i$ with importance weights probability proportional to the (normalized) squared norm of its gradient. But  computing this probability naively is as costly as computing $\\nabla F$. \n\nThis paper proposes a more sophisticated approach to computing importance weights, in the setting where the functions $f_i$ have a special form (essentially polynomials in the inner product with a fixed vector). There is a matrix $A$ where each row $a_i$ defines a function $f_i(x) = P(<a_i,x> )$ for some univariate polynomial $P(t)$, so that its gradient is of the form $\\nabla f_i = P'(<a_i,x>)a_i$. The punchline of this paper is that it can compute a reasonable approximation to the true importance weights in time linear in $O(nnz(A))$ where $nnz(A)$ is the number of non-zero entries in $A$. \n \n#### Pros\n\n* The general problem is clearly important and interesting. Conditioned on the class of functions  (more about that later), the result they get is strong. \n\n* The individual technical ingredients such as the  polynomial inner product sampler and the leverage score sampler are interesting in their own right. \n\n* The paper seems mathematically sound and comfortably above bar at a technical level. \n\n#### Cons\n\nMy main concerns are with the presentation. \n\n* Let us start with the class of functions that can be handled. The paper defines this class as functions such that\n$$\\nabla f_i(x) = p(<a_i, x>)a_i$$\nTo my knowledge, this is just polynomails in the inner product:\n$$f_i(x) = P(<a_i, x>) \\ \\text{where} \\ P'(t) = p(t).$$\nIf so, stating it in the form of a gradient equation is a little strange. \n\n* There needs to be some discussion of why this is an interesting class of functions in this setting, what it covers, and what it doesn't. . The authors mention applications to neural nets at some point in the Introduction. It is unclear how their results apply to neural nets (or even to logistic regression, on which they report experiments), given $P$ is polynomial in the inner product.  \n\n* In general, the paper is very dense. I was unable to get a sense of what the technical novelty of the paper is, and where the authors  are putting together known ideas from previous work.  (See suggestions below). \n\n#### Suggestions\n\n* It might be worth pointing out that in the simplest case when $f_i = < a_i, x>$ is trivial since you only need the sum of the rows. For the case $f_i = <a_i, x>^2$, it seems related to SVD decomposition. So is it the degree $3$ case the simplest case you where you need to do something new? Is there a simple explanation for how your algorithm works in this case? Is there a connection to the $L_p$ sampling problem?\n\n#### Summary of Recommendation\n\nAt this point, my concerns are mainly with the presentation of the paper, not its content. Whereas the problem it tackles is of broad interest, in its current form, the paper will not appeal to a broad ML audience,   \n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "B5sxqWhA4b", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Here is the review of the article: ``Adaptive Single-Pass Stochastic Gradient Descent in Input Sparsity Time''.\n\n**Summary**\n\nThe authors proposed a method consisting in proposing the best-sampling schemes to reduce the variance of the Stochastic Gradient Descent algorithm in the finite sum setting.\n\nMore precisely, they resort to importance sampling techniques to lower (significantly) the variance of the SGD procedure. However, as a perfectly tuned variance reduced scheme would lead to the same complexity as Gradient descent, they try to estimate for each sample the importance of each gradient. \n\nThe main results consists in giving a proof of how works this sampling techniques and why it avoids a $T \\cdot nnz(A)$ running time (the one of full GD). Roughly speaking, the algorithm rests on the following building blocks:\n\n(i) The ability to estimate $\\sum{\\|f(\\langle a_i,x \\rangle)\\cdot a_i\\|_2^2}$ efficiently;\n\n(ii) The ability to solve the problem at each step efficiently in $nnz(A)$;\n\n(iii) The ability to circumvent the use of the (ii) $T$-times by separating the data set into $T\\cdot d$ buckets and apply a coarse-grained version of (ii).\n\n\n\n**Clarity**\n\nDespite the relative complexity of the procedure the paper is relatively comprehensive. The goals are clearly stated and the paper reads rather well when we have understood the purpose of each block. However, it will be very useful for the reader to have an overview understanding of how the algorithm works at the end of the introduction. I cannot yet imagine how the authors can modified the current version to make it more comprehensive, but it is sure that the complexity of the procedure is a real barrier for the reader.\n\n**Quality**\n\nAll the results are clearly stated, and I have nothing to say about the veracity of the theorems that seem to be properly referenced and proved (yet I did not have a precise look at the proofs of the article).\n\n**Originality**\n\nI am not a specialist of this importance sampling literature. However, this approach of trying to estimate the best sampling procedure seems quite new. I will come back to possible and proper comparison with the literature in the {\\bfseries Comments} section.\n\n**Comments**\n\n*General Comments.* Below the general comments and ways to improve the quality of the article.\n\n\n\n**Improve the clarity.** First, some comments regarding the clarity and the paper have been given earlier in the properly named section. As a rather technical and complex article, it is very important to tell to the reader a clear story of how the different blocks build the algorithm. Moreover, I really wonder if the target of such a article is really the ICLR conference. 8 pages do not seem enough to explain the rich ideas of your paper and detail its technical ideas.\n\n**Proper comparison and randomized Kaczmarz algorithm.** Second, and once again, I am not a specialist of the literature, but the referencing on importance sampling SGD does not seem very complete. I would have very liked the article to be compared to what I know from the field, and in my mind: A randomized Kaczmarz algorithm with exponential convergence (Thomas Strohmer, Roman Vershynin) and Stochastic Gradient Descent, Weighted Sampling, and the Randomized Kaczmarz algorithm (Deanna Needell, Nathan Srebro and Rachel Ward) are interesting comparisons to have in mind when it comes to importance sampling for SGD. How does you algorithm compare to them? Is this a computational explanation of how to perform SGD properly without estimating the magnitude of the ligns of the matrix to invert ? \n\n**More experiments.** I know that it is difficult to give a new algorithm, prove that it converges to the right distribution and extensively test it. However, I really think that I will be convinced only if the experiments on real or synthetic examples show a nice (regarding time!) behavior of your algorithm. Indeed, as shown in your example there is a trade-off between the fact that you need more time than SGD on the one hand but lower the variance on the other hand. How can I see this in practice is very important and for now, I am not totally convinced by your simulations. Comparing this to invert largely overparametrized linear system could be a very good benchmark !\n\n\n\n*Minor Comments*\n\n\nThe title in itself does not seem to reflect the content on the paper. Indeed, the word \"adaptive\" is quite confusing (even if I see that it is adaptive to the magnitude of the gradient). Moreover, I would really stress on the fact that it is an importance sampling technique for finite sum SGD.\n\nI really would like that the authors give an intuition for the theorem 2.1 which is the principal building block of the paper. Is this new ? How is it constructed ?\n\nIt seems that at the end of page 1 and beginning of page 2 there is a mistake in the variance. I would expect a $n^2$ in front of $\\|\\nabla F(x_t)\\|^2$ in the $\\sigma_t^2$ expressions.\n\nTypo: there are several confusion with $p_{i,t}$ and $p_{i_{t}}$ and $p_{i_{t}, t}$ \n\n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: While this paper was received pretty well, especially after the revision, reviewers still find it borderline and request further revisions which we cannot check in this short review cycle. Therefore, we encourage the authors to improve the paper and resubmit to a future venue. In particular, please take into account the reviewers' comment to improve the clarity of the paper. Particularly it is critical to clarify the function class you are working with (essentially polynomials) more clearly than what you currently do (i.e., your current gradient definition). It would be helpful for future work to clearly state that this function class is a shortcoming of your work, and that an interesting direction is to extend this to natural function classes in ML (e.g., logistic loss).
Decision: Reject