ZONEOUT: REGULARIZING RNNS BY RANDOMLY PRESERVING HIDDEN ACTIVATIONS
We propose zoneout, a novel method for regularizing RNNs. At each timestep, zoneout stochastically forces some hidden units to maintain their previous values. Like dropout, zoneout uses random noise to train a pseudo-ensemble, improving generalization. But by preserving instead of dropping hidden units, gradient information and state information are more readily propagated through time, as in feedforward stochastic depth networks. We perform an empirical investigation of various RNN regularizers, and find that zoneout gives significant performance improvements across tasks. We achieve competitive results with relatively simple models in characterand word-level language modelling on the Penn Treebank and Text8 datasets, and combining with recurrent batch normalization (Cooijmans et al., 2016) yields state-of-the-art results on permuted sequential MNIST.
1 INTRODUCTION
Regularizing neural nets can significantly improve performance, as indicated by the widespread use of early stopping, and success of regularization methods such as dropout and its recurrent variants (Hinton et al., 2012; Srivastava et al., 2014; Zaremba et al., 2014; Gal, 2015). In this paper, we address the issue of regularization in recurrent neural networks (RNNs) with a novel method called zoneout.
RNNs sequentially construct fixed-length representations of arbitrary-length sequences by folding new observations into their hidden state using an input-dependent transition operator. The repeated application of the same transition operator at the different time steps of the sequence, however, can make the dynamics of an RNN sensitive to minor perturbations in the hidden state; the transition dynamics can magnify components of these perturbations exponentially. Zoneout aims to improve RNNs’ robustness to perturbations in the hidden state in order to regularize transition dynamics.
Like dropout, zoneout injects noise during training. But instead of setting some units’ activations to 0 as in dropout, zoneout randomly replaces some units’ activations with their activations from the previous timestep. As in dropout, we use the expectation of the random noise at test time. This results in a simple regularization approach which can be applied through time for any RNN architecture, and can be conceptually extended to any model whose state varies over time.
Compared with dropout, zoneout is appealing because it preserves information flow forwards and backwards through the network. This helps combat the vanishing gradient problem (Hochreiter, 1991; Bengio et al., 1994), as we observe experimentally.
We also empirically evaluate zoneout on classification using the permuted sequential MNIST dataset, and on language modelling using the Penn Treebank and Text8 datasets, demonstrating competitive or state of the art performance across tasks. In particular, we show that zoneout performs competitively with other proposed regularization methods for RNNs, including recently-proposed dropout variants. Code for replicating all experiments can be found at: http://github.com/teganmaharaj/zoneout
2 RELATED WORK
2.1 RELATIONSHIP TO DROPOUT
2.2 DROPOUT IN RNNS
Initially successful applications of dropout in RNNs (Pham et al., 2013; Zaremba et al., 2014) only applied dropout to feed-forward connections (“up the stack”), and not recurrent connections (“forward through time”), but several recent works (Semeniuta et al., 2016; Moon et al., 2015; Gal, 2015) propose methods that are not limited in this way. Bayer et al. (2013) successfully apply fast dropout (Wang & Manning, 2013), a deterministic approximation of dropout, to RNNs.
Semeniuta et al. (2016) apply recurrent dropout to the updates to LSTM memory cells (or GRU states), i.e. they drop out the input/update gate in LSTM/GRU. Like zoneout, their approach prevents the loss of long-term memories built up in the states/cells of GRUs/LSTMS, but zoneout does this by preserving units’ activations exactly. This difference is most salient when zoning out the hidden states (not the memory cells) of an LSTM, for which there is no analogue in recurrent dropout. Whereas saturated output gates or output nonlinearities would cause recurrent dropout to suffer from vanishing gradients (Bengio et al., 1994), zoned-out units still propagate gradients effectively in this situation. Furthermore, while the recurrent dropout method is specific to LSTMs and GRUs, zoneout generalizes to any model that sequentially builds distributed representations of its input, including vanilla RNNs.
Also motivated by preventing memory loss, Moon et al. (2015) propose rnnDrop. This technique amounts to using the same dropout mask at every timestep, which the authors show results in improved performance on speech recognition in their experiments. Semeniuta et al. (2016) show, however, that past states’ influence vanishes exponentially as a function of dropout probability when taking the expectation at test time in rnnDrop; this is problematic for tasks involving longer-term dependencies.
Gal (2015) propose another technique which uses the same mask at each timestep. Motivated by variational inference, they drop out the rows of weight matrices in the input and output embeddings and LSTM gates, instead of dropping units’ activations. The proposed variational RNN technique achieves single-model state-of-the-art test perplexity of 73.4 on word-level language modelling of Penn Treebank.
2.3 RELATIONSHIP TO STOCHASTIC DEPTH
Zoneout can also be viewed as a per-unit version of stochastic depth (Huang et al., 2016), which randomly drops entire layers of feed-forward residual networks (ResNets (He et al., 2015)). This is
equivalent to zoning out all of the units of a layer at the same time. In a typical RNN, there is a new input at each timestep, causing issues for a naive implementation of stochastic depth. Zoning out an entire layer in an RNN means the input at the corresponding timestep is completely ignored, whereas zoning out individual units allows the RNN to take each element of its input sequence into account. We also found that using residual connections in recurrent nets led to instability, presumably due to the parameter sharing in RNNs. Concurrent with our work, Singh et al. (2016) propose zoneout for ResNets, calling it SkipForward. In their experiments, zoneout is outperformed by stochastic depth, dropout, and their proposed Swapout technique, which randomly drops either or both of the identity or residual connections. Unlike Singh et al. (2016), we apply zoneout to RNNs, and find it outperforms stochastic depth and recurrent dropout.
2.4 SELECTIVELY UPDATING HIDDEN UNITS
Like zoneout, clockwork RNNs (Koutnik et al., 2014) and hierarchical RNNs (Hihi & Bengio, 1996) update only some units’ activations at every timestep, but their updates are periodic, whereas zoneout’s are stochastic. Inspired by clockwork RNNs, we experimented with zoneout variants that target different update rates or schedules for different units, but did not find any performance benefit. Hierarchical multiscale LSTMs (Chung et al., 2016) learn update probabilities for different units using the straight-through estimator (Bengio et al., 2013; Courbariaux et al., 2015), and combined with recently-proposed Layer Normalization (Ba et al., 2016), achieve competitive results on a variety of tasks. As the authors note, their method can be interpreted as an input-dependent form of adaptive zoneout.
In recent work, Ha et al. (2016) use a hypernetwork to dynamically rescale the row-weights of a primary LSTM network, achieving state-of-the-art 1.21 BPC on character-level Penn Treebank when combined with layer normalization (Ba et al., 2016) in a two-layer network. This scaling can be viewed as an adaptive, differentiable version of the variational LSTM (Gal, 2015), and could similarly be used to create an adaptive, differentiable version of zoneout. Very recent work conditions zoneout probabilities on suprisal (a measure of the discrepancy between the predicted and actual state), and sets a new state of the art on enwik8 (Rocki et al., 2016).
3 ZONEOUT AND PRELIMINARIES
We now explain zoneout in full detail, and compare with other forms of dropout in RNNs. We start by reviewing recurrent neural networks (RNNs).
3.1 RECURRENT NEURAL NETWORKS
Recurrent neural networks process data x1, x2, . . . , xT sequentially, constructing a corresponding sequence of representations, h1, h2, . . . , hT . Each hidden state is trained (implicitly) to remember and emphasize all task-relevant aspects of the preceding inputs, and to incorporate new inputs via a transition operator, T , which converts the present hidden state and input into a new hidden state: ht = T (ht−1, xt). Zoneout modifies these dynamics by mixing the original transition operator T̃ with the identity operator (as opposed to the null operator used in dropout), according to a vector of Bernoulli masks, dt:
Zoneout: T = dt T̃ + (1− dt) 1 Dropout: T = dt T̃ + (1− dt) 0
3.2 LONG SHORT-TERM MEMORY
In long short-term memory RNNs (LSTMs) (Hochreiter & Schmidhuber, 1997), the hidden state is divided into memory cell ct, intended for internal long-term storage, and hidden state ht, used as a transient representation of state at timestep t. In the most widely used formulation of an LSTM (Gers et al., 2000), ct and ht are computed via a set of four “gates”, including the forget gate, ft, which directly connects ct to the memories of the previous timestep ct−1, via an element-wise multiplication. Large values of the forget gate cause the cell to remember most (not all) of its previous value. The other gates control the flow of information in (it, gt) and out (ot) of the cell. Each gate has a weight matrix and bias vector; for example the forget gate has Wxf , Whf , and bf . For brevity, we will write these as Wx,Wh, b.
An LSTM is defined as follows:
it, ft, ot = σ(Wxxt +Whht−1 + b)
gt = tanh(Wxgxt +Whght−1 + bg)
ct = ft ct−1 + it gt ht = ot tanh(ct)
A naive application of dropout in LSTMs would zero-mask either or both of the memory cells and hidden states, without changing the computation of the gates (i, f, o, g). Dropping memory cells, for example, changes the computation of ct as follows:
ct = dt (ft ct−1 + it gt)
Alternatives abound, however; masks can be applied to any subset of the gates, cells, and states. Semeniuta et al. (2016), for instance, zero-mask the input gate:
ct = (ft ct−1 + dt it gt)
When the input gate is masked like this, there is no additive contribution from the input or hidden state, and the value of the memory cell simply decays according to the forget gate.
In zoneout, the values of the hidden state and memory cell randomly either maintain their previous value or are updated as usual. This introduces stochastic identity connections between subsequent time steps:
ct = d c t ct−1 + (1− dct) ( ft ct−1 + it gt ) ht = d h t ht−1 + (1− dht ) ( ot tanh ( ft ct−1 + it gt
)) We usually use different zoneout masks for cells and hiddens. We also experiment with a variant of recurrent dropout that reuses the input dropout mask to zoneout the corresponding output gates:
ct = (ft ct−1 + dt it gt) ht = ((1− dt) ot + dt ot−1) tanh(ct)
The motivation for this variant is to prevent the network from being forced (by the output gate) to expose a memory cell which has not been updated, and hence may contain misleading information.
4 EXPERIMENTS AND DISCUSSION
We evaluate zoneout’s performance on the following tasks: (1) Character-level language modelling on the Penn Treebank corpus (Marcus et al., 1993); (2) Word-level language modelling on the Penn Treebank corpus (Marcus et al., 1993); (3) Character-level language modelling on the Text8 corpus (Mahoney, 2011); (4) Classification of hand-written digits on permuted sequential MNIST (pMNIST) (Le et al., 2015). We also investigate the gradient flow to past hidden states, using pMNIST.
4.1 PENN TREEBANK LANGUAGE MODELLING DATASET
The Penn Treebank language model corpus contains 1 million words. The model is trained to predict the next word (evaluated on perplexity) or character (evaluated on BPC: bits per character) in a sequence. 1
4.1.1 CHARACTER-LEVEL
For the character-level task, we train networks with one layer of 1000 hidden units. We train LSTMs with a learning rate of 0.002 on overlapping sequences of 100 in batches of 32, optimize using Adam, and clip gradients with threshold 1. These settings match those used in Cooijmans et al. (2016). We also train GRUs and tanh-RNNs with the same parameters as above, except sequences are nonoverlapping and we use learning rates of 0.001, and 0.0003 for GRUs and tanh-RNNs respectively. Small values (0.1, 0.05) of zoneout significantly improve generalization performance for all three models. Intriguingly, we find zoneout increases training time for GRU and tanh-RNN, but decreases training time for LSTMs.
We focus our investigation on LSTM units, where the dynamics of zoning out states, cells, or both provide interesting insight into zoneout’s behaviour. Figure 3 shows our exploration of zoneout in LSTMs, for various zoneout probabilities of cells and/or hiddens. Zoneout on cells with probability 0.5 or zoneout on states with probability 0.05 both outperform the best-performing recurrent dropout (p = 0.25). Combining zc = 0.5 and zh = 0.05 leads to our best-performing model, which achieves 1.27 BPC, competitive with recent state-of-the-art set by (Ha et al., 2016). We compare zoneout to recurrent dropout (for p ∈ {0.05, 0.2, 0.25, 0.5, 0.7}), weight noise (σ = 0.075), norm stabilizer (β = 50) (Krueger & Memisevic, 2015), and explore stochastic depth (Huang et al., 2016) in a recurrent setting (analagous to zoning out an entire timestep). We also tried a shared-mask variant of zoneout as used in pMNIST experiments, where the same mask is used for both cells and hiddens. Neither stochastic depth or shared-mask zoneout performed as well as separate masks, sampled per unit. Figure 3 shows the best performance achieved with each regularizer, as well as an unregularized LSTM baseline. Results are reported in Table 1, and learning curves shown in Figure 4.
Low zoneout probabilities (0.05-0.25) also improve over baseline in GRUs and tanh-RNNs, reducing BPC from 1.53 to 1.41 for GRU and 1.67 to 1.52 for tanh-RNN. Similarly, low zoneout probabilities work best on the hidden states of LSTMs. For memory cells in LSTMs, however, higher probabilities (around 0.5) work well, perhaps because large forget-gate values approximate the effect of cells zoning out. We conjecture that best performance is achieved with zoneout LSTMs because of the stability of having both state and cell. The probability that both will be zoned out is very low, but having one or the other zoned out carries information from the previous timestep forward, while having the other react ’normally’ to new information.
4.1.2 WORD-LEVEL
For the word-level task, we replicate settings from Zaremba et al. (2014)’s best single-model performance. This network has 2 layers of 1500 units, with weights initialized uniformly [-0.04, +0.04]. The model is trained for 14 epochs with learning rate 1, after which the learning rate is reduced by a factor of 1.15 after each epoch. Gradient norms are clipped at 10.
With no dropout on the non-recurrent connections (i.e. zoneout as the only regularization), we do not achieve competitive results. We did not perform any search over models, and conjecture that the large model size requires regularization of the feed-forward connections. Adding zoneout (zc = 0.25 and zh = 0.025) on the recurrent connections to the model optimized for dropout on the non-recurrent connections however, we are able to improve test perplexity from 78.4 to 77.4. We report the best performance achieved with a given technique in Table 1.
1 These metrics are deterministic functions of negative log-likelihood (NLL). Specifically, perplexity is exponentiated NLL, and BPC (entropy) is NLL divided by the natural logarithm of 2.
4.2 TEXT8
Enwik8 is a corpus made from the first 109 bytes of Wikipedia dumped on Mar. 3, 2006. Text8 is a "clean text" version of this corpus; with html tags removed, numbers spelled out, symbols converted to spaces, all lower-cased. Both datasets were created and are hosted by Mahoney (2011).
We use a single-layer network of 2000 units, initialized orthogonally, with batch size 128, learning rate 0.001, and sequence length 180. We optimize with Adam (Kingma & Ba, 2014), clip gradients to a maximum norm of 1 (Pascanu et al., 2012), and use early stopping, again matching the settings of Cooijmans et al. (2016). Results are reported in Table 1, and Figure 4 shows training and validation learning curves for zoneout (zc = 0.5, zh = 0.05) compared to an unregularized LSTM and to recurrent dropout.
4.3 PERMUTED SEQUENTIAL MNIST
In sequential MNIST, pixels of an image representing a number [0-9] are presented one at a time, left to right, top to bottom. The task is to classify the number shown in the image. In pMNIST , the pixels are presented in a (fixed) random order.
We compare recurrent dropout and zoneout to an unregularized LSTM baseline. All models have a single layer of 100 units, and are trained for 150 epochs using RMSProp (Tieleman & Hinton, 2012) with a decay rate of 0.5 for the moving average of gradient norms. The learning rate is set to 0.001 and the gradients are clipped to a maximum norm of 1 (Pascanu et al., 2012).
As shown in Figure 5 and Table 2, zoneout gives a significant performance boost compared to the LSTM baseline and outperforms recurrent dropout (Semeniuta et al., 2016), although recurrent batch normalization (Cooijmans et al., 2016) outperforms all three. However, by adding zoneout to the recurrent batch normalized LSTM, we achieve state of the art performance. For this setting, the zoneout mask is shared between cells and states, and the recurrent dropout probability and zoneout probabilities are both set to 0.15.
4.4 GRADIENT FLOW
We investigate the hypothesis that identity connections introduced by zoneout facilitate gradient flow to earlier timesteps. Vanishing gradients are a perennial issue in RNNs. As effective as many techniques are for mitigating vanishing gradients (notably the LSTM architecture Hochreiter & Schmidhuber (1997)), we can always imagine a longer sequence to train on, or a longer-term dependence we want to capture.
We compare gradient flow in an unregularized LSTM to zoning out (stochastic identity-mapping) and dropping out (stochastic zero-mapping) the recurrent connections after one epoch of training on pMNIST. We compute the average gradient norms ‖ ∂L∂ct ‖ of loss L with respect to cell activations ct at each timestep t, and for each method, normalize the average gradient norms by the sum of average gradient norms for all timesteps.
Figure 6 shows that zoneout propagates gradient information to early timesteps much more effectively than dropout on the recurrent connections, and even more effectively than an unregularized LSTM. The same effect was observed for hidden states ht.
zoneout (zc = 0.5), dropout (zc = 0.5), and an unregularized LSTM on one epoch of pMNIST .
5 CONCLUSION
We have introduced zoneout, a novel and simple regularizer for RNNs, which stochastically preserves hidden units’ activations. Zoneout improves performance across tasks, outperforming many alternative regularizers to achieve results competitive with state of the art on the Penn Treebank and Text8 datasets, and state of the art results on pMNIST. While searching over zoneout probabilites allows us to tune zoneout to each task, low zoneout probabilities (0.05 - 0.2) on states reliably improve performance of existing models.
We perform no hyperparameter search to achieve these results, simply using settings from the previous state of the art. Results on pMNIST and word-level Penn Treebank suggest that Zoneout works well in combination with other regularizers, such as recurrent batch normalization, and dropout on feedforward/embedding layers. We conjecture that the benefits of zoneout arise from two main factors: (1) Introducing stochasticity makes the network more robust to changes in the hidden state; (2) The identity connections improve the flow of information forward and backward through the network.
ACKNOWLEDGMENTS
We are grateful to Hugo Larochelle, Jan Chorowski, and students at MILA, especially Çağlar Gülçehre, Marcin Moczulski, Chiheb Trabelsi, and Christopher Beckham, for helpful feedback and discussions. We thank the developers of Theano (Theano Development Team, 2016), Fuel, and Blocks (van Merriënboer et al., 2015). We acknowledge the computing resources provided by ComputeCanada and CalculQuebec. We also thank IBM and Samsung for their support. We would also like to acknowledge the work of Pranav Shyam on learning RNN hierarchies. This research was developed with funding from the Defense Advanced Research Projects Agency (DARPA) and the Air
Force Research Laboratory (AFRL). The views, opinions and/or findings expressed are those of the authors and should not be interpreted as representing the official views or policies of the Department of Defense or the U.S. Government.
6 APPENDIX
6.1 STATIC IDENTITY CONNECTIONS EXPERIMENT
This experiment was suggested by AnonReviewer2 during the ICLR review process with the goal of disentangling the effects zoneout has (1) through noise injection in the training process and (2) through identity connections. Based on these results, we observe that noise injection is essential for obtaining the regularization benefits of zoneout.
In this experiment, one zoneout mask is sampled at the beginning of training, and used for all examples. This means the identity connections introduced are static across training examples (but still different for each timestep). Using static identity connections resulted in slightly lower training (but not validation) error than zoneout, but worse performance than an unregularized LSTM on both train and validation sets, as shown in Figure 7.

