V ARIANCE R EDUCTION WITH S PARSE G RADIENTS
Anonymous authors
Paper under double-blind review
A BSTRACT
Variance reduction methods which use a mixture of large and small batch gradients, such as SVRG (Johnson & Zhang, 2013) and SpiderBoost (Wang et al., 2018), require signiﬁcantly more computational resources per update than SGD (Robbins & Monro, 1951).
We reduce the computational cost per update of variance reduction methods by introducing a sparse gradient operator blending the top-K operator (Stich et al., 2018; Aji & Heaﬁeld, 2017) and the randomized coordinate descent operator.
While the computational cost of computing the derivative of a model parameter is constant, we make the observation that the gains in variance reduction are proportional to the magnitude of the derivative.
In this paper, we show that a sparse gradient based on the magnitude of past gradients reduces the computational cost of model updates without a signiﬁcant loss in variance reduction.
Theoretically, our algorithm is at least as good as the best available algorithm (e.g. SpiderBoost) under appropriate settings of parameters and can be much more efﬁcient if our algorithm succeeds in capturing the sparsity of the gradients.
Empirically, our algorithm consistently outperforms SpiderBoost using various models to solve various image classiﬁcation tasks.
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.
I NTRODUCTION
Optimization tools for machine learning applications seek to minimize the ﬁnite sum objective
[Equation] [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.
[Equation]
The batch size |I| in batch SGD directly impacts the stochastic variance and gradient query complexity of each iteration of the update rule.
Lower variance improves convergence rate without any changes to learning rate, but the step-size in the convergence analysis of SGD decreases with variance (Robbins & Monro, 1951), which suggests that learning rates can be increased when stochastic variance is decreased to further improve the convergence rate of gradient-based machine learning optimization algorithms.
This is generally observed behavior in practice (Smith et al., 2018; Hoffer et al., 2017).
In recent years, new 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.
This leads to the critique by Defazio (2019), which questions the utility of variance reduction techniques in modern machine learning problems.
While their position is reasonable, their results are on basic implementations of variance reduction methods, which does not rule out the possibility that variance reduction can be effective when combined with new ideas.
In this paper, we take a step in this direction by signiﬁcantly improving the computational complexity of variance reduction methods.
In this work, we reduce the query complexity of variance reduction methods by applying a sparse gradient operator, which has been used to reduce the communication complexity of distributed optimization (Stich et al., 2018; Aji & Heaﬁeld, 2017) applications.
Stich et al.
(2018) addressed the communication complexity of distributed SGD by transmitting sparse gradients between computer nodes over a low bandwidth communication channel.
Aji & Heaﬁeld (2017) expands on this work by using a memory vector to record dense gradient updates.
This is done by adding a dense gradient to the memory vector at each iteration and transmitting to other nodes on the network the gradient coordinates which rank among the top k in terms of their magnitude.
The transmitted coordinates are set to zero, which allows coordinates with smaller derivatives to accumulate enough magnitude to eventually be transmitted.
Like the memory vector in Aji & Heaﬁeld (2017), we also use a vector to rank the variance of the additional gradient computations introduced by variance reduction methods, and compute both the derivative of model parameters which rank among the top k 1 of the stored memory vector, and additionally compute the gradient of k 2 randomly selected model parameters from the remaining d − k 1 parameters. 1. C ONTRIBUTIONS
We introduce a novel way to reduce the computational complexity of SVRG-style variance reduction methods using gradient sparsity estimates.
Concretely, we deﬁne an algorithm which applies these ideas to SpiderBoost (Wang
et al., We provide a complete theoretical complexity analysis of our algorithm, which shows algorithmic improvements in the presence of gradient sparsity structure. We experimentally show the presence of sparsity structure for some deep neural networks, which is an important assumption of our algorithm. Our experiments show that, for those deep neural networks, sparse gradients improve the empirical convergence rate by reducing both variance and computational complexity.
We include additional experiments on natural language processing and sparse matrix factorization, and compare our algorithms to two different SGD baselines.
These experiments demonstrate different ways in which variance reduction methods can be adapted to obtain competitive performance on challenging optimization tasks.
The rest of the paper is organized as follows.
We begin by providing a sparse variance reduction algorithm based on a combination of SCSG (Lei et al., 2017) and SpiderBoost (Wang et al., 2018).
We then explain how to perform sparse back-propagation in order to realize the beneﬁts of sparsity.
We prove both that our algorithm is as good as SpiderBoost, and under reasonable assumptions, has better complexity than SpiderBoost.
Finally, we present our experimental results which include an empirical analysis of the sparsity of various image classiﬁcation problems, and a comparison between our algorithm and SpiderBoost.
S TOCHASTIC V ARIANCE R EDUCTION WITH S PARSE G RADIENTS
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 .
2. S PIDER B OOST WITH S PARSE G RADIENTS
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 deﬁne an operator which takes vectors x, y and outputs y (cid:48) , where y (cid:48) 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 .
The kcoordinate indices and k 2 coordinate indices are disjoint.
Formally, the operator rtop k 1 ,k 2 : R d →
R d is deﬁned for x, y ∈ R d as
[Equation]
where | x | denotes 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 { (cid:96) : | x | (cid:96) <
| x | ( k 1 ) } .
For instance, if x = (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 , 4 y 2 , 0 , 0 , y 5 )
=
(0 , − 96 , 0 , 0 , 11) .
If k 1 + k 2 =
d , rtop k 1 ,k 2 ( x, y )
=
y .
On the other hand, if k 1 = 0 , rtop 0 ,k 2 ( 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 k 1 ,k 2 ( x, y ) is linear in y .
The following Lemma shows that rtop k 1 ,k 2 ( 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 ,
[Equation]
where E is taken over the random subset S involved in the rtop k 1 ,k 2 operator and
[Equation] if k 1 > 0 and | x | (cid:96) < | x | ( k 1 ) otherwise.
Our algorithm is detailed as below.
Algorithm 1: 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 .
1 I 0 ∼ Unif( { 1 , . . . , n } ) with |I
|
=
B 2 M 0 :=
|∇ f I 0 ( x 0 ) | 3 for j = 1 , ..., T do x ( j ) 0 :=
x j − 1 , M ( j ):=
M j −
I j ∼ Unif( { 1 , . . . , n } ) with |I
j | =
B ν ( j ) 0 := ∇ f I j ( x ( j ) 0 )
N j := m (for implementation) or N j ∼ Geom( 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 k 1 ,k 2 (cid:16) M ( j ) t , ∇ f I ( j ) t ( x ( j ) t +1 )
− ∇ f I ( j ) t ( x ( j ) t ) (cid:17)
M ( j ) t +1 :
= α | ν ( j ) t +1 | + (1 − α )
M
( j ) t x j := x ( j ) N j , M j :=
M ( j ) N j
Output: x out = x T (for implementation) or x out = x T (cid:48) where T (cid:48) ∼ 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 simpliﬁes analysis (e.g. Lei et al., 2017;
Allen-Zhu, 2018a).
In practice, as observed by Lei et al.
(2017), it does not make a difference if N j is simply set to be m .
For this reason, we apply ”geometrization” in theory to make arguments clean and readable.
On the other hand, in theory the output is taken as uniformly random elements from the set of last iterates in each outer loop.
This is a generic strategy for nonconvex optimization, as an analogue of the average iterates for convex optimization, proposed by Nemirovski et al.
In practice, we simply use the last iterate as convention.
Similar to Aji & Heaﬁeld (2017), we maintain a memory vector at each iteration of our algorithm.
We assume the optimization procedure is taking place locally and thus do not transmit and zero out any components.
Instead, we maintain an exponential moving average M ( j ) t of the magnitudes of each coordinate of our gradient estimate ν ( j ) t .
We then use M ( j ) t as an approximation to the variance of each gradient coordinate in our rtop k 1 ,k 2 operator.
With M ( j ) t as input, the rtop k 1 ,k 2 operator targets k 1 high variance gradient coordinates in addition to the k 2 randomly selected coordinates.
The cost of invoking rtop k 1 ,k 2 is dominated by the algorithm for selecting the top k coordinates, which has linear worst case complexity when using the introselect algorithm.
Algorithmic implementation details for a sparse back-propagation algorithm can be found in appendix B.
G RADIENT Q UERY C OMPLEXITY
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 k 1 ,k 2 ( y, ∇ f i ( x )) incur ( k 1 + k 2 )
/d units of cost.
Note that calculating rtop k 1 ,k 2 ( y, ∇ 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]
T HEORETICAL C OMPLEXITY A NALYSIS 3. N OTATION AND ASSUMPTIONS
Denote by (cid:107) · (cid:107) the Euclidean norm and by a ∧ b the minimum of a and b .
For a random vector
Y ∈ R d ,
Var
[Equation]
Var [Equation]
We say a random variable N has a geometric distribution, N ∼ Geom( m ) , if N is supported on the non-negative integers with
[Equation]
for some γ such that E N = m .
Here we allow N to be zero to facilitate the analysis.
Assumption A 1 on the smoothness of individual functions will be made throughout the paper.
A 1 f i is differentiable with
[Equation]
for some L < ∞ and for all i ∈ { 1 , . . . , n } .
As a direct consequence of assumption A 1, it holds for any x, y ∈ R d that
[Equation]
To formulate our complexity bounds, we deﬁne
[Equation]
Further we deﬁne σ 2 as an upper bound on the variance of the stochastic gradients:
[Equation]
[Equation]
3. W ORST - CASE GUARANTEE
Theorem 1.
Under the following setting of parameters
[Equation]
For any T ≥ T ( (cid:15) ) (cid:44) 4∆ f /ηm(cid:15) 2 , E
[Equation]
f [Equation] x out [Equation]
If we further set
[Equation]
the complexity to achieve the above condition is
[Equation]
Recall that the complexity of SpiderBoost (Wang et al., 2018) is
[Equation]
Thus as long as b =
O (1) , k
=
O ( k 2 ) , our algorithm has the same complexity as SpiderBoost under appropriate settings.
The penalty term O ( (cid:112) b ( k 1 + k 2 ) /k 2 ) is due to the information loss by sparsiﬁcation.
3. D ATA ADAPTIVE ANALYSIS
Let[-g ( j ) t = (cid:107) top − k 
 ( M ( j ) t , ∇ f ( x ( j ) t +1 )-]
[Equation]
and
[Equation]
By Cauchy-Schwarz inequality and the linearity of top − k 1 , it is easy to see that g ( j ) t ≤ G ( j ) t .
If our algorithm succeeds in capturing the sparsity, both g ( j ) t and G ( j ) t will be small.
In this subsection we will analyze the complexity under this case.
Further deﬁne 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
[Equation]
(cid:15)
[Equation]
For any T ≥ T ( (cid:15) ) (cid:44) 6∆ f /ηm(cid:15) 2 , E (cid:107)∇ f ( x out ) (cid:107) 2 ≤ 2 (cid:15) 2 3 + ( d − k 1 − k 2 ) m k
E ¯ R T , where ¯ R T = 1 T
T (cid:88)
j =
R j .
If E ¯ R T ≤ (cid:15)k 2 3( d − k 1 − k 2 ) m , then
E (cid:107)∇ f ( x out ) (cid:107) ≤ (cid:15).
If we further set
[Equation]
the complexity to achieve the above condition is
[Equation]
In practice, m is usually much larger than b .
As a result, the complexity of our algorithm is
O ( (cid:112) ( k 1 + k 2 ) /d ) smaller than that of SpiderBoost if our algorithm captures the sparsity as desired.
Although this type of data adaptive analysis is not as clean as the worst-case guarantee (Theorem 1), it can reveal the potentially superior performance of the algorithm.
Similar analyses have been done for various other algorithms, including AdaGrad (Duchi et al., 2011) and Adam (Kingma &
Ba, 2014).
E XPERIMENTS
We ran a variety of experiments to demonstrate the performance of Sparse SpiderBoost, as well as to illustrate the potential of sparsity as a way to improve the gradient query complexity of variance reduction methods.
We include performance on image classiﬁcation to further illustrate the performance of SpiderBoost with and without sparsity.
We also provide additional experiments, including a natural language processing task and sparse matrix factorization to evaluate our algorithm on a variety of tasks.
For all experiments, unless otherwise speciﬁed, 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
= 10% the sparsity of the optimization algorithm.
4. G RADIENT S PARSITY AND I MAGE C LASSIFICATION
Our experiments in this section test a number of image classiﬁcation 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 C, 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 for the other models we include in this work.
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 over the magnitude of the derivative of the model parameters.
In Algorithm 1, the following term updates our estimate of the variance of each coordinate’s derivative:
[Equation]
Consider the entropy of the following probability vector p =
[Equation]
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
=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 SpiderBoost on a variety of datasets and model architectures.
The results of running this experiment are summarized in the following table.
The entries of tables 1a and 1b correspond to the entropy of the memory vector before and after training.
For each model, the entropy at the beginning of training is almost maximal.
For example, [Table] maximum entropy of the convolutional model, which consists of 62 , 006 parameters, is 15 .
This is mainly 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.
Hyperparameter tuning of sparsity values 0 .
1 , 0 . 2 , 0 . 4 , and 0 .
8 for Sparse SpiderBoost showed that a value of 0 .
1 performs best.
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 N has a gradient query cost of N .
For a model with d parameters, using a single sample to update kd model parameters has a gradient query cost of kd , etc.
Our results of ﬁtting the convolutional neural network to MNIST show that sparsity provides a signiﬁcant 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.
Results for the rest of these experiments can be found in appendix C.
4. N ATURAL L ANGUAGE P ROCESSING AND S PARSE M ATRIX F ACTORIZATION
To further evaluate SpiderBoost with sparse gradients, as well as variance reduction methods in general, we test our algorithm on an LSTM (Hochreiter & Schmidhuber, 1997) generative language model, and Bayesian Personalized Ranking (Rendle et al., 2009).
The LSTM model is trained on the Penn Treebank (Marcus et al., 1994) corpus.
The details of the training algorithm can be found in appendix C. For the NLP experiment, 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 . The matrix factorization model is trained on the MovieLens database (Harper & Konstan, 2015) with a latent dimension of 20 . Further details can be found in C. We run SpiderBoost and Sparse SpiderBoost with a large-batch size B = 1030 , small-batch size b = 103 , inner loop length of m = 10 .
For this experiment, we run SpiderBoost with a a learning rate schedule that interpolates from η = 1 . 0 to η = 0 .
1 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 it is 0 .
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 ﬁgure 2.
We see that for both tasks, SpiderBoost is slightly worse than SGD, and sparsity provides a slight improvement over SGD.
C ONCLUSION
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 explicitly 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 SpiderBoost validates this assumption, and our entropy experiment empirically supports the hypothesis that gradient sparsity does exist.
The results of our entropy experiment also support the results in Aji & Heaﬁeld (2017), which show that the top k operator generally outperforms the 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.
The results of our experiments on natural language processing and matrix factorization demonstrate that, with extra engineering effort, variance reduction methods can be competitive with SGD baselines.
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 T ECHNICAL P ROOFS A. P REPARATORY 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
[Equation]
D N
[Equation]
D N
[Equation]
1 ) =
[Equation]
Remark 1.
The requirement E | D N | < ∞ is essential.
A useful sufﬁcient condition if | D t | =
O (Poly( t ))
because a geometric random variable has ﬁnite 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
[Equation]
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
[Equation]
As a result,
[Equation]
and [Equation]
> k 1 ) .
Therefore,
Var
[Equation]
A. A NALYSIS OF A SINGLE INNER LOOP
Lemma 4.
For any j, t ,
[Equation]
and
[Equation]
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 k 1 ,k 2 operator.
Proof.
By deﬁnition,
[Equation]
Let S be the random subset involved in rtop k 1 ,k 2 .
Then S is independent of ( I ( j ) t , M ( j ) t , x ( j ) t +1 , x ( j ) t ) .
By Lemma 1,
[Equation]
and
[Equation]
Since I ( j ) t is independent of ( M ( j ) t , x ( j ) t +1 , x ( j ) t ) , the tower property of conditional expectation and variance implies that
[Equation]
and
[Equation]
To bound the ﬁrst term
[Equation]
we note that top − k 1 is linear in y and thus
[Equation]
where the last inequality uses Lemma 3.
To bound the second term of equation 6, by Lemma 3,
[Equation]
where (i) uses assumption A 1 and (ii) uses the deﬁnition that x ( j ) t +
= x ( j ) t − ην
( j ) t .
Lemma 5.
For any j, t ,
[Equation]
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 k 1 ,k 2 operator.
Proof.
By Lemma 4, we have
[Equation]
Since I ( j ) t is independent of ( ν ( j ) t , x ( j ) t ) ,
[Equation]
As a result
[Equation]
The proof is then completed by Lemma 4.
Lemma 6.
For any j ,
[Equation]
where E j is taken over all randomness in j -th outer loop (line 4-13 of Algorithm 1).
Proof.
By deﬁnition,
[Equation]
As a result,
[Equation]
Thus,
[Equation]
This implies that we can apply Lemma 2 on the sequence D t
= (cid:107) ν
( j ) t − ∇ f ( x
[Equation]
j ) t ) (cid:107) 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 deﬁnition that x j − 1 = x ( j ) 0 , x j = x ( j ) N j .
By Lemma 3,
[Equation]
The proof is completed by putting equation 9, equation 10 and equation 11 together.
Lemma 7.
For any j, t ,
[Equation]
Proof.
By equation
[Equation]
The proof is then completed.
Lemma 8.
For any j ,
[Equation]
where E j is taken over all randomness in j -th outer loop (line 4-13 of Algorithm 1).
Proof.
Since (cid:107)∇ f ( x ) (cid:107) ≤ σ for any x ,
[Equation]
This implies that
[Equation]
As shown in equation 8, (cid:107) ν
( j ) t (cid:107) =
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
[Equation]
By Lemma 2,
[Equation]
The proof is then completed.
Combining Lemma 6 and Lemma 8, we arrive at the following key result on one inner loop
[Equation]
Theorem 3.
For any j ,
[Equation]
A. C OMPLEXITY ANALYSIS
Proof of Theorem 1 .
By deﬁnition equation 5 of R j and the smoothness assumption A 1,
[Equation]
By Theorem 3,
[Equation]
Since η = (cid:112) k 2 / 6 dm ,
[Equation]
As a result
[Equation]
Since x out = x T (cid:48) where T (cid:48) ∼ Unif([ T ])
[Equation]
we have
[Equation]
The setting of T and B guarantees that
[Equation]
Therefore,
[Equation]
By Cauchy-Schwarz inequality,
[Equation]
In this case
[Equation]
the average computation cost is
[Equation]
The proof is then proved by the setting of B
[Equation]
Proof of Theorem 2 .
Under the setting of η ,
[Equation]
By Theorem 3,
[Equation]
By deﬁnition of x out ,
[Equation]
Under the settings of T and B ,
[Equation]
This proves the ﬁrst result.
The second result follows directly.
For the computation cost, similar to the proof of Theorem 1, we have
[Equation]
The proof is then completed by trivial algebra.
B S PARSE B ACK -P ROPAGATION
A weakness of our method is the technical difﬁculty of implementing a sparse backpropagation algorithm in modern machine learning libraries, such as Tensorﬂow (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.
This makes the required sparse backpropagation trivial and suggests that our algorithm will perform best on libraries that 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 θ (cid:96) can be rewritten as a constrained optimization problem as follows:
[Equation]
In this form, z L +1 i is the model estimate for data point i .
Consider φ (cid:96) ( x ; θ (cid:96) )
=
σ ( x T θ (cid:96) ) for 1 ≤ (cid:96) <
L , φ L be the output layer, and σ be some subdifferentiable activation function.
If we apply the rtop k 1 ,k 2 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 k 1 ,k 2 ( v, x ) T rtop k 1 ,k 2 ( v, θ (cid:96) )) .
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).
C E XPERIMENTS C. D ESCRIPTION OF S IMPLE C ONVOLUTIONAL N EURAL N ETWORK
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 ×(side is the size of the second dimension of the input), followed by a fully connected layer of size 120 × 84 , followed by a ﬁnal fully connected layer of size 84 × the output dimension. C. [Figure] N ATURAL L ANGUAGE P ROCESSING
The natural language processing model consists of a word embedding of dimension 128 of 1 ,tokens, which is jointly learned with the task.
The LSTM has a hidden and cell state dimension of 1024 .
The variance reduction training algorithm for this type of model is given below.
{+do if mod ( t, (cid:96) 2 B ) = 
+}
then s t =
S t
{+= if mod ( t, m ) = 
+}
then ν t
[Equation]
S t +1 := ∇ f J mod ( t,(cid:96) 2) ( x t , S t ) s t +1 =
s t else
[Equation]
) g
[Equation]
) ν t :=
ν t
[Equation]
+ ( g c − g b )
S t
[Equation]
=
S t x t +1 := x
[Equation]
− ην t
Output: x T
In Algorithm 2, for a sequence dataset D of length | D | , we choose B = mb
{+and | D | = (cid:96) 2 B = (cid:96) 1 b , so if mod ( t, (cid:96) 2 B ) = 
+}
then mod ( t, (cid:96) 1 b ) =
mod ( t, m ) = 0 .
Algorithm 2 can be made clear by considering the batch gradient of generative language models.
A dataset D consists of a sequence of tokens of length | D | .
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 model M can be thought of as a classiﬁer with cross entropy loss L and additional dependence on s i .
The batch gradient objective can therefore be formulated by considering the full sequence of predictions from i = 0 to i
=
| D | − 1 , generating for each step i the output
ˆ D
i +1 , s i +1 .
Each token is one-hot encoded, so the empirical risk is given by
[Equation]
In this setting, a dataset of length | D | is split into b contiguous sequences of length | D | /b and stored in a matrix Z b ∈ R b × ( | D | /b ) .
Taking a pass over Z b requires maintaining a state s i for each entry in a batch, which is reset before every pass over Z b .
To deal with maintaining state for batches at different time scales, we deﬁne a different matrix Z B ∈ R b × ( | D | /B ) which maintains a different set of states S i for each entry of batch size B .