POINTER SENTINEL MIXTURE MODELS
Recent neural network sequence models with softmax classifiers have achieved their best language modeling performance only with very large hidden states and large vocabularies. Even then they struggle to predict rare or unseen words even if the context makes the prediction unambiguous. We introduce the pointer sentinel mixture architecture for neural sequence models which has the ability to either reproduce a word from the recent context or produce a word from a standard softmax classifier. We explore applying the pointer sentinel mixture model to the LSTM, a standard recurrent neural network building block. Utilizing an LSTM that achieves 80.6 perplexity on the Penn Treebank, the pointer sentinel-LSTM model pushes perplexity down to 70.9 while using far fewer parameters than an LSTM that achieves similar results. In order to evaluate how well language models can exploit longer contexts and deal with more realistic vocabularies and corpora we also introduce the freely available WikiText corpus.1
Recent neural network sequence models with softmax classifiers have achieved their best language modeling performance only with very large hidden states and large vocabularies. Even then they struggle to predict rare or unseen words even if the context makes the prediction unambiguous. We introduce the pointer sentinel mixture architecture for neural sequence models which has the ability to either reproduce a word from the recent context or produce a word from a standard softmax classifier. We explore applying the pointer sentinel mixture model to the LSTM, a standard recurrent neural network building block. Utilizing an LSTM that achieves 80.6 perplexity on the Penn Treebank, the pointer sentinel-LSTM model pushes perplexity down to 70.9 while using far fewer parameters than an LSTM that achieves similar results. In order to evaluate how well language models can exploit longer contexts and deal with more realistic vocabularies and corpora we also introduce the freely available WikiText corpus.1
1 INTRODUCTION
A major difficulty in language modeling is learning when to predict specific words from the immediate context. For instance, imagine a new person is introduced and two paragraphs later the context would allow one to very accurately predict this person’s name as the next word. For standard neural sequence models to predict this name, they would have to encode the name, store it for many time steps in their hidden state, and then decode it when appropriate. As the hidden state is limited in capacity and the optimization of such models suffer from the vanishing gradient problem, this is a lossy operation when performed over many timesteps. This is especially true for rare words.
Models with soft attention or memory components have been proposed to help deal with this challenge, aiming to allow for the retrieval and use of relevant previous hidden states, in effect increasing hidden state capacity and providing a path for gradients not tied to timesteps. Even with attention, the standard softmax classifier that is being used in these models often struggles to correctly predict rare or previously unknown words.
Pointer networks (Vinyals et al., 2015) provide one potential solution for rare and out of vocabulary (OoV) words as a pointer network uses attention to select an element from the input as output. This allows it to produce previously unseen input tokens. While pointer networks improve performance on rare words and long-term dependencies they are unable to select words that do not exist in the input.
We introduce a mixture model, illustrated in Fig. 1, that combines the advantages of standard softmax classifiers with those of a pointer component for effective and efficient language modeling. Rather than relying on the RNN hidden state to decide when to use the pointer, as in the recent work of Gülçehre et al. (2016), we allow the pointer component itself to decide when to use the softmax vocabulary through a sentinel. The model improves the state of the art perplexity on the Penn Treebank. Since this commonly used dataset is small and no other freely available alternative exists that allows for learning long range dependencies, we also introduce a new benchmark dataset for language modeling called WikiText.
1Available for download at the WikiText dataset site
2 THE POINTER SENTINEL FOR LANGUAGE MODELING
Given a sequence of words w1, . . . , wN−1, our task is to predict the next word wN .
2.1 THE SOFTMAX-RNN COMPONENT
Recurrent neural networks (RNNs) have seen widespread use for language modeling (Mikolov et al., 2010) due to their ability to, at least in theory, retain long term dependencies. RNNs employ the chain rule to factorize the joint probabilities over a sequence of tokens: p(w1, . . . , wN ) =∏N
i=1 p(wi|w1, . . . , wi−1). More precisely, at each time step i, we compute the RNN hidden state hi according to the previous hidden state hi−1 and the input xi such that hi = RNN(xi, hi−1). When all the N − 1 words have been processed by the RNN, the final state hN−1 is fed into a softmax layer which computes the probability over a vocabulary of possible words: pvocab(w) = softmax(UhN−1), where pvocab ∈ RV , U ∈ RV×H , H is the hidden size, and V the vocabulary size. RNNs can suffer from the vanishing gradient problem. The LSTM (Hochreiter & Schmidhuber, 1997) architecture has been proposed to deal with this by updating the hidden state according to a set of gates. Our work focuses on the LSTM but can be applied to any RNN architecture that ends in a vocabulary softmax.
2.2 THE POINTER NETWORK COMPONENT
In this section, we propose a modification to pointer networks for language modeling. To predict the next word in the sequence, a pointer network would select the member of the input sequence p(w1, . . . , wN−1) with the maximal attention score as the output.
The simplest way to compute an attention score for a specific hidden state is an inner product with all the past hidden states h, with each hidden state hi ∈ RH . However, if we want to compute such a score for the most recent word (since this word may be repeated), we need to include the last hidden state itself in this inner product. Taking the inner product of a vector with itself results in the vector’s magnitude squared, meaning the attention scores would be strongly biased towards the most recent word. Hence we project the current hidden state to a query vector q first. To produce the query q we compute q = tanh(WhN−1 + b), where W ∈ RH×H , b ∈ RH , and q ∈ RH . To generate the pointer attention scores, we compute the match between the previous RNN output states hi and the query q by taking the inner product, followed by a softmax activation function to obtain a probability distribution:
zi = q Thi, (1)
a = softmax(z), (2) where z ∈ RL, a ∈ RL, and L is the total number of hidden states. The probability mass assigned to a given word is the sum of the probability mass given to all token positions where the given word appears:
pptr(w) = ∑
i∈I(w,x) ai, (3)
where I(w, x) results in all positions of the word w in the input x and pptr ∈ RV . This technique, referred to as pointer sum attention, has been used for question answering (Kadlec et al., 2016).
Given the length of the documents used in language modeling, it may not be feasible for the pointer network to evaluate an attention score for all the words back to the beginning of the dataset. Instead, we may elect to maintain only a window of the L most recent words for the pointer to match against. The length L of the window is a hyperparameter that can be tuned on a held out dataset or by empirically analyzing how frequently a word at position t appears within the last L words.
To illustrate the advantages of this approach, consider a long article featuring two sentences President Obama discussed the economy and President Obama then flew to Prague. If the query was Which President is the article about?, probability mass could be applied to Obama in either sentence. If the question was instead Who flew to Prague?, only the latter occurrence of Obama provides the proper context. The attention sum model ensures that, as long as the entire attention probability mass is distributed on the occurrences of Obama, the pointer network can achieve zero loss. This flexibility provides supervision without forcing the model to put mass on supervision signals that may be incorrect or lack proper context.
2.3 THE POINTER SENTINEL MIXTURE MODEL
While pointer networks have proven to be effective, they cannot predict output words that are not present in the input, a common scenario in language modeling. We propose to resolve this by using a mixture model that combines a standard softmax with a pointer.
Our mixture model has two base distributions: the softmax vocabulary of the RNN output and the positional vocabulary of the pointer model. We refer to these as the RNN component and the pointer component respectively. To combine the two base distributions, we use a gating function g = p(zi = k|xi) where zi is the latent variable stating which base distribution the data point belongs to. As we only have two base distributions, g can produce a scalar in the range [0, 1]. A value of 0 implies that only the pointer is used and 1 means only the softmax-RNN is used.
p(yi|xi) = g pvocab(yi|xi) + (1− g) pptr(yi|xi). (4) While the models could be entirely separate, we re-use many of the parameters for the softmaxRNN and pointer components. This sharing minimizes the total number of parameters in the model and capitalizes on the pointer network’s supervision for the RNN component.
2.4 DETAILS OF THE GATING FUNCTION
To compute the new pointer sentinel gate g, we modify the pointer component. In particular, we add an additional element to z, the vector of attention scores as defined in Eq. 1. This element is
computed using an inner product between the query and the sentinel2 vector s ∈ RH . This change can be summarized by changing Eq. 2 to a = softmax ([ z; qT s ]) . We define a ∈ RV+1 to be the attention distribution over both the words in the pointer window as well as the sentinel state. We interpret the last element of this vector to be the gate value: g = a[V + 1].
Any probability mass assigned to g is given to the standard softmax vocabulary of the RNN. The final updated, normalized pointer probability over the vocabulary in the window then becomes:
pptr(yi|xi) = 1
1− g a[1 : V ], (5)
where we denoted [1 : V ] to mean the first V elements of the vector. The final mixture model is the same as Eq. 4 but with the updated Eq. 5 for the pointer probability.
This setup encourages the model to have both components compete: use pointers whenever possible and back-off to the standard softmax otherwise. By integrating the gating function into the pointer computation, it is influenced by both the RNN hidden state and the pointer window’s hidden states.
2.5 MOTIVATION FOR THE SENTINEL AS GATING FUNCTION
To make the best decision possible regarding which component to use the gating function must have as much context as possible. As we increase the window of words for the pointer component to consider, the RNN hidden state by itself isn’t guaranteed to accurately recall the identity or order of words it has recently seen (Adi et al., 2016). This is an obvious limitation of encoding a variable length sequence into a fixed dimensionality vector.
If we want a pointer window where the length L is in the hundreds, accurately modeling all of this information within the RNN hidden state is impractical. The position of specific words is also a vital feature as relevant words eventually fall out of the pointer component’s window. To correctly model this would require the RNN hidden state to store both the identity and position of each word in the pointer window. This is far beyond the capability of the fixed dimensionality RNN hidden state.
For this reason, we integrate the gating function directly into the pointer network by use of the sentinel. The decision to back-off to the softmax vocabulary is then informed by both the query q, generated using the RNN hidden state hN−1, and from the contents of the hidden states in the pointer window itself. This allows the model to accurately query what hidden states are contained in the pointer window and avoid maintaining state for words that may have fallen out of the window.
2.6 POINTER SENTINEL LOSS FUNCTION
We minimize the cross-entropy loss of −∑j ŷij log p(yij |xi), where ŷi is a one hot encoding of the correct output. During training, as ŷi is one hot, only a single mixed probability p(yij) must be computed for calculating the loss. This can result in a far more efficient GPU implementation. At prediction time, when we want all values for p(yi|xi), a maximum of L word probabilities must be mixed, as there is a maximum of L unique words in the pointer window of length L. This mixing can occur on the CPU where random access indexing is more efficient than the GPU.
Following the pointer sum attention network, the aim is to place probability mass from the attention mechanism on the correct output ŷi if it exists in the input. In the case of our mixture model the pointer loss instead becomes − log ( g + ∑ i∈I(y,x) ai ) , where I(y, x) results in all positions of the correct output y in the input x. The gate g may be assigned all probability mass if, for instance, the correct output ŷi exists only in the softmax-RNN vocabulary. There is no penalty if the model places the entire probability mass on any of the instances of the correct word in the input window. If the pointer component places the entirety of the probability mass on the gate g, the pointer network incurs no penalty and the loss is entirely determined by the loss of the softmax-RNN component.
2A sentinel value is inserted at the end of a search space in order to ensure a search algorithm terminates if no matching item is found. Our sentinel value terminates the pointer search space and distributes the rest of the probability mass to the RNN vocabulary.
2.7 PARAMETERS AND COMPUTATION TIME
The pointer sentinel-LSTM mixture model results in a relatively minor increase in parameters and computation time, especially when compared to the model size required to achieve similar performance using a standard LSTM. The only two additional parameters required by the model are those required for computing q, specifically W ∈ RH×H and b ∈ RH , and the sentinel vector embedding, s ∈ RH . This is independent of the depth of the RNN as the pointer component only interacts with the output of the final RNN layer. The additional H2 + 2H parameters are minor compared to a single LSTM layer’s 8H2 + 4H parameters. Most models also use multiple LSTM layers.
In terms of additional computation, a pointer sentinel-LSTM of window size L only requires computing the query q (a linear layer with tanh activation), a total of L parallelizable inner product calculations, and the attention scores for the L resulting scalars via the softmax function.
3 RELATED WORK
Considerable research has been dedicated to the task of language modeling, from traditional machine learning techniques such as n-grams to deep neural sequence models.
Mixture models composed of various knowledge sources have been proposed in the past for language modeling. Rosenfeld (1996) uses a maximum entropy model to combine a variety of information sources to improve language modeling on news text and speech. These information sources include complex overlapping n-gram distributions and n-gram caches that aim to capture rare words.
Beyond n-grams, neural sequence models such as recurrent neural networks have been shown to achieve state of the art results (Mikolov et al., 2010). A variety of RNN regularization methods have been explored, including a number of dropout variations (Zaremba et al., 2014; Gal, 2015) which prevent overfitting of complex LSTM language models. Other work has modified the RNN architecture to better handle increased recurrence depth (Zilly et al., 2016).
In order to increase capacity and minimize the impact of vanishing gradients, some language and translation models have also added a soft attention or memory component (Bahdanau et al., 2015; Sukhbaatar et al., 2015; Cheng et al., 2016; Kumar et al., 2016; Xiong et al., 2016; Ahn et al., 2016). These mechanisms allow for the retrieval and use of relevant previous hidden states. Soft attention mechanisms need to first encode the relevant word into a state vector and then decode it again, even if the output word is identical to the input word used to compute that hidden state or memory. A drawback to soft attention is that if, for instance, January and March are both equally attended candidates, the attention mechanism may blend the two vectors, resulting in a context vector closest to February (Kadlec et al., 2016). Even with attention, the standard softmax classifier being used in these models often struggles to correctly predict rare or previously unknown words.
Attention-based pointer mechanisms were introduced in Vinyals et al. (2015) where the pointer network is able to select elements from the input as output. In the above example, only January or March would be available as options, as February does not appear in the input. The use of pointer networks have been shown to help with geometric problems (Vinyals et al., 2015), code generation (Ling et al., 2016), summarization (Gu et al., 2016; Gülçehre et al., 2016), question answering (Kadlec et al., 2016). While pointer networks improve performance on rare words and long-term dependencies they are unable to select words that do not exist in the input.
Gülçehre et al. (2016) introduce a pointer softmax model that can generate output from either the vocabulary softmax of an RNN or the location softmax of the pointer network. Not only does this allow for producing OoV words which are not in the input, the pointer softmax model is able to better deal with rare and unknown words than a model only featuring an RNN softmax. Rather than constructing a mixture model as in our work, they use a switching network to decide which component to use. For neural machine translation, the switching network is conditioned on the representation of the context of the source text and the hidden state of the decoder. The pointer network is not used as a source of information for the switching network as in our model. The pointer and RNN softmax are scaled according to the switching network and the word or location with the highest final attention score is selected for output. Although this approach uses both a pointer and RNN component, it is not a mixture model and does not combine the probabilities for a word if it occurs in both the pointer location softmax and the RNN vocabulary softmax. In our model the word probability is a mix of both the RNN and pointer components, allowing for better predictions when the context may be ambiguous.
Extending this concept further, the latent predictor network (Ling et al., 2016) generates an output sequence conditioned on an arbitrary number of base models where each base model may have differing granularity. In their task of code generation, the output could be produced one character at a time using a standard softmax or instead copy entire words from referenced text fields using a pointer network. As opposed to Gülçehre et al. (2016), all states which produce the same output are merged by summing their probabilities. The model requires a complex training process involving the forward-backward algorithm for Semi-Markov models to prevent an exponential path explosion.
4 WIKITEXT - A BENCHMARK FOR LANGUAGE MODELING
We first describe the most commonly used language modeling dataset and its pre-processing in order to then motivate the need for a new benchmark dataset.
4.1 PENN TREEBANK
In order to compare our model to the many recent neural language models, we conduct word-level prediction experiments on the Penn Treebank (PTB) dataset (Marcus et al., 1993), pre-processed by Mikolov et al. (2010). The dataset consists of 929k training, 73k validation, and 82k test words. As part of the pre-processing performed by Mikolov et al. (2010), words were lower-cased, numbers were replaced with N, newlines were replaced with 〈eos〉, and all other punctuation was removed. The vocabulary is the most frequent 10k words with OoV tokens replaced by an 〈unk〉 token. For full statistics, refer to Table 1.
4.2 REASONS FOR A NEW DATASET
While the processed version of the PTB above has been frequently used for language modeling, it has many limitations. The tokens in PTB are all lower case, stripped of any punctuation, and limited to a vocabulary of only 10k words. These limitations mean that the PTB is unrealistic for real language use, especially when far larger vocabularies with many rare words are involved. The appendix contains a graph illustrating this using a Zipfian plot over the training partition of the PTB, with the curve stopping abruptly at the 10k limit. Given that accurately predicting rare words, such as named entities, is an important task for many applications, the lack of a long tail is problematic.
Other larger scale language modeling datasets exist. Unfortunately, they either have restrictive licensing which prevents widespread use or have randomized sentence ordering (Chelba et al., 2013) which is unrealistic for most language use and prevents the effective learning and evaluation of longer term dependencies. Hence, we constructed a language modeling dataset using text extracted from Wikipedia and have made this available to the community.
4.3 CONSTRUCTION AND PRE-PROCESSING
We selected articles only fitting the Good or Featured article criteria specified by editors on Wikipedia. These articles have been reviewed by humans and are considered well written, factually accurate, broad in coverage, neutral in point of view, and stable. This resulted in 23,805 Good
articles and 4,790 Featured articles. The text for each article was extracted using the Wikipedia API. Extracting text from Wikipedia mark-up is nontrivial due to the large number of macros in use, used for metric conversions, abbreviations, language notation, and date handling.
Once extracted, specific sections which primarily featured lists were removed by default. Other minor bugs, such as sort keys and Edit buttons that leaked in from the HTML, were also removed. Mathematical formulae and LATEX code were replaced with 〈formula〉 tokens. Normalization and tokenization were performed using the Moses tokenizer (Koehn et al., 2007), slightly augmented to further split numbers (8,600 → 8 @,@ 600) and with some additional minor fixes. Following Chelba et al. (2013) a vocabulary was constructed by discarding all words with a count below 3. Words outside of the vocabulary were mapped to the 〈unk〉 token, also a part of the vocabulary.
4.4 STATISTICS
The full WikiText dataset is over 103 million words in size, a hundred times larger than the PTB. It is also a tenth the size of the One Billion Word Benchmark (Chelba et al., 2013), one of the largest publicly available language modeling benchmarks, whilst consisting of articles that allow for the capture and usage of longer term dependencies as might be found in many real world tasks.
The dataset is available in two different sizes: WikiText-2 and WikiText-103. Both feature punctuation, original casing, a larger vocabulary, and numbers. WikiText-2 is two times the size of the Penn Treebank dataset. WikiText-103 features all extracted articles. Both datasets use the same articles for validation and testing, only differing in the vocabularies. For full statistics, refer to Table 1.
5 EXPERIMENTS
5.1 TRAINING DETAILS
As the pointer sentinel mixture model uses the outputs of the RNN from up to L timesteps back, this presents a challenge for training. If we do not regenerate the stale historical outputs of the RNN when we update the gradients, backpropagation through these stale outputs may result in incorrect gradient updates. If we do regenerate all stale outputs of the RNN, the training process is far slower. As we can make no theoretical guarantees on the impact of stale outputs on gradient updates, we opt to regenerate the window of RNN outputs used by the pointer component after each gradient update.
We also use truncated backpropagation through time (BPTT) in a different manner to many other RNN language models. Truncated BPTT allows for practical time-efficient training of RNN models but has fundamental trade-offs that are rarely discussed. For running truncated BPTT, BPTT is run for k2 timesteps once every k1 timesteps. For many RNN language modeling training schemes, k1 = k2, meaning that every k timesteps truncated BPTT is performed for the k previous timesteps. This results in only a single RNN output receiving backpropagation for k timesteps, with the other extreme being that the first token receives backpropagation for 0 timesteps. As such, most words in the training data will never experience a full backpropagation for k timesteps.
In our task, the pointer component always looks L timesteps into the past if L past timesteps are available. We select k1 = 1 and k2 = L such that for each timestep we perform backpropagation for L timesteps and advance one timestep at a time. Only the loss for the final predicted word is used for backpropagation through the window.
5.2 MODEL DETAILS
Our experimental setup reflects that of Zaremba et al. (2014) and Gal (2015). We increased the number of timesteps used during training from 35 to 100, matching the length of the window L. Batch size was increased to 32 from 20. We also halve the learning rate when validation perplexity is worse than the previous iteration, stopping training when validation perplexity fails to improve for three epochs or when 64 epochs are reached. The gradients are rescaled if their global norm exceeds 1 (Pascanu et al., 2013b).3 We evaluate the medium model configuration which features a two layer
3The highly aggressive clipping is likely due to the increased BPTT length. Even with such clipping early batches may experience excessively high perplexity, though this settles rapidly.
LSTM of hidden size 650. We compare against the large model configuration which features a two layer LSTM of hidden size 1500.
We produce results for two model types, an LSTM model that uses dropout regularization and the pointer sentinel-LSTM model. The variants of dropout used were zoneout (Krueger et al., 2016) and variational inference based dropout (Gal, 2015). Zoneout, which stochastically forces some recurrent units to maintain their previous values, was used for the recurrent connections within the LSTM. Variational inference based dropout, where the dropout mask for a layer is locked across timesteps, was used on the input to each RNN layer and also on the output of the final RNN layer. We used a value of 0.5 for both dropout connections.
5.3 COMPARISON OVER PENN TREEBANK
Table 2 compares the pointer sentinel-LSTM to a variety of other models on the Penn Treebank dataset. The pointer sentinel-LSTM achieves the lowest perplexity, followed by the recent Recurrent Highway Networks (Zilly et al., 2016). The medium pointer sentinel-LSTM model also achieves lower perplexity than the large LSTM models. Note that the best performing large variational LSTM model uses computationally intensive Monte Carlo (MC) dropout averaging. Monte Carlo dropout averaging is a general improvement for any sequence model that uses dropout but comes at a greatly increased test time cost. In Gal (2015) it requires rerunning the test model with 1000 different dropout masks. The pointer sentinel-LSTM is able to achieve these results with far fewer parameters than other models with comparable performance, specifically with less than a third the parameters used in the large variational LSTM models.
We also test a variational LSTM that uses zoneout, which serves as the RNN component of our pointer sentinel-LSTM mixture. This variational LSTM model performs BPTT for the same length L as the pointer sentinel-LSTM, where L = 100 timesteps. The results for this model ablation are worse than that of Gal (2015)’s variational LSTM without Monte Carlo dropout averaging.
5.4 COMPARISON OVER WIKITEXT-2
As WikiText-2 is being introduced in this dataset, there are no existing baselines. We provide two baselines to compare the pointer sentinel-LSTM against: our variational LSTM using zoneout and
the medium variational LSTM used in Gal (2015).4 Attempts to run the Gal (2015) large model variant, a two layer LSTM with hidden size 1500, resulted in out of memory errors on a 12GB K80 GPU, likely due to the increased vocabulary size. We chose the best hyperparameters from PTB experiments for all models. Table 3 shows a similar gain made by the pointer sentinel-LSTM over the variational LSTM models. The variational LSTM from Gal (2015) again beats out the variational LSTM used as a base for our experiments.
6 ANALYSIS
6.1 IMPACT ON RARE WORDS
A hypothesis as to why the pointer sentinel-LSTM can outperform an LSTM is that the pointer component allows the model to effectively reproduce rare words. The RNN may better use hidden state capacity by relying on the pointer component. The pointer component may also allow for a sharper selection of a single word than may be possible using only the softmax.
The appendix contains a graph which shows the improvement of perplexity when comparing the LSTM to the pointer sentinel-LSTM. Words are split across buckets according to frequency. As the words become rarer, the pointer sentinel-LSTM has stronger improvements in perplexity. Even on the Penn Treebank, where there is a relative absence of rare words due to only selecting the most frequent 10k words, we can see the pointer sentinel-LSTM mixture model provides a direct benefit.
While the improvements are largest on rare words, we can see the pointer sentinel-LSTM is still helpful on relatively frequent words. This may be the pointer component directly selecting the word or through the pointer supervision signal improving the RNN by allowing gradients to flow directly to other occurrences of the word in that window.
6.2 QUALITATIVE ANALYSIS OF POINTER USAGE
In a qualitative analysis, we visualized the gate use and pointer attention for a variety of examples in the validation set, focusing on predictions where the gate primarily used the pointer component. These visualizations are available in the appendix.
As expected, the pointer component is heavily used for rare names such as Seidman (23 times in training), Iverson (7 times in training), and Rosenthal (3 times in training). The pointer component was also heavily used when it came to other named entity names such as companies like Honeywell (8 times in training) and Integrated (41 times in training, though due to lowercasing of words this includes integrated circuits, fully integrated, and other generic usage). Surprisingly, the pointer component was also used for many frequent tokens. For selecting units of measurement (tons, kilograms, . . . ) or the short scale of numbers (thousands, millions, billions, . . . ), the pointer would refer to recent usage. This is to be expected, especially when phrases are of the form increased from N tons to N tons. The model can even be found relying on a mixture of the softmax and the pointer for predicting frequent verbs such as said.
Finally, the pointer component can be seen pointing to words at the very end of the 100 word window (position 97), a far longer horizon than the 35 steps that most language models truncate their backpropagation training to. This illustrates why the gating function must be integrated into the pointer component. If the gating function could only use the RNN hidden state, it would need to be wary of words that were near the tail of the pointer, especially if it was not able to accurately
4https://github.com/yaringal/BayesianRNN
track exactly how long it was since seeing a word. By integrating the gating function into the pointer component, we avoid the RNN hidden state having to maintain this intensive bookkeeping.
7 CONCLUSION
We introduced the pointer sentinel mixture model and the WikiText language modeling dataset. The pointer sentinel mixture model can be applied to any classifier that ends in a softmax, including various recurrent neural network building blocks. When applied to a standard LSTM, the pointer sentinel-LSTM achieves state of the art results in language modeling over the Penn Treebank while using few additional parameters and little additional computational complexity at prediction time.
We have also motivated the need to move from Penn Treebank to a new language modeling dataset for long range dependencies, providing WikiText-2 and WikiText-103 as potential options. We hope these new datasets can serve as a platform to improve handling of rare words and the usage of long term dependencies in language modeling.