The authors propose a conceptually simple method for regularisation of recurrent neural networks. The idea is related to dropout, but instead of zeroing out units, they are instead set to their respective values at the preceding time step element-wise with a certain probability.
Overall, the paper is well written. The method is clearly represented up to issues raised by reviewers during the pre-review question phase. The related work is complete and probably the best currently available on the matter of regularising RNNs.
The experimental section focuses on comparing the method with the current SOTA on a set of NLP benchmarks and a synthetic problem. All of the experiments focus on sequences over discrete values. An additional experiment also shows that the sequential Jacobian is far higher for long-term dependencies than in the dropout case.
Overall, the paper bears great potential. However, I do see some points.
1) As raised during the pre-review questions, I would like to see the results of experiments that feature a complete hyper parameter search. I.e. a proper model selection process,as it should be standard in the community. I do not see why this was not done, especially as the author count seems to indicate that the necessary resources are available.
I want to repeat at this point that Table 2 of the paper shows that validation error is not a reliable estimator for testing error in the respective data set. Thus, overfitting the model selection process is a serious concern here.
Zoneout does not seem to improve that much in the other tasks.
2) Zoneout is not investigated well mathematically. E.g. an analysis of the of the form of gradients from unit K at time step T to unit K’ at time step T-R would have been interesting, especially as these are not necessarily non-zero for dropout. Also, the question whether zoneout has a variational interpretation in the spirit of Yarin Gal’s work is an obvious one. I can see that it is if we treat zoneout in a resnet framework and dropout on the incremental parts. Overall, little effort is done answering the question *why* zoneout works well, even though the literature bears plenty of starting points for such analysis.
3) The data sets used are only symbolic. It would have been great if more ground was covered, i.e. continuous data such as from dynamical systems. To me it is not obvious whether it will transfer right away.
An extreme amount of “tricks” is being published currently for improved RNN training. How does zoneout stand out? It is a nice idea, and simple to implement. However, the paper under delivers: the experiments do not convince me (see 1) and 3)). There authors do not provide convincing theoretical insights either. (2)
Consequently, the paper reduces to a “epsilon improvement, great text, mediocre experimental evaluation, little theoretical insight”.


