Title: VARIANCE REDUCTION WITH SPARSE GRADIENTS

Abstract: Variance reduction methods such as SVRG (Johnson & Zhang, 2013) and Spi-derBoost (Wang et al., 2018) use a mixture of large and small batch gradients to reduce the variance of stochastic gradients. Compared to SGD (Robbins & Monro, 1951) , these methods require at least double the number of operations per update to model parameters. To reduce the computational cost of these methods, we introduce a new sparsity operator: The random-top-k operator. Our operator reduces computational complexity by estimating gradient sparsity exhibited in a variety of applications by combining the top-k operator (Stich et al., 2018; Aji & Heafield, 2017) and the randomized coordinate descent operator. With this operator, large batch gradients offer an extra benefit beyond variance reduction: A reliable estimate of gradient sparsity. Theoretically, our algorithm is at least as good as the best algorithm (SpiderBoost), and further excels in performance whenever the random-top-k operator captures gradient sparsity. Empirically, our algorithm consistently outperforms SpiderBoost using various models on various tasks including image classification, natural language processing, and sparse matrix factorization. We also provide empirical evidence to support the intuition behind our algorithm via a simple gradient entropy computation, which serves to quantify gradient sparsity at every iteration.

1. INTRODUCTION:

Optimization tools for machine learning applications seek to minimize the finite sum objective

EQUATION

where x is a vector of parameters, and f i : R d → R is the loss associated with sample i. Batch SGD serves as the prototype for modern stochastic gradient methods. It updates the iterate x with x -η∇f I (x), where η is the learning rate and ∇f I (x) is the batch stochastic gradient, i.e.

∇f I (x) = 1 |I| i∈I ∇f i (x).

The batch size |I| in batch SGD directly impacts the stochastic variance and gradient query complexity of each iteration of the update rule.

In recent years, variance reduction techniques have been proposed by carefully blending large and small batch gradients (e.g. Roux et al., 2012; Johnson & Zhang, 2013; Defazio et al., 2014; Xiao & Zhang, 2014; Allen-Zhu & Yuan, 2016; Allen-Zhu & Hazan, 2016; Reddi et al., 2016a; b; Allen-Zhu, 2017; Lei & Jordan, 2017; Lei et al., 2017; Allen-Zhu, 2018b; Fang et al., 2018; Zhou et al., 2018; Wang et al., 2018; Pham et al., 2019; Nguyen et al., 2019; Lei & Jordan, 2019) . They are alternatives to batch SGD and are provably better than SGD in various settings. While these methods allow for greater learning rates than batch SGD and have appealing theoretical guarantees, they require a per-iteration query complexity which is more than double than that of batch SGD. Defazio (2019) questions the utility of variance reduction techniques in modern machine learning problems, empirically identifying query complexity as one issue. In this paper, we show that gradient sparsity (Aji & Heafield, 2017 ) can be used to significantly reduce the query complexity of variance reduction methods. Our work is motivated by the observation that gradients tend to be "sparse," having only a small fraction of large coordinates. Specifically, if the indices of large gradient coordinates (measured in absolute value) are known before updating model parameters, we compute the derivative of only those coordinates while setting the remaining gradient coordinates to zero. In principle, if sparsity is exhibited, using large gradient coordinates will not effect performance and will significantly reduce the number of operations required to update model parameters. Nevertheless, this heuristic alone has three issues: (1) bias is introduced by setting other entries to zero; (2) the locations of large coordinates are typically unknown;

(3) accessing a subset of coordinates may not be easily implemented for some problems like deep neural networks.

We provide solutions for all three issues. First, we introduce a new sparse gradient operator: The random-top-k operator. The random-top-k operator is a composition of the randomized coordinate descent operator and the top-k operator. In prior work, the top-k operator has been used to reduce the communication complexity of distributed optimization (Stich et al., 2018; Aji & Heafield, 2017) applications. The random-top-k operator has two phases: Given a stochastic gradient and a pair of integers (k 1 , k 2 ) that sum to k, the operator retains k 1 coordinates which are most "promising" in terms of their "likelihood" to be large on average, then randomly selects k 2 of the remaining coordinates with appropriate rescaling. The first phase captures sparsity patterns while the second phase eliminates bias. Second, we make use of large batch gradients in variance reduction methods to estimate sparsity patterns. Inspired by the use of a memory vector in Aji & Heafield (2017), the algorithm maintains a memory vector initialized with the absolute value of the large batch gradient at the beginning of each outer loop and updated by taking an exponential moving average over subsequent stochastic gradients. Coordinates with large values in the memory vector are more "promising," and the random-top-k operator will pick the top k 1 coordinate indices based on the memory vector. Since larger batch gradients have lower variance, the initial estimate is quite accurate. Finally, for software that supports dynamic computation graphs, we provide a cost-effective way (sparse back-propagation) to implement the random-top-k operator.

In this work we apply the random-top-k operator to SpiderBoost (Wang et al., 2018) , a recent variance reduction method that achieves optimal query complexity, with a slight modification based on the "geometrization" technique introduced by Lei & Jordan (2019) . Theoretically, we show that our algorithm is never worse than SpiderBoost and can strictly outperform it when the random-top-k operator captures gradient sparsity. Empirically, we demonstrate the improvements in computation for various tasks including image classification, natural language processing, and sparse matrix factorization.