This paper proposes augmenting RNN-based language models with a pointer network in order to deal better with rare words. The pointer network can point to words in the recent context, and hence the prediction for each time step is a mixture between the usual softmax output and the pointer distribution over the recent words. The paper also introduces a new language modelling dataset, which overcomes some of the shortcomings of previous datasets.
The reason for the score I gave for this paper is that I find the proposed model a direct application of the previous work Gulcehre et al., which follows a similar approach but for machine translation and summarization. The main differences I find is that Gulcehre et al. use an encoder-decoder architecture, and use the attention weights of the encoder to point to locations of words in the input, while here an RNN is used and a pointer network produces a distribution over the full vocabulary (by summing the softmax probabilities of words in the recent context). The context (query) vector for the pointing network is also different, but this is also a direct consequence of having a different application.
While the paper describes the differences between the proposed approach and Gulcehre et al.’s approach, I find some of the claims either wrong or not that significant. For example, quoting from Section 1:
“Rather than relying on the RNN hidden state to decide when to use the pointer, as in the recent work of Gulcehre et al. (2016), we allow the pointer component itself to decide when to use the softmax vocabulary through a sentinel.”
As far as I can tell, your model also uses the recent hidden state to form a query vector,  which is matched by the pointer network to previous words. Can you please clarify what you mean here?
In addition, quoting from section 3 which describes the model of Gulcehre et al.:
“Rather than constructing a mixture model as in our work, they use a switching network to decide which component to use”
This is not correct. The model of Gulcehre is also a mixture model, where an MLP with sigmoid output (switching network) is used to form a mixture between softmax prediction and locations of the input text.
Finally, in the following quote, also from section 3:
“The pointer network is not used as a source of information for the switching network as in our model.”
It is not clear what the authors mean by “source of information” here. Is it the fact that the switching probability is part of the pointer softmax? I am wondering how significant this difference is.
With regards to the proposed dataset, there are also other datasets typically used for language modelling, including The Hutter Prize Wikipedia (enwik8) dataset (Hutter, 2012) and e Text8 dataset (Mahoney, 2009). Can you please comment on the differences between your dataset and those as well?
I would be happy to discuss with the authors the points I raised, and I am open to changing my vote if there is any misunderstanding on my part.