DEEP PROBABILISTIC PROGRAMMING
We propose Edward, a Turing-complete probabilistic programming language. Edward defines two compositional representations—random variables and inference. By treating inference as a first class citizen, on a par with modeling, we show that probabilistic programming can be as flexible and computationally efficient as traditional deep learning. For flexibility, Edward makes it easy to fit the same model using a variety of composable inference methods, ranging from point estimation to variational inference to MCMC. In addition, Edward can reuse the modeling representation as part of inference, facilitating the design of rich variational models and generative adversarial networks. For efficiency, Edward is integrated into TensorFlow, providing significant speedups over existing probabilistic systems. For example, we show on a benchmark logistic regression task that Edward is at least 35x faster than Stan and 6x faster than PyMC3. Further, Edward incurs no runtime overhead: it is as fast as handwritten TensorFlow.
1 INTRODUCTION
The nature of deep neural networks is compositional. Users can connect layers in creative ways, without having to worry about how to perform testing (forward propagation) or inference (gradientbased optimization, with back propagation and automatic differentiation).
In this paper, we design compositional representations for probabilistic programming. Probabilistic programming lets users specify generative probabilistic models as programs and then “compile” those models down into inference procedures. Probabilistic models are also compositional in nature, and much work has enabled rich probabilistic programs via compositions of random variables (Goodman et al., 2012; Ghahramani, 2015; Lake et al., 2016).
Less work, however, has considered an analogous compositionality for inference. Rather, many existing probabilistic programming languages treat the inference engine as a black box, abstracted away from the model. These cannot capture probabilistic inferences that reuse the model’s representation—a key idea in recent advances in variational inference (Kingma & Welling, 2014; Rezende & Mohamed, 2015; Tran et al., 2016b), generative adversarial networks (Goodfellow et al., 2014), and also in more classic inferences (Dayan et al., 1995; Gutmann & Hyvärinen, 2010).
We propose Edward1, a Turing-complete probabilistic programming language which builds on two compositional representations—one for random variables and one for inference. By treating inference as a first class citizen, on a par with modeling, we show that probabilistic programming can be as flexible and computationally efficient as traditional deep learning. For flexibility, we show how Edward makes it easy to fit the same model using a variety of composable inference methods, ranging from point estimation to variational inference to MCMC. For efficiency, we show how to integrate Edward into existing computational graph frameworks such as TensorFlow (Abadi et al., 2016). Frameworks like TensorFlow provide computational benefits like distributed training, parallelism, vectorization, and GPU support “for free.” For example, we show on a benchmark task that Edward’s Hamiltonian Monte Carlo is many times faster than existing software. Further, Edward incurs no runtime overhead: it is as fast as handwritten TensorFlow.
1See Tran et al. (2016a) for details of the API. A companion webpage for this paper is available at http: //edwardlib.org/iclr2017. It contains more complete examples with runnable code.
2 RELATED WORK
Probabilistic programming languages (PPLs) typically trade off the expressiveness of the language with the computational efficiency of inference. On one side, there are languages which emphasize expressiveness (Pfeffer, 2001; Milch et al., 2005; Pfeffer, 2009; Goodman et al., 2012), representing a rich class beyond graphical models. Each employs a generic inference engine, but scales poorly with respect to model and data size. On the other side, there are languages which emphasize efficiency (Spiegelhalter et al., 1995; Murphy, 2001; Plummer, 2003; Salvatier et al., 2015; Carpenter et al., 2016). The PPL is restricted to a specific class of models, and inference algorithms are optimized to be efficient for this class. For example, Infer.NET enables fast message passing for graphical models (Minka et al., 2014), and Augur enables data parallelism with GPUs for Gibbs sampling in Bayesian networks (Tristan et al., 2014). Edward bridges this gap. It is Turing complete—it supports any computable probability distribution—and it supports efficient algorithms, such as those that leverage model structure and those that scale to massive data.
There has been some prior research on efficient algorithms in Turing-complete languages. Venture and Anglican design inference as a collection of local inference problems, defined over program fragments (Mansinghka et al., 2014; Wood et al., 2014). This produces fast program-specific inference code, which we build on. Neither system supports inference methods such as programmable posterior approximations, inference models, or data subsampling. Concurrent with our work, WebPPL features amortized inference (Ritchie et al., 2016). Unlike Edward, WebPPL does not reuse the model’s representation; rather, it annotates the original program and leverages helper functions, which is a less flexible strategy. Finally, inference is designed as program transformations in Kiselyov & Shan (2009); Ścibior et al. (2015); Zinkov & Shan (2016). This enables the flexibility of composing inference inside other probabilistic programs. Edward builds on this idea to compose not only inference within modeling but also modeling within inference (e.g., variational models).
3 COMPOSITIONAL REPRESENTATIONS FOR PROBABILISTIC MODELS
We first develop compositional representations for probabilistic models. We desire two criteria: (a) integration with computational graphs, an efficient framework where nodes represent operations on data and edges represent data communicated between them (Culler, 1986); and (b) invariance of the representation under the graph, that is, the representation can be reused during inference.
Edward defines random variables as the key compositional representation. They are class objects with methods, for example, to compute the log density and to sample. Further, each random variable x is associated to a tensor (multi-dimensional array) x∗, which represents a single sample x∗ ∼ p(x). This association embeds the random variable onto a computational graph on tensors.
The design’s simplicity makes it easy to develop probabilistic programs in a computational graph framework. Importantly, all computation is represented on the graph. This enables one to compose random variables with complex deterministic structure such as deep neural networks, a diverse set of math operations, and third party libraries that build on the same framework. The design also enables compositions of random variables to capture complex stochastic structure.
As an illustration, we use a Beta-Bernoulli model, p(x, θ) = Beta(θ | 1, 1) ∏50
n=1 Bernoulli(xn | θ), where θ is a latent probability shared across the 50 data points x ∈ {0, 1}50. The random variable x is 50-dimensional, parameterized by the random tensor θ∗. Fetching the object x runs the graph: it simulates from the generative process and outputs a binary vector of 50 elements.
All computation is registered symbolically on random variables and not over their execution. Symbolic representations do not require reifying the full model, which leads to unreasonable memory
consumption for large models (Tristan et al., 2014). Moreover, it enables us to simplify both deterministic and stochastic operations in the graph, before executing any code (Ścibior et al., 2015; Zinkov & Shan, 2016).
With computational graphs, it is also natural to build mutable states within the probabilistic program. As a typical use of computational graphs, such states can define model parameters; in TensorFlow, this is given by a tf.Variable. Another use case is for building discriminative models p(y |x), where x are features that are input as training or test data. The program can be written independent of the data, using a mutable state (tf.placeholder) for x in its graph. During training and testing, we feed the placeholder the appropriate values.
In Appendix A, we provide examples of a Bayesian neural network for classification (A.1), latent Dirichlet allocation (A.2), and Gaussian matrix factorization (A.3). We present others below.
3.1 EXAMPLE: VARIATIONAL AUTO-ENCODER
Figure 2 implements a variational auto-encoder (VAE) (Kingma & Welling, 2014; Rezende et al., 2014) in Edward. It comprises a probabilistic model over data and a variational model designed to approximate the former’s posterior. Here we use random variables to construct both the probabilistic model and the variational model; they are fit during inference (more details in Section 4).
There are N data points xn ∈ {0, 1}28·28 each with d latent variables, zn ∈ Rd. The program uses Keras (Chollet, 2015) to define neural networks. The probabilistic model is parameterized by a 2-layer neural network, with 256 hidden units (and ReLU activation), and generates 28× 28 pixel images. The variational model is parameterized by a 2-layer inference network, with 256 hidden units and outputs parameters of a normal posterior approximation.
The probabilistic program is concise. Core elements of the VAE—such as its distributional assumptions and neural net architectures—are all extensible. With model compositionality, we can embed it into more complicated models (Gregor et al., 2015; Rezende et al., 2016) and for other learning tasks (Kingma et al., 2014). With inference compositionality (which we discuss in Section 4), we can embed it into more complicated algorithms, such as with expressive variational approximations (Rezende & Mohamed, 2015; Tran et al., 2016b; Kingma et al., 2016) and alternative objectives (Ranganath et al., 2016a; Li & Turner, 2016; Dieng et al., 2016).
3.2 EXAMPLE: BAYESIAN RECURRENT NEURAL NETWORK WITH VARIABLE LENGTH
Random variables can also be composed with control flow operations. As an example, Figure 3 implements a Bayesian recurrent neural network (RNN) with variable length. The data is a sequence of inputs {x1, . . . ,xT } and outputs {y1, . . . , yT } of length T with xt ∈ RD and yt ∈ R per time step. For t = 1, . . . , T , a RNN applies the update
ht = tanh(Whht−1 +Wxxt + bh),
where the previous hidden state is ht−1 ∈ RH . We feed each hidden state into the output’s likelihood, yt ∼ Normal(Wyht + by, 1), and we place a standard normal prior over all parameters {Wh ∈ RH×H ,Wx ∈ RD×H ,Wy ∈ RH×1,bh ∈ RH ,by ∈ R}. Our implementation is dynamic: it differs from a RNN with fixed length, which pads and unrolls the computation.
3.3 STOCHASTIC CONTROL FLOW AND MODEL PARALLELISM
Random variables can also be placed in the control flow itself, enabling probabilistic programs with stochastic control flow. Stochastic control flow defines dynamic conditional dependencies, known in the literature as contingent or existential dependencies (Mansinghka et al., 2014; Wu et al., 2016). See Figure 4, where x may or may not depend on a for a given execution. In Appendix A.4, we use stochastic control flow to implement a Dirichlet process mixture model. Tensors with stochastic shape are also possible: for example, tf.zeros(Poisson(lam=5.0)) defines a vector of zeros with length given by a Poisson draw with rate 5.0.
Stochastic control flow produces difficulties for algorithms that use the graph structure because the relationship of conditional dependencies changes across execution traces. The computational graph, however, provides an elegant way of teasing out static conditional dependence structure (p) from dynamic dependence structure (a). We can perform model parallelism (parallel computation across components of the model) over the static structure with GPUs and batch training. We can use more generic computations to handle the dynamic structure.
4 COMPOSITIONAL REPRESENTATIONS FOR INFERENCE
We described random variables as a representation for building rich probabilistic programs over computational graphs. We now describe a compositional representation for inference. We desire two criteria: (a) support for many classes of inference, where the form of the inferred posterior depends on the algorithm; and (b) invariance of inference under the computational graph, that is, the posterior can be further composed as part of another model.
To explain our approach, we will use a simple hierarchical model as a running example. Figure 5 displays a joint distribution p(x, z, β) of data x, local variables z, and global variables β. The ideas here extend to more expressive programs.
4.1 INFERENCE AS STOCHASTIC GRAPH OPTIMIZATION
The goal of inference is to calculate the posterior distribution p(z, β | xtrain;θ) given data xtrain, where θ are any model parameters that we will compute point estimates for.2 We formalize this as
2For example, we could replace x’s sigma argument with tf.exp(tf.Variable(0.0))*tf.ones([N, D]). This defines a model parameter initialized at 0 and positive-constrained.
the following optimization problem:
min λ,θ L(p(z, β | xtrain;θ), q(z, β;λ)), (1)
where q(z, β;λ) is an approximation to the posterior p(z, β |xtrain;θ), and L is a loss function with respect to p and q.
The choice of approximation q, loss L, and rules to update parameters {θ,λ} are specified by an inference algorithm. (Note q can be nonparametric, such as a point or a collection of samples.)
In Edward, we write this problem as follows: 1 inference = ed.Inference({beta: qbeta, z: qz}, data={x: x_train})
Inference is an abstract class which takes two inputs. The first is a collection of latent random variables beta and z, associated to their “posterior variables” qbeta and qz respectively. The second is a collection of observed random variables x, which is associated to their realizations x_train.
The idea is that Inference defines and solves the optimization in Equation 1. It adjusts parameters of the distribution of qbeta and qz (and any model parameters) to be close to the posterior.
Class methods are available to finely control the inference. Calling inference.initialize() builds a computational graph to update {θ,λ}. Calling inference.update() runs this computation once to update {θ,λ}; we call the method in a loop until convergence. Importantly, no efficiency is lost in Edward’s language: the computational graph is the same as if it were handwritten for a specific model. This means the runtime is the same; also see our experiments in Section 5.2.
A key concept in Edward is that there is no distinct “model” or “inference” block. A model is simply a collection of random variables, and inference is a way of modifying parameters in that collection subject to another. This reductionism offers significant flexibility. For example, we can infer only parts of a model (e.g., layer-wise training (Hinton et al., 2006)), infer parts used in multiple models (e.g., multi-task learning), or plug in a posterior into a new model (e.g., Bayesian updating).
4.2 CLASSES OF INFERENCE
The design of Inference is very general. We describe subclasses to represent many algorithms below: variational inference, Monte Carlo, and generative adversarial networks.
Variational inference posits a family of approximating distributions and finds the closest member in the family to the posterior (Jordan et al., 1999). In Edward, we build the variational family in the graph; see Figure 6 (left). For our running example, the family has mutable variables as parameters λ = {π, µ, σ}, where q(β;µ, σ) = Normal(β;µ, σ) and q(z;π) = Categorical(z;π). Specific variational algorithms inherit from the VariationalInference class. Each defines its own methods, such as a loss function and gradient. For example, we represent maximum a posteriori (MAP) estimation with an approximating family (qbeta and qz) of PointMass random variables, i.e., with all probability mass concentrated at a point. MAP inherits from VariationalInference and defines the negative log joint density as the loss function; it uses existing optimizers inside TensorFlow. In Section 5.1, we experiment with multiple gradient estimators for black box variational inference (Ranganath et al., 2014). Each estimator implements the same loss (an objective proportional to the divergence KL(q ‖ p)) and a different update rule (stochastic gradient).
Monte Carlo approximates the posterior using samples (Robert & Casella, 1999). Monte Carlo is an inference where the approximating family is an empirical distribution, q(β; {β(t)}) = 1 T ∑T t=1 δ(β, β (t)) and q(z; {z(t)}) = 1T ∑T t=1 δ(z, z (t)). The parameters are λ = {β(t), z(t)}. See Figure 6 (right). Monte Carlo algorithms proceed by updating one sample β(t), z(t) at a time in the empirical approximation. Specific MC samplers determine the update rules: they can use gradients such as in Hamiltonian Monte Carlo (Neal, 2011) and graph structure such as in sequential Monte Carlo (Doucet et al., 2001).
Edward also supports non-Bayesian methods such as generative adversarial networks (GANs) (Goodfellow et al., 2014). See Figure 7. The model posits random noise eps over N data points, each with d dimensions; this random noise feeds into a generative_network function, a neural network that outputs real-valued data x. In addition, there is a discriminative_network which takes data as input and outputs the probability that the data is real (in logit parameterization). We build GANInference; running it optimizes parameters inside the two neural network functions. This approach extends to many advances in GANs (e.g., Denton et al. (2015); Li et al. (2015)).
Finally, one can design algorithms that would otherwise require tedious algebraic manipulation. With symbolic algebra on nodes of the computational graph, we can uncover conjugacy relationships between random variables. Users can then integrate out variables to automatically derive classical Gibbs (Gelfand & Smith, 1990), mean-field updates (Bishop, 2006), and exact inference. These algorithms are being currently developed in Edward.
4.3 COMPOSING INFERENCES
Core to Edward’s design is that inference can be written as a collection of separate inference programs. Below we demonstrate variational EM, with an (approximate) E-step over local variables and an M-step over global variables. We instantiate two algorithms, each of which conditions on inferences from the other, and we alternate with one update of each (Neal & Hinton, 1993),
1 qbeta = PointMass(params=tf.Variable(tf.zeros([K, D]))) 2 qz = Categorical(logits=tf.Variable(tf.zeros([N, K]))) 3 4 inference_e = ed.VariationalInference({z: qz}, data={x: x_train, beta: qbeta}) 5 inference_m = ed.MAP({beta: qbeta}, data={x: x_train, z: qz}) 6 ...
7 for _ in range(10000): 8 inference_e.update() 9 inference_m.update()
This extends to many other cases such as exact EM for exponential families, contrastive divergence (Hinton, 2002), pseudo-marginal methods (Andrieu & Roberts, 2009), and Gibbs sampling within variational inference (Wang & Blei, 2012; Hoffman & Blei, 2015). We can also write message passing algorithms, which solve a collection of local inference problems (Koller & Friedman, 2009). For example, classical message passing uses exact local inference and expectation propagation locally minimizes the Kullback-Leibler divergence, KL(p ‖ q) (Minka, 2001).
4.4 DATA SUBSAMPLING
Stochastic optimization (Bottou, 2010) scales inference to massive data and is key to algorithms such as stochastic gradient Langevin dynamics (Welling & Teh, 2011) and stochastic variational inference (Hoffman et al., 2013). The idea is to cheaply estimate the model’s log joint density in an unbiased way. At each step, one subsamples a data set {xm} of size M and then scales densities with respect to local variables,
log p(x, z, β) = log p(β) + N∑ n=1 [ log p(xn | zn, β) + log p(zn |β) ] ≈ log p(β) + N
M M∑ m=1 [ log p(xm | zm, β) + log p(zm |β) ] .
To support stochastic optimization, we represent only a subgraph of the full model. This prevents reifying the full model, which can lead to unreasonable memory consumption (Tristan et al., 2014). During initialization, we pass in a dictionary to properly scale the arguments. See Figure 8.
Conceptually, the scale argument represents scaling for each random variable’s plate, as if we had seen that random variable N/M as many times. As an example, Appendix B shows how to implement stochastic variational inference in Edward. The approach extends naturally to streaming data (Doucet et al., 2000; Broderick et al., 2013; McInerney et al., 2015), dynamic batch sizes, and data structures in which working on a subgraph does not immediately apply (Binder et al., 1997; Johnson & Willsky, 2014; Foti et al., 2014).
5 EXPERIMENTS
In this section, we illustrate two main benefits of Edward: flexibility and efficiency. For the former, we show how it is easy to compare different inference algorithms on the same model. For the latter, we show how it is easy to get significant speedups by exploiting computational graphs.
5.1 RECENT METHODS IN VARIATIONAL INFERENCE
We demonstrate Edward’s flexibility for experimenting with complex inference algorithms. We consider the VAE setup from Figure 2 and the binarized MNIST data set (Salakhutdinov & Murray,
2008). We use d = 50 latent variables per data point and optimize using ADAM. We study different components of the VAE setup using different methods; Appendix C.1 is a complete script. After training we evaluate held-out log likelihoods, which are lower bounds on the true value.
Table 1 shows the results. The first method uses the VAE from Figure 2. The next three methods use the same VAE but apply different gradient estimators: reparameterization gradient without an analytic KL; reparameterization gradient with an analytic entropy; and the score function gradient (Paisley et al., 2012; Ranganath et al., 2014). This typically leads to the same optima but at different convergence rates. The score function gradient was slowest. Gradients with an analytic entropy produced difficulties around convergence: we switched to stochastic estimates of the entropy as it approached an optima. We also use hierarchical variational models (HVMs) (Ranganath et al., 2016b) with a normalizing flow prior; it produced similar results as a normalizing flow on the latent variable space (Rezende & Mohamed, 2015), and better than importance-weighted auto-encoders (IWAEs) (Burda et al., 2016).
We also study novel combinations, such as HVMs with the IWAE objective, GAN-based optimization on the decoder (with pixel intensity-valued data), and Rényi divergence on the decoder. GAN-based optimization does not enable calculation of the log-likelihood; Rényi divergence does not directly optimize for log-likelihood so it does not perform well. The key point is that Edward is a convenient research platform: they are all easy modifications of a given script.
5.2 GPU-ACCELERATED HAMILTONIAN MONTE CARLO
We benchmark runtimes for a fixed number of Hamiltonian Monte Carlo (HMC; Neal, 2011) iterations on modern hardware: a 12-core Intel i7-5930K CPU at 3.50GHz and an NVIDIA Titan X (Maxwell) GPU. We apply logistic regression on the Covertype dataset (N = 581012, D = 54; responses were binarized) using Edward, Stan (with PyStan) (Carpenter et al., 2016), and PyMC3 (Salvatier et al., 2015). We ran 100 HMC iterations, with 10 leapfrog updates per iteration, a step size of 0.5/N , and single precision. Figure 9 illustrates the program in Edward.
Table 2 displays the runtimes.3 Edward (GPU) features a dramatic 35x speedup over Stan (1 CPU) and 6x speedup over PyMC3 (12 CPU). This showcases the value of building a PPL on top of com-
3In a previous version of this paper, we reported PyMC3 took 361s. This was caused by a bug preventing PyMC3 from correctly handling single-precision floating point. (PyMC3 with double precision is roughly 14x
putational graphs. The speedup stems from fast matrix multiplication when calculating the model’s log-likelihood; GPUs can efficiently parallelize this computation. We expect similar speedups for models whose bottleneck is also matrix multiplication, such as deep neural networks.
There are various reasons for the speedup. Stan only used 1 CPU as it leverages multiple cores by running HMC chains in parallel. Stan also used double-precision floating point as it does not allow single-precision. For PyMC3, we note Edward’s speedup is not a result of PyMC3’s Theano backend compared to Edward’s TensorFlow. Rather, PyMC3 does not use Theano for all its computation, so it experiences communication overhead with NumPy. (PyMC3 was actually slower when using the GPU.) We predict that porting Edward’s design to Theano would feature similar speedups.
In addition to these speedups, we highlight that Edward has no runtime overhead: it is as fast as handwritten TensorFlow. Following Section 4.1, this is because the computational graphs for inference are in fact the same for Edward and the handwritten code.
5.3 PROBABILITY ZOO
In addition to Edward, we also release the Probability Zoo, a community repository for pre-trained probability models and their posteriors.4 It is inspired by the model zoo in Caffe (Jia et al., 2014), which provides many pre-trained discriminative neural networks, and which has been key to making large-scale deep learning more transparent and accessible. It is also inspired by Forest (Stuhlmüller, 2012), which provides examples of probabilistic programs.
6 DISCUSSION: CHALLENGES & EXTENSIONS
We described Edward, a Turing-complete PPL with compositional representations for probabilistic models and inference. Edward expands the scope of probabilistic programming to be as flexible and computationally efficient as traditional deep learning. For flexibility, we showed how Edward can use a variety of composable inference methods, capture recent advances in variational inference and generative adversarial networks, and finely control the inference algorithms. For efficiency, we showed how Edward leverages computational graphs to achieve fast, parallelizable computation, scales to massive data, and incurs no runtime overhead over handwritten code.
In present work, we are applying Edward as a research platform for developing new probabilistic models (Rudolph et al., 2016; Tran et al., 2017) and new inference algorithms (Dieng et al., 2016). As with any language design, Edward makes tradeoffs in pursuit of its flexibility and speed for research. For example, an open challenge in Edward is to better facilitate programs with complex control flow and recursion. While possible to represent, it is unknown how to enable their flexible inference strategies. In addition, it is open how to expand Edward’s design to dynamic computational graph frameworks—which provide more flexibility in their programming paradigm—but may sacrifice performance. A crucial next step for probabilistic programming is to leverage dynamic computational graphs while maintaining the flexibility and efficiency that Edward offers.
slower than Edward (GPU).) This has been fixed after discussion with Thomas Wiecki. The reported numbers also exclude compilation time, which is significant for Stan.
4The Probability Zoo is available at http://edwardlib.org/zoo. It includes model parameters and inferred posterior factors, such as local and global variables during training and any inference networks.
ACKNOWLEDGEMENTS
We thank the probabilistic programming community—for sharing our enthusiasm and motivating further work—including developers of Church, Venture, Gamalon, Hakaru, and WebPPL. We also thank Stan developers for providing extensive feedback as we developed the language, as well as Thomas Wiecki for experimental details. We thank the Google BayesFlow team—Joshua Dillon, Ian Langmore, Ryan Sepassi, and Srinivas Vasudevan—as well as Amr Ahmed, Matthew Johnson, Hung Bui, Rajesh Ranganath, Maja Rudolph, and Francisco Ruiz for their helpful feedback. This work is supported by NSF IIS-1247664, ONR N00014-11-1-0651, DARPA FA8750-14-2-0009, DARPA N66001-15-C-4032, Adobe, Google, NSERC PGS-D, and the Sloan Foundation.
A MODEL EXAMPLES
There are many examples available at http://edwardlib.org, including models, inference methods, and complete scripts. Below we describe several model examples; Appendix B describes an inference example (stochastic variational inference); Appendix C describes complete scripts. All examples in this paper are comprehensive, only leaving out import statements and fixed values. See the companion webpage for this paper (http://edwardlib.org/iclr2017) for examples in a machine-readable format with runnable code.
A.1 BAYESIAN NEURAL NETWORK FOR CLASSIFICATION
A Bayesian neural network is a neural network with a prior distribution on its weights.
Define the likelihood of an observation (xn, yn) with binary label yn ∈ {0, 1} as
p(yn |W0,b0,W1,b1 ; xn) = Bernoulli(yn |NN(xn ; W0,b0,W1,b1)),
where NN is a 2-layer neural network whose weights and biases form the latent variables W0,b0,W1,b1. Define the prior on the weights and biases to be the standard normal. See Figure 10. There are N data points, D features, and H hidden units.
A.2 LATENT DIRICHLET ALLOCATION
See Figure 11. Note that the program is written for illustration. We recommend vectorization in practice: instead of storing scalar random variables in lists of lists, one should prefer to represent few random variables, each which have many dimensions.
A.3 GAUSSIAN MATRIX FACTORIZATIONN
See Figure 12.
1 N = 10 2 M = 10 3 K = 5 # latent dimension 4 5 U = Normal(mu=tf.zeros([M, K]), sigma=tf.ones([M, K])) 6 V = Normal(mu=tf.zeros([N, K]), sigma=tf.ones([N, K])) 7 Y = Normal(mu=tf.matmul(U, V, transpose_b=True), sigma=tf.ones([N, M]))
A.4 DIRICHLET PROCESS MIXTURE MODEL
See Figure 13.
A Dirichlet process mixture model is written as follows: 1 mu = DirichletProcess(alpha=0.1, base_cls=Normal, mu=tf.zeros(D), sigma=tf.ones(D), sample_n=N) 2 x = Normal(mu=mu, sigma=tf.ones([N, D]))
where mu has shape (N, D). The DirichletProcess random variable returns sample_n=N draws, each with shape given by the base distribution Normal(mu, sigma). The essential component defining the DirichletProcess random variable is a stochastic while loop. We define it below. See Edward’s code base for a more involved version with a base distribution.
B INFERENCE EXAMPLE: STOCHASTIC VARIATIONAL INFERENCE
In the subgraph setting, we do data subsampling while working with a subgraph of the full model. This setting is necessary when the data and model do not fit in memory. It is scalable in that both the algorithm’s computational complexity (per iteration) and memory complexity are independent of the data set size.
For the code, we use the running example, a mixture model described in Figure 5.
1 N = 10000000 # data set size 2 D = 2 # data dimension 3 K = 5 # number of clusters
The model is
p(x, z, β) = p(β) N∏ n=1 p(zn | β)p(xn | zn, β).
To avoid memory issues, we work on only a subgraph of the model,
p(x, z, β) = p(β) M∏ m=1 p(zm | β)p(xm | zm, β)
1 M = 128 # mini-batch size 2 3 beta = Normal(mu=tf.zeros([K, D]), sigma=tf.ones([K, D])) 4 z = Categorical(logits=tf.zeros([M, K])) 5 x = Normal(mu=tf.gather(beta, z), sigma=tf.ones([M, D]))
Assume the variational model is
q(z, β) = q(β;λ) N∏ n=1 q(zn | β; γn),
parameterized by {λ, {γn}}. Again, we work on only a subgraph of the model,
q(z, β) = q(β;λ) M∏ m=1 q(zm | β; γm).
parameterized by {λ, {γm}}. Importantly, onlyM parameters are stored in memory for {γm} rather than N .
1 qbeta = Normal(mu=tf.Variable(tf.zeros([K, D])), 2 sigma=tf.nn.softplus(tf.Variable(tf.zeros[K, D]))) 3 qz_variables = tf.Variable(tf.zeros([M, K])) 4 qz = Categorical(logits=qz_variables)
We use KLqp, a variational method that minimizes the divergence measure KL(q ‖ p) (Jordan et al., 1999). We instantiate two algorithms: a global inference over β given the subset of z and a local inference over the subset of z given β. We also pass in a TensorFlow placeholder x_ph for the data, so we can change the data at each step.
1 x_ph = tf.placeholder(tf.float32, [M]) 2 inference_global = ed.KLqp({beta: qbeta}, data={x: x_ph, z: qz}) 3 inference_local = ed.KLqp({z: qz}, data={x: x_ph, beta: qbeta})
We initialize the algorithms with the scale argument, so that computation on z and x will be scaled appropriately. This enables unbiased estimates for stochastic gradients.
1 inference_global.initialize(scale={x: float(N) / M, z: float(N) / M}) 2 inference_local.initialize(scale={x: float(N) / M, z: float(N) / M})
We now run the algorithm, assuming there is a next_batch function which provides the next batch of data.
1 qz_init = tf.initialize_variables([qz_variables]) 2 for _ in range(1000): 3 x_batch = next_batch(size=M) 4 for _ in range(10): # make local inferences 5 inference_local.update(feed_dict={x_ph: x_batch}) 6 7 # update global parameters 8 inference_global.update(feed_dict={x_ph: x_batch}) 9 # reinitialize the local factors
10 qz_init.run()
After each iteration, we also reinitialize the parameters for q(z | β); this is because we do inference on a new set of local variational factors for each batch. This demo readily applies to other inference algorithms such as SGLD (stochastic gradient Langevin dynamics): simply replace qbeta and qz with Empirical random variables; then call ed.SGLD instead of ed.KLqp.
Note that if the data and model fit in memory but you’d still like to perform data subsampling for fast inference, we recommend not defining subgraphs. You can reify the full model, and simply index the local variables with a placeholder. The placeholder is fed at runtime to determine which of the local variables to update at a time. (For more details, see the website’s API.)
C COMPLETE EXAMPLES
C.1 VARIATIONAL AUTO-ENCODER
See Figure 14.
C.2 PROBABILISTIC MODEL FOR WORD EMBEDDINGS
See Figure 15. This example uses data subsampling (Section 4.4). The priors and conditional likelihoods are defined only for a minibatch of data. Similarly the variational model only models the embeddings used in a given minibatch. TensorFlow variables contain the embedding vectors for the entire vocabulary. TensorFlow placeholders ensure that the correct embedding vectors are used as variational parameters for a given minibatch.
The Bernoulli variables y_pos and y_neg are fixed to be 1’s and 0’s respectively. They model whether a word is indeed the target word for a given context window or has been drawn as a negative sample. Without regularization (via priors), the objective we optimize is identical to negative sampling.