The rest of the paper is organized as follows. In Section 2, we define the random-top-k operator, our optimization algorithm, and a description of sparse backpropagation. The theoretical analyses are presented in Section 3, followed by experimental results in Section 4. All technical proofs are relegated to Appendix A, and additional experimental details can be found in Appendix B.

2. STOCHASTIC VARIANCE REDUCTION WITH SPARSE GRADIENTS:

Generally, variance reduction methods reduce the variance of stochastic gradients by taking a snapshot ∇f (y) of the gradient ∇f (x) every m steps of optimization, and use the gradient information in this snapshot to reduce the variance of subsequent smaller batch gradients ∇f I (x) (Johnson & Zhang, 2013; Wang et al., 2018) . Methods such as SCSG (Lei & Jordan, 2017) utilize a large batch gradient, which is typically some multiple in size of the small batch gradient b, which is much more practical and is what we do in this paper. To reduce the cost of computing additional gradients, we use sparsity by only computing a subset k of the total gradients d, where y ∈ R d .

For d, k, k 1 , k 2 ∈ Z + , let k = k 1 + k 2 , where 1 ≤ k ≤ d for a parametric model of dimension d.

In what follows, we define an operator which takes vectors x, y and outputs y , where y retains only k of the entries in y, k 1 of which are selected according to the coordinates in x which have the k 1 largest absolute values, and the remaining k 2 entries are randomly selected from y. (11, 12, 13, -14, -15) , y = (-25, -24, 13, 12, 11) and k 1 = k 2 = 1, then S is a singleton uniformly drawn from {1, 2, 3, 4}. Suppose S = {2}, then rtop 1,1 (x, y) = (0, 4y 2 , 0, 0, y 5 ) = (0, -96, 0, 0, 11). If k 1 + k 2 = d, rtop k1,k2 (x, y) = y. On the other hand, if k 1 = 0, rtop 0,k2 (x, y) does not depend on x and returns a rescaled random subset of y. This is the operator used in coordinate descent methods. Finally, rtop k1,k2 (x, y) is linear in y. The following Lemma shows that rtop k1,k2 (x, y) is an unbiased estimator of y, which is a crucial property in our later analysis. Lemma 1. Given any x, y ∈ R d ,

E rtop k1,k2 (x, y) = y, Var rtop k1,k2 (x, y) = d -k 1 -k 2 k 2 top -k1 (x, y) 2 ,

where E is taken over the random subset S involved in the rtop k1,k2 operator and

(top -k1 (x, y)) = y if k 1 > 0 and |x| < |x| (k1) 0 otherwise.

Our algorithm is detailed as below.

Algorithm 1: SpiderBoost with Sparse Gradients.

Input: Learning rate η, inner loop size m, outer loop size T , large batch size B, small batch size b, initial iterate x 0 , memory decay factor α, sparsity parameters k 1 , k 2 .

I 0 ∼ Unif({1, . . . , n}) with |I 0 | = B M 0 := |∇f I0 (x 0 )| for j = 1, ..., T do x (j) 0 := x j-1 , M (j) 0 := M j-1 I j ∼ Unif({1, . . . , n}) with |I j | = B ν (j) 0 := ∇f Ij (x (j) 0 )

N j := m (for implementation) or N j ∼ geometric distribution with mean m (for theory)

for t = 0, . . . , N j -1 do x (j) t+1 := x (j) t -ην (j) t I (j) t ∼ Unif([n]) with |I (j) t | = b ν (j) t+1 := ν (j) t + rtop k1,k2 M (j) t , ∇f I (j) t (x (j) t+1 ) -∇f I (j) t (x (j) t ) M (j) t+1 := α|ν (j) t+1 | + (1 -α)M (j) t x j := x (j) Nj , M j := M (j) Nj Output: x out = x T (for implementation) or x out = x T where T ∼ Unif([T ]) (for theory)

The algorithm includes an outer-loop and an inner-loop. In the theoretical analysis, we generate N j as Geometric random variables. This trick is called "geometrization", proposed by Lei & Jordan (2017) and dubbed by Lei & Jordan (2019) . It greatly simplifies analysis (e.g. Lei et al., 2017; Allen-Zhu, 2018a) . In practice, as observed by Lei et al. (2017) , setting N j to m does not impact performance in any significant way. We only use "geometrization" in our theoretical analysis for clarity. Similarly, for our theoretical analysis, the output of our algorithm is selected uniformly at random from the set of outer loop iterations. Like the use of average iterates in convex optimization, this is a common technique for nonconvex optimization proposed by Nemirovski et al. (2009) . In practice, we simply use the last iterate.

Similar to Aji & Heafield (2017), we maintain a memory vector M (j) t at each iteration of our algorithm. The memory vector is initialized to the large batch gradient computed before every pass through the inner loop, which provides a relatively accurate gradient sparsity estimate of x (j) 0 . The exponential moving average gradually incorporates information from subsequent small batch gradients to account for changes to gradient sparsity. We then use M (j) t as an approximation to the variance of each gradient coordinate in our rtop k1,k2 operator. With M (j) t as input, the rtop k1,k2 operator targets k 1 high variance gradient coordinates in addition to k 2 randomly selected coordinates.

The cost of invoking rtop k1,k2 is dominated by the algorithm for selecting the top k coordinates, which has linear worst case complexity when using the introselect algorithm (Musser, 1997).

2.1. SPARSE BACK-PROPAGATION:

A weakness of our method is the technical difficulty of implementing a sparse backpropagation algorithm in modern machine learning libraries, such as Tensorflow (Abadi et al., 2015) and Pytorch (Paszke et al., 2017) . Models implemented in these libraries generally assume dense structured parameters. The optimal implementation of our algorithm makes use of a sparse forward pass and assumes a sparse computation graph upon which backpropagation is executed. Libraries that support dynamic computation graphs, such as Pytorch, will construct the sparse computation graph in the forward pass, which makes the required sparse backpropagation trivial. We therefore expect our algorithm to perform quite well on libraries which support dynamic computation graphs.

Consider the forward pass of a deep neural network, where φ is a deep composition of parametric functions,

EQUATION

The unconstrained problem of minimizing over the θ can be rewritten as a constrained optimization problem as follows:

min θ 1 n n i=1 loss(z (L+1) i , y i ) s.t. z (L+1) i = φ L (z (L) i ; θ L ) . . . z ( +1) i = φ (z ( ) i ; θ ) . . . z (1) i = φ 0 (x i ; θ 0 ).

(3)

In this form, z L+1 i is the model estimate for data point i. Consider φ (x; θ ) = σ(x T θ ) for 1 ≤ < L, φ L be the output layer, and σ be some subdifferentiable activation function. If we apply the rtop k1,k2 operator per-layer in the forward-pass, with appropriate scaling of k 1 and k 2 to account for depth, we see that the number of multiplications in the forward pass is reduced to

k 1 + k 2 : σ(rtop k1,k2 (v, x) T rtop k1,k2 (v, θ )).

A sparse forward-pass yields a computation graph for a (k 1 + k 2 )-parameter model, and back-propagation will compute the gradient of the objective with respect to model parameters in linear time (Chauvin & Rumelhart, 1995) .

3. THEORETICAL COMPLEXITY ANALYSIS:



3.1. NOTATION AND ASSUMPTIONS:

Denote by • the Euclidean norm and by a ∧ b the minimum of a and b. For a random vector

Y ∈ R d , Var(Y ) = d i=1 Var(Y i ).

We say a random variable N has a geometric distribution, N ∼ Geom(m), if N is supported on the non-negative integers with

P(N = k) = γ k (1 -γ), ∀k = 0, 1, . . . ,

for some γ such that EN = m. Here we allow N to be zero to facilitate the analysis.

Assumption A1 on the smoothness of individual functions will be made throughout the paper.

A1 f i is differentiable with ∇f i (x) -∇f i (y) ≤ L x -y ,

for some L < ∞ and for all i ∈ {1, . . . , n}.

As a direct consequence of assumption A1, it holds for any x, y ∈ R d that

EQUATION

)

To formulate our complexity bounds, we define

f * = inf x f (x), ∆ f = f (x 0 ) -f * .

Further we define σ2 as an upper bound on the expected norm of the stochastic gradients:

EQUATION

)

By Cauchy-Schwarz inequality, it is easy to see that σ 2 is also a uniform bound of ∇f (x) 2 .

Finally, we assume that sampling an index i and accessing the pair ∇f i (x) incur a unit of cost and accessing the truncated version rtop k1,k2 (m,

∇f i (x)) incur (k 1 + k 2 )/d units of cost. Note that calculating rtop k1,k2 (m, ∇f I (x)) incurs |I|(k 1 + k 2 )/d units of computational cost.

Given our framework, the theoretical complexity of the algorithm is

EQUATION

3.2. WORST-CASE GUARANTEE:

Theorem 1. Under the following setting of parameters

ηL = k 2 6dm , B = 2σ 2 2 ∧ n For any T ≥ T ( ) 4∆ f /ηm 2 , E ∇f (x out ) ≤ . If we further set m = Bd b(k 1 + k 2 )

, the complexity to achieve the above condition is

EC comp ( ) = O   σ 3 ∧ √ n 2 L∆ f b(k 1 + k 2 ) k 2   .

Recall that the complexity of SpiderBoost (Wang et al., 2018) is

O σ 3 ∧ √ n

Published as a conference paper at ICLR 2020

3.3 DATA ADAPTIVE ANALYSIS Let g (j) t = top -k1 (M (j) t , ∇f (x (j) t+1 ) -∇f (x (j) t )) 2

, and

G (j) t = 1 n n i=1 top -k1 (M (j) t , ∇f i (x (j) t+1 ) -∇f i (x (j) t )) 2 .

By Cauchy-Schwarz inequality and the linearity of top -k1 , it is easy to see that g

(j) t ≤ G (j)

t . If our algorithm succeeds in capturing sparsity, both g (j)

t and G (j) t will be small. In this subsection we will analyze the complexity under this case. Further define R j as

EQUATION

where E j is taken over all randomness in j-th outer loop (line 4-13 of Algorithm 1). Theorem 2. Under the following setting of parameters

ηL = b ∧ m 3m , B = 3σ 2 2 ∧ n For any T ≥ T ( ) 6∆ f /ηm 2 , E ∇f (x out ) 2 ≤ 2 2 3 + (d -k 1 -k 2 )m k 2 E RT ,

where

RT = 1 T T j=1 R j . If E RT ≤ 2 k2 3(d-k1-k2)m , then E ∇f (x out ) ≤ . If we further set m = Bd b(k 1 + k 2 )

, the complexity to achieve the above condition is