Representation learning seeks to expose certain aspects of observed data in a learned representation that’s amenable to downstream tasks like classification. For instance, a good representation for 2D images might be one that describes only global structure and discards information about detailed texture. In this paper, we present a simple but principled method to learn such global representations by combining Variational Autoencoder (VAE) with neural autoregressive models such as RNN, MADE and PixelRNN/CNN. Our proposed VAE model allows us to have control over what the global latent code can learn and by designing the architecture accordingly, we can force the global latent code to discard irrelevant information such as texture in 2D images, and hence the VAE only “autoencodes” data in a lossy fashion. In addition, by leveraging autoregressive models as both prior distribution p(z) and decoding distribution p(x|z), we can greatly improve generative modeling performance of VAEs, achieving new state-of-the-art results on MNIST, OMNIGLOT and Caltech-101 Silhouettes density estimation tasks as well as competitive results on CIFAR10.
1 INTRODUCTION
A key goal of representation learning is to identify and disentangle the underlying causal factors of the data, so that it becomes easier to understand the data, to classify it, or to perform other tasks (Bengio et al., 2013). For image data this often means that we are interested in uncovering the “global structure” that captures the content of an image (for example, the identity of objects present in the image) and its “style”, but that we are typically less interested in the local and high frequency sources of variation such as the specific textures or white noise patterns.
A popular approach for learning representations is to fit a probabilistic latent variable model, an approach also known as analysis-by-synthesis (Yuille & Kersten, 2006; Nair et al., 2008). By learning a generative model of the data with the appropriate hierarchical structure of latent variables, it is hoped that the model will somehow uncover and untangle those causal sources of variations that we happen to be interested in. However, without further assumptions, representation learning via generative modeling is ill-posed: there are many different possible generative models with different (or no) kinds of latent variables that all encode the same probability density function on our observed data. Thus, the results we empirically get using this approach are highly dependent on the specific architectural and modeling choices that are made. Moreover, the objective that we optimize is often completely disconnected from the goal of learning a good representation: An autoregressive model of the data may achieve the same log-likelihood as a variational autoencoder (VAE) (Kingma & Welling, 2013), but the structure learned by the two models is completely different: the latter typically has a clear hierarchy of latent variables, while the autoregressive model has no stochastic latent variables at all (although it is conceivable that the deterministic hidden units of the autoregressive models will have meaningful and useful representations). For this reason, autoregressive models have thus far not been popular for the purpose of learning representations, even though they are extremely powerful as generative models (see e.g. van den Oord et al., 2016a).
A natural question becomes: is it possible to have a model that is a powerful density estimator and at the same time has the right hierarchical structure for representation learning? A potential solution would be to use a hybrid model that has both the latent variable structure of a VAE, as
well as the powerful recurrence of an autoregressive model. However, earlier attempts at combining these two kinds of models have run into the problem that the autoregressive part of the model ends up explaining all structure in the data, while the latent variables are not used (Fabius & van Amersfoort, 2014; Chung et al., 2015; Bowman et al., 2015; Serban et al., 2016; Fraccaro et al., 2016; Xu & Sun, 2016). Bowman et al. (2015) noted that weakening the autoregressive part of the model by, for example, dropout can encourage the latent variables to be used. We analyze why weakening is necessary, and we propose a principled solution that takes advantage of this property to control what kind of information goes into latent variables. The model we propose performs well as a density estimator, as evidenced by state-of-the-art log-likelihood results on MNIST, OMNIGLOT and Caltech-101, and also has a structure that is uniquely suited for learning interesting global representations of data.
2 VAES DO NOT AUTOENCODE IN GENERAL
A VAE is frequently interpreted as a regularized autoencoder (Kingma & Welling, 2013; Zhang et al., 2016), but the conditions under which it is guaranteed to autoencode (reconstruction being close to original datapoint) are not discussed. In this section, we discuss the often-neglected fact that VAEs do not always autoencode and give explicit reasons why previous attempts to apply VAE in sequence modeling found that the latent code is generally not used unless the decoder is weakened (Bowman et al., 2015; Serban et al., 2016; Fraccaro et al., 2016). The understanding of when VAE does autoencode will be an essential building piece for VLAE.
2.1 TECHNICAL BACKGROUND
Let x be observed variables, z latent variables and let p(x, z) be the parametric model of their joint distribution, called the generative model defined over the variables. Given a dataset X = {x1, ...,xN} we wish to perform maximum likelihood learning of its parameters:
log p(X) = N∑ i=1 log p(x(i)), (1)
but in general this marginal likelihood is intractable to compute or differentiate directly for flexible generative models that have high-dimensional latent variables and flexible priors and likelihoods. A solution is to introduce q(z|x), a parametric inference model defined over the latent variables, and optimize the variational lower bound on the marginal log-likelihood of each observation x:
log p(x) ≥ Eq(z|x) [log p(x, z)− log q(z|x)] = L(x; θ) (2) where θ indicates the parameters of p and q models.
There are various ways to optimize the lower bound L(x; θ); for continuous z it can be done efficiently through a re-parameterization of q(z|x) (Kingma & Welling, 2013; Rezende et al., 2014). This way of optimizing the variational lower bound with a parametric inference network and reparameterization of continuous latent variables is usually called VAE. The “autoencoding” terminology comes from the fact that the lower bound L(x; θ) can be re-arranged:
L(x; θ) = Eq(z|x) [log p(x, z)− log q(z|x)] (3) = Eq(z|x) [log p(x|z)]−DKL(q(z|x)||p(z)) (4)
where the first term can be seen as the expectation of negative reconstruction error and the KL divergence term can be seen as a regularizer, which as a whole could be seen as a regularized autoencoder loss with q(z|x) being the encoder and p(x|z) being the decoder. In the context of 2D images modeling, the decoding distribution p(x|z) is usually chosen to be a simple factorized distribution, i.e. p(x|z) = ∏ i p(xi|z), and this setup often yields a sharp decoding distribution p(x|z) that tends to reconstruct original datapoint x exactly.
2.2 BITS-BACK CODING AND INFORMATION PREFERENCE
It’s straightforward to see that having a more powerful p(x|z) will make VAE’s marginal generative distribution p(x) = ∫ z p(z)p(x|z)dz more expressive. This idea has been explored extensively
in previous work applying VAE to sequence modeling (Fabius & van Amersfoort, 2014; Chung et al., 2015; Bowman et al., 2015; Serban et al., 2016; Fraccaro et al., 2016; Xu & Sun, 2016), where the decoding distribution is a powerful RNN with autoregressive dependency, i.e., p(x|z) =∏ i p(xi|z,x<i). Since RNNs are universal function approximators and any joint distribution over x admits an autoregressive factorization, the RNN autoregressive decoding distribution can in theory represent any probability distribution even without dependence on z.
However, previous attempts have found it hard to benefit from VAE when using an expressive decoding distribution p(x|z). Indeed it’s documented in detail by Bowman et al. (2015) that in most cases when an RNN autoregressive decoding distribution is used, the latent code z is completely ignored and the model regresses to be a standard unconditional RNN autoregressive distribution that doesn’t depend on the latent code. This phenomenon is commonly attributed to “optimization challenges” of VAE in the literature (Bowman et al., 2015; Serban et al., 2016; Kaae Sønderby et al., 2016) because early in the training the approximate posterior q(z|x) carries little information about datapoint x and hence it’s easy for the model to just set the approximate posterior to be the prior to avoid paying any regularization cost DKL(q(z|x)||p(z)). Here we present a simple but often-neglected observation that this phenomenon arises not just due to optimization challenges and instead even if we can solve the optimization problems exactly, the latent code should still be ignored at optimum for most practical instances of VAE that have intractable true posterior distributions and sufficiently powerful decoders. It is easiest to understand this observation from a Bits-Back Coding perspective of VAE.
It is well-known that Bits-Back Coding is an information-theoretic view of Variational Inference (Hinton & Van Camp, 1993; Honkela & Valpola, 2004) and specific links have been established between Bits-Back Coding and the Helmholtz Machine/VAE (Hinton & Zemel, 1994; Gregor et al., 2013). Here we briefly relate VAE to Bits-Back Coding for self-containedness:
First recall that the goal of designing an efficient coding protocol is to minimize the expected code length of communicating x. To explain Bits-Back Coding, let’s first consider a more naive coding scheme. VAE can be seen as a way to encode data in a two-part code: p(z) and p(x|z), where z can be seen as the essence/structure of a datum and is encoded first and then the modeling error (deviation from z’s structure) is encoded next. The expected code length under this naive coding scheme for a given data distribution is hence:
Cnaive(x) = Ex∼data,z∼q(z|x) [− log p(z)− log p(x|z)] (5)
This coding scheme is, however, inefficient. Bits-Back Coding improves on it by noticing that the encoder distribution q(z|x) can be used to transmit additional information, up to H(q(z|x)) expected nats, as long as the receiver also has access to q(z|x). The decoding scheme works as follows: a receiver first decodes z from p(z), then decodes x from p(x|z) and, by running the same approximate posterior that the sender is using, decodes a secondary message from q(z|x). Hence, to properly measure the code length of VAE’s two-part code, we need to subtract the extra information from q(z|x). Using Bit-Back Coding, the expected code length equates to the negative variational lower bound or the so-called Helmholtz variational free energy, which means minimizing code length is equivalent to maximizing the variational lower bound:
CBitsBack(x) = Ex∼data,z∼q(z|x) [log q(z|x)− log p(z)− log p(x|z)] (6) = Ex∼data [−L(x)] (7)
Casting the problem of optimizing VAE into designing an efficient coding scheme easily allows us to reason when the latent code z will be used: the latent code z will be used when the two-part code is an efficient code. Recalling that the lower-bound of expected code length for data is given by the Shannon entropy of data generation distribution: H(data) = Ex∼data [− log pdata(x)], we can analyze VAE’s coding efficiency:
CBitsBack(x) = Ex∼data,z∼q(z|x) [log q(z|x)− log p(z)− log p(x|z)] (8) = Ex∼data [− log p(x) +DKL(q(z|x)||p(z|x))] (9) ≥ Ex∼data [− log pdata(x) +DKL(q(z|x)||p(z|x))] (10) = H(data) + Ex∼data [DKL(q(z|x)||p(z|x))] (11)
Since Kullback Leibler divergence is always non-negative, we know that using the two-part code derived from VAE suffers at least an extra code length of DKL(q(z|x)||p(z|x)) nats for using a posterior that’s not precise. Many previous works in Variational Inference have designed flexible approximate posteriors to better approximate true posterior (Salimans et al., 2014; Rezende & Mohamed, 2015; Tran et al., 2015; Kingma et al., 2016). Improved posterior approximations have shown to be effective in improving variational inference but none of the existing methods are able to completely close the gap between approximate posterior and true posterior. This leads us to believe that for most practical models, at least in the near future, the extra coding costDKL(q(z|x)||p(z|x)) will exist and will not be negligible.
Once we understand the inefficiency of the Bits-Back Coding mechanism, it’s simple to realize why sometimes the latent code z is not used: if the p(x|z) could model pdata(x) without using information from z, then it will not use z, in which case the true posterior p(z|x) is simply the prior p(z) and it’s usually easy to set q(z|x) to be p(z) to avoid incurring an extra cost DKL(q(z|x)||p(z|x)). And it’s exactly the case when a powerful decoding distribution is used like an RNN autoregressive distribution, which given enough capacity is able to model arbitrarily complex distributions. Hence there exists a preference of information when a VAE is optimized: information that can be modeled locally by decoding distribution p(x|z) without access to z will be encoded locally and only the remainder will be encoded in z.
We note that one common way to encourage putting information into the code is to use a factorized decoder p(x|z) = ∏ i p(xi|z) but so long as there is one dimension xj that’s independent of all other dimensions for true data distribution, pdata(x) = pdata(xj)pdata(x 6=j), then the latent code doesn’t contain all the information about x since at least xj will be modeled locally by factorized p(x|z). This kind of independence structure rarely exists in images so common VAEs that have factorized decoder autoencode almost exactly. Other techniques to encourage the usage of the latent code include annealing the relative weight of of DKL(q(z|x)||p(z)) in the variational lower bound (Bowman et al., 2015; Kaae Sønderby et al., 2016) or the use of free bits (Kingma et al., 2016), which can serve the dual purpose of smoothing the optimization landscape and canceling out part of the Bits-Back Code inefficiency DKL(q(z|x)||p(z|x)).
3 VARIATIONAL LOSSY AUTOENCODER
The discussion in Section 2.2 suggests that autoregressive models cannot be combined with VAE since information will be preferred to be modeled by autoregressive models. Nevertheless, in this section, we present two complementary classes of improvements to VAE that utilize autoregressive models fruitfully to explicitly control representation learning and improve density estimation.
3.1 LOSSY CODE VIA EXPLICIT INFORMATION PLACEMENT
Even though the information preference property of VAE might suggest that one should always use the full autoregressive models to achieve a better code length/log-likelihood, especially when slow data generation is not a concern, we argue that this information preference property can be exploited to turn the VAE into a powerful representation learning method that gives us fine-grained control over the kind of information that gets included in the learned representation.
When we try to learn a lossy compression/representation of data, we can simply construct a decoding distribution that’s capable of modeling the part of information that we don’t want the lossy representation to capture, but, critically, that’s incapable of modelling the information that we do want the lossy representation to capture.
For instance, if we are interested in learning a global representation for 2D images that doesn’t encode information about detailed texture, we can construct a specific factorization of the autoregressive distribution such that it has a small local receptive field as decoding distribution, e.g., plocal(x|z) = ∏ i p(xi|z,xWindowAround(i)). Notice that, as long as xWindowAround(i) is smaller than x<i, plocal(x|z) won’t be able to represent arbitrarily complex distribution over x without dependence on z since the receptive field is limited such that not all distributions over x admit such factorizations. In particular, the receptive field window can be a small rectangle adjacent to a pixel xi and in this case long-range dependency will be encoded in the latent code z. On the other hand, if the true data distribution admits such factorization for a given datum x and dimension i, i.e.
pdata(xi|xWindowAround(i)) = pdata(xi|x<i), then the information preference property discussed in Section 2.2 will apply here, which means that all the information will be encoded in local autoregressive distribution for xi. Local statistics of 2D images like texture will likely be modeled completely by a small local window, whereas global structural information of an images like shapes of objects is long-range dependency that can only be communicated through latent code z. Therefore we have given an example VAE that will produce a lossy compression of 2D images carrying exclusively global information that can’t be modeled locally.
Notice that a global representation is only one of many possible lossy representations that we can construct using this information preference property. For instance, the conditional of an autoregressive distribution might depend on a heavily down-sampled receptive field so that it can only model long-range pattern whereas local high-frequency statistics need to be encoded into the latent code. Hence we have demonstrated that we can achieve explicit placement of information by constraining the receptive field/factorization of an autoregressive distribution that’s used as decoding distribution.
We want to additionally emphasize the information preference property is an asymptotic view in a sense that it only holds when the variational lowerbound can be optimized well. Thus, we are not proposing an alternative to techniques like free bits Kingma et al. (2016) or KL annealing, and indeed they are still useful methods to smooth the optimization problem and used in this paper’s experiments.
3.2 LEARNED PRIOR WITH AUTOREGRESSIVE FLOW
Inefficiency in Bits-Back Coding, i.e., the mismatch between approximate posterior and true posterior, can be exploited to construct a lossy code but it’s still important to minimize such inefficiency to improve overall modeling performance/coding efficiency. We propose to parametrize the prior distribution p(z; θ) with an autoregressive model and show that a type of autoregressive latent code can in theory reduce inefficiency in Bits-Back coding.
It is well-known that limited approximate posteriors impede learning and therefore various expressive posterior approximations have been proposed to improve VAE’s density estimation performance (Turner et al., 2008; Mnih & Gregor, 2014; Salimans et al., 2014; Rezende & Mohamed, 2015; Kingma et al., 2016). One such class of approximate posteriors that has been shown to attain good empirical performance is based on the idea of Normalizing Flow, which is to apply an invertible mapping to a simple random variable, for example a factorized Gaussian as commonly used for q(z|x), in order to obtain a complicated random variable. For an invertible transformation between a simple distribution y and a more flexible z, we know from the change-of-variable technique that log q(z|x) = log q(y|x) − log det dzdy and using q(z|x) as approximate posterior will decrease the coding efficiency gap DKL(q(z|x)||p(z|x)) provided the transformation is sufficiently expressive. Kingma et al. (2016) introduced Inverse Autoregressive Flow, which is a powerful class of such invertible mappings that have simple determinant: zi =
yi−µi(y1:i−1) σi(y1:i−1) , where µi(.) ∈ R, σi(.) ∈ R+ are general functions that can be parametrized by expressive neural networks, such as MADE and PixelCNN variants (Germain et al., 2015; van den Oord et al., 2016a). Inverse autoregressive flow is the inverse/whitening of autoregressive flow: yi = ziσi(y1:i−1) + µi(y1:i−1). We refer interested readers to (Rezende & Mohamed, 2015; Kingma et al., 2016) for in-depth discussions on related topics.
In this paper, we propose to parametrize our learnable prior as an autoregressive flow from some simple noise source like spherical Gaussian. Next, we show that using latent code transformed by autoregressive flow (AF) is equivalent to using inverse autoregressive flow (IAF) approximate posterior, which explains why it can similarly improve Bits-Back Coding efficiency. Moreover, compared with an IAF posterior, an AF prior has a more expressive generative model that essentially “comes for free”.
For an autoregressive flow f , some continuous noise source is transformed into latent code z: z = f( ). Assuming the density function for noise source is u( ), we similarly know that log p(z) = log u( ) + log det d dz .
Simply re-arranging the variational lowerbound for using AF prior reveals that having an AF latent code z is equivalent to using an IAF posterior for that we can interpret as the new latent code:
L(x; θ) = Ez∼q(z|x) [log p(x|z) + log p(z)− log q(z|x)] (12) = Ez∼q(z|x), =f−1(z) [ log p(x|f( )) + log u( ) + log det d
dz − log q(z|x)
] (13)
= Ez∼q(z|x), =f−1(z) log p(x|f( )) + log u( )− (log q(z|x)− log det d dz
)︸ ︷︷ ︸ IAF Posterior  (14) AF prior is the same as IAF posterior along the encoder path, f−1(q(z|x)), but differs along the decoder/generator path: IAF posterior has a shorter decoder path p(x|z) whereas AF prior has a deeper decoder path p(x|f( )). The crucial observation is that AF prior and IAF posterior have the same computation cost under the expectation of z ∼ q(z|x), so using AF prior makes the model more expressive at no training time cost.
4 EXPERIMENTS
In this paper, we evaluate VLAE on 2D images and leave extensions to other forms of data to future work. For the rest of the section, we define a VLAE model as a VAE that uses AF prior and autoregressive decoder. We choose to implement conditional distribution p(x|z) with a smallreceptive-field PixelCNN (van den Oord et al., 2016a), which has been proved to be a scalable autoregressive model.
For evaluation, we use binary image datasets that are commonly used for density estimation tasks: MNIST (LeCun et al., 1998) (both statically binarized 1 and dynamically binarized version (Burda et al., 2015a)), OMNIGLOT (Lake et al., 2013; Burda et al., 2015a) and Caltech-101 Silhouettes (Marlin et al., 2010). All datasets uniformly consist of 28x28 binary images, which allow us to use a unified architecture. VAE networks used in binary image datasets are simple variants of ResNet VAEs described in (Salimans et al., 2014; Kingma et al., 2016). For the decoder, we use a variant of PixelCNN that has 6 layers of masked convolution with filter size 3, which means the window of dependency, xWindowAround(i), is limited to a small local patch. During training, ”free bits” (Kingma et al., 2016) is used improve optimization stability. Experimental setup and hyperparameters are detailed in the appendix. Reported marginal NLL is estimated using Importance Sampling with 4096 samples.
We designed experiments to answer the following questions:
• Can VLAE learn lossy codes that encode global statistics? • Does using AF priors improves upon using IAF posteriors as predicted by theory? • Does using autoregressive decoding distributions improve density estimation performance?
4.1 LOSSY COMPRESSION
First we are interested in whether VLAE can learn a lossy representation/compression of data by using the PixelCNN decoder to model local statistics. We trained VLAE model on Statically Binarized MNIST and the converged model has E[DKL(q(z|x)||p(z))] = 13.3 nats = 19.2 bits, which is the number of bits it uses on average to encode/compress one MNIST image. By comparison, an identical VAE model with factorized decoding distribution will uses on average 37.3 bits in latent code, and this thus indicates that VLAE can learn a lossier compression than a VAE with regular factorized conditional distribution.
The next question is whether VLAE’s lossy compression encodes global statistics and discards local statistics. In Fig 1a, we visualize original images xdata and one random “decompression” xdecompressed from VLAE: z ∼ q(z|xdata),xdecompressed ∼ p(x|z). We observe that none of the
1We use the version provided by Hugo Larochelle.
decompressions is an exact reconstruction of the original image but instead the global structure of the image was encoded in the lossy code z and regenerated. Also worth noting is that local statistics are not preserved but a new set of likely local statistics are generated in the decompressed images: the binary masks are usually different and local styles like stroke width are sometimes slightly different.
However, we remark that the lossy code z doesn’t always capture the kind of global information that we care about and it’s dependent on the type of constraint we put on the decoder. For instance, in Fig 4b, we show decompressions for OMNIGLOT dataset, which has more meaningful variations in small patches than MNIST, and we can observe that semantics are not preserved in some cases. This highlights the need to specify the type of statistics we care about in a representation, which will be different across tasks and datasets, and design decoding distribution accordingly.
4.2 DENSITY ESTIMATION
Next we investigate whether leveraging autoregressive models as latent distribution p(z) and as decoding distribution p(x|z) would improve density estimation performance. To verify whether AF prior is able to improve upon IAF posterior alone, it’s desirable to test this model without using autoregressive decoder but instead using the conventional independent Bernoulli distribution for p(x|z). Hence we use the best performing model from Kingma et al.
(2016) on statically binarized MNIST and make the single modification of replacing the original IAF posterior with an equivalent AF prior, removing the context. As seen in Table 1, VAE with AF prior is outperforming VAE with an equivalent IAF posterior, indicating that the deeper generative model from AF prior is beneficial. A similar gain carries over when an autoregressive decoder is used: on statically binarized MNIST, using AF prior instead of IAF posterior reduces train NLL by 0.8 nat and test NLL by 0.6 nat.
Next we evaluate whether using autoregressive decoding distribution can improve performance and we show in Table 1 that a VLAE model, with AF prior and PixelCNN conditional, is able to outperform a VAE with just AF prior and achieves new state-of-the-art results on statically binarized MNIST.
In addition, we hypothesize that the separation of different types of information, the modeling global structure in latent code and local statistics in PixelCNN, likely has some form of good inductive biases for 2D images. In order to evaluate if VLAE is an expressive density estimator with good inductive biases, we will test a single VLAE model, with the same network architecture, on all binary datasets. We choose hyperparameters manually on statically binarized MNIST and use the same hyperparameters to evaluate on dynamically binarized MNIST, OMNIGLOT and Caltech-101 Silhouettes. We also note that better performance can be obtained if we individually tune hyperparameters for each dataset. As a concrete demonstration, we report the performance of a fine-tuned VLAE on OMNIGLOT dataset in Table 3.
As seen in Table 2,3,4, with the same set of hyperparameters tuned on statically binarized MNIST, VLAE is able to perform well on the rest of datasets, significantly exceeding previous state-ofthe-art results on dynamically binarized MNIST and Caltech-101 Silhouettes and tying statistically with best previous result on OMNIGLOT. In order to isolate the effect of expressive PixelCNN as decoder, we also report performance of the same PixelCNN trained without VAE part under the name “Unconditional Decoder”.
4.3 NATURAL IMAGES: CIFAR10
In addition to binary image datasets, we have applied VLAE to the CIFAR10 dataset of natural images. Density estimation of CIFAR10 images has been a challenging benchmark problem used by many recent generative models and hence is great task to position VLAE among existing methods.
We investigated using ResNet (He et al., 2016) and DenseNet (Huang et al., 2016) as building blocks for VAE networks and observed that DenseNet reduces overfitting. We also propose a new optimization technique that blends the advantages of KL annealing (Serban et al., 2016) and ”free bits” (Kingma et al., 2016) to stabilize learning on this challenging dataset. Detailed experimental setup is described in Appendix.
VLAE is compared to other methods on CIFAR10 in Table 5. We show that VLAE models attain new state-of-the-art performance among other variationally trained latent-variable models. DenseNet VLAE model also outperforms most other tractable likelihood models including Gated PixelCNN and PixelRNN and has results only slightly worse than currently unarchived state-of-the-art PixelCNN++.
We also investigate learning lossy codes on CIFAR10 images. To illustrate how does the receptive field size of PixelCNN decoder influence properties of learned latent codes, we show visualizations of similar VLAE models with receptive fields of different sizes. Specifically we say a receptive field, xWindowAround(i), has size AxB when a pixel xi can depend on the rectangle block of size AxB immediately on top of xi as well as the ⌈ A−1 2 ⌉ pixels immediately to the left of xi. We use this notation to refer to different types of PixelCNN decoders in Figure 3.
From (a)-(c) in Figure 3, we can see that larger receptive fields progressively make autoregressive decoders capture more structural information. In (a), a smaller receptive field tends to preserve rather detailed shape information in the lossy code whereas the latent code only retains rough shape in (c) with a larger receptive field.
It’s interesting to also note that in (a)-(c), oftentimes color information is partially omitted from latent codes and one explanation can be that color is very predictable locally. However, color information can be important to preserve if our task is, for example, object classification. To demonstrate how we can encode color information in the lossy code, we can choose to make PixelCNN decoder depend only on images’ grayscale versions. In other words, instead of choosing the decoder to be plocal(x|z) = ∏ i p(xi|z,xWindowAround(i)), we use a decoder of the form
plocal(x|z) = ∏ i p(xi|z,Grayscale(xWindowAround(i))). In (d) of Figure 3, we visualize lossy codes for a VLAE that has the same receptive field size as (c) but uses a “grayscale receptive field”. We note that the lossy codes in (d) encode roughly the same structural information as those in (c) but generally generate objects that are more recognizable due to the preservation of color information. This serves as one example of how we can design the lossy latent code carefully to encode what’s important and what’s not.
5 RELATED WORK
We investigate a fusion between variational autoencoders with continuous latent variables (Kingma & Welling, 2013; Rezende et al., 2014) and neural autoregressive models. For autoregression, we specifically apply a novel type of architecture where autoregression is realised through a carefully
constructed deep convolutional network, introduced in the PixelCNN model for images (van den Oord et al., 2016a,b). These family of convolutional autoregressive models was further explored, and extended, for audio in WaveNet (Oord et al., 2016), video in Video Pixel Networks (Kalchbrenner et al., 2016b) and language in ByteNet (Kalchbrenner et al., 2016a).
The combination of latent variables with expressive decoder was previously explored using recurrent networks mainly in the context of language modeling (Chung et al., 2015; Bowman et al., 2015; Serban et al., 2016; Fraccaro et al., 2016; Xu & Sun, 2016). Bowman et al. (2015) has also proposed to weaken an otherwise too expressive decoder by dropout to force some information into latent codes.
Concurrent with our work, PixelVAE (Gulrajani et al., 2016) also explored using conditional PixelCNN as a VAE’s decoder and has obtained impressive density modeling results through the use of multiple levels of stochastic units.
Using autoregressive model on latent code was explored in the context of discrete latent variables in DARN (Gregor et al., 2013). Kingma et al. (2016), Kaae Sønderby et al. (2016), Gregor et al. (2016) and Salimans (2016) explored VAE architecture with an explicitly deep autoregressive prior for continuous latent variables, but the autoregressive data likelihood is intractable in those architectures and needs to inferred variationally. In contrast, we use multiple steps of autoregressive flows that has exact likelihood and analyze the effect of using expressive latent code.
Optimization challenges for using (all levels of) continuous latent code were discussed before and practical solutions were proposed (Bowman et al., 2015; Kaae Sønderby et al., 2016; Kingma et al., 2016). In this paper, we present a complementary perspective on when/how should the latent code be used by appealing to a Bits-Back interpretation of VAE.
Learning a lossy compressor with latent variable model has been investigated with ConvDRAW (Gregor et al., 2016). It learns a hierarchy of latent variables and just using high-level latent variables will result in a lossy compression that performs similarly to JPEG. Our model similarly learns a lossy compressor but it uses an autoregressive model to explicitly control what kind of information should be lost in compression.
6 CONCLUSION
In this paper, we analyze the condition under which the latent code in VAE should be used, i.e. when does VAE autoencode, and use this observation to design a VAE model that’s a lossy compressor of observed data. At modeling level, we propose two complementary improvements to VAE that are shown to have good empirical performance.
VLAE has the appealing properties of controllable representation learning and improved density estimation performance but these properties come at a cost: compared with VAE models that have simple prior and decoder, VLAE is slower at generation due to the sequential nature of autoregressive model.
Moving forward, we believe it’s exciting to extend this principle of learning lossy codes to other forms of data, in particular those that have a temporal aspect like audio and video. Another promising direction is to design representations that contain only information for downstream tasks and utilize those representations to improve semi-supervised learning.
A DETAILED EXPERIMENT SETUP FOR BINARY IMAGES
For VAE’s encoder and decoder, we use the same ResNet (He et al., 2015) VAE architecture as the one used in IAF MNIST experiment (Kingma et al., 2016). The only difference is that the decoder network now, instead of outputing a 28x28x1 spatial feature map to specify the mean of a factorized bernoulli distribution, outputs a 28x28x4 spatial feature map that’s concatenated with the original binary image channel-wise, forming a 28x28x5 feature map that’s then fed through a typical masked PixelCNN (van den Oord et al., 2016a). As such even though the PixelCNN conditions on the latent code, we don’t call it a Conditional PixelCNN because it doesn’t use the specific architecture that was proposed in van den Oord et al. (2016b). For the PixelCNN, it has 6 masked convolution layers with 12 3x3 filters organized in ResNet blocks, and it has 4 additional 1x1 convolution ResNet block between every other masked convolution layer to increase processing capacity since it employs fewer masked convolutions than usual. All the masked convolution layer have their weights tied to reduce overfitting on statically binarized MNIST, and untying the weights will increase performance for other datasets. Experiments are tuned on the validation set and then final experiment was run with train and validation set, with performance evaluated with test set. Exponential Linear Units (Clevert et al., 2015) are used as activation functions in both VAE network and PixelCNN network. Weight normalization is everywhere with data-dependent initialization (Salimans & Kingma, 2016).
A latent code of dimension 64 was used. For AF prior, it’s implemented with MADE (Germain et al., 2015) as detailed in Kingma et al. (2016). We used 4 steps of autoregressive flow and each flow is implemented by a 3-layer MADE that has 640 hidden units and uses Relu (Nair & Hinton, 2010) as activation functions. Differing from the practice of Kingma et al. (2016), we use mean-only autoregressive flow, which we found to be more numerically stable.
In terms of training, Adamax (Kingma & Ba, 2014) was used with a learning rate of 0.002. 0.01 nats/data-dim free bits (Kingma et al., 2016) was found to be effective in dealing with the problem of all the latent code being ignored early in training. Polyak averaging (Polyak & Juditsky, 1992) was used to compute the final parameters, with α = 0.998.
All experiments are implemented using TensorFlow (Abadi et al., 2016).
B ADDITIONAL EXPERIMENT SETUP FOR CIFAR10
Latent codes are represented by 16 feature maps of size 8x8, and this choice of spatial stochastic units are inspired by ResNet IAF VAE (Kingma et al., 2016). Prior distribution is factorized Gaussian noise transformed by 6 autoregressive flows, each of which is implemented by a PixelCNN (van den Oord et al., 2016a) with 2 hidden layers and 128 feature maps. Between every other autoregressive flow, the ordering of stochastic units is reversed.
ResNet VLAE has the following structure for encoder: 2 ResNet blocks, Conv w/ stride=2, 2 ResNet blocks, Conv w/ stride=2, 3 ResNet blocks, 1x1 convolution and has a symmetric decoder. Channel size = 48 for 32x32 feature maps and 96 for other feature maps. DenseNet VLAE follows a similar structure: replacing 2 ResNet blocks with one DenseNet block of 3 steps and each step produces a certain number of feature maps such that at the end of a block, the concatenated feature maps is slightly more than the ResNet VLAE at the same stage.
Conditional PixelCNN++ (Salimans et al., 2017) is used as the decoder. Specifically the channelautoregressive variant is used to ensure there is sufficient capacity even when the receptive field is small. Specifically, the decoder PixelCNN has 4 blocks of 64 feature maps where each block is conditioned on previous blocks with Gated ResNet connections and hence the PixelCNN decoders we use are shallow but very wide. For 4x2 receptive field experiment, we use 1 layer of vertical stack convolutions and 2 layers of horizontal stack convolutions; for 5x3 receptive field experiment, we use 2 layers of vertical stack convolutions and 2 layers of horizontal stack convolutions; For 5x3 receptive field experiment, we use 2 layers of vertical stack convolutions and 2 layers of horizontal stack convolutions; For 7x4 receptive field experiment, we use 3 layers of vertical stack convolutions and 3 layers of horizontal stack convolutions; for 7x4 Grayscale experiment, we transform RGB
images into gray-scale images via this specific transformation: (0.299∗R)+(0.587G)+(0.114B). Best density estimation result is obtained with 7x4 receptive field experiments.
C SOFT FREE BITS
”Free bits” was a technique proposed in (Kingma et al., 2016) where K groups of stochastic units are encouraged to be used through the following surrogate objective:
L̃λ = Ex∼M [ Eq(z|x) [log p(x|z)] ] − K∑ j=1 maximum(λ,Ex∼M [DKL(q(zj |x)||p(zj))])
This technique is easy to use since it’s usually easy to determine the minimum number of bits/nats, λ, stochastic units need to encode. Choosing λ is hence easier than setting a fixed KL annealing schedule (Serban et al., 2016).
On the other hand, Kl annealing has the benefit of the surrogate objective will smoothly become the true objective, the variational lower bound where as ”free bits” has a sharp transition at the boundary. Therefore, we propose to still use λ as hyperparameter to specify at least λ nats should be used but try to change the optimization objective as slowly as possible:
LSoftFreeBits(x; θ) =Eq(z|x) [log p(x|z)]− γDKL(q(z|x)||p(z)) where 0 < γ ≤ 1. And we make the optimization smoother by changing γ slowly online to make sure at least λ nats are used: when Kl is too much higher than λ (we experimented wide range of thresholds from 3% to 30%, all of which yield improved results, and we tend to use 5% us a threshold), γ is increased, and when Kl lower than λ, γ is decreased to encourage information flow.
We found it sufficient to increase/decrease at 10% increment and didn’t further tune this parameter.
D AUTOREGRESSIVE DECODER WITHOUT AUTOREGRESSIVE PRIOR
In this section, we investigate the scenario of just using an autoregressive decoder without using an autoregressive prior. We compare the exact same model in three configurations: 1) using smallreceptive-field PixelCNN as an unconditional density estimator; 2) using small-receptive-field as a decoder in a VAE with Gaussian latent variables; 3) replacing Gaussian latent variables with autoregressive flow latent variables in 2).
In Table 1, we can observe that each step of modification improves density estimation performance. In addition, using an autoregressive latent code makes the latent code transmit more information as shown in the difference of E[DKL(q(z|x)||p(z))].
E CIFAR10 GENERATED SAMPLES