The authors propose a new software package for probabilistic programming, taking advantage of recent successful tools used in the deep learning community. The software looks very promising and has the potential to transform the way we work in the probabilistic modelling community, allowing us to perform rapid-prototyping to iterate through ideas quickly. The composability principles are used insightfully, and the extension of inference to HMC for example, going beyond VI inference (which is simple to implement using existing deep learning tools), makes the software even more compelling.
However, the most important factor of any PPL is whether it is practical for real-world use cases. This was not demonstrated sufficiently in the submission. There are many example code snippets given in the paper, but most are not evaluated. The Dirichlet process mixture model example (Figure 12) is an important one: do the proposed black-box inference tools really work for this snippet? and will the GAN example (Figure 7) converge when optimised with real data? To convince the community of the practicality of the package it will be necessary to demonstrate these empirically. Currently the only evaluated model is a VAE with various inference techniques, which are not difficult to implement using pure TF.
Presentation:
* Paper presentation could be improved. For example the authors could use more signalling for what is about to be explained. On page 5 qbeta and qz are used without explanation - the authors could mention that an example will be given thereafter.
* I would also suggest to the authors to explain in the preface how the layers are implemented, and how the KL is handled in VI for example.
It will be useful to discuss what values are optimised and what values change as inference is performed (even before section 4.4). This was not clear for the majority of the paper.
Experiments:
* Why is the run time not reported in table 1?
* What are the "difficulties around convergence" encountered with the analytical entropies? inference issues become more difficult to diagnose as inference is automated. Are there tools to diagnose these with the provided toolbox?
* Did HMC give sensible results in the experiment at the bottom of page 8? only run time is reported.
* How difficult is it to get the inference to work (eg HMC) when we don't have full control over the computational graph structure and sampler?
* It would be extremely insightful to give a table comparing the performance (run time, predictive log likelihood, etc) of the various inference tools on more models.
* What benchmarks do you intend to use in the Model Zoo? the difficulty with probabilistic modelling is that there are no set benchmarks over which we can evaluate and compare many models. Model zoo is sensible for the Caffe ecosystem because there exist few benchmarks a large portion of the community was working on (ImageNet for example). What datasets would you use to compare the DPMM on for example?
Minor comments:
* Table 1: I would suggest to compare to Li & Turner with alpha=0.5 (the equivalent of Hellinger distance) as they concluded this value performs best. I'm not sure why alpha=-1 was chosen here.
* How do you handle discrete distributions (eg Figure 5)?
* x_real is not defined in Figure 7.
* I would suggest highlighting M in Figure 8.
* Comma instead of period after "rized), In" on page 8.
In conclusion I would say that the software developments presented here are quite exciting, and I'm glad the authors are pushing towards practical and accessible "inference for all". In its current form though I am forced to give the submission itself a score of 5.