EC comp ( ) = O σ 3 ∧ √ n 2 L∆ f k 1 + k 2 d b b ∧ m .

In practice, m is usually much larger than b. As a result, the complexity of our algorithm is O( (k 1 + k 2 )/d) smaller than that of SpiderBoost if our algorithm captures gradient sparsity. Although this type of data adaptive analyses is not as clean as the worst-case guarantee (Theorem 1), it reveals the potentially superior performance of our algorithm. Similar analyses have been done for various other algorithms, including AdaGrad (Duchi et al., 2011) and Adam (Kingma & Ba, 2014) .

4. EXPERIMENTS:

In this section, we present a variety of experiments to illustrate gradient sparsity and demonstrate the performance of Sparse SpiderBoost. By computing the entropy of the empirical distribution of the absolute value of stochastic gradient coordinates, we show that certain models exhibit gradient sparsity during optimization. To evaluate the performance of variance reduction with sparse gradients, we compute the loss over gradient queries per epoch of Sparse Spiderboost and SpiderBoost for a number of image classification problems. We also compare Sparse SpiderBoost, SpiderBoost, and SGD on a natural language processing task and sparse matrix factorization.

For all experiments, unless otherwise specified, we run SpiderBoost and Sparse SpiderBoost with a learning rate η = 0.1, large-batch size B = 1000, small-batch size b = 100, inner loop length of m = 10, memory decay factor of α = 0.5, and k 1 and k 2 both set to 5% of the total number of model parameters. We call the sum k 1 + k 2 = k = 10% the sparsity of the optimization algorithm.

4.1. GRADIENT SPARSITY AND IMAGE CLASSIFICATION:

Our experiments in this section test a number of image classification tasks for gradient sparsity, and plot the learning curves of some of these tasks. We test a 2-layer fully connected neural network with hidden layers of width 100, a simple convolutional neural net which we describe in detail in Appendix B, and Resnet-18 (He et al., 2015) . All models use ReLu activations. For datasets, we use CIFAR-10 (Krizhevsky et al.), SVHN (Netzer et al., 2011), and MNIST (LeCun & Cortes, 2010) . None of our experiments include Resnet-18 on MNIST as MNIST is an easier dataset; it is included primarily to provide variety.

Our method relies partially on the assumption that the magnitude of the derivative of some model parameters are greater than others. To measure this, we compute the entropy of the empirical distribution of the absolute value of stochastic gradient coordinates. In Algorithm 1, the following term updates our estimate of the variance of each coordinate's derivative:

M (j) t+1 := α|ν (j) t+1 | + (1 -α)M (j) t .

Consider the entropy of the following probability vector p

(j) t = M (j) t / M (j) t

1 . The entropy of p provides us with a measure of how much structure there is in our gradients. To see this, consider the hypothetical scenario where p i = 1/d. In this scenario we have no structure; the top k 1 component of our sparsity operator is providing no value and entropy is maximized. On the other hand, if a single entry p i = 1 and all other entries p j = 0, then the top k 1 component of our sparsity operator is effectively identifying the only relevant model parameter.

To measure the potential of our sparsity operator, we compute the entropy of p while running Spi-derBoost on a variety of datasets and model architectures. The results of running this experiment are summarized in the following table. 1 provides the maximum entropy as well as the entropy of the memory vector before and after training for 150 epochs, for each dataset and each model. For each model, the entropy at the beginning of training is almost maximal. This is due to random initialization of model parameters.

After 150 epochs, the entropy of M t for the convolutional model drops to approximately 3, which suggests a substantial amount of gradient structure. Note that for the datasets that we tested, the gradient structure depends primarily on the model and not the dataset. In particular, for Resnet-18, the entropy appears to vary minimally after 150 epochs. Figure 1 compares SpiderBoost alone to SpiderBoost with 10% sparsity (10% of parameter derivatives). All experiments in this section are run for 50 epochs. In our comparison to SpiderBoost, we measure the number of gradient queries over the size of the dataset N . A single gradient query is taken to be the cost of computing a gradient for a single data point. If i is the index of a single sample, then ∇f i (x) is a single gradient query. Using the batch gradient to update model parameters for a dataset of size B has a gradient query cost of B. For a model with d parameters, using a single sample to update k model parameters has a gradient query cost of k/d, etc.

Our results of fitting the convolutional neural network to MNIST show that sparsity provides a significant advantage compared to using SpiderBoost alone. We only show 2 epochs of this experiment since the MNIST dataset is fairly simple and convergence is rapidly achieved. The results of training Resnet-18 on CIFAR-10 suggests that our sparsity algorithm works well on large neural networks, and non-trivial datasets. We believe Resnet-18 on CIFAR-10 does not do as well due to the gradient density we observe for Resnet-18 in general. Sparsity here not only has the additional benefit of reducing gradient query complexity, but also provides a dampening effect on variance due to the additional covariates in SpiderBoost's update to model parameters. Results for the rest of these experiments can be found in Appendix B.

4.2. NATURAL LANGUAGE PROCESSING:

We evaluate Sparse SpiderBoost's performance on an LSTM-based (Hochreiter & Schmidhuber, 1997) generative language model. We compare Sparse SpiderBoost, SpiderBoost, and SGD. We train our LSTM model on the Penn Treebank (Marcus et al., 1994) 

η(a, b, t) = b + (a -b) m -t m ,