This paper introduces the notion of a "variational lossy autoencoder", where a powerful autoregressive conditional distribution on the inputs x given the latent code z is crippled in a way that forces it to use z in a meaningful way. Its three main contributions are:
(1) It gives an interesting information-theoretical insight as to why VAE-type models don't tend to take advantage of their latent representation when the conditional distribution on x given z is powerful enough.
(2) It shows that this insight can be used to efficiently train VAEs with powerful autoregressive conditional distributions such that they make use of the latent code.
(3) It presents a powerful way to parametrize the prior in the form of an autoregressive flow transformation which is equivalent to using an inverse autoregressive flow transformation on the approximate posterior.
By itself, I think the information-theoretical explanation of why VAEs do not use their latent code when the conditional distribution on x given z is powerful enough constitutes an excellent addition to our understanding of VAE-related approaches.
However, the way this intuition is empirically evaluated is a bit weak. The "crippling" method used feels hand-crafted and very task-dependent, and the qualitative evaluation of the "lossyness" of the learned representation is carried out on three datasets (MNIST, OMNIGLOT and Caltech-101 Silhouettes) which feature black-and-white images with little-to-no texture. Figures 1a and 2a do show that reconstructions discard low-level information, as observed in the slight variations in strokes between the input and the reconstruction, but such an analysis would have been more compelling with more complex image datasets. Have the authors tried applying VLAE to such datasets?
I think the Caltech101 Silhouettes benchmark should be treated with caution, as no comparison is made against other competitive approaches like IAF VAE, PixelRNN and Conv DRAW. This means that VLAE significantly outperforms the state-of-the-art in only one of the four settings examined.
A question which is very relevant to this paper is "Does a latent representation on top of an autoregressive model help improve the density modeling performance?" The paper touches this question, but very briefly: the only setting in which VLAE is compared against recent autoregressive approaches shows that it wins against PixelRNN by a small margin.
The proposal to transform the latent code with an autoregressive flow which is equivalent to parametrizing the approximate posterior with an inverse autoregressive flow transformation is also interesting. There is, however, one important distinction to be made between the two approaches: in the former, the prior over the latent code can potentially be very complex whereas in the latter the prior is limited to be a simple, factorized distribution.
It is not clear to me that having a very powerful prior is necessarily a good thing from a representation learning point of view: oftentimes we are interested in learning a representation of the data distribution which is untangled and composed of roughly independent factors of variation. The degree to which this can be achieved using something as simple as a spherical gaussian prior is up for discussion, but finding a good balance between the ability of the prior to fit the data and its usefulness as a high-level representation certainly warrants some thought. I would be interested in hearing the authors' opinion on this.
Overall, the paper introduces interesting ideas despite the flaws outlined above, but weaknesses in the empirical evaluation prevent me from recommending its acceptance.
UPDATE: The rating has been revised to a 7 following the authors' reply.