F LIPOUT : E FFICIENT P SEUDO -I NDEPENDENT W EIGHT P ERTURBATIONS ON M INI -B ATCHES
Anonymous authors
Paper under double-blind review
A BSTRACT
Stochastic neural net weights are used in a variety of contexts, including regularization, Bayesian neural nets, exploration in reinforcement learning, and evolution strategies.
Unfortunately, due to the large number of weights, all the examples in a mini-batch typically share the same weight perturbation, thereby limiting the variance reduction effect of large mini-batches.
We introduce ﬂipout, an efﬁcient method for decorrelating the gradients within a mini-batch by implicitly sampling pseudo-independent weight perturbations for each example.
Empirically, ﬂipout achieves the ideal linear variance reduction for fully connected networks, convolutional networks, and RNNs.
We ﬁnd signiﬁcant speedups in training neural networks with multiplicative Gaussian perturbations.
Flipout also enables us to vectorize evolution strategies: a single GPU with ﬂipout can handle the same throughput at least as 40 CPU cores under existing methods, equivalent to a factorof- 4 cost reduction on Amazon Web Services.
I NTRODUCTION
Stochasticity is a key component of many modern neural net architectures and training algorithms.
The most widely used regularization methods are based on randomly perturbing a network’s computations (Srivastava et al., 2014; Ioffe & Szegedy, 2015).
Bayesian neural nets can be trained with variational inference by perturbing the weights (Graves, 2011; Blundell et al., 2015).
Weight noise was found to aid exploration in reinforcement learning (Plappert et al., 2017; Fortunato et al., 2017).
Evolution strategies minimizes a black-box objective by evaluating many weight perturbations in parallel, with impressive performance on robotic control tasks (Salimans et al., 2017).
Some methods perturb a network’s activations (Srivastava et al., 2014; Ioffe & Szegedy, 2015), while others perturb its weights (Graves, 2011; Blundell et al., 2015; Plappert et al., 2017;
Fortunato et al., 2017; Salimans et al., 2017).
Stochastic weights are appealing in the context of regularization or exploration because they can be viewed as a form of posterior uncertainty about the parameters.
However, compared with stochastic activations, they have a serious drawback: because a network typically has many more weights than units, it is very expensive to compute and store separate weight perturbations for every example in a mini-batch.
Therefore, stochastic weight methods are typically done with a single sample per mini-batch.
In contrast, activations are easy to sample independently for different training examples within a mini-batch.
This allows the training algorithm to see orders of magnitude more perturbations in a given amount of time, and the variance of the stochastic gradients decays as 1 /N , where N is the mini-batch size.
We believe this is the main reason stochastic activations are far more prevalent than stochastic weights for neural net regularization.
In other settings such as Bayesian neural nets and evolution strategies, one is forced to use weight perturbations and live with the resulting inefﬁciency.
In order to achieve the ideal 1 /N variance reduction, the gradients within a mini-batch need not be independent but merely uncorrelated.
In this paper, we present ﬂipout, an efﬁcient method for decorrelating the gradients between different examples while still giving unbiased gradient estimates.
It applies to any perturbation distribution which factorizes by weight and is symmetric around 0— including DropConnect, multiplicative Gaussian perturbations, evolution strategies, and variational
Bayesian neural nets—and to many architectures, including fully connected nets, convolutional nets, and RNNs.
In Section 3, we introduce ﬂipout, show that it gives unbiased stochastic gradients, and discuss its efﬁcient vectorized implementation which incurs only a factor-of-2 computational overhead compared with shared perturbations.
We then analyze the asymptotics of gradient variance with and without ﬂipout, demonstrating strictly reduced variance.
In Section 4, we measure the variance reduction effects on a variety of architectures.
Empirically, ﬂipout gives the ideal 1 /N variance reduction in all architectures we have investigated, just as if the perturbations were done fully independently for each training example.
We demonstrate that Flipout achieves competitive performance in regularizing LSTMs, compared to several methods that use dropout.
We demonstrate signiﬁcant speedups in training time in a large batch regime.
Finally, we apply ﬂipout to vectorize evolution strategies (Salimans et al., 2017), allowing a single GPU to handle the same throughput under 40 CPU cores using existing approaches; this corresponds to a factor-of-4 cost reduction on Amazon Web Services.
B ACKGROUND 2. W EIGHT PERTURBATIONS
We use the term “weight perturbation” to refer to a class of methods which sample the weights of a neural network stochastically at training time.
More precisely, let f ( x, W ) denote the output of a network with weights W on input x .
The weights are sampled from some distribution q θ parameterized by θ .
We aim to minimize the expected loss E ( x,y ) ∼D ,W ∼ q θ [ L ( f ( x, W ) , y )] , where
L is a loss function, and D denotes the data distribution.
The distribution q θ can often be described in terms of perturbations: W = W + ∆ W , where W are the mean weights (typically represented explicitly as part of θ ) and ∆ W is a stochastic perturbation.
We now give some speciﬁc examples of weight perturbations.
Gaussian perturbations.
If the entries ∆ W ij are sampled independently from Gaussian distributions with variance σ 2 ij , this gives the distribution q ( W ij ) =
N ( W ij ; W ij , σ 2 ij ) .
Using the reparameterization trick (Kingma et al., 2015), this can be rewritten as W ij =
W ij + σ ij (cid:15) ij , where (cid:15) ij ∼ N (0 , 1) ; this representation allows the gradients to be computed using backprop.
A variant on this is multiplicative Gaussian perturbation, where the perturbations are scaled according to the weights:
W ij ∼ N
( W ij , σ 2 ij W 2 ij ) , or W ij =
W ij (1 + σ ij (cid:15) ij ) , where again (cid:15) ij ∼ N (0 , 1) .
Multiplicative perturbations can be more effective than additive ones because the information content of the weights is the same, regardless of their scale.
DropConnect.
DropConnect (Wan et al., 2013) is a regularization method inspired by dropout (Srivastava et al., 2014) which randomly zeros out a random subset of the weights.
In the case of a 50% drop rate, this can be thought of as weight perturbation where W = W/
and each entry ∆ W ij is sampled uniformly from ± W ij .
Variational Bayesian neural nets.
Rather than ﬁtting a point estimate of a neural net’s weights, one can adopt a Bayesian approach of putting a prior distribution p ( W ) over the weights and approximating the posterior distribution p ( W |D )
∝ p ( W ) p ( D| W ) , where D denotes the observed data.
Graves (2011) observed that one could ﬁt an approximation q θ ( W )
≈ p ( W |D )
using variational inference; in particular, one could maximize the evidence lower bound (ELBO) with respect to θ :
[Equation]
The negation of the second term can be viewed as the description length of the data, and the negation of the ﬁrst term can be viewed as the description length of the weights (Hinton & Van Camp, 1993).
Graves (2011) observed that if q is chosen to be a factorial Gaussian, it can be thought of as Gaussian weight perturbation, except where the variance is adapted to maximize F .
(Note that the ﬁrst term is proportional to the negative expected loss if we deﬁne L = − log p ( y | x ) .)
Blundell et al.
(2015) later combined this insight with the reparameterization trick (Kingma & Welling, 2014) to derive unbiased stochastic estimates of the gradient of F .
Evolution strategies.
Evolution strategies (ES) (Eigen, 1973) is a family of black box optimization algorithms which involves a heavy use of weight perturbation to search for model parameters.
It doesn’t scale well to large dimensional space so it has limited use in neural networks until recently proposed as an alternative reinforcement learning algorithm
(Schmidhuber et al., 2007; Salimans et al., 2017).
ES generates a collection of weight perturbations as candidates and evaluates them by some objective functions.
The gradient of the parameters can be estimated given the evaluation result.
ES is a highly parallelizable algorithm, since generating and evaluating perturbations can be done independently.
Suppose N is the number of workers, σ is the perturbation standard deviation and W is the model parameter.
The algorithm is trying to maximize E ∆ W (cid:2) F (cid:0) W + σ ∆ W (cid:1)(cid:3) , where F ( · ) is the objective function.
The gradient of the objective function and the update rule can be given as:
[Equation]
where α is the learning rate, F ( · ) is the objective function, and ∆ W i is the Gaussian noise generated at worker i .
2. R EPARAMETERIZATION GRADIENT
As mentioned in Section 2.1, the reparameterization trick allows the gradients to be computed using backprop (Kingma & Welling, 2014).
It is important to contrast this with regularizers such as dropout (Srivastava et al., 2014), which perturb activations rather than weights.
Because activations can be perturbed independently for every training example within a mini-batch, the variance of the stochastic gradients decreases linearly with the size of the mini-batch.
Unfortunately, it would be extremely expensive to explicitly sample an independent ∆ W for every training example — this would require storing tens or hundreds of copies of the parameters in memory.
Therefore, weight perturbation gradients are typically computed using only a single perturbation per mini-batch.
The variance reduction effect of mini-batches is one of the key advantages of activation perturbations over weight perturbations.
It is possible to derive efﬁcient independent weight perturbation methods in particular cases; for instance, independent DropConnect masks could be stored since they require only one bit per weight Wan et al.
Kingma et al.
(2015) showed that in the case of fully connected networks with no weight sharing, unbiased stochastic gradients could be computed without explicit weight perturbations using the local reparameterization trick (LRT).
Speciﬁcally, suppose X is the input mini-batch, W is the weight matrix and B = XW is the activation.
The
LRT samples activation B rather than the weights W .
In the case of a Gaussian posterior, the local reparameterization trick is given as:
[Equation]
Notice that the local reparameterization trick gives an unbiased gradient estimate when the activations are independent.
However, when weight sharing is involved, such as in convolutional neural networks and recurrent neural networks, the LRT does not give an unbiased gradient estimate.
2. O THER RELATED WORK
Control variates are another general class of strategies for variance reduction, both for black-box optimization (Williams, 1992; Ranganath et al., 2014; Mnih & Gregor, 2014) and for gradient-based optimization (Roeder et al., 2016; Miller et al., 2017;
Louizos et al., 2017).
Control variates are complementary to ﬂipout, so one could potentially combine these techniques to achieve a larger variance reduction.
We also note that the fastfood transform (Le et al., 2013) is based on similar mathematical techniques.
However, whereas fastfood is used to approximately multiply by a large
Gaussian matrix, ﬂipout preserves the random matrix’s distribution and instead decorrelates the gradients between different samples.
M ETHODS
As described above, weight perturbation algorithms suffer from a high variance of the gradient estimates because all training examples in a mini-batch share the same perturbation.
More precisely, sharing the perturbation induces correlations between the gradients, implying that the variance can’t be eliminated by averaging.
(This argument is detailed in Section 3.2.)
In this section, we propose ﬂipout, an efﬁcient way to perturb the weights quasi-independently within a mini-batch.
We show that this almost fully decorrelates the gradients within a mini-batch, allowing us to achieve the ideal variance reduction as if the gradients were independent.
3. F LIPOUT
We make two assumptions about the weight distribution q θ : (1) the perturbations of different weights are independent; and (2) the perturbation distribution is symmetric around zero.
These are nontrivial constraints, but they encompass important applications: independent Gaussian perturbations (e.g. as used in variational BNNs and evolution strategies) and DropConnect with drop probability 0.5.
We observe that, under these assumptions, the perturbation distribution is invariant to elementwise multiplication by a random sign matrix (i.e. matrix whose entries are ± 1 ).
In the following, we denote elementwise multiplication by ◦ .
Proposition 1.
Let q θ be a perturbation distribution that satisﬁes the above assumptions, and let (cid:100) ∆ W ∼ q θ .
Let E be a random sign matrix that is independent of (cid:100) ∆ W .
Then ∆ W = (cid:100) ∆ W ◦
E is identically distributed to (cid:100) ∆ W .
Furthermore, the loss derivatives computed using ∆ W are identically distributed to those computed using (cid:100) ∆ W .
Flipout exploits this fact by using a base perturbation (cid:100) ∆ W shared by all examples in the mini-batch, and multiplies it by a different rank-one sign matrix for each example:
[Equation]
[Equation]
where the superscript denotes the index within the mini-batch, and e ( i ) 1 and e ( i ) 2 are random vectors whose entries are sampled uniformly from ± 1 .
According to Theorem 1, the marginal distribution over gradients computed for individual training examples will be identical to the distributions computed using shared weight perturbations.
In particular, ﬂipout yields an unbiased gradient estimator for the original loss function.
However, by decorrelating the gradients between different training examples, we can achieve much lower variance updates when averaging over a mini-batch.
Vectorization.
The advantage of ﬂipout over explicit perturbations is that computations on a minibatch can be written in terms of matrix multiplications.
This enables efﬁcient implementations on
GPUs and modern accelerators such as the Tensor Processing Unit (TPU) (Jouppi et al., 2017).
Let x denote the activations in one layer of a neural net.
The next layer’s activations are given by:
[Equation]
where φ denotes the activation function.
To vectorize these computations, we deﬁne matrices Eand E 2 whose rows correspond to the random sign vectors e ( i ) 1 and e ( i ) 2 for all examples in the batch.
The above equation is vectorized as:
[Equation]
This deﬁnes the forward pass.
Because E 1 and E 2 are sampled independently of W and (cid:100) ∆ W , we can also backpropagate through Eqn.
4 to obtain derivatives with respect to W , (cid:100) ∆ W , and X .
Computational cost.
In general, the most expensive operation in the forward pass is matrix multiplication.
Flipout’s forward pass requires two matrix multiplications instead of one, and therefore should be roughly twice as expensive as a forward pass with a single shared perturbation when done in sequence.
However, note the two matrix multiplications are independent and can be parallelized.
As for the backward pass, a general rule of thumb for neural nets is that it requires roughly twice as many FLOPs as the forward pass.
This suggests that each update using ﬂipout ought to be about twice as expensive as an update with a single shared perturbation (again, if not parallelized); this is consistent with our experience.
(Note this is the same overhead for the local reparameterization trick (Kingma et al., 2015).)
Evolution strategies.
Evolution Strategies (ES) is a highly parallelizable algorithm.
However, most ES systems are engineered to run on multi-core CPU machines and are not able to take full advantage of GPU parallelism.
Flipout enables the ES to run more efﬁciently on a GPU because it allows each worker to evaluate a batch of quasi-independent perturbations rather than only one.
To applying ﬂipout to ES on reinforcement learning tasks, we can simply replicate the starting state by the number of ﬂipout perturbations M , at each worker.
Instead of Eqn.
1, the update rule on using
N workers becomes:
[Equation]
[Equation] where e ( i ) 1 ,j and e ( i ) 2 ,j are the i th random ﬂipout sign vectors sampled at the worker j , and F ij is the reward evaluated with i
th perturbation at the worker j .
The same ﬂipout algorithm can also be applied to supervised learning tasks.
As a result, it can accelerate the training process of the current
ES algorithm in terms of the wall-clock time with the same amount of resources.
3. V ARIANCE ANALYSIS
In this section, we analyze the variance of stochastic gradients under the ﬂipout perturbation.
We show that ﬂipout is guaranteed to reduce the variance of the gradient estimates compared to the na ¨ ıve shared perturbation.
Without loss of generality, let G ( x, ∆ W ) denote the stochastic gradient under the perturbation ∆ W for a single training example x w.r.t.
the weight W , i.e. G ( x, ∆ W ) = ∇ W L ( y, f ( x, W, ∆ W )) .
We deﬁne the mini-batch gradient G B (∆ W ) as the averaged gradient under perturbation over a minibatch B of size N , G B (∆ W ) = 1 N (cid:80) x ∈B G ( x, ∆ W ) , where the subscript is used to denote the random mini-batch sampled from the whole training set.
Assume that the training examples in the mini-batch are sampled independently from the data distribution.
Using the Law of Total Variance, we can expand the variance of the mini-batch gradient
Var( G B (∆ W )) into the gradient variance between mini-batches (data term) and the estimation variance (estimation term) under perturbation noise:
[Equation]
Notice that the data term in the above equation is not related to the perturbation noise ∆ W and decays with the mini-batch size N .
The estimation term can be further divided and manipulated so that we can reach the following conclusion.
Theorem 1 (Variance Decomposition Theorem) .
Let ∆ W be the weight perturbations sampled independently given the base perturbations on from Eqn. 3.
Deﬁne α , β and γ to be
[Equation]
Under the assumption of Proposition
[Equation]
the variance of the gradients under the shared and the ﬂipout perturbation can be written in terms of α , β and γ as following:
[Equation]
Proof.
Details of the proof are provided in Appendix A.
We can interpret α , β and γ in the above theorem.
The γ term is the variance from sampling different mini-batch and perturbation, α + β is the expected covariance of the gradients among the training examples in the mini-batch.
When a single perturbation is used for whole mini-batch, it introduces additional covariance among the gradients contributing to the total variance of the minibatch gradient.
The quasi-independent perturbation from ﬂipout is guaranteed to reduce the gradient covariance (the α term) linearly in the number of training examples.
Therefore, ﬂipout achieves a lower-variance gradient estimate than methods based on a shared perturbation.
We further examine the contribution of the gradient covariance to the variance of the mini-batch gradient.
In the shared weight perturbation case, if the mini-batch size N is large enough, Eqn.is dominated by the ﬁrst two covariance terms, i.e. α + β , which cannot be reduced by increasing the size of the mini-batch N .
This suggest the gradient variance should level off after a certain mini-batch size for the shared perturbation.
In contrast, the gradient variance under ﬂipout should be considerably lower given that the α term decays with the mini-batch size.
Moreover, we argue that if the β term is insigniﬁcant in practice, then the total ﬂipout variance should decrease linearly in terms of N .
Empirically, we found that the β term is indeed close to 0 for a various neural network architectures as discussed in Section 4.1.
Thus, in practice, ﬂipout obtains the ideal variance reduction in the weight perturbation.
EXPERIMENTS
We ﬁrst empirically verify the variance reduction effect using ﬂipout predicted by Theorem 1 in Section 4.1.
We study the variance of the gradients under different perturbations for a wide variety of neural network architectures and batch sizes.
In Section 4.2, we show that ﬂipout is effective at regularizing LSTM networks.
Furthermore, we demonstrate that the ﬂipout perturbation outperforms the shared perturbations method, which we call “non-ﬂipout”, when training with large mini-batches in
Section 4.3.
Lastly, we present the experimental results of Evolution Strategies with ﬂipout in both supervised learning and reinforcement learning tasks in Section 4.4.
4. VARIANCE REDUCTION
Since the main effect of ﬂipout is intended to be variance reduction of the gradients, we ﬁrst estimated the gradient variances of several architectures with mini-batch sizes ranging from 1 to 8196.
We experimented with three perturbation methods: a single shared perturbation per mini-batch, the local reparameterization trick (LRT) of Kingma et al.
(2015), and ﬂipout.
We computed Monte
Carlo estimates of the gradient variance, including both the data and estimation terms in Eqn. 6.
(The estimation variance dominated data variance in all of our results except for LRT.)
Conﬁdence intervals are based on 50 independent runs of the estimator.
We considered four architectures: a fully connected MNIST classiﬁer, a LeNet-like MNIST conv net (LeCun et al., 1998), a VGG-like CIFAR conv net (Simonyan & Zisserman, 2014), and an LSTM language model trained on Penn Treebank (Marcus et al., 1993).
These models are henceforth referred to as FC , Conv1 , Conv2 , and LSTM , respectively.
For each of these networks, we froze a partially trained network to use for all variance estimations, and we used multiplicative Gaussian perturbations with σ 2 = 1 .
Details of architectures and training methods are given in Appendix B.
Estimation variance plots for all of these models are given in Figure 1.
We note that the analysis of Section 3.2 makes strong predictions about the shapes of these curves.
In particular, by Theorem 1, the variance curves for ﬂipout and shared each have the form a + b/N , where N is the mini-batch size.
On a log-log plot, this functional form appears as a linear regime with slope -1, a constant regime, and a smooth phase transition in between.
Also, because the distribution of individual gradients is identical with and without ﬂipout, the curves must agree for N = 1 .
Our plots are consistent with both of these predictions.
We observe that for shared perturbations, the phase transition consistently occurs for batch sizes somewhere between 100 and 1000.
In contrast, ﬂipout gives the ideal linear variance reduction throughout the range of mini-batch sizes we investigated, i.e., its behavior is indistinguishable from fully independent perturbations.
In the following experiments, we measure the variance of the stochastic gradients in two settings: (1) non-ﬂipout , with per-batch
[Figure]
Gaussian weight perturbations (shared for all examples in a mini-batch), and (2) with ﬂipout applied to yield quasi-independent perturbations for each example
As analyzed by Kingma et al.
(2015) the LRT gradients are fully independent within a mini-batch, therefore are guaranteed to achieve the ideal 1 /N variance reduction.
Furthermore, they reduce the variance below that of explicit weight perturbations, so we would expect them to achieve smaller variance than ﬂipout.
This is borne out by our plots.
However, ﬂipout is applicable to a wider variety of architectures, including convolutional nets and RNNs.
4. R EGULARIZATION FOR L ANGUAGE M ODELING
We evaluate the regularization performance of ﬂipout on both the character-level and word-level language modeling tasks with the Penn Treebank corpus (PTB) (Marcus et al., 1993).
We compare ﬂipout to several other methods for regularizing RNNs: na ¨ ıve dropout (Zaremba et al., 2014), variational dropout (Gal & Ghahramani, 2016), recurrent dropout (Semeniuta et al., 2016), zoneout (Krueger et al., 2016), and DropConnect (Merity et al., 2017).
Zaremba et al.
(2014) apply dropout only to the feed-forward connections (to the input, output, and connections between layers).
The other methods regularize the recurrent connections as well: Semeniuta et al.
(2016) apply dropout to the cell update vector, with masks sampled either per step or per sequence; Gal & Ghahramani (2016) propose to apply dropout to the forward and recurrent connections, with all dropout masks sampled per sequence. Merity et al.
(2017) use DropConnect to regularize the hidden-to-hidden weight matrices, with a single DropConnect mask shared between examples in a minibatch. We denote their model WD (for weight-dropped LSTM).
Character-Level.
For our character-level experiments, we use a single-layer LSTM with 1000 hidden units, sequences of 100 characters in batches of size 32, and train for 55 epochs using Adam with a learning rate of 0.002.
Here, we apply ﬂipout to all input-to-hidden and hidden-to-hidden weight matrices, as well as to the hidden-to-output layer, providing feed-forward and recurrent regularization.
Model hyperparameters are given in Appendix D.
The results, measured in bits-per-character (BPC) for the validation and test sequences of PTB are shown in Table 1.
In the table, non-ﬂipout and ﬂipout (with Gaussian noise sampling) are denoted by Mult.
Gauss and Mult.
Gauss + Flipout, respectively.
We see that ﬂipout achieves results competitive with the other methods.
Word-Level.
For word-level language modeling, we use a 2-layer LSTM with 650 hidden units per layer and a 650-dimensional word embedding layer.
We train on sequences of length 35 in batches of size 40, for 100 epochs.
We use SGD with initial learning rate 30, and decay the learning rate by a factor of 4 based on the nonmonotonic criterion introduced by Merity et al.
We use ﬂipout to implement DropConnect, as described in Section 2.1, and call this WD+Flipout.
We apply
WD+Flipout to the hidden-to-hidden weight matrices for recurrent regularization, and use the same hyperparameters as (Merity et al., 2017).
We use embedding dropout (setting rows of the embedding matrix to 0) with probability 0.1 for all models except Gal, where we use probability 0.2 as speciﬁed in their paper.
More hyperparameter details are given in Appendix D. We show in Table 2 that WD+Flipout outperforms the other methods with respect to both validation and test perplexity.
In
Appendix E.4, we show that WD+Flipout yields signiﬁcant variance reduction for large minibatches, and that when training with batch size 8192, it converges faster than WD.
[Table] [Table]
4. LARGE BATCH TRAINING WITH FLIPOUT
Theorem 1 and Fig.1 suggest that the variance reduction is more effective when the model is trained with a large mini-batch size.
In this section, we train a Bayesian neural network with mini-batch size 8192 and show that ﬂipout can speed up the training process in terms of the number of iterations, due to its variance reduction effect.
We ran the large mini-batch size training experiment with the Bayes by Backprop algorithm (Blundell et al., 2015), under a 3-layer fully connected network and a LeNet-like convolutional neural network (FC and Conv1 respectively).
The purpose of this section is to show the speed up of the training process due to the variance reduction.
Thus, we focus on the training loss in this section and present other training details in Appendix E.
Fig.2a shows the training loss of the large batch experiment.
In the fully connected network (FC), we compare the local reparameterization trick, ﬂipout, and non-ﬂipout.
Only ﬂipout and non-ﬂipout are compared in the convolutional network (Conv1) experiment since the local reparameterization trick does not provide an unbiased estimator.
We found that ﬂipout trained in about 3 times fewer iterations than non-ﬂipout in both fully connected and convolutional networks, while achieving comparable performance with the local reparameterization trick.
In current experiments, ﬂipout is roughly two times slower than non-ﬂipout in FLOPs, which leads to a 1.5 times faster convergence overall.
This slowdown is due to an additional matrix multiplication (second half of Eqn. 4).
However, this computation is independent from the ﬁrst half and can be parallelized.
In preliminary experiments, we have found this is possible using modern accelerators via the second-generation Tensor Processing Unit (TPU) (Jouppi et al., 2017).
A key challenge in utilizing the TPU is to prevent a bottleneck in the input pipeline.
In general, large batch sizes, as well as deeper and wider nets, improve utilization where an initial enqueing phase is most expensive.
In this regime, weight variance is the biggest concern and ﬂipout provides a promising solution.
4. EVOLUTION STRATEGIES
Evolution Strategies (ES) is a highly parallelizable algorithm that typically runs on multiple CPU cores.
The limitation that only one perturbation can be evaluated locally restricts its ability to run efﬁciently on GPU.
In Section 3.1, we showed that ES with ﬂipout allows each local worker to evaluate a batch of perturbations, which enables the algorithm to run on a single GPU more efﬁciently.
In this section, we show that ﬂipout accelerates the ES algorithm and narrows the computational gap between ES and back-propagation.
The gradient approximated by Eqn. 1 has high variance.
Therefore, a large number of samples are generally needed before applying the update Eqn. 1.
We empirically found that 5 , 000 samples are needed to achieve a stable performance in the supervised learning tasks.
The na¨ıve way to run ES on GPU would be iterating the forward pass 5 , 000 times to perform the update.
FlipES is designed to reduce the number of iterations, since it can evaluate a batch of perturbations in one single forward pass.
Throughout the experiments, we set the ﬂip size to be 40 (i.e. M = 40 in
Eqn. 5).
Na¨ıve ES uses fully independent noise during each forward pass while ﬂipout uses pseudoindependent noise.
We compared these two methods with a fully connected network (FC) on the
MNIST dataset.
We only sampled 1 , 600 times per update during training, because the purpose of this experiment is to show no performance loss with pseudo-independent noise, as shown in Fig.2b.
Next, we compared FlipES and cpuES (40 cores) in terms of the per update time with respect to the model size.
The result (Appendix E.3) shows cpuES has an increasing computational cost as the model size increases, while FlipES scales better because it runs on GPU.
Finally, we compare
FlipES and the back propagation algorithm in both a fully connected network (FC) and a LeNet-like neural network (Conv1).
Fig.2c and Fig.2d show that FlipES achieves a data efﬁciency comparable with the back propagation algorithm.
The na¨ıve ES has a much higher computational cost than
[Figure]
the back propagation algorithm, due to the large number of forward passes.
FlipES can narrow the computational gap between them.
Additionally, FlipES is still a highly parallelizable algorithm since there is no dependency in each iteration.
Although evolution strategies is more expensive than back propagation, it applies to models which are not entirely differentiable, such as a model with discrete loss (accuracy or BLEU score) or with stochastic units.
Moreover, evolution strategies naturally induces the variance of gradient, which makes it easier to perform second-order optimization.
We leave these as future work.
C ONCLUSIONS
We have introduced ﬂipout, an efﬁcient method for decorrelating the weight gradients between different examples in a mini-batch.
We showed that ﬂipout is guaranteed to reduce the variance compared with shared perturbations.
Empirically, we demonstrated signiﬁcant variance reduction in the large-batch setting for a variety of network architectures, as well as signiﬁcant speedups in training time.
Flipout also makes it practical to apply GPUs to evolution strategies, resulting in substantially increased throughput for a given computational cost.
We believe ﬂipout will make weight perturbations practical in the large batch setting favored by modern accelerators such as Tensor Processing
Units (Jouppi et al., 2017).
A PROOF OF T HEOREM 1(V ARIANCE DECOMPOSITION THEOREM )
In this section, we provide the proof to Theorem 1 (Variance Decomposition Theorem).
Proof.
We use the notations in 3.2.
Speciﬁcally, let G ( x, ∆ W ) denote the stochastic gradient under the perturbation ∆ W for a single training example x w.r.t.
the weight W , i.e. G ( x, ∆ W ) = ∇ W L ( y, f ( x, W, ∆ W )) .
We deﬁne the mini-batch gradient G B (∆ W ) as the averaged gradient under perturbation over a mini-batch B of size N , G B (∆ W ) = 1 N (cid:80) x ∈B G ( x, ∆ W ) , where the subscript is used to denote the random mini-batch sampled from the whole training set.
Let x, x (cid:48) be two arbitrarily chosen training examples from the mini-batch, ∆ W, ∆ W (cid:48) be the two weight perturbations they received, and 1 Flip be the indicator function for using ﬂipout or not.
Following Eqn. 6, the variance can be split into a data term and an estimation term.
We simplify the estimation term by replacing the averaged mini-batch gradient with the single training example gradient :
[Equation]
We can further simplify the last term in Eqn.
11 for ﬂipout and shared perturbation cases.
In the shared perturbation case
[Equation]
it leads to quadratic N 2 number of the variance terms:
Shared:
[Equation]
[Equation] With ﬂipout, ∆ W are conditionally independent given on , which leads to linear number of variance terms.
Flipout:
[Equation]
Substituting the above expansions back into Eqn. 6, the total variance of the gradient estimates can be expressed as: [Equation]
[Equation]
Therefore, the conditionally independence perturbation help reduce the ﬁrst term in the gradient variance in Eqn.
B N ETWORK CONFIGURATION IN THE EXPERIMENTS
In this section, we provide the detailed network conﬁgurations to reproduce the results in the experiments Section 4.
[Table]
The (Shallow) convolutional neural networks (Conv1) is a LeNet-like network (LeCun et al., 1998).
In speciﬁc, the ﬁrst two layers are convolutional layers with 32 and 64 ﬁlters and relu non-linearity.
Both ﬁlters are of size [5 , 5] .
A 2 × 2 max pooling layer is followed after each convolutional layer.
Dimension reduction only takes place in the pooling layer, which means the stride of the pooling is two and padding is used in the convolutional layers to keep the dimension.
Two fully connected layer with 1024 and 10 are attached to produce the classiﬁcation result.
The (Deep) convolutional neural networks (Conv2) is a VGG16-like network (Simonyan & Zisserman, 2014).
We only modiﬁed the last fully connected layer in order to run experiments on CIFAR10 dataset.
We didn’t use the batch normalization in the variance reduction experiment since it introduces extra stochasticity.
Also, multiple GPUs are needed when running the experiment with large mini-batch size (such as 4096 and 8192 ).
In such case, the algorithm naturally induces multiple (depends on the number of GPUs) independent noises due to the distributed computation.
In
Section 4.1, we ﬁxed the noise among all GPUs to produce the plot.
We show in the Appendix E that ﬂipout still results in smaller variance even in the case of distributed computation.
[Table]
network is a 3-layer fully connected network.
The number of hidden units are 512 ,and 10 without extra mentioning.
C V ARIANCE REDUCTION EXPERIMENT DETAILS
Given a network architecture and a speciﬁc algorithm, we compute the empirical stochastic gradient update variance as follows.
Starting from a moderately pre-trained model, such as a network with 85% training accuracy on MNIST.
Without updating the parameters, we obtain the gradient of each weight in each layer by performing the feed-forward pass with sampling (cid:15) , E 1 and E 2 and back propagation.
The gradient variance of each weight is computed by repeating the above procedure 200 times in the experiments.
Denote the variance estimation of layer l , the weight j as (cid:93)
Var lj , the above can be written in the following form,
[Equation]
where g ilj is the gradient received by the weight j in the layer l .
Then we can use the average variance over weights 1 | J | (cid:80) j (cid:93)
Var lj as the variance estimation of the gradients in the layer l .
In order to compute the conﬁdence interval of the gradient variance estimation, we repeat the above procedure 50 times.
Suppose we obtain the average variance estimation sequence (cid:102) V 1 , . . .
, (cid:102)
V 50 .
Assuming the sequence follows an i.i.d.
normal distribution with unknown mean and unknown variance, it follows that V − µ S/ √ 50 ∼ t 49 , where V is the mean of the sequence and S is the sample standard deviation.
From this, we can compute the 90% conﬁdence interval of the mean of the variance estimation and it leads to the error plot in Fig.1.
To measure the variance reduction effect of ﬂipout in LSTMs, we train a ﬂipout LSTM with the two-layer architecture described in Section 4.2, on the word-level Penn Treebank dataset (Marcus et al., 1993).
We train on sequences of length 35 for 3 epochs using plain SGD with learning rate 20.
We then measure the variance for each of the hidden-to-hidden weight matrices in the ﬁrst layer ( W f , W i , W c , and W o ).
We split large mini-batches (size 128 and higher) into multiple sub-batches of size 64; we sample the noise (cid:15) once for the set of sub-batches ( (cid:15) is shared among sub-batches), and we sample new E 1 and E 2 matrices for each sub-batch.
We see that as the mini-batch size increases, the variance computed with the ﬂipout model continues to decrease, while the non-ﬂipout model converges to a constant variance.
D LSTM E XPERIMENT D ETAILS
Long Short-Term Memory networks (LSTMs) are deﬁned by the following equations:
[Equation]
where i t , f t , and o t are the input, forget, and output gates, respectively, g t is the candidate update, and ◦ denotes elementwise multiplication.
Na ¨ ıve application of dropout on the hidden state of an LSTM is not effective, because it leads to signiﬁcant memory loss over long sequences.
Several approaches have been proposed to regularize the recurrent connections, based on applying dropout to speciﬁc terms in the LSTM equations.
Semeniuta et al.
(2016) propose to drop the cell update vector, with a dropout mask d t sampled either per-step or per-sequence: c t =
f t ◦ c t − 1 + i t ◦ ( d t ◦ g t ) .
Gal & Ghahramani (2016) apply dropout to the input and hidden state at each time step, x t ◦ d x and h t − 1 ◦ d h , with dropout masks d x and d h sampled once per sequence (and repeated in each time step).
Krueger et al.
(2016) proposes to zone out units rather than dropping them; in the LSTM, the hidden state and cell values are either stochastically updated or maintain their previous value: c t =
d ct ◦ c t − 1 +(1 − d ct ) ◦
( f t ◦ c t − 1 + i t ◦ g t ) and h t =
d ht ◦ h t − 1 +(1 − d ht ) ◦ ( o t ◦ tanh( f t ◦ c t −
+ i t ◦ g t )) , with zoneout masks
d ht and d ct sampled per step.
D. H YPERPARAMETER D ETAILS
For the word-level models, we use the following hyperparameters:
• For (Gal & Ghahramani, 2016), we use the parameters given in their paper: 0.35 dropout probability on inputs and outputs, 0.2 hidden state dropout, and 0.2 embedding dropout.
• For (Semeniuta et al., 2016), we use 0.1 embedding dropout, 0.5 dropout on inputs and outputs, and 0.3 dropout on cell updates, with per-step mask sampling.
• For (Krueger et al., 2016), we use 0.1 embedding dropout, 0.5 dropout on inputs and outputs, and cell and hidden state zoneout probabilities of 0.25 and 0.025, respectively.
• For (Merity et al., 2017), we use the parameters given in their paper: 0.1 embedding dropout, 0.4 dropout probability on inputs and outputs, and 0.3 dropout probability on the output between layers (the same masks are used for each step of a sequence). We use 0.5 probability for DropConnect applied to the hidden-to-hidden weight matrices. • For WD+Flipout, we use the same parameters as (Merity et al., 2017), given above, but we regularize the hidden-to-hidden weight matrices with the variant of ﬂipout described in Section 2.1, which implements DropConnect with probability 0.5.
For the character-level models, we use the following hyperparameters.
We do not use input or output dropout for any of these models
• For (Semeniuta et al., 2016), we use 0.25 dropout probability on the cell state, and per-step mask sampling.
• For Zoneout (Krueger et al., 2016), we use 0.5 and 0.05 for the cell and hidden state zoneout probabilities, respectively.
• For the variational LSTM (Gal & Ghahramani, 2016), we use 0.25 recurrent dropout.
• For the ﬂipout and non-ﬂipout LSTMs, we use Gaussian noise with σ = 0 . 3 .
E SUPPLEMENTARY PLOTS E. V ARIANCE REDUCTION
As discussed in Appendix B
[Figure]
training on multiple GPUs naturally induces independent noise for every sub-batch.
Fig.3 below shows that ﬂipout still has lower variance in such cases
[Figure] [Figure]
E. L ARGE BATCH TRAINING WITH FLIPOUT
Fig. 4 is the plot of the training and test accuracy in the large batch training experiments in Section 4.3.
The fully connected network and convolutional network we used below have exactly the same architecture as explained in Appendix B.
We used an Adam optimizer with a 0 . 003 learning rate in both networks.
We downscaled the KL term by a factor of 10 to achieve a higher accuracy.
[Equation]
[Figure] Fig. 5 shows that cpuES has an increasing computational cost as the model size increases, while
FlipES scales better because it runs on GPU.
E. L ARGE BATCH LSTM TRAINING
The variance reduction offered by ﬂipout allows us to use DropConnect (Wan et al., 2013) efﬁciently in a large-minibatch setting.
Here, we use ﬂipout to implement DropConnect as described
[Figure]
# Hidden Units
FlipES cpuES 0.12s 0.51s 0.13s 1.22s 0.18s 5.10s 1.86s 38.0s # Filters
FlipES cpuES 0.
2.3s 16s 0.
5.48s 46s 1.
7.12s 67s 1.
11.77s 132s
[Figure]
[Figure]
in Section 2.1, and use it to regularize an LSTM word-level language model.
We use the LSTM architecture proposed by Merity et al.
(2017), which has 400-dimensional word embedddings and three layers with hidden dimension 1150.
Following Merity et al.
, we tie the weights of the embedding layer and the decoder layer.
Merity et al.
(2017) use DropConnect to regularize the hidden-to-hidden weight matrices, with a single mask shared for all examples in a batch.
We use ﬂipout to achieve a different DropConnect weight matrix per example.
We apply WD+Flipout to both the hidden-to-hidden (h2h) and input-to-hidden (i2h) weight matrices, and compare to the model from Merity et al.
(2017), which we call the WD -LSTM (for weight-dropped LSTM), with
DropConnect applied to both h2h and i2h.
Both models use embedding dropout 0.1, output dropout 0.4, and have DropConnect probability 0.5 for the i2h and h2h weights.
Both models are trained using Adam with learning rate 0.001.
Fig.
6 compares the variance of the gradients of the ﬁrst-layer hidden-to-hidden weights between
WD-LSTM and WD+Flipout, and shows that ﬂipout achieves signiﬁcant variance reduction for minibatch sizes larger than 256.
Fig.
7 shows the training curves of both models with batch size 8192.
We see that WD+Flipout converges faster than WD, and achieves a lower training perplexity, showcasing the optimization beneﬁts of ﬂipout in large-minibatch settings.