where a = 1.0 and b = 0.1. The schedule interpolates from a to b as the algorithm progresses through the inner loop. For instance, within the inner loop, at iteration 0 the learning rate is 1.0, and at iteration m the learning rate is 0.1. We believe this is a natural way to utilize the low variance at the beginning of the inner loop, and is a fair comparison to an exponential decay learning rate schedule for SGD. Details of the SGD baselines are provided in Figure 2 . We see SpiderBoost is slightly worse than SGD, and sparsity provides a slight improvement over SGD, especially in the first few epochs.

5. CONCLUSION:

In this paper, we show how sparse gradients with memory can be used to improve the gradient query complexity of SVRG-type variance reduction algorithms. While we provide a concrete sparse variance reduction algorithm for SpiderBoost, the techniques developed in this paper can be adapted to other variance reduction algorithms.

We show that our algorithm provides a way to explicitly control the gradient query complexity of variance reduction methods, a problem which has thus far not been addressed. Assuming our algorithm captures the sparsity structure of the optimization problem, we also prove that the complexity of our algorithm is an improvement over SpiderBoost. The results of our comparison to Spider-Boost validates this assumption, and entropy measures provided in Table 1 empirically support our hypothesis that gradient sparsity exists.

Table 1 also supports the results in Aji & Heafield (2017) , which shows that the top-k operator generally outperforms the random-k operator. Our random-top-k operator takes advantage of the superior performance of the top-k operator while eliminating bias via a secondary random-k operator. Not every problem we tested exhibited sparsity structure. While this is true, our analysis proves that our algorithm performs no worse than SpiderBoost in these settings. Even when there is no structure, our algorithm reduces to a random sampling of k 1 + k 2 coordinates, which is essentially a randomized coordinate descent analogue of SpiderBoost. Empirically, we see that Sparse Spider-Boost outperforms SpiderBoost when no sparsity structure is present. We believe this is due to the variance introduced by additional covariates in the SpiderBoost update, which is mitigated in Sparse SpiderBoost by our random-top-k operator.

The results of our experiments on natural language processing and matrix factorization demonstrate that, with additional effort, variance reduction methods are competitive with SGD. While we view this as progress toward improving the practical viability of variance reduction algorithms, we believe further improvements can be made, such as better utilization of reduced variance during training, and better control over increased variance in very high dimensional models such as dense net (Defazio, 2019) . We recognize these issues and hope to make progress on them in future work.

A TECHNICAL PROOFS:

A.1 PREPARATORY RESULTS

Lemma 2 (Lemma 3.1 of Lei & Jordan (2019) ). Let N ∼ Geom(m). Then for any sequence

D 0 , D 1 , . . . with E|D N | < ∞, E(D N -D N +1 ) = 1 m (D 0 -ED N ) . Remark 1. The requirement E|D N | < ∞ is essential. A useful sufficient condition if |D t | = O(Poly(t)

) because a geometric random variable has finite moments of any order.

Lemma 3 (Lemma B.2 of Lei & Jordan (2019) ). Let z 1 , . . . , z M ∈ R d be an arbitrary population and J be a uniform random subset of [M ] with size m. Then

Var   1 m j∈J z j   ≤ I(m < M ) m • 1 M M j=1 z j 2 2 . Proof of Lemma 1. WLOG, assume that |x 1 | ≥ |x 2 | ≥ . . . ≥ |x d |.

Let S be a random subset of {k 1 + 1, . . . , d} with size k 2 . Then

rtop k1,k2 (x, y) = y I( ≤ k 1 ) + d -k 1 k 2 I( ∈ S) .

As a result,

EQUATION

Var rtop k1,k2 (x, y) = d -k 1 k 2 2 y 2 I( > k 1 )P ( ∈ S)(1 -P ( ∈ S)) = d -k 1 -k 2 k 2 y 2 I( > k 1 ).

Therefore,

Var rtop k1,k2 (x, y) = d -k 1 -k 2 k 2 >k1 y 2 = d -k 1 -k 2 k 2 top -k1 (x, y) 2 .

A.2 ANALYSIS OF A SINGLE INNER LOOP:

Lemma 4. For any j, t,

EQUATION

t ) and Var j,t (ν

(j) t+1 -ν (j) t ) ≤ η 2 L 2 b ν (j) t 2 + d -k 1 -k 2 k 2 g (j) t + G (j) t b ,

where E j,t and Var j,t are taken over the randomness of I (j) t and the random subset S involved in the rtop k1,k2 operator.

Proof. By definition,

ν (j) t+1 -ν (j) t = rtop k1,k2 M (j) t , ∇f I (j) t (x (j) t+1 ) -∇f I (j) t (x (j) t ) .

Let S be the random subset involved in rtop k1,k2 . Then S is independent of (I

EQUATION

t ), the tower property of conditional expectation and variance implies that

E j,t ν (j) t+1 -ν (j) t = E I (j) t ∇f I (j) t (x (j) t+1 ) -∇f I (j) t (x (j) t ) = ∇f (x (j) t+1 ) -∇f (x (j) t ),

and

Var j,t ν (j) t+1 -ν (j) t = E I (j) t Var S ν (j) t+1 -ν (j) t + Var I (j) t E S ν (j) t+1 -ν (j) t . (8)

To bound the first term, we note that top -k1 is linear in y and thus

EQUATION

where the last inequality uses Lemma 3. To bound the second term of (8), by Lemma 3,