RECURRENT BATCH NORMALIZATION
We propose a reparameterization of LSTM that brings the benefits of batch normalization to recurrent neural networks. Whereas previous works only apply batch normalization to the input-to-hidden transformation of RNNs, we demonstrate that it is both possible and beneficial to batch-normalize the hidden-to-hidden transition, thereby reducing internal covariate shift between time steps. We evaluate our proposal on various sequential problems such as sequence classification, language modeling and question answering. Our empirical results show that our batch-normalized LSTM consistently leads to faster convergence and improved generalization.
1 INTRODUCTION
Recurrent neural network architectures such as LSTM (Hochreiter & Schmidhuber, 1997) and GRU (Cho et al., 2014) have recently exhibited state-of-the-art performance on a wide range of complex sequential problems including speech recognition Amodei et al. (2015), machine translation (Bahdanau et al., 2015) and image and video captioning (Xu et al., 2015; Yao et al., 2015). Top-performing models, however, are based on very high-capacity networks that are computationally intensive and costly to train. Effective optimization of recurrent neural networks is thus an active area of study (Pascanu et al., 2012; Martens & Sutskever, 2011; Ollivier, 2013).
It is well-known that for deep feed-forward neural networks, covariate shift (Shimodaira, 2000; Ioffe & Szegedy, 2015) degrades the efficiency of training. Covariate shift is a change in the distribution of the inputs to a model. This occurs continuously during training of feed-forward neural networks, where changing the parameters of a layer affects the distribution of the inputs to all layers above it. As a result, the upper layers are continually adapting to the shifting input distribution and unable to learn effectively. This internal covariate shift (Ioffe & Szegedy, 2015) may play an especially important role in recurrent neural networks, which resemble very deep feed-forward networks.
Batch normalization (Ioffe & Szegedy, 2015) is a recently proposed technique for controlling the distributions of feed-forward neural network activations, thereby reducing internal covariate shift. It involves standardizing the activations going into each layer, enforcing their means and variances to be invariant to changes in the parameters of the underlying layers. This effectively decouples each layer’s parameters from those of other layers, leading to a better-conditioned optimization problem. Indeed, deep neural networks trained with batch normalization converge significantly faster and generalize better.
Although batch normalization has demonstrated significant training speed-ups and generalization benefits in feed-forward networks, it is proven to be difficult to apply in recurrent architectures (Laurent et al., 2016; Amodei et al., 2015). It has found limited use in stacked RNNs, where the normalization is applied “vertically”, i.e. to the input of each RNN, but not “horizontally” between timesteps. RNNs are deeper in the time direction, and as such batch normalization would be most beneficial when applied horizontally. However, Laurent et al. (2016) hypothesized that applying batch normalization in this way hurts training because of exploding gradients due to repeated rescaling.
Our findings run counter to this hypothesis. We show that it is both possible and highly beneficial to apply batch normalization in the hidden-to-hidden transition of recurrent models. In particular, we describe a reparameterization of LSTM (Section 3) that involves batch normalization and demonstrate that it is easier to optimize and generalizes better. In addition, we empirically analyze the
gradient backpropagation and show that proper initialization of the batch normalization parameters is crucial to avoiding vanishing gradient (Section 4). We evaluate our proposal on several sequential problems and show (Section 5) that our LSTM reparameterization consistently outperforms the LSTM baseline across tasks, in terms of both time to convergence and performance.
Liao & Poggio (2016) simultaneously investigated batch normalization in recurrent neural networks, albeit only for very short sequences (10 steps). Ba et al. (2016) independently developed a variant of batch normalization that is also applicable to recurrent neural networks and delivers similar improvements as our method.
2 PREREQUISITES
2.1 LSTM
Long Short-Term Memory (LSTM) networks are an instance of a more general class of recurrent neural networks (RNNs), which we review briefly in this paper. Given an input sequence X = (x1,x2, . . . ,xT ), an RNN defines a sequence of hidden states ht according to
ht = φ(Whht−1 +Wxxt + b), (1)
where Wh ∈ Rdh×dh ,Wx ∈ Rdx×dh ,b ∈ Rdh and the initial state h0 ∈ Rdh are model parameters. A popular choice for the activation function φ( · ) is tanh. RNNs are popular in sequence modeling thanks to their natural ability to process variable-length sequences. However, training RNNs using first-order stochastic gradient descent (SGD) is notoriously difficult due to the well-known problem of exploding/vanishing gradients (Bengio et al., 1994; Hochreiter, 1991; Pascanu et al., 2012). Gradient vanishing occurs when states ht are not influenced by small changes in much earlier states hτ , t τ , preventing learning of long-term dependencies in the input data. Although learning long-term dependencies is fundamentally difficult (Bengio et al., 1994), its effects can be mitigated through architectural variations such as LSTM (Hochreiter & Schmidhuber, 1997), GRU (Cho et al., 2014) and iRNN/uRNN (Le et al., 2015; Arjovsky et al., 2015).
In what follows, we focus on the LSTM architecture (Hochreiter & Schmidhuber, 1997) with recurrent transition given by
f̃t ĩt õt g̃t  = Whht−1 +Wxxt + b (2) ct = σ(f̃t) ct−1 + σ(̃it) tanh(g̃t) (3) ht = σ(õt) tanh(ct), (4)
where Wh ∈ Rdh×4dh ,WxRdx×4dh ,b ∈ R4dh and the initial states h0 ∈ Rdh , c0 ∈ Rdh are model parameters. σ is the logistic sigmoid function, and the operator denotes the Hadamard product.
The LSTM differs from simple RNNs in that it has an additional memory cell ct whose update is nearly linear which allows the gradient to flow back through time more easily. In addition, unlike the RNN which overwrites its content at each timestep, the update of the LSTM cell is regulated by a set of gates. The forget gate ft determines the extent to which information is carried over from the previous timestep, and the input gate it controls the flow of information from the current input xt. The output gate ot allows the model to read from the cell. This carefully controlled interaction with the cell is what allows the LSTM to robustly retain information for long periods of time.
2.2 BATCH NORMALIZATION
Covariate shift (Shimodaira, 2000) is a phenomenon in machine learning where the features presented to a model change in distribution. In order for learning to succeed in the presence of covariate shift, the model’s parameters must be adjusted not just to learn the concept at hand but also to adapt to the changing distribution of the inputs. In deep neural networks, this problem manifests as
internal covariate shift (Ioffe & Szegedy, 2015), where changing the parameters of a layer affects the distribution of the inputs to all layers above it.
Batch Normalization (Ioffe & Szegedy, 2015) is a recently proposed network reparameterization which aims to reduce internal covariate shift. It does so by standardizing the activations using empirical estimates of their means and standard deviations. However, it does not decorrelate the activations due to the computationally costly matrix inversion. The batch normalizing transform is as follows:
BN(h; γ, β) = β + γ h− Ê[h]√ V̂ar[h] +
(5)
where h ∈ Rd is the vector of (pre)activations to be normalized, γ ∈ Rd, β ∈ Rd are model parameters that determine the mean and standard deviation of the normalized activation, and ∈ R is a regularization hyperparameter. The division should be understood to proceed elementwise.
At training time, the statistics E[h] and Var[h] are estimated by the sample mean and sample variance of the current minibatch. This allows for backpropagation through the statistics, preserving the convergence properties of stochastic gradient descent. During inference, the statistics are typically estimated based on the entire training set, so as to produce a deterministic prediction.
3 BATCH-NORMALIZED LSTM
This section introduces a reparameterization of LSTM that takes advantage of batch normalization. Contrary to Laurent et al. (2016); Amodei et al. (2015), we leverage batch normalization in both the input-to-hidden and the hidden-to-hidden transformations. We introduce the batch-normalizing transform BN( · ; γ, β) into the LSTM as follows:
 f̃t ĩt õt g̃t  = BN(Whht−1; γh, βh) + BN(Wxxt; γx, βx) + b (6) ct = σ(f̃t) ct−1 + σ(̃it) tanh(g̃t) (7) ht = σ(õt) tanh(BN(ct; γc, βc)) (8)