Var I (j) t E S ν (j) t+1 -ν (j) t = Var I (j) t ∇f I (j) t (x (j) t+1 ) -∇f I (j) t (x (j) t ) ≤ 1 b 1 n n i=1 ∇f i (x (j) t+1 ) -∇f i (x (j) t ) 2 (i) ≤ L 2 b x (j) t+1 -x (j) t 2 (ii) = η 2 L 2 b ν (j) t 2 ,

where (i) uses assumption A1 and (ii) uses the definition that x

EQUATION

t . Lemma 5. For any j, t,

E j,t ν (j) t+1 -∇f (x (j) t+1 ) 2 ≤ ν (j) t -∇f (x (j) t ) 2 + η 2 L 2 b ν (j) t 2 + d -k 1 -k 2 k 2 g (j) t + G (j) t b ,

where E j,t and Var j,t are taken over the randomness of I (j) t

and the random subset S involved in the rtop k1,k2 operator.

Proof. By Lemma 4, we have

ν (j) t+1 -∇f (x (j) t+1 ) = ν (j) t -∇f (x (j) t ) + ν (j) t+1 -ν (j) t -E j,t (ν (j) t+1 -ν (j) t ) . Since I (j) t is independent of (ν (j) t , x (j) t ), Cov j,t ν (j) t -∇f (x (j) t ), ν (j) t+1 -ν (j) t = 0.

As a result,

E j,t ν (j) t+1 -∇f (x (j) t+1 ) 2 = ν (j) t -∇f (x (j) t ) 2 + Var j,t (ν (j) t+1 -ν (j) t ).

The proof is then completed by Lemma 4. Lemma 6. For any j,

E j ν (j) Nj -∇f (x (j) Nj ) 2 ≤ mη 2 L 2 b E j ν (j) Nj 2 + σ 2 I(B < n) B + (d -k 1 -k 2 )m k 2 R j ,

where E j is taken over all randomness in j-th outer loop (line 4-13 of Algorithm 1). 4.

Proof. By definition,

ν (j) t+1 ≤ ν (j) t + rtop k1,k2 M (j) t , ∇f I (j) t (x (j) t+1 ) -∇f I (j) t (x (j) t ) ≤ ν (j) t + ∇f I (j) t (x (j) t+1 ) -∇f I (j) t (x (j) t ) ≤ ν (j) t + 1 b i∈I (j) t ∇f i (x (j) t+1 ) -∇f i (x (j) t ) ≤ ν (j) t + 1 b i∈I (j) t ∇f i (x (j) t+1 ) -∇f i (x (j) t ) 2 ≤ ν (j) t + 2 b    i∈I (j) t ∇f i (x (j) t+1 ) 2 + i∈I (j) t ∇f i (x (j) t ) 2    ≤ ν (j) t + 2n b 1 n n i=1 ∇f i (x (j) t+1 ) 2 + 1 n n i=1 ∇f i (x (j) t ) 2 ≤ ν (j) t + √ 2nσ

As a result, ν

EQUATION

t -∇f (x

EQUATION

t ) 2 = Poly(t). This implies that we can apply Lemma 2 on the sequence D t = ν (j) t -∇f (x (j) t ) 2 . Letting j = N j in Lemma 5 and taking expectation over all randomness in E j , we have

EQUATION

By Lemma 2,

EQUATION

where the last line uses the definition that x j-1 = x (j) 0 , x j = x (j) Nj . By Lemma 3,

EQUATION

)

The proof is completed by putting ( 11), ( 12) and ( 13) together.

Lemma 7. For any j, t, f (x

(j) t+1 ) ≤ f (x (j) t ) + η 2 ν (j) t -∇f (x (j) t ) 2 - η 2 ∇f (x (j) t ) 2 - η 2 (1 -ηL) ν (j) t 2 .

Proof. By (4),

f (x (j) t+1 ) ≤ f (x (j) t ) + ∇f (x (j) t ), x (j) t+1 -x (j) t + L 2 x (j) t -x (j) t+1 2 = f (x (j) t ) -η ∇f (x (j) t ), ν (j) t + η 2 L 2 ν (j) t 2 = f (x (j) t ) + η 2 ν (j) t -∇f (x (j) t ) 2 - η 2 ∇f (x (j) t ) 2 - η 2 ν (j) t 2 + η 2 L 2 ν (j) t 2 .

The proof is then completed.

Lemma 8. For any j,

E j ∇f (x j ) 2 ≤ 2 ηm E j (f (x j-1 ) -f (x j )) + E j ν (j) Nj -∇f (x j ) 2 -(1 -ηL)E j ν (j) Nj 2 ,

where E j is taken over all randomness in j-th outer loop (line 4-13 of Algorithm 1).

Proof. Since ∇f (x) ≤ σ for any x,

|f (x (j) t+1 ) -f (x (j) t )| ≤ σ ν (j) t . This implies that |f (x (j) t )| ≤ σ t k=0 ν (j) t + |f (x (j) 0 )|.

As shown in (10), ν

= Poly(t) and thus |f (x (j) t )| = Poly(t). This implies that we can apply Lemma 2 on the sequence D t = f (x (j) t ). Letting j = N j in Lemma 7 and taking expectation over all randomness in E j , we have

E j f (x (j) Nj +1 ) ≤ E j f (x (j) Nj ) + η 2 ν (j) Nj -∇f (x (j) Nj ) 2 - η 2 ∇f (x (j) Nj ) 2 - η 2 (1 -ηL) ν (j) Nj 2 .

By Lemma 2,

E j f (x (j) Nj ) -E j f (x (j) Nj +1 ) = 1 m E j (f (x (j) 0 ) -f (x (j) Nj )) = 1 m E j (f (x j-1 ) -f (x j )).

The proof is then completed.

Combining Lemma 6 and Lemma 8, we arrive at the following key result on one inner loop.

Theorem 3. For any j,

E ∇f (x j ) 2 ≤ 2 ηm E j (f (x j-1 ) -f (x j )) + σ 2 I(B < n) B + (d -k 1 -k 2 )m k 2 R j -1 -ηL - mη 2 L 2 b E j ν (j) Nj 2 .

A.3 COMPLEXITY ANALYSIS:

Proof of Theorem 1. By definition (7) of R j and the smoothness assumption A1,

ER j ≤ b + 1 b L 2 E x (j) Nj +1 -x (j) Nj 2 ≤ 2η 2 L 2 E ν (j) Nj 2 .

By Theorem 3,

E ∇f (x j ) 2 ≤ 2 ηm E j (f (x j-1 ) -f (x j )) + σ 2 I(B < n) B -1 -ηL - mη 2 L 2 b - 2(d -k 1 -k 2 )mη 2 L 2 k 2 E j ν (j) Nj 2 . Since ηL = k 2 /6dm, ηL + mη 2 L 2 b + 2(d -k 1 -k 2 )mη 2 L 2 k 2 ≤ 1 √ 6 + 1 6 + 1 3 ≤ 1.

As a result,

E ∇f (x j ) 2 ≤ 2 ηm E j (f (x j-1 ) -f (x j )) + σ 2 I(B < n) B .

Since

x out = x T where T ∼ Unif([T ]), we have E ∇f (x out ) 2 ≤ 2 ηmT E(f (x 0 ) -f (x T +1 )) + σ 2 I(B < n) B ≤ 2∆ f ηmT + σ 2 I(B < n) B .

The setting of T and B guarantees that

2∆ f ηmT ≤ 2 2 , σ 2 I(B < n) B ≤ 2 2 .

Therefore, E ∇f (x out ) 2 ≤ 2 . By Cauchy-Schwarz inequality, E ∇f (x out ) ≤ E ∇f (x out ) 2 ≤ . In this case, the average computation cost is

EC comp ( ) = T ( ) B + 2(k 1 + k 2 ) d bm = 3BT ( ) = O B∆ f ηm 2 = O √ BbL∆ f 2 k 1 + k 2 k 2 .

The proof is then proved by the setting of B.

Proof of Theorem 2. Under the setting of η,

ηL + mη 2 L 2 b ≤ 1 √ 3 + 1 3 ≤ 1.

By Theorem 3,

E ∇f (x j ) 2 ≤ 2 ηm E j (f (x j-1 ) -f (x j )) + σ 2 I(B < n) B + d -k 1 -k 2 k 2 R j .

By definition of x out ,

E ∇f (x out ) 2 ≤ 2∆ f ηmT + σ 2 I(B < n) B + (d -k 1 -k 2 )m k 2 E RT .

Under the settings of T and B,

2∆ f ηmT ≤ 2 3 , σ 2 I(B < n) B ≤ 2 3 .

This proves the first result. The second result follows directly. For the computation cost, similar to the proof of Theorem 1, we have

EC comp ( ) = O(BT ) = O L∆ f 2 B m(b ∧ m) .

The proof is then completed by trivial algebra. The simple convolutional neural network used in the experiments consists of a convolutional layer with a kernel size of 5, followed by a max pool layer with kernel size 2, followed by another convolutional layer with kernel size 5, followed by a fully connected layer of input size 16 × side 2 × 120 (side is the size of the second dimension of the input), followed by a fully connected layer of size 120 × 84, followed by a final fully connected layer of size 84× the output dimension.

B.2 NATURAL LANGUAGE PROCESSING:

The natural language processing model consists of a word embedding of dimension 128 of 1000 tokens, which is jointly learned with the task. The LSTM has a hidden and cell state dimension of 1024.

Algorithm 2: SpiderBoost for Natural Language Processing.

Input: Learning rate η, inner loop size m, number of iterations T , large batch matrix Z 2 with 2 batches of size B, small batch matrix Z 1 with 1 batches of size b, initial iterate x 0 , initial states s 0 and S 0 . for t = 0, ..., T -1 do i = mod(t, 1 ) j = mod(t, 2 ) if i = 0 then s t = 0 if j = 0 then S t = 0 if mod(t, m) = 0 then ν t , S t+1 := ∇f Z2j (x t , S t ) s t+1 = s t else g p := ∇f Z1i (x t-1 , s t-1 ) g c , s t+1 := ∇f Z1i (x t , s t ) ν t := ν t-1 + (g c -g b ) S t+1 = S t x t+1 := x t -ην t Output: x T Before describing Algorithm 2, let us derive the full batch gradient of a generative language model. We encode the vocabulary of our dataset of length N so that D ∈ N N is a sequence of integers corresponding to one-hot encodings of each token. We model the transition p(D i+1 |D i , s i ) using an RNN model M as M (D i , s i ) = D i+1 , s i+1 , where s i is the sequential model state at step i. The