In our formulation, we normalize the recurrent term Whht−1 and the input term Wxxt separately. Normalizing these terms individually gives the model better control over the relative contribution of the terms using the γh and γx parameters. We set βh = βx = 0 to avoid unnecessary redundancy, instead relying on the pre-existing parameter vector b to account for both biases. In order to leave the LSTM dynamics intact and preserve the gradient flow through ct, we do not apply batch normalization in the cell update.
The batch normalization transform relies on batch statistics to standardize the LSTM activations. It would seem natural to share the statistics that are used for normalization across time, just as recurrent neural networks share their parameters over time. However, we find that simply averaging statistics over time severely degrades performance. Although LSTM activations do converge to a stationary distribution, we observe that their statistics during the initial transient differ significantly (see Figure 5 in Appendix A). Consequently, we recommend using separate statistics for each timestep to preserve information of the initial transient phase in the activations.1
Generalizing the model to sequences longer than those seen during training is straightforward thanks to the rapid convergence of the activations to their steady-state distributions (cf. Figure 5). For our experiments we estimate the population statistics separately for each timestep 1, . . . , Tmax where
1 Note that we separate only the statistics over time and not the γ and β parameters.
Tmax is the length of the longest training sequence. When at test time we need to generalize beyond Tmax, we use the population statistic of time Tmax for all time steps beyond it.
During training we estimate the statistics across the minibatch, independently for each timestep. At test time we use estimates obtained by averaging the minibatch estimates over the training set.
4 INITIALIZING γ FOR GRADIENT FLOW
Although batch normalization allows for easy control of the pre-activation variance through the γ parameters, common practice is to normalize to unit variance. We suspect that the previous difficulties with recurrent batch normalization reported in Laurent et al. (2016); Amodei et al. (2015) are largely due to improper initialization of the batch normalization parameters, and γ in particular. In this section we demonstrate the impact of γ on gradient flow.
In Figure 1(a), we show how the pre-activation variance impacts gradient propagation in a simple RNN on the sequential MNIST task described in Section 5.1. Since backpropagation operates in reverse, the plot is best read from right to left. The quantity plotted is the norm of the gradient of the loss with respect to the hidden state at different time steps. For large values of γ, the norm quickly goes to zero as gradient is propagated back in time. For small values of γ the norm is nearly constant.
To demonstrate what we think is the cause of this vanishing, we drew samples x from a set of centered Gaussian distributions with standard deviation ranging from 0 to 1, and computed the derivative tanh′(x) = 1− tanh2(x) ∈ [0, 1] for each. Figure 1(b) shows the empirical distribution of the derivative as a function of standard deviation. When the input standard deviation is low, the input tends to be close to the origin where the derivative is close to 1. As the standard deviation increases, the expected derivative decreases as the input is more likely to be in the saturation regime. At unit standard deviation, the expected derivative is much smaller than 1.
We conjecture that this is what causes the gradient to vanish, and recommend initializing γ to a small value. In our trials we found that values of 0.01 or lower caused instabilities during training. Our choice of 0.1 seems to work well across different tasks.
5 EXPERIMENTS
This section presents an empirical evaluation of the proposed batch-normalized LSTM on four different tasks. Note that for all the experiments, we initialize the batch normalization scale and shift parameters γ and β to 0.1 and 0 respectively.
5.1 SEQUENTIAL MNIST
We evaluate our batch-normalized LSTM on a sequential version of the MNIST classification task (Le et al., 2015). The model processes each image one pixel at a time and finally predicts the label. We consider both sequential MNIST tasks, MNIST and permuted MNIST (pMNIST). In MNIST, the pixels are processed in scanline order. In pMNIST the pixels are processed in a fixed random order.
Our baseline consists of an LSTM with 100 hidden units, with a softmax classifier to produce a prediction from the final hidden state. We use orthogonal initialization for all weight matrices, except for the hidden-to-hidden weight matrix which we initialize to be the identity matrix, as this yields better generalization performance on this task for both models. The model is trained using RMSProp (Tieleman & Hinton, 2012) with learning rate of 10−3 and 0.9 momentum. We apply gradient clipping at 1 to avoid exploding gradients.
The in-order MNIST task poses a unique problem for our model: the input for the first hundred or so timesteps is constant across examples since the upper pixels are almost always black. This causes the variance of the hidden states to be exactly zero for a long period of time. Normalizing these zerovariance activations involves dividing zero by a small number at many timesteps, which does not affect the forward-propagated activations but causes the back-propagated gradient to explode. We work around this by adding Gaussian noise to the initial hidden states. Although the normalization amplifies the noise to signal level, we find that it does not hurt performance compared to datadependent ways of initializing the hidden states.
In Figure 2 we show the validation accuracy while training for both LSTM and batch-normalized LSTM (BN-LSTM). BN-LSTM converges faster than LSTM on both tasks. Additionally, we observe that BN-LSTM generalizes significantly better on pMNIST. It has been highlighted in Arjovsky et al. (2015) that pMNIST contains many longer term dependencies across pixels than in the original pixel ordering, where a lot of structure is local. A recurrent network therefore needs to
characterize dependencies across varying time scales in order to solve this task. Our results suggest that BN-LSTM is better able to capture these long-term dependencies.
Table 1 reports the test set accuracy of the early stop model for LSTM and BN-LSTM using the population statistics. Recurrent batch normalization leads to a better test score, especially for pMNIST where models have to leverage long-term temporal depencies. In addition, Table 1 shows that our batch-normalized LSTM achieves state of the art on both MNIST and pMNIST.
5.2 CHARACTER-LEVEL PENN TREEBANK
We evaluate our model on the task of character-level language modeling on the Penn Treebank corpus (Marcus et al., 1993) according to the train/valid/test partition of Mikolov et al. (2012). For training, we segment the training sequence into examples of length 100. The training sequence does not cleanly divide by 100, so for each epoch we randomly crop a subsequence that does and segment that instead.
Our baseline is an LSTM with 1000 units, trained to predict the next character using a softmax classifier on the hidden state ht. We use stochastic gradient descent on minibatches of size 64, with gradient clipping at 1.0 and step rule determined by Adam (Kingma & Ba, 2014) with learning rate 0.002. We use orthogonal initialization for all weight matrices. The setup for the batch-normalized LSTM is the same in all respects except for the introduction of batch normalization as detailed in 3.
We show the learning curves in Figure 3(a). BN-LSTM converges faster and generalizes better than the LSTM baseline. Figure 3(b) shows the generalization of our model to longer sequences. We observe that using the population statistics improves generalization performance, which confirms that repeating the last population statistic (cf. Section 3) is a viable strategy. In table 2 we report the performance of our best models (early-stopped on validation performance) on the Penn Treebank test sequence. Follow up works havd since improved the state of the art (Krueger et al., 2016; Chung et al., 2016; Ha et al., 2016).
5.3 TEXT8
We evaluate our model on a second character-level language modeling task on the much larger text8 dataset (Mahoney, 2009). This dataset is derived from Wikipedia and consists of a sequence of 100M characters including only alphabetical characters and spaces. We follow Mikolov et al. (2012); Zhang et al. (2016) and use the first 90M characters for training, the next 5M for validation and the final 5M characters for testing. We train on nonoverlapping sequences of length 180.
Both our baseline and batch-normalized models are LSTMs with 2000 units, trained to predict the next character using a softmax classifier on the hidden state ht. We use stochastic gradient descent on minibatches of size 128, with gradient clipping at 1.0 and step rule determined by Adam (Kingma & Ba, 2014) with learning rate 0.001. All weight matrices were initialized to be orthogonal.
We early-stop on validation performance and report the test performance of the resulting model in table 3. We observe that BN-LSTM obtains a significant performance improvement over the LSTM baseline. Chung et al. (2016) has since improved on our performance.
5.4 TEACHING MACHINES TO READ AND COMPREHEND
Recently, Hermann et al. (2015) introduced a set of challenging benchmarks for natural language processing, along with neural network architectures to address them. The tasks involve reading real news articles and answering questions about their content. Their principal model, the Attentive Reader, is a recurrent neural network that invokes an attention mechanism to locate relevant information in the document. Such models are notoriously hard to optimize and yet increasingly popular.
To demonstrate the generality and practical applicability of our proposal, we apply batch normalization in the Attentive Reader model and show that this drastically improves training.
We evaluate several variants. The first variant, referred to as BN-LSTM, consists of the vanilla Attentive Reader model with the LSTM simply replaced by our BN-LSTM reparameterization. The second variant, termed BN-everywhere, is exactly like the first, except that we also introduce batch normalization into the attention computations, normalizing each term going into the tanh nonlinearities.
Our third variant, BN-e*, is like BN-everywhere, but improved to more carefully handle variablelength sequences. Throughout this experiment we followed the common practice of padding each batch of variable-length data with zeros. However, this biases the batch mean and variance of xt toward zero. We address this effect using sequencewise normalization of the inputs as proposed by Laurent et al. (2016); Amodei et al. (2015). That is, we share statistics over time for normalization
of the input terms Wxxt, but not for the recurrent terms Whht or the cell output ct. Doing so avoids many issues involving degenerate statistics due to input sequence padding.
Our fourth and final variant BN-e** is like BN-e* but bidirectional. The main difficulty in adapting to bidirectional models also involves padding. Padding poses no problem as long as it is properly ignored (by not updating the hidden states based on padded regions of the input). However to perform the reverse application of a bidirectional model, it is common to simply reverse the padded sequences, thus moving the padding to the front. This causes similar problems as were observed on the sequential MNIST task (Section 5.1): the hidden states will not diverge during the initial timesteps and hence their variance will be severely underestimated. To get around this, we reverse only the unpadded portion of the input sequences and leave the padding in place.
See Appendix C for hyperparameters and task details.
Figure 4(a) shows the learning curves for the different variants of the attentive reader. BN-LSTM trains dramatically faster than the LSTM baseline. BN-everywhere in turn shows a significant improvement over BN-LSTM. In addition, both BN-LSTM and BN-everywhere show a generalization benefit over the baseline. The validation curves have minima of 50.3%, 49.5% and 50.0% for the baseline, BN-LSTM and BN-everywhere respectively. We emphasize that these results were obtained without any tweaking – all we did was to introduce batch normalization.
BN-e* and BN-e** converge faster yet, and reach lower minima: 47.1% and 43.9% respectively.
We train and evaluate our best model, BN-e**, on the full task from (Hermann et al., 2015). On this dataset we had to reduce the number of hidden units to 120 to avoid severe overfitting. Training curves for BN-e** and a vanilla LSTM are shown in Figure 4(b). Table 4 reports performances of the early-stopped models.
6 CONCLUSION
Contrary to previous findings by Laurent et al. (2016); Amodei et al. (2015), we have demonstrated that batch-normalizing the hidden states of recurrent neural networks greatly improves optimization. Indeed, doing so yields benefits similar to those of batch normalization in feed-forward neural networks: our proposed BN-LSTM trains faster and generalizes better on a variety of tasks including language modeling and question-answering. We have argued that proper initialization of the batch normalization parameters is crucial, and suggest that previous difficulties (Laurent et al., 2016; Amodei et al., 2015) were due in large part to improper initialization. Finally, we have shown our model to apply to complex settings involving variable-length data, bidirectionality and highly nonlinear attention mechanisms.
ACKNOWLEDGEMENTS
The authors would like to acknowledge the following agencies for research funding and computing support: the Nuance Foundation, Samsung, NSERC, Calcul Québec, Compute Canada, the Canada Research Chairs and CIFAR. Experiments were carried out using the Theano (Team et al., 2016) and the Blocks and Fuel (van Merriënboer et al., 2015) libraries for scientific computing. We thank David Krueger, Saizheng Zhang, Ishmael Belghazi and Yoshua Bengio for discussions and suggestions.
C TEACHING MACHINES TO READ AND COMPREHEND: TASK SETUP
We evaluate the models on the question answering task using the CNN corpus (Hermann et al., 2015), with placeholders for the named entities. We follow a similar preprocessing pipeline as Hermann et al. (2015). During training, we randomly sample the examples with replacement and shuffle the order of the placeholders in each text inside the minibatch. We use a vocabulary of 65829 words.
We deviate from Hermann et al. (2015) in order to save computation: we use only the 4 most relevant sentences from the description, as identified by a string matching procedure. Both the training and validation sets are preprocessed in this way. Due to imprecision this heuristic sometimes strips the
answers from the passage, putting an upper bound of 57% on the validation accuracy that can be achieved.
For the reported performances, the first three models (LSTM, BN-LSTM and BN-everywhere) are trained using the exact same hyperparameters, which were chosen because they work well for the baseline. The hidden state is composed of 240 units. We use stochastic gradient descent on minibatches of size 64, with gradient clipping at 10 and step rule determined by Adam (Kingma & Ba, 2014) with learning rate 8× 10−5. For BN-e* and BN-e**, we use the same hyperparameters except that we reduce the learning rate to 8× 10−4 and the minibatch size to 40.
D HYPERPARAMETER SEARCHES
Table 5 reports hyperparameter values that were tried in the experiments.
For MNIST and pMNIST, the hyperparameters were varied independently. For Penn Treebank, we performed a full grid search on learning rate and hidden state size, and later performed a sensitivity
analysis on the batch size and initial γ. For the text8 task and the experiments with the Attentive Reader, we carried out a grid search on the learning rate and hidden state size.
The same values were tried for both the baseline and our BN-LSTM. In each case, our reported results are those of the model with the best validation performance.