:

L∆ f . Thus as long as b = O(1), k 1 = O(k 2 ), our algorithm has the same complexity as SpiderBoost under appropriate settings. The penalty term O( b(k 1 + k 2 )/k 2 ) is due to the information loss by sparsification.

Figures and Tables:

Figure 1: SpiderBoost with 10% sparsity (k = 0.1d) compared to SpiderBoost without sparsity. Left figure compares the two algorithms using Resnet-18 on Cifar-10. Right figure compares the two algorithms using a convolutional neural network trained on MNIST. The x-axis measures gradient queries over N , where N is the size of the respective datasets. Plots are in log-scale.

corpus. The natural language processing model consists of a word embedding of dimension 128 of 1000 tokens, which is jointly learned with the task. The LSTM has a hidden and cell state dimension of 1024. All three optimization algorithms operate on this model. The variance reduction training algorithm for this type of model can be found in Appendix B. We run SpiderBoost and Sparse SpiderBoost with a learning rate η = 0.2, large-batch size B = 40, small-batch size b = 20, inner loop length of m = 2. We run SGD with learning rate 0.2 and batch size is 20. Figure2shows SpiderBoost is slightly worse than SGD, and sparsity provides a noticeable improvement over SGD.

Figure 2: (a): SGD learning rate is 0.2 and batch size is 20. (b): SGD batch size is 103 and learning rate schedule is 0.1 for epochs 0 -10, 0.01 for epochs 10 -20, and 0.001 for epochs 20 -40. The x-axis measures gradient queries over N , where N is the size of the respective datasets. Plots are in log-scale.

Figure 3: SpiderBoost with various values of sparsity, where (sparsity=k/d) corresponds to Spider-Boost with sparsity k/d. Both figures use MNIST. The x-axis measures gradient queries over N , where N is the size of the respective datasets. Plots are in log-scale.

a vector of absolute values, |x| (1) ≥ |x| (2) ≥ . . . ≥ |x| (d) denotes the order statistics of coordinates of x in absolute values, and S denotes a random subset with size k 2 that is uniformly drawn from the set { : |x| < |x| (k1) }. For instance, if x =

Table content: <table><tr><td>where |x| denotes</td><td/><td/></tr><tr><td>rtop k1,k2 (x, y) =</td><td>  </td><td colspan="2">y (d-k1) k2 y if ∈ S if k 1 &gt; 0 and |x| ≥ |x| (k1)</td></tr><tr><td/><td> </td><td>0</td><td>otherwise,</td></tr></table>

Entropy of Memory Vectors

Table content: <table><tr><td/><td/><td>FC NN</td><td/><td/><td>Conv NN</td><td/><td>Resnet-18</td></tr><tr><td/><td>Max</td><td colspan="2">Before After</td><td>Max</td><td colspan="2">Before After</td><td>Max</td><td>Before After</td></tr><tr><td colspan="3">CIFAR-10 18.234 16.41</td><td colspan="3">8.09 15.920 13.38</td><td colspan="2">2.66 23.414 22.59 21.70</td></tr><tr><td>SVHN</td><td colspan="2">18.234 15.36</td><td colspan="3">8.05 15.920 13.00</td><td colspan="2">2.97 23.414 22.62 21.31</td></tr><tr><td>MNIST</td><td colspan="2">18.234 14.29</td><td colspan="3">9.77 15.920 14.21</td><td>2.77</td><td>-</td><td>-</td><td>-</td></tr></table>



Table content: <table/>

Peer review (iclr2020):
This paper aims at improving the computational cost of variance reduction methods while preserving their benefits regarding the fast provable convergence. The existing variance reduction based methods suffer from higher per-iteration gradient query complexity as compared to the vanilla mini-batch SGD, which limits their utility in many practical settings. This paper notices that, for many models, as the training progresses the gradient vectors start exhibiting structure in the sense that only a small number of coordinates have large magnitude. Based on this observation, the paper proposes a modified variance reduction method (by modifying the SpiderBoost method), where a *memory vector* keeps track of the coordinates of the gradient vectors with large variance. Let be the size of the model parameter. During each iteration, one computes the gradient for coordinates with the highest variance (according to the memory vector) and an additional random coordinates. The paper shows that in the worst case, the proposed method has the same gradient query complexity as the SpiderBoost variance reduction method. Assuming that the proposed method can track the sparsity of the gradient vector, the proposed method achieves a gradient query complexity which is times that of the SpiderBoost method. The paper demonstrates the gradient query complexity improvement over the SpiderBoost method on MNIST and CIFAR-10 data set. The paper presents novel results by utilizing the ideas from the field of communication-efficient distributed optimization. As far as the reviewer can tell, the results in the paper are correct. That said, there is quite a bit of room for improvement in terms of the writing of the paper. The paper appears to have way too many typos. For example, In Section 2.1: - Why is introduced? - denotes a random subset with size ---> ? - drawn from the set ----> ? - --> In Lemma 1: - while defining , *.... if |x_{ell}| >= |x_{(k_1)}|* ----> *.... if |x_{ell}| <= |x_{(k_1)}|*? In Section 2.2: - What are ? Shouldn*t these be ? In A1: - right after (5), what is in the definition of ? The authors may also consider making the empirical evaluation more comprehensive by considering tasks from the NLP domain, e.g., language modeling. This would further help asses the utility of the proposed method.