Contributions
The paper presents an adaptation of batch normalization for RNNs in the case of LSTMs, along the horizontal depth. Contrary to previous work from (Laurent 2015; Amodei 2016), the work demonstrates that batch-normalizing the hidden states of RNNs can improve optimization, and argues with quantitative experiments that the key factor to making this work is proper initialization of parameters, in particular gamma. Experiments show some gain in performance over vanilla LSTMs on Sequential MNIST, PTB, Text8 and CNN Question-Answering.
Novelty+Significance
Batch normalization has been key for training deeper and deeper networks (e.g. ResNets) and it seems natural that we would want to extend it to RNNs.  The paper shows that it is possible to do so with proper initialization of parameters, contrary to previous work from (Laurent 2015; Amodei 2016). Novelty comes from where to batch norm (i.e. not in the cell update) and in the per-time step statistics.
Adding batch normalization to LSTMs incurs additional computational cost and bookkeeping; for training speed comparisons (e.g. Figure 2) the paper only compares LSTM and BN-LSTM by iteration count; given the additional complexity of the BN-LSTM I would have also liked to see a wall-clock comparison.
As RNNs are used across many tasks, this work is of interest to many.  However, the results gains are generally minor and require several tricks to work in practice. Also, this work doesn’t address a question about batch normalization that it seems natural that it helps with faster training, but why would it also improve generalization?
Clarity
The paper is overall very clear and well-motivated. The model is well described and easy to understand, and the plots illustrate the points clearly.
Summary
Interesting though relatively incremental adaptation, but shows batch normalization to work for RNNs where previous works have not succeeded. Comprehensive set of experiments though it is questionable if the empirical gains are significant enough to justify the increased model complexity as well as computational overhead.
Pros
- Shows batch normalization to work for RNNs where previous works have not succeeded
- Good empirical analysis of hyper-parameter choices and of the activations
- Experiments on multiple tasks
- Clarity
Cons
- Relatively incremental
- Several ‘hacks’ for the method (per-time step statistics, adding noise for exploding variance, sequence-wise normalization)
- No mention of computational overhead
- Only character or pixel-level tasks, what about word-level?