STICK-BREAKING VARIATIONAL AUTOENCODERS
We extend Stochastic Gradient Variational Bayes to perform posterior inference for the weights of Stick-Breaking processes. This development allows us to define a Stick-Breaking Variational Autoencoder (SB-VAE), a Bayesian nonparametric version of the variational autoencoder that has a latent representation with stochastic dimensionality. We experimentally demonstrate that the SB-VAE, and a semisupervised variant, learn highly discriminative latent representations that often outperform the Gaussian VAE’s.
1 INTRODUCTION
Deep generative models trained via Stochastic Gradient Variational Bayes (SGVB) (Kingma & Welling, 2014a; Rezende et al., 2014) efficiently couple the expressiveness of deep neural networks with the robustness to uncertainty of probabilistic latent variables. This combination has lead to their success in tasks ranging from image generation (Gregor et al., 2015; Rezende et al., 2016) to semi-supervised learning (Kingma et al., 2014; Maaløe et al., 2016) to language modeling (Bowman et al., 2016). Various extensions to SGVB have been proposed (Burda et al., 2016; Maaløe et al., 2016; Salimans et al., 2015), but one conspicuous absence is an extension to Bayesian nonparametric processes. Using SGVB to perform inference for nonparametric distributions is quite attractive. For instance, SGVB allows for a broad class of non-conjugate approximate posteriors and thus has the potential to expand Bayesian nonparametric models beyond the exponential family distributions to which they are usually confined. Moreover, coupling nonparametric processes with neural network inference models equips the networks with automatic model selection properties such as a selfdetermined width, which we explore in this paper.
We make progress on this problem by first describing how to use SGVB for posterior inference for the weights of Stick-Breaking processes (Ishwaran & James, 2001). This is not a straightforward task as the Beta distribution, the natural choice for an approximate posterior, does not have the differentiable non-centered parametrization that SGVB requires. We bypass this obstacle by using the little-known Kumaraswamy distribution (Kumaraswamy, 1980).
Using the Kumaraswamy as an approximate posterior, we then reformulate two popular deep generative models—the Variational Autoencoder (Kingma & Welling, 2014a) and its semi-supervised variant (model M2 proposed by Kingma et al. (2014))—into their nonparametric analogs. These models perform automatic model selection via an infinite capacity hidden layer that employs as many stick segments (latent variables) as the data requires. We experimentally show that, for datasets of natural images, stick-breaking priors improve upon previously proposed deep generative models by having a latent representation that better preserves class boundaries and provides beneficial regularization for semi-supervised learning.
2 BACKGROUND
We begin by reviewing the relevant background material on Variational Autoencoders (Kingma & Welling, 2014a), Stochastic Gradient Variational Bayes (also known as Stochastic Backpropagation) (Kingma & Welling, 2014a; Rezende et al., 2014), and Stick-Breaking Processes (Ishwaran & James, 2001).
2.1 VARIATIONAL AUTOENCODERS
A Variational Autoencoder (VAE) is model comprised of two multilayer perceptrons: one acts as a density network (MacKay & Gibbs, 1999) mapping a latent variable zi to an observed datapoint xi, and the other acts as an inference model (Salimans & Knowles, 2013) performing the reverse mapping from xi to zi. Together the two form a computational pipeline that resembles an unsupervised autoencoder (Hinton & Salakhutdinov, 2006). The generative process can be written mathematically as
zi ∼ p(z) , xi ∼ pθ(x|zi) (1) where p(z) is the prior and pθ(x|zi) is the density network with parameters θ. The approximate posterior of this generative process, call it qφ(z|xi), is then parametrized by the inference network (with parameters φ). In previous work (Kingma & Welling, 2014a; Rezende & Mohamed, 2015; Burda et al., 2016; Li & Turner, 2016), the prior p(z) and variational posterior have been marginally Gaussian.
2.2 STOCHASTIC GRADIENT VARIATIONAL BAYES
The VAE’s generative and variational parameters are estimated by Stochastic Gradient Variational Bayes (SGVB). SGVB is distinguished from classical variational Bayes by it’s use of differentiable Monte Carlo (MC) expectations. To elaborate, consider SGVB’s approximation of the usual evidence lowerbound (ELBO) (Jordan et al., 1999):
L̃(θ,φ;xi) = 1
S S∑ s=1 log pθ(xi|ẑi,s)−KL(qφ(zi|xi)||p(z)) (2)
for S samples of zi and where KL is the Kullback-Leibler divergence. An essential requirement of SGVB is that the latent variable be represented in a differentiable, non-centered parametrization (DNCP) (Kingma & Welling, 2014b); this is what allows the gradients to be taken through the MC expectation, i.e.:
∂
∂φ S∑ s=1 log pθ(xi|ẑi,s) = S∑ s=1 ∂ ∂ẑi,s log pθ(xi|ẑi,s) ∂ẑi,s ∂φ .
In other words, z must have a functional form that deterministically exposes the variational distribution’s parameters and allows the randomness to come from draws from some fixed distribution. Location-scale representations and inverse cumulative distribution functions are two examples of DNCPs. For instance, the VAE’s Gaussian latent variable (with diagonal covariance matrix) is represented as ẑi = µ + σ where ∼ N(0,1).
2.3 STICK-BREAKING PROCESSES
Lastly, we define stick-breaking processes with the ultimate goal of using their weights for the VAE’s prior p(z). A random measure is referred to as a stick-breaking prior (SBP) (Ishwaran & James, 2001) if it is of the form G(·) = ∑∞ k=1 πkδζk where δζk is a discrete measure concentrated at ζk ∼ G0, a draw from the base distribution G0 (Ishwaran & James, 2001). The πks are random weights independent of G0, chosen such that 0 ≤ πk ≤ 1, and ∑ k πk = 1 almost surely. SBPs have been termed as such because of their constructive definition known as the stick-breaking process (Sethuraman, 1994). Mathematically, this definition implies that the weights can be drawn according to the following iterative procedure:
πk = { v1 if k = 1 vk ∏ j<k(1− vj) for k > 1
(3)
where vk ∼ Beta(α, β). When vk ∼ Beta(1, α0), then we have the stick-breaking construction for the Dirichlet Process (Ferguson, 1973). In this case, the name for the joint distribution over the infinite sequence of stick-breaking weights is the Griffiths, Engen and McCloskey distribution with concentration parameter α0 (Pitman, 2002): (π1, π2, . . .) ∼ GEM(α0).
3 SGVB FOR GEM RANDOM VARIABLES
Having covered the relevant background material, we now discuss the first contribution of this paper, using Stochastic Gradient Variational Bayes for the weights of a stick-breaking process. Inference for the random measure G(·) is an open problem that we leave to future work. We focus on performing inference for just the series of stick-breaking weights, which we will refer to as GEM random variables after their joint distribution.
3.1 COMPOSITION OF GAMMA RANDOM VARIABLES
In the original SGVB paper, Kingma & Welling (2014a) suggest representing the Beta distribution as a composition of Gamma random variables by using the fact v ∼ Beta(α, β) can be sampled by drawing Gamma variables x ∼ Gamma(α, 1), y ∼ Gamma(β, 1) and composing them as v = x/(x + y). However, this representation still does not admit a DNCP as the Gamma distribution does not have one with respect to its shape parameter. Knowles (2015) suggests that when the shape parameter is near zero, the following asymptotic approximation of the inverse CDF is a suitable DNCP:
F−1(û) ≈ (ûaΓ(a)) 1 a
b (4)
for û ∼ Uniform(0, 1), shape parameter a, and scale parameter b. This approximation becomes poor as a increases, however, and Knowles recommends a finite difference approximation of the inverse CDF when a ≥ 1.
3.2 THE KUMARASWAMY DISTRIBUTION
Another candidate posterior is the little-known Kumaraswamy distribution (Kumaraswamy, 1980). It is a two-parameter continuous distribution also on the unit interval with a density function defined as
Kumaraswamy(x; a, b) = abxa−1(1− xa)b−1 (5)
for x ∈ (0, 1) and a, b > 0. In fact, if a = 1 or b = 1 or both, the Kumaraswamy and Beta are equivalent, and for equivalent parameter settings, the Kumaraswamy resembles the Beta albeit with higher entropy. The DNCP we desire is the Kumaraswamy’s closed-form inverse CDF. Samples can be drawn via the inverse transform:
x ∼ (1− u 1b ) 1a where u ∼ Uniform(0, 1). (6)
Not only does the Kumaraswamy make sampling easy, its KL-divergence from the Beta can be closely approximated in closed-form (for ELBO computation).
3.2.1 GAUSS-LOGIT PARAMETRIZATION
Another promising parametrization is inspired by the Probit Stick-Breaking Process (Rodriguez & Dunson, 2011). In a two-step process, we can draw a Gaussian and then use a squashing function to map it on (0, 1):
v̂k = g(µk + σk ) (7)
where ∼ N(0, 1). In the Probit SBP, g(·) is taken to be the Gaussian CDF, and it is chosen as such for posterior sampling considerations. This choice is impractical for our purposes, however, since the Gaussian CDF does not have a closed form. Instead, we use the logistic function g(x) = 1/(1+e−x).
4 STICK-BREAKING VARIATIONAL AUTOENCODERS
Given the discussion above, we now propose the following novel modification to the VAE. Instead of drawing the latent variables from a Gaussian distribution, we draw them from the GEM distribution, making the hidden representation an infinite sequence of stick-breaking weights. We term this model a Stick-Breaking Variational Autoencoder (SB-VAE) and below detail the generative and inference processes implemented in the decoding and encoding models respectively.
4.1 GENERATIVE PROCESS
The generative process is nearly identical to previous VAE formulations. The crucial difference is that we draw the latent variable from a stochastic process, the GEM distribution. Mathematically, the hierarchical formulation is written as
πi ∼ GEM(α0) , xi ∼ pθ(xi|πi) (8)
where πi is the vector of stick-breaking weights and α0 is the concentration parameter of the GEM distribution. The likelihood model pθ(xi|πi) is a density network just as described in Section 2.1.
4.2 INFERENCE
The inference process—how to draw πi ∼ qφ(πi|zi)—requires modification from the standard VAE’s in order to sample from the GEM’s stick-breaking construction. Firstly, an inference network computes the parameters of K fraction distributions and samples values vi,k according to one of the parametrizations in Section 3. Next, a linear-time operation composes the stick segments from the sampled fractions:
πi = (πi,1, πi,2, . . . , πi,K) = vi,1, vi,2(1− vi,1), . . . ,K−1∏ j=1 (1− vi,j)  . (9) The computation path is summarized in Figure 1 (c) with arrows denoting the direction of feedforward computation. The gray blocks represent any deterministic function that can be trained with gradient descent—i.e. one or more neural network layers. Optimization of the SB-VAE is done just as for the VAE, by optimizing Equation 2 w.r.t. φ and θ. The KL divergence term can be computed (or closely approximated) in closed-form for all three parametrizations under consideration; the Kumaraswamy-to-Beta KL divergence is given in the appendix.
An important detail is that the Kth fraction vi,K is always set to one to ensure the stick segments sum to one. This truncation of the variational posterior does not imply that we are using a finite dimensional prior. As explained by Blei & Jordan (2006), the truncation level is a variational parameter and not part of the prior model specification. Truncation-free posteriors have been proposed, but these methods use split-and-merge steps (Hughes et al., 2015) or collapsed Gibbs sampling, both of which are not applicable to the models we consider. Nonetheless, because SGVB imposes few limitations on the inference model, it is possible to have an untruncated posterior. We conducted exploratory experiments using a truncation-free posterior by adding extra variational parameters in an on-line fashion, initializing new weights if more than 1% of the stick remained unbroken. However, we found this made optimization slower without any increase in performance.
5 SEMI-SUPERVISED MODEL
We also propose an analogous approach for the semi-supervised relative of the VAE, the M2 model described by Kingma et al. (2014). A second latent variable yi is introduced that represents a class label. Its distribution is the categorical one: qφ(yi|xi) = Cat(y|gy(xi)) where gy is a non-linear function of the inference network. Although y’s distribution is written as independent of z, the two share parameters within the inference network and thus act to regularize one another. We assume the same factorization of the posterior and use the same objectives as in the finite dimensional version (Kingma et al., 2014). Since yi is present for some but not all observations, semi-supervised DGMs need to be trained with different objectives depending on whether the label is present or not. If the label is present, following Kingma et al. (2014) we optimize
J̃ (θ,φ;xi, yi) = 1
S S∑ s=1 log pθ(xi|πi,s, yi)−KL(qφ(πi|xi)||p(πi;α0)) + log qφ(yi|xi) (10)
where log qφ(yi|xi) is the log-likelihood of the label. And if the label is missing, we optimize
J̃ (θ,φ;xi) = 1
S S∑ s=1 ∑ yj qφ(yj |xi) [log pθ(xi|πi,s, yj)] + H[qφ(y|xi)]
−KL(qφ(πi|xi)||p(πi;α0)) (11)
where H[qφ(yi|xi)] is the entropy of y’s variational distribution.
6 RELATED WORK
To the best of our knowledge, neither SGVB nor any of the other recently proposed amortized VI methods (Kingma & Welling, 2014b; Rezende & Mohamed, 2015; Rezende et al., 2014; Tran et al., 2016) have been used in conjunction with BNP priors. There has been work on using nonparametric posterior approximations—in particular, the Variational Gaussian Process (Tran et al., 2016)—but in that work the variational distribution is nonparametric, not the generative model. Moreover, we are not aware of prior work that uses SGVB for Beta (or Beta-like) random variables1.
In regards to the autoencoder implementations we describe, they are closely related to the existing work on representation learning with adaptive latent factors—i.e. where the number of latent dimensions grows as the data necessitates. The best known model of this kind is the infinite binary latent feature model defined by the Indian Buffet Process (Ghahramani & Griffiths, 2005); but its discrete latent variables prevent this model from admitting fully differentiable inference. Recent work that is much closer in spirit is the Infinite Restricted Boltzmann Machine (iRBM) (Côté & Larochelle, 2016), which has gradient-based learning, expands its capacity by adding hidden units, and induces a similar ordering on latent factors. The most significant difference between our SB-VAE and the iRBM is that the latter’s nonparametric behavior arises from a particular definition of the energy function of the Gibbs distribution, not from an infinite dimensional Bayesian prior. Lastly, our training procedure bears some semblance to Nested Dropout (Rippel et al., 2014), which removes all hidden units at an index lower than some threshold index. The SB-VAE can be seen as performing soft nested dropout since the latent variable values decrease as their index increases.
7 EXPERIMENTS
We analyze the behavior of the three parametrizations of the SB-VAE and examine how they compare to the Gaussian VAE. We do this by examining their ability to reconstruct the data (i.e. density estimation) and to preserve class structure. Following the original DGM papers (Kingma et al., 2014; Kingma & Welling, 2014a; Rezende et al., 2014), we performed unsupervised and semi-supervised
1During preparation of this draft, the work of Ruiz et al. (2016) on the Generalized Reparametrization Gradient was released (on 10/7/16), which can be used for Beta random variables. We plan to compare their technique to our proposed use of the Kumaraswamy in subsequent drafts.
tasks on the following image datasets: Frey Faces2, MNIST, MNIST+rot, and Street View House Numbers3 (SVHN). MNIST+rot is a dataset we created by combining MNIST and rotated MNIST4 for the purpose of testing the latent representation under the conjecture that the rotated digits should use more latent variables than the non-rotated ones.
Complete implementation and optimization details can be found in the appendix and code repository5. In all experiments, to best isolate the effects of Gaussian versus stick-breaking latent variables, the same architecture and optimization hyperparameters were used for each model. The only difference was in the prior: p(z) = N(0,1) for Gaussian latent variables and p(v) = Beta(1, α0) (Dirichlet process) for stick-breaking latent variables. We cross-validated the concentration parameter over the range α0 ∈ {1, 3, 5, 8}. The Gaussian model’s performance potentially could have been improved by cross validating its prior variance. However, the standard Normal prior is widely used as a default choice (Bowman et al., 2016; Gregor et al., 2015; Kingma et al., 2014; Kingma & Welling, 2014a; Rezende et al., 2014; Salimans et al., 2015), and our goal is to experimentally demonstrate a stick-breaking prior is a competitive alternative.
7.1 UNSUPERVISED
We first performed unsupervised experiments testing each model’s ability to recreate the data as well as preserve class structure (without having access to labels). The inference and generative models both contained one hidden layer of 200 units for Frey Faces and 500 units for MNIST and MNIST+rot. For Frey Faces, the Gauss VAE had a 25 dimensional (factorized) distribution, and we set the truncation level of the SB-VAE also to K = 25, so the SB-VAE could use only as many latent variables as the Gauss VAE. For the MNIST datasets, the latent dimensionality/truncation-level was set at 50. Cross-validation chose α0 = 1 for Frey Faces and α0 = 5 for both MNISTs.
Density Estimation. In order to show each model’s optimization progress, Figure 2 (a), (b), and (c) report test expected reconstruction error (i.e. the first term in the ELBO) vs training progress (epochs) for Frey Faces, MNIST, and MNIST+rot respectively. Optimization proceeds much the same in both models except that the SB-VAE learns at a slightly slower pace for all parametrizations. This is not too surprising since the recursive definition of the latent variables likely causes coupled gradients.
We compare the final converged models in Table 1, reporting the marginal likelihood of each model via the MC approximation log p(xi) ≈ log 1S ∑ s p(xi|ẑi,s)p(ẑi,s)/q(ẑi,s) using 100 samples. The Gaussian VAE has a better likelihood than all stick-breaking implementations (∼ 96 vs ∼ 98). Between the stick-breaking parametrizations, the Kumaraswamy outperforms both the Gamma and Gauss-Logit on both datasets, which is not surprising given the others’ flaws (i.e. the Gamma is approximate, the Gauss-Logit is restricted). Given this result, we used the Kumaraswamy parametriza-
2Available at http://www.cs.nyu.edu/~roweis/data.html 3Available at http://ufldl.stanford.edu/housenumbers/ 4Available at http://www.iro.umontreal.ca/~lisa/twiki/bin/view.cgi/Public/MnistVariations 5Theano implementations available at https://github.com/enalisnick/stick-breaking_
dgms
tion for all subsequently reported experiments. Note that the likelihoods reported are worse than the ones reported by Burda et al. (2016) because our training set consisted of 50k examples whereas theirs contained 60k (training and validation).
We also investigated whether the SB-VAE is using its adaptive capacity in the manner we expect, i.e., the SB-VAE should use a larger latent dimensionality for the rotated images in MNIST+rot than it does for the non-rotated ones. We examined if this is the case by tracking how many ‘breaks’ it took the model to deconstruct 99% of the stick. On average, the rotated images in the training set were represented by 28.7 dimensions and the non-rotated by 27.4. Furthermore, the rotated images used more latent variables in eight out of ten classes. Although the difference is not as large as we were expecting, it is statistically significant. Moreover, the difference is made smaller by the non-rotated one digits, which use 32 dimensions on average, the most for any class. The non-rotated average decreases to 26.3 when ones are excluded.
Figure 3 (a) shows MNIST digits drawn from the SB-VAE by sampling from the prior—i.e. vk ∼ Beta(1, 5), and Figure 3 (b) shows Gauss VAE samples for comparison. SB-VAE samples using all fifty dimensions of the truncated posterior are shown in the bottom block. Samples from Dirichlets constrained to a subset of the dimensions are shown in the two columns in order to test that the latent features are concentrating onto lower-dimensional simplices. This is indeed the case: adding a latent variable results in markedly different but still coherent samples. For instance, the second and third dimensions seem to capture the 7-class, the fourth and fifth the 6-class, and the eighth the 5-class. The seventh dimension seems to model notably thick digits.
Discriminative Qualities. The discriminative qualities of the models’ latent spaces are assessed by running a k-Nearest Neighbors classifier on (sampled) MNIST latent variables. Results are shown in the table in Figure 4 (a). The SB-VAE exhibits conspicuously better performance than the Gauss VAE at all choices of k, which suggests that although the Gauss VAE converges to a better likelihood, the SB-VAE’s latent space better captures class structure. We also report results for two Gaussian mixture VAEs: Dilokthanakul et al. (2016)’s Gaussian mixture Variational Autoencoder (GMVAE)
k=3 k=5 k=10 SB-VAE 9.34 8.65 8.90 DLGMM 9.14 8.38 8.42 Gauss VAE 28.4 20.96 15.33 Raw Pixels 2.95 3.12 3.35 GMVAE6 — 8.96 —
(a) MNIST: Test error for kNN on latent space
t-SNE EmEeddLng of StLcN-BreaNLng VAE'V Latent SSace
(b) MNIST SB-VAE
t-SNE EmEeddLng of StLcN-BreaNLng VAE'V Latent SSace
(c) MNIST Gauss VAE
Figure 4: Subfigure (a) shows results of a kNN classifier trained on the latent representations produced by each model. Subfigures (b) and (c) show t-SNE projections of the latent representations learned by the SB-VAE and Gauss VAE respectively.
(a) Gauss VAE (b) Stick-Breaking VAE
Figure 5: Sparsity in the latent representation vs sparsity in the decoder network. The Gaussian VAE ‘turns off’ unused latent dimensions by setting the outgoing weights to zero (in order to dispel the sampled noise). The SB VAE, on the other hand, also has sparse representations but without decay of the associated decoder weights.
and Nalisnick et al. (2016)’s Deep Latent Gaussian Mixture Model (DLGMM). The GMVAE6 has sixteen mixture components and the DLGMM has five, and hence both have many more parameters than the SB-VAE. Despite the SB-VAE’s lower capacity, we see that its performance is competitive to the mixture VAEs’ (8.65 vs 8.38/8.96).
The discriminative qualities of the SB-VAE’s latent space are further supported by Figures 4 (b) and (c). t-SNE was used to embed the Gaussian (c) and stick-breaking (b) latent MNIST representations into two dimensions. Digit classes (denoted by color) in the stick-breaking latent space are clustered with noticeably more cohesion and separation.
Combating Decoder Pruning. The ‘component collapsing’ behavior of the variational autoencoder has been well noted (Maaløe et al., 2016): the model will set to zero the outgoing weights of latent variables that remain near the prior. Figure 5 (a) depicts this phenomenon for the Gauss VAE by plotting the KL divergence from the prior and outgoing decoder weight norm for each latent dimension. We see the weights are only nonzero in the dimensions in which there is posterior deviation. Ostensibly the model receives only sampling noise from the dimensions that remain at the prior, and setting the decoder weights to zero quells this variance. While the behavior of the Gauss VAE is not necessarily improper, all examples are restricted to pass through the same latent variables. A sparse-coded representation—one having few active components per example (like the Gauss VAE) but diversity of activations across examples–would likely be better.
We compare the activation patterns against the sparsity of the decoder for the SB-VAE in Figure 5 (b). Since KL-divergence doesn’t directly correspond to sparsity in stick-breaking latent variables like it does for Gaussian ones, the black lines denote the average activation value per dimension. Similarly to (a), blue lines denoted the decoder weight norms, but they had to be down-scaled by a factor of 100 so they could be visualized on the same plot. The SB-VAE does not seem to have any component collapsing, which is not too surprising since the model can set latent variables to zero to deactivate decoder weights without being in the heart of the prior. We conjecture that this increased capacity is
6The GMVAE’s evaluation is different from performing kNN. Rather, test images are assigned to clusters and whole clusters are given a label. Thus results are not strictly comparable but the ultimate goal of unsupervised MNIST classification is the same.
the reason stick-breaking variables demonstrate better discriminative performance in many of our experiments.
7.2 SEMI-SUPERVISED
We also performed semi-supervised classification, replicating and extending the experiments in the original semi-supervised DGMs paper (Kingma et al., 2014). We used the MNIST, MNIST+rot, and SVHN datasets and reduced the number of labeled training examples to 10%, 5%, and 1% of the total training set size. Labels were removed completely at random and as a result, class imbalance was all but certainly introduced. Similarly to the unsupervised setting, we compared DGMs with stick-breaking (SB-DGM) and Gaussian (Gauss-DGM) latent variables against one another and a baseline k-Nearest Neighbors classifier (k=5). We used 50 for the latent variable dimensionality / truncation level. The MNIST networks use one hidden layer of 500 hidden units. The MNIST+rot and SVHN networks use four hidden layers of 500 units in each. The last three hidden layers have identity function skip-connections. Cross-validation chose α0 = 5 for MNISTs and α0 = 8 for SVHN.
Quantitative Evaluation. Table 2 shows percent error on a test set when training with the specified percentage of labeled examples. We see the the SB-DGM performs markedly better across almost all experiments. The Gauss-DGM achieves a superior error rate only on the easiest tasks: MNIST with 10% and 5% of the data labeled.
8 CONCLUSIONS
We have described how to employ the Kumaraswamy distribution to extend Stochastic Gradient Variational Bayes to the weights of stick-breaking Bayesian nonparametric priors. Using this development we then defined deep generative models with infinite dimensional latent variables and showed that their latent representations are more discriminative than those of the popular Gaussian variant. Moreover, the only extra computational cost is in assembling the stick segments, a linear operation on the order of the truncation size. Not only are the ideas herein immediately useful as presented, they are an important first-step to integrating black box variational inference and Bayesian nonparametrics, resulting in scalable models that have differentiable control of their capacity. In particular, applying SGVB to full Dirichlet processes with non-trivial base measures is an interesting next step. Furthermore, differentiable stick-breaking has the potential to increase the dynamism and adaptivity of neural networks, a subject of recent interest (Graves, 2016), in a probabilistically principled way.
ACKNOWLEDGEMENTS
Many thanks to Marc-Alexandre Côté and Hugo Larochelle for helpful discussions. This work was supported in part by NSF award number IIS-1320527.

This paper presents an approach which modifies the variational auto-encoder (VAE) framework so as to use stochastic latent dimensionality. This is achieved by using an inherently infinite prior, the stick-breaking process. This is coupled with inference tailored to this model, specifically the Kumaraswamy distribution as an approximate variational posterior. The resulting model is named the SB-VAE which also has a semi-supervised extension, in similar vein to the original VAE paper.
There's a lot of interest in VAEs these days; many lines of work seek to achieve automatic "black-box" inference in these models. For example, the authors themselves mention parallel work by Blei's lab (also others) towards this direction. However, there's a lot of merit in investigating more bespoke solutions to new models, which is what the authors are doing in this paper. Indeed, a (useful) side-effect of providing efficient inference for the SB-VAE is drawing attention to the use of the Kumaraswamy distribution which hasn't been popular in ML.
Although the paper is in general well structured, I found it confusing at parts. I think the major source of confusion comes from the fact that the model specification and model inference are discussed in a somehow mixed manner. The pre-review questions clarified most parts.
I have two main concerns regarding the methodology and motivation of this paper. Firstly, conditioning the model directly on the stick-breaking weights seems a little odd. I initially thought that there was some mixture probabilistic model involved, but this is not the case. To be fair, the authors discuss about this issue (which became clearer to me after the pre-review questions), and explain that they're investigating the apparently challenging problem of using a base distribution G_0. The question is whether their relaxation is still useful. From the experiments it seems that the method is at least competitive, so the answer is yes. Hopefully an extension will come in the future, as the authors mention.
The second concern is about the motivation of this method. It seems that the paper fails to clearly explain in a convincing way why it is beneficial to reformulate the VAE as a SB-VAE. I understand that the non-parametric property induced by the prior might result in better capacity control, however I feel that this advantage (and potentially others which are still unclear to me) is not sufficiently explained and demonstrated. Perhaps some comparison with a dropout approach or a more thorough discussion related to dropout would make this clearer.
Overall, I found this to be an interesting paper, it would be a good fit for ICLR.


IMPROVING NEURAL CONVERSATION MODELS
We study reinforcement learning of chatbots with recurrent neural network architectures when the rewards are noisy and expensive to obtain. For instance, a chatbot used in automated customer service support can be scored by quality assurance agents, but this process can be expensive, time consuming and noisy. Previous reinforcement learning work for natural language processing uses on-policy updates and/or is designed for on-line learning settings. We demonstrate empirically that such strategies are not appropriate for this setting and develop an off-policy batch policy gradient method (BPG). We demonstrate the efficacy of our method via a series of synthetic experiments and an Amazon Mechanical Turk experiment on a restaurant recommendations dataset.
1 INTRODUCTION
Chatbots are one of the classical applications of artificial intelligence and are now ubiquitous in technology, business and everyday life. Many corporate entities are now increasingly using chatbots to either replace or assist humans in customer service contexts. For example, Microsoft is currently actively building a chat bot to optimise and streamline its technical support service.
In these scenarios, there is usually an abundance of historical data since past conversations between customers and human customer service agents are usually recorded by organisations. An apparently straightforward solution would be to train chatbots to reproduce the responses by human agents using standard techniques such as maximum likelihood. While this seems natural, it is far from desirable for several reasons. It has been observed that such procedures have a tendency to produce very generic responses (Sordoni et al., 2015). For instance, when we trained chatbots via maximum likelihood on a restaurant recommendations dataset, they repeatedly output responses to the effect of How large is your group?, What is your budget? etc. Further, they also produce responses such as Let me look that up. or Give me a second. which, although permissible for a human agent to say, are not appropriate for a chatbot. Although there are ways to increase the diversity of responses (Li et al., 2015), our focus is on encouraging the bot to meaningfully advance the conversation. One way to address this problem is to provide some form of weak supervision for responses generated by a chatbot. For example, a human labeller, such as a quality assurance agent, could score each response generated by a chatbot in a conversation with a customer. This brings us to the reinforcement learning (RL) paradigm where these rewards (scores) are to be used to train a good chatbot. In this paper we will use the terms score, label, and reward interchangeably. Labelled data will mean conversations which have been assigned a reward of some form as explained above.
Nonetheless, there are some important differences in the above scenario when compared to the more popular approaches for RL.
• Noisy and expensive rewards: Obtaining labels for each conversation can be time consuming and economically expensive. As a result, there is a limited amount of labelled data available. Moreover, labels produced by humans are invariably noisy due to human error and subjectivity.
• Off-line evaluations: Unlike conventional RL settings, such as games, where we try to find the optimal policy while interacting with the system, the rewards here are not immediately available. Previous conversations are collected, labelled by human experts, and then given to an algorithm which has to manage with the data it has.
• Unlabelled Data: While labelled data is limited, a large amount of unlabelled data is available.
If labelled data is in short supply, reinforcement learning could be hopeless. However, if unlabelled data can be used to train a decent initial bot, say via maximum likelihood, we can use policy iteration techniques to refine this bot by making local improvements using the labelled data (Bellman, 1956). Besides chatbots, this framework also finds applications in tasks such as question answering (Ferrucci et al., 2010; Hermann et al., 2015; Sachan et al., 2016), generating image descriptions (Karpathy & Fei-Fei, 2015) and machine translation (Bahdanau et al., 2014) where a human labeller can provide weak supervision in the form of a score to a sentence generated by a bot.
To contextualise the work in this paper, we make two important distinctions in policy iteration methods in reinforcement learning. The first is on-policy vs off-policy. In on-policy settings, the goal is to improve the current policy on the fly while exploring the space. On-policy methods are used in applications where it is necessary to be competitive (achieve high rewards) while simultaneously exploring the environment. In off-policy, the environment is explored using a behaviour policy, but the goal is to improve a different target policy. The second distinction is on-line vs batch (off-line). In on-line settings one can interact with the environment. In batch methods, which is the setting for this work, one is given past exploration data from possibly several behaviour policies and the goal is to improve a target policy using this data. On-line methods can be either on-policy or off-policy whereas batch methods are necessarily off-policy.
In this paper, we study reinforcement learning in batch settings, for improving chat bots with Seq2Seq recurrent neural network (RNN) architectures. One of the challenges when compared to on-line learning is that we do not have interactive control over the environment. We can only hope to do as well as our data permits us to. On the other hand, the batch setting affords us some luxuries. We can reuse existing data and use standard techniques for hyper-parameter tuning based on cross validation. Further, in on-line policy updates, we have to be able to “guess” how an episode will play out, i.e. actions the behaviour/target policies would take in the future and corresponding rewards. However, in batch learning, the future actions and rewards are directly available in the data. This enables us to make more informed choices when updating our policy.
RELATED WORK
Recently there has been a surge of interest in deep learning approaches to reinforcement learning, many of them adopting Q-learning, e.g. (He et al., 2015; Mnih et al., 2013; Narasimhan et al., 2015). In Q-learning, the goal is to estimate the optimal action value function Q∗. Then, when an agent is at a given state, it chooses the best greedy action according to Q∗. While Q-learning has been successful in several applications, it is challenging in the settings we consider since estimating Q∗ over large action and state spaces will require a vast number of samples. In this context, policy iteration methods are more promising since we can start with an initial policy and make incremental local improvements using the data we have. This is especially true given that we can use maximum likelihood techniques to estimate a good initial bot using unlabelled data.
Policy gradient methods, which fall within the paradigm of policy iteration, make changes to the parameters of a policy along the gradient of a desired objective (Sutton et al., 1999). Recently, the natural language processing (NLP) literature has turned its attention to policy gradient methods for improving language models. Ranzato et al. (2015) present a method based on the classical REINFORCE algorithm (Williams, 1992) for improving machine translation after preliminary training with maximum likelihood objectives. Bahdanau et al. (2016) present an actor-critic method also for machine translation. In both cases, as the reward, the authors use the BLEU (bilingual evaluation understudy) score of the output and the translation in the training dataset. This setting, where the rewards are deterministic and cheaply computable, does not reflect difficulties inherent to training chatbots where labels are noisy and expensive. Li et al. (2016) develop a policy gradient method bot for chatbots. However, they use user defined rewards (based on some simple rules) which, once again, are cheaply obtained and deterministic. Perhaps the closest to our work is that of Williams & Zweig (2016) who use a REINFORCE based method for chat bots. We discuss the differences of
this and other methods in greater detail in Section 3. The crucial difference between all of the above efforts and ours is that they use on-policy and/or on-line updates in their methods.
The remainder of this manuscript is organised as follows. In Section 2 we review Seq2Seq models and Markov decision processes (MDP) and describe our framework for batch reinforcement learning. Section 3 presents our method BPG and compares it with prior work in the RL and NLP literature. Section 4 presents experiments on a synthetic task and a customer service dataset for restaurant recommendations.
2 PRELIMINARIES
2.1 A REVIEW OF SEQ2SEQ MODELS
The goal of a Seq2Seq model in natural language processing is to produce an output sequence y = [a1, a2, . . . , aT ] given an input sequence x (Cho et al., 2014; Kalchbrenner & Blunsom, 2013; Sutskever et al., 2014). Here ai ∈ A where A is a vocabulary of words. For example, in machine translation from French to English, x is the input sequence in French, and y is its translation in English. In customer service chatbots, x is the conversation history until the customer’s last query and y is the response by an agent/chatbot. In a Seq2Seq model, we use an encoder network to represent the input sequence as a euclidean vector and then a decoder network to convert this vector to an output sequence. Typically, both the encoder and decoder networks are recurrent neural networks (RNN) (Mikolov et al., 2010) where the recurrent unit processes each word in the input/output sequences one at a time. In this work, we will use the LSTM (long short term memory) (Hochreiter & Schmidhuber, 1997) as our recurrent unit due to its empirical success in several applications.
In its most basic form, the decoder RNN can be interpreted as assigning a probability distribution over A given the current “state”. At time t, the state st is the input sequence x and the words yt−1 = [a1, . . . , at−1] produced by the decoder thus far, i.e. st = (x, yt−1). We sample the next word at from this probability distribution π(·|st), then update our state st+1 = (x, yt) where yt = [yt−1, at], and proceed in a similar fashion. The vocabulary A contains an end-of-statement token <EOS>. If we sample <EOS> at time T + 1, we terminate the sequence and output yT .
2.2 A REVIEW OF MARKOV DECISION PROCESSES (MDP)
We present a formalism for MDPs simplified to our setting. In an MDP, an agent takes an action a in a state s and transitions to a state s′. An episode refers to a sequence of transitions s1 → a1 → s2 → a2 → · · · → aT → sT+1 until the agent reaches a terminal state sT+1. At a terminal state, the agent receives a reward. Formally, an MDP is the triplet (S,A, R). Here, S is a set of states andA is a set of actions. When we take an action a at state s we transition to a new state s′ = s′(s, a) which, in this work, will be deterministic. A will be a finite but large discrete set and S will be discrete but potentially infinite. R : S → R is the expected reward function such that when we receive a reward r at state s ∈ S, E[r] = R(s). Let S0 ⊂ S be a set of terminal states. When we transition to any s ∈ S0, the episode ends. In this work, we will assume that the rewards are received only at a terminal state, i.e R(s) is nonzero only on S0. A policy π is a rule to select an action at a given state. We will be focusing on stochastic policies π : A×S → R+ where π(a|s) denotes the probability an agent will execute action a at state s. We define the value function V π : S → R of policy π, where V (s) is the expected reward at the end of the episode when we follow policy π from state s. For any terminal state s ∈ S0, V π(s) = R(s) regardless of π. We will also find it useful to define the action-value function Qπ : S × A :→ R, where Qπ(s, a) is the expected reward of taking action a at state s and then following policy π. With deterministic state transitions this is simply Qπ(s, a) = V π(s′(s, a)). It can be verified that V π(s) = Ea∼π(·|s) [Qπ(s, a)] (Sutton & Barto, 1998).
2.3 SET UP
We now frame our learning from labels scenario for RNN chatbots as an MDP. The treatment has similarities to some recent RL work in the NLP literature discussed above.
Let x be the input and yt−1 = [a1, . . . , at−1] be the words output by the decoder until time t. The state of our MDP at time t of the current episode will be st = (x, yt−1). Therefore, the set of states S will be all possible pairs of inputs and partial output sequences. The actions A will be the vocabulary. The terminal states S0 will be (x, y) such that the last literal of y is <EOS>. The stochastic policy π will be a Seq2Seq RNN which produces a distribution overA given state st. When we wish to make the dependence of the policy on the RNN parameters θ explicit, we will write πθ. When we sample an action at ∼ π(·|st), we deterministically transition to state (x, [yt−1, at]). If we sample aT+1 = <EOS> at time T + 1, the episode terminates and we observe a stochastic reward.
We are given a dataset of input-output-reward triples {(x(i), y(i), r(i))}ni=1 where y(i) = (a
(i) 1 , . . . , a (i) Ti , <EOS>) is the sequence of output words. This data was collected from possibly multiple behaviour policies which output y(i) for the given input x(i). In the above customer service example, the behaviour policies could be chatbots, or even humans, which were used for conversations with a customer. The rewards ri are scores assigned by a human quality assurance agent to each response of the chatbot. Our goal is to use this data to improve a given target policy πθ. We will use q to denote the distribution of the data. q(s) is the distribution of the states in the dataset, q(a|s) is the conditional distribution of an action given a state, and q(s, a) = q(s)q(a|s) is the joint distribution over states and actions. q will be determined by the initial distribution of the inputs x(i) and the behaviour policies used to collect the training data. Our aim is to find a policy that does well with respect to q. Specifically, we wish to maximise the following objective,
J(θ) = ∑ s∈S q(s)V πθ (s). (1)
Here, the value function V πθ is not available to us but has to be estimated from the data. This is similar to objectives used in on-line off-policy policy gradient literature where q is replaced by the limiting distribution of the behaviour policy (Degris et al., 2012). In the derivation of our algorithm, we will need to know q(a|s) to compute the gradient of our objective. In off-policy reinforcement learning settings this is given by the behaviour policy which is readily available. If the behaviour policy if available to us, then we can use this too. Otherwise, a simple alternative is to “learn” a behaviour policy. For example, in our experiments we used an RNN trained using the unlabelled data to obtain values for q(a|s). As long as this learned policy can capture the semantics of natural language (for example, the word apple is more likely than car when the current state is (x, I ate an)), then it can be expected to do reasonably well. In the following section, we will derive a stochastic gradient descent (SGD) procedure that will approximately minimise (1).
Before we proceed, we note that it is customary in the RL literature to assume stochastic transitions between states and use rewards at all time steps instead of the terminal step. Further, the future rewards are usually discounted by a discount factor γ < 1. While we use the above formalism to simplify the exposition, the ideas presented here extend naturally to more conventional settings.
3 BATCH POLICY GRADIENT
Our derivation follows the blueprint in Degris et al. (2012) who derive an off-policy on-line actor critic algorithm. Following standard policy gradient methods, we will aim to update the policy by taking steps along the gradient of the objective∇J(θ).
∇J(θ) = ∇Es∼q [∑ a∈A πθ(a|s)Qπθ (s, a) ] = Es∼q [∑ a∈A ∇πθ(a|s)Qπθ (s, a) + πθ(a|s)∇Qπθ (s, a) ] .
The latter term inside the above summation is difficult to work with, so the first step is to ignore it and work with the approximate gradient g(θ) = Es∼q[ ∑ a∈A∇πθ(a|s)Qπθ (s, a)] ≈ ∇J(θ). Degris et al. (2012) provide theoretical justification for this approximation in off policy settings by establishing that J(θ) ≤ J(θ + αg(θ)) for all small enough α. Expanding on g(θ), we obtain:
g(θ) = Es∼q [∑ a∈A πθ(a|s) ∇πθ(a|s) πθ(a|s) Qπθ (s, a) ] = E s∼q a∼q(·|s) [ ρ(s, a)ψ(a, s)Qπθ (s, a) ] = E(st,at)∼q(·,·) [ρ(st, at)ψ(at, st)(Q πθ (st, at)− V πθ (st))] . (2)
Here ψ(a, s) = ∇πθ(a|s)πθ(a|s) = ∇ log πθ(a|s) is the score function of the policy and ρ(s, a) = πθ(a|s)/q(a|s) is the importance sampling coefficient. In the last step, we have used the fact that E[π(a|s)ψ(a|s)h(s)] = 0 for any function h : S → R of the current state (Szepesvári, 2010). The purpose of introducing the value function V πθ is to reduce the variance of the SGD updates – we want to assess how good/bad action at is relative to how well πθ will do at state st in expectation. If at is a good action (Qπθ (st, at) is large relative to V πθ (st)), the coefficient of the score function is positive and it will change θ so as to assign a higher probability to action at at state st.
The Qπθ , V πθ functions are not available to us so we will replace them with estimates. For V πθ (st) we will use an estimate V̂ (st) – we will discuss choices for this shortly. However, the action value function is usually not estimated in RL policy gradient settings to avoid the high sample complexity. A sensible stochastic approximation for Qπθ (st, at) is to use the sum of future rewards from the current state (Sutton & Barto, 1998)1. If we receive reward r at the end of the episode, we can then use Qπθ (st, at) ≈ r for all time steps t in the episode. However, since q(at|st) is different from πθ(at|st) we will need to re-weight future rewards via importance sampling r ∏T i=t ρ(si, ai). This is to account for the fact that an action a given s may have been more likely under the policy πθ(·|s) than it was under q(·|s) or vice versa. Instead of directly using the re-weighted rewards, we will use the so called λ–return which is a convex combination of the re-weighted rewards and the value function (Sutton, 1988; 1984). In our setting, they are defined recursively from the end of the episode t = T + 1 to t = 1 as follows. For λ ∈ (0, 1],
r̃λT+1 = r, r̃ λ t = (1− λ)V πθ (st+1) + λρ(st, at)r̃λt+1 for t = T, . . . , 1. (3)
The purpose of introducing λ is to reduce the variance of using the future rewards alone as an estimate for Qπθ (st, at). This is primarily useful when rewards are noisy. If the rewards are deterministic, λ = 1 which ignores the value function is the best choice. In noisy settings, it is recommended to use λ < 1 (see Sec 3.1 of (Szepesvári, 2010)). In our algorithm, we will replace r̃λt with r λ t where V πθ is replaced with the estimate V̂ . Putting it all together, and letting α denote the step size, we have the following update rule for the parameters θ of our policy:
θ ← θ + αρ(st, at)ψ(st, at)(rλt − V̂ (st)). In Algorithm 1, we have summarised the procedure where the updates are performed after an entire pass through the dataset. In practice, we perform the updates in mini-batches.
An Estimator for the Value Function: All that is left to do is to specify an estimator V̂ for the value function. We first need to acknowledge that this is a difficult problem: S is quite large and for typical applications for this work there might not be enough data since labels are expensive. That said, the purpose of V̂ in (2), (3) is to reduce the variance of our SGD updates and speed up convergence so it is not critical that this be precise – even a bad estimator will converge eventually. Secondly, standard methods for estimating the value function based on minimising the projected Bellman error require the second derivatives, which might be intractable for highly nonlinear parametrisations of V̂ (Maei, 2011). For these two statistical and computational reasons, we resort to simple estimators for V πθ . We will study two options. The first is a simple heuristic used previously in the RL literature, namely a constant estimator for V̂ which is equal to the mean of all rewards in the dataset (Williams, 1992). The second uses the parametrisation V̂ (s) = σ(ξ>φ(s)) where σ is the logistic function and φ(s) ∈ Rd is a Euclidean representation of the state. For V̂ (s) of the above form, the Hessian ∇2ξV̂ (s) can be computed in O(d) time. To estimate this value function, we use the GTD(λ) estimator from Maei (2011). As φ(s) we will be using the hidden state of the LSTM. The rationale for this is as follows. In an LSTM trained using maximum likelihood, the hidden state contains useful information about the objective. If there is overlap between the maximum likelihood and reinforcement learning objectives, we can expect the hidden state to also carry useful information about the RL objective. Therefore, we can use the hidden state to estimate the value function whose expectation is the RL objective. We have described our implementation of GTD(λ) in Appendix A and specified some implementation details in Section 4.
1 Note Qπθ (st, at) = V πθ (st+1) for deterministic transitions. However, it is important not to interpret the term in (2) as the difference in the value function between successive states. Conditioned on the current time step, V πθ (st) is deterministic, while V πθ (st+1) is stochastic. In particular, while a crude estimate suffices for the former, the latter is critical and should reflect the rewards received during the remainder of the episode.
Algorithm 1 Batch Policy Gradient (BPG) Given: Data {(xi, yi, ri)}ni=1, step size α, return coefficient λ, initial θ0.
– Set θ ← θ0. – For each epoch k = 1, 2, . . .
I Set ∆θ ← 0 I For each episode i = 1, . . . , n • rλT+1 ← ri • ρt ← πθ(a(i)t |s (i) t )/q(a (i) t |s (i) t ) for t = 1, . . . , T
(i). • For each time step in reverse t = T (i), . . . , 1
(i) rλt ← (1− λ)V̂ (s (i) t+1) + λρtr λ t+1
(ii) ∆θ ← ∆θ + 1 T (i) ρtψ(s (i) t , a (i) t )(r λ t − V̂ (s (i) t ))
(iii) Compute updates for the value function estimate V̂ . I Update the policy θ ← θ + α∆θ I Update the value function estimate V̂ .
COMPARISON WITH OTHER RL APPROACHES IN NLP
Policy gradient methods have been studied extensively in on policy settings where the goal is to improve the current policy on the fly (Amari, 1998; Williams, 1992). To our knowledge, all RL approaches in Seq2Seq models have also adopted on-policy policy gradient updates (Bahdanau et al., 2016; Li et al., 2016; Ranzato et al., 2015; Williams & Zweig, 2016). However, on policy methods break down in off-policy settings, because any update must account for the probability of the action under the target policy. For example, suppose the behaviour policy took action a at state s and received a low reward. Then we should modify the target policy θ so as to reduce πθ(a|s). However, if the target policy is already assigning low probability to a|s then we should not be as aggressive when making the updates. The re-weighting ρ(s, a) via importance sampling does precisely this.
A second difference is that we study batch RL. Standard on-line methods are designed for settings where we have to continually improve the target while exploring using the behaviour policy. Critical to such methods are the estimation of future rewards at the current state and the future actions that will be taken by both the behaviour and target policies. In order to tackle this, previous research either ignore future rewards altogether (Williams, 1992), resort to heuristics to distribute a delayed reward to previous time steps (Bahdanau et al., 2016; Williams & Zweig, 2016), or make additional assumptions about the distribution of the states such as stationarity of the Markov process (Degris et al., 2012; Maei, 2011). However, in batch settings, the λ-return from a given time step can be computed directly (3) since the future action and rewards are available in the dataset. Access to this information provides a crucial advantage over techniques designed for on-line settings.
4 EXPERIMENTS
Implementation Details: We implement our methods using Chainer (Tokui et al., 2015), and group sentences of the same length together in the same batch to make use of GPU parallelisation. Since different batches could be of different length, we do not normalise the gradients by the batch size as we should take larger steps after seeing more data. However, we normalise by the length of the output sequence to allocate equal weight to all sentences. We truncate all output sequences to length 64 and use a maximum batch size of 32. We found it necessary to use a very small step size (10−5), otherwise the algorithm has a tendency to get stuck at bad parameter values. While importance reweighting is necessary in off-policy settings, it can increase the variance of the updates, especially when q(at|st) is very small. A common technique to alleviate this problem is to clip the ρ(st, at) value (Swaminathan & Joachims, 2015). In addition to single ρ(st, at) values, our procedure has a product of ρ(st, at) values when computing the future rewards (3). The effect of large ρ values is a large weight ρt(rλt −V̂ (st)) for the score function in step (ii) of Algorithm 1. In our implementation,
LSTM LSTM LSTM
LSTM LSTM LSTM
softmax softmax softmax
LSTM LSTM
LSTM LSTM
Encoder Decoder
we clip this weight at 5 which controls the variance of the updates and ensures that a single example does not disproportionately affect the gradient.
RNN Design: In both experiments we use deep LSTMs with two layers for the encoder and decoder RNNs. The output of the bottom layer is fed to the top layer and in the decoder RNN, the output of the top layer is fed to a softmax layer of size |A|. When we implement GTD(λ) to estimate V πθ we use the hidden state of the bottom LSTM as φ(s). When performing our policy updates, we only change the parameters of the top LSTM and the softmax layer in our decoder RNN. If we were to change the bottom LSTM too, then the state representation φ(s) would also change as the policy changes. This violates the MDP framework. In other words, we treat the bottom layer as part of the environment in our MDP. To facilitate a fair comparison, we only modify the top LSTM and softmax layers in all methods. We have illustrated this set up in Fig. 1. We note that if one is content with using the constant estimator, then one can change all parameters of the RNN.
4.1 SOME SYNTHETIC EXPERIMENTS ON THE EUROPARL DATASET
To convey the main intuitions of our method, we compare our methods against other baselines on a synthetic task on the European parliament proceedings corpus (Koehn, 2005). We describe the experimental set up briefly, deferring details to Appendix B.1. The input sequence to the RNN was each sentence in the dataset. Given an input, the goal was to reproduce the words in the input without repeating words in a list of forbidden words. The RL algorithm does not explicitly know either goal of the objective but has to infer it from the stochastic rewards assigned to input output sequences in the dataset. We used a training set of 500 input-output-reward triplets for the RL methods.
We initialised all methods by maximum likelihood training on 6000 input output sequences where the output sequence was the reverse of the input sequence. The maximum likelihood objective captures part of the RL objective. This set up reflects naturally occurring practical scenarios for the algorithm where a large amount unlabelled data can be used to bootstrap a policy if the maximum likelihood and reinforcement learning objectives are at least partially aligned. We trained the RL algorithms for 200 epochs on the training set. At the end of each epoch, we generated outputs from the policy on test set of 500 inputs and scored them according to our criterion. We plot the test set error against the number of epochs for various methods in Fig. 2.
Fig. 2(a) compares 3 methods: BPG with and without maximum likelihood initialisation and a version of BPG which does not use importance sampling. Clearly, bootstrapping an RL algorithm with ML can be advantageous especially if data is abundantly available for ML training. Further, without importance sampling, the algorithm is not as competitive for reasons described in Section 3. In all 3 cases, we used a constant estimator for V̂ and λ = 0.5. The dashed line indicates the performance of ML training alone. BPG-NIS is similar to the algorithms of Ranzato et al. (2015); Williams & Zweig (2016) except that there, their methods implicitly use λ = 1.
Fig. 2(b) compares 4 methods: BPG and its on-line version OPG with constant (CONST) and GTD(λ) estimators for V̂ . The on-line versions of the algorithms are a direct implementation of the method in Degris et al. (2012) which do not use the future rewards as we do. The first observation is that while GTD(λ) is slightly better in the early iterations, it performs roughly the same as using a constant estimator in the long run. Next, BPG performs significantly better than OPG. We believe this is due to the following two reasons. First, the online updates assume stationarity of the MDP. When this does not hold, such as in limited data instances like ours, the SGD updates can be
very noisy. Secondly, the value function estimate plays a critical role in the online version. While obtaining a reliable estimate V̂ is reasonable in on-line settings where we can explore indefinitely to collect a large number of samples, it is difficult when one only has a limited number of labelled samples. Finally, we compare BPG with different choices for λ in Fig. 2(c). As noted previously, λ < 1 is useful with stochastic rewards, but choosing too small a value is detrimental. The optimal λ value may depend on the problem.
4.2 RESTAURANT RECOMMENDATIONS
We use data from an on-line restaurant recommendation service. Customers log into the service and chat with a human agent asking recommendations for restaurants. The agents ask a series of questions such as food preferences, group size etc. before recommending a restaurant. The goal is to train a chatbot (policy) which can replace or assist the agent. For reasons explained in Section 1, maximum likelihood training alone will not be adequate. By obtaining reward labels for responses produced by various other bots, we hope to improve on a bot initialised using maximum likelihood.
Data Collection: We collected data for RL as follows. We trained five different RNN chatbots with different LSTM parameters via maximum likelihood on a dataset of 6000 conversations from this dataset. The bots were trained to reproduce what the human agent said (output y) given the past conversation history (input x). While the dataset is relatively small, we can still expect our bots to do reasonably well since we work in a restricted domain. Next, we generated responses from these bots on 1216 separate conversations and had them scored by workers on Amazon Mechanical Turk (AMT). For each response by the bots in each conversation, the workers were shown the history before the particular response and asked to score (label) each response on a scale of 0− 1− 2. We collected scores from three different workers for each response and used the mean as the reward.
Policies and RL Application: Next, we initialised 2 bots via maximum likelihood and then used BPG to improve them using the labels collected from AMT. For the 2 bots we used the following LSTM hidden state size H , word embedding size E and BPG parameters. These parameters were chosen arbitrarily and are different from those of the bots used in data collection described above.
• Bot-1: H = 512, E = 256. BPG: λ = 0.5, GTD(λ) estimator for V̂ .
• Bot-2: H = 400, E = 400. BPG: λ = 0.5, constant estimator for V̂ .
Testing: We used a separate test set of 500 conversations which had a total of more than 3500 inputoutput (conversation history - response) pairs. For each Bot-1 and Bot-2 we generated responses before and after applying BPG, totalling 4 responses per input. We then had them scored by workers on AMT using the same set up described above. The same worker labels the before-BPG and afterBPG responses from the same bot. This controls spurious noise effects and allows us to conduct a paired test. We collected 16, 808 before and after label pairs each for Bot-1 and Bot-2 and compare them using a paired t-test and a Wilcoxon signed rank test.
Results: The results are shown in Table 1. The improvements on Bot-2 are statistically significant at the 10% level on both tests, while Bot-1 is significant on the Wilcoxon test. The large p-values for Bot-1 are due to the noisy nature of AMT experiments and we believe that we can attain significance if we collect more labels which will reduce the standard error in both tests. In Appendix B.2 we present some examples of conversation histories and the responses generated by the bots before and after applying BPG. We qualitatively discuss specific kinds of issues that we were able to overcome via reinforcement learning.
5 CONCLUSION
We presented a policy gradient method for batch reinforcement learning to train chatbots. The data to this algorithm are input-output sequences generated using other chatbots/humans and stochastic rewards for each output in the dataset. This setting arises in many applications, such as customer service systems, where there is usually an abundance of unlabelled data, but labels (rewards) are expensive to obtain and can be noisy. Our algorithm is able to efficiently use minimal labelled data to improve chatbots previously trained through maximum likelihood on unlabelled data. While our method draws its ideas from previous policy gradient work in the RL and NLP literature, there are some important distinctions that contribute to its success in the settings of interest for this work. Via importance sampling we ensure that the probability of an action is properly accounted for in off-policy updates. By explicitly working in the batch setting, we are able to use knowledge of future actions and rewards to converge faster to the optimum. Further, we use the unlabelled data to initialise our method and also learn a reasonable behaviour policy. Our method outperforms baselines on a series of synthetic and real experiments.
The ideas presented in this work extend beyond chatbots. They can be used in applications such as question answering, generating image descriptions and machine translation where an output sentence generated by a policy is scored by a human labeller to provide a weak supervision signal.
ACKNOWLEDGEMENTS
We would like to thank Christoph Dann for the helpful conversations and Michael Armstrong for helping us with the Amazon Mechanical Turk experiments.
APPENDIX
A IMPLEMENTATION OF GTD(λ)
We present the details of the GTD(λ) algorithm (Maei, 2011) to estimate a value function in Algorithm 2. However, while Maei (2011) give an on-line version we present the batch version here where the future rewards of an episode are known. We use a parametrisation of the form V̂ (s) = V̂ξ(s) = σ(ξ
>φ(s)) where ξ ∈ Rd is the parameter to be estimated. σ(z) = 1/(1 + e−z) is the logistic function.
The algorithm requires two step sizes α′, α′′ below for the updates to ξ and the ancillary parameter w. Following the recommendations in Borkar (1997), we use α′′ α. In our implementations, we used α′ = 10−5 and α′′ = 10−6. When we run BPG, we perform steps (a)-(f) of Algorithm 2 in step (iii) of Algorithm 1 and the last two update steps of Algorithm 2 in the last update step of Algorithm 1.
The gradient and Hessian of V̂ξ have the following forms,
∇ξV̂ξ(s) = V̂ξ(s)(1− V̂ξ(s))φ(s), ∇2ξV̂ξ(s) = V̂ξ(s)(1− V̂ξ(s))(1− 2V̂ξ(s))φ(s)φ(s)>.
The Hessian product in step (d) of Algorithm 2 can be computed in O(d) time via, ∇2ξV̂ξ(s) · w = [ V̂ξ(s)(1− V̂ξ(s))(1− 2V̂ξ(s))(φ(s)>w) ] φ(s).
Algorithm 2 GTD(λ) Given: Data {(xi, yi, ri)}ni=1, step sizes α′, α′′, return coefficient λ, initial ξ0.
– Set ξ ← ξ0, w ← 0. – For each epoch k = 1, 2, . . .
I Set ∆ξ ← 0, ∆w ← 0. I For each episode i = 1, . . . , n • Set rλT+1 ← ri, gλT+1 ← 0, qλT+1 ← 0 • ρt ← πθ(a(i)t |s (i) t )/q(a (i) t |s (i) t ) for t = 1, . . . , T
(i). • For each time step in reverse t = T (i), . . . , 1:
(a) gλt ← ρt ( (1− λ)V̂ξ(s(i)t+1) + λρtrλt+1 )
(b) qλt ← ρt ( (1− λ)∇ξV̂ξ(s(i)t+1) + λqλt+1 )
(c) δt ← gλt − V̂ξ(s (i) t ) (d) ht ← ( δt − w>∇ξV̂ξ(s(i)t ) ) ∇2ξV̂ξ(s (i) t ) · w
(e) ∆w ← ∆w + 1 T (i) ( δt − w>∇ξV̂ξ(s(i)t ) ) ∇ξV̂ξ(s(i)t )
(f) ∆ξ ← ∆ξ + 1 T (i) ( δt∇ξV̂ξ(s(i)t )− qλt w>∇ξV̂ξ(s (i) t )− ht ) I w ← w + α′′∆w. I ξ ← ξ + α′∆ξ.
B ADDENDUM TO EXPERIMENTS
B.1 DETAILS OF THE SYNTHETIC EXPERIMENT SET UP
Given an input and output sequence, we used the average of five Bernoulli rewards Bern(r), where the parameter r was r = 0.75 × rr + 0.25 × rf. Here rr was the fraction of common words in the input and output sequences while rf = 0.01pf where pf is the fraction of forbidden words in the dataset. As the forbidden words, we used the 50 most common words in the dataset. So if an input
had 10 words of which 2 were forbidden, an output sequence repeating 7 of the allowed words and 1 forbidden word would receive an expected score of 0.75× (8/10) + 0.25× 0.01(1/8) = 0.7406. The training and testing set for reinforcement learning were obtained as follows. We trained 4 bots using maximum likelihood on 6000 input output sequences as indicated in Section 4.1. The LSTM hidden state size H and word embedding size E for the 4 bots were, (H,E) = (256, 128), (128, 64), (64, 32), (32, 16). The vocabulary size was |A| = 12000. We used these bots to generate outputs for 500 different input sequences each. This collection of input and output pairs was scored stochastically as described above to produce a pool of 2000 input-output-score triplets. From this pool we use a fixed set of 500 triplets for testing across all our experiments. From the remaining 1500 data points, we randomly select 500 for training for each execution of an algorithm. For all RL algorithms, we used an LSTM with 16 layers and 16 dimensional word embeddings.
B.2 ADDENDUM TO THE AMT RESTAURANT RECOMMENDATIONS EXPERIMENT
MORE DETAILS ON THE EXPERIMENTAL SET UP
We collected the initial batch of training data for RL as follows: We trained, via maximum likelihood on 6000 conversations, five RNN bots whose LSTM hidden sizeH and word embedding sizeE were (H,E) = (512, 512), (256, 256), (128, 128), (512, 256), (256, 64). The inputs x were all words from the history of the conversation truncated at length 64, i.e. the most recent 64 words in the conversation history. The outputs were the actual responses of the agent which were truncated to length 64. As the vocabulary we use the |A| = 4000 most commonly occurring words in the dataset and replace the rest with an <UNK> token.
Using the bots trained this way we generate responses on 1216 separate conversations. This data was sent to AMT workers who were asked to label the conversations on the following scale.
• 2: The response is coherent and appropriate given the history and advances the conversation forward.
• 1: The response has some minor flaws but is discernible and appropriate. • 0: The response is either completely incoherent or inappropriate and fails to advance the
conversation forward.
SOME QUALITATIVE RESULTS
In Tables 2 and 3 we have presented some examples. The text in black/grey shows the conversation history, the response in blue is by the bot trained via maximum likelihood (ML) alone and in red is the bot after improvement using our BPG reinforcement learning algorithm. The first two examples of Table 2 present examples where the ML algorithm repeated generic questions (on budget, group size etc.) even though they had already been answered previously. After applying BPG, we are able to correct such issues, even though there are some grammatical errors. In the second, third and fourth example, we see that the ML+BPG bot is able to take context into consideration well when responding. For example, the customer asks for oriental/Mexican/Italian food. While the ML bot doesn’t take this into consideration, the ML+BPG bot is able to provide relevant answers. However, in the third example, the name of the restaurant suggests that the food might be Indian and not Mexican. In the final example of Table 2 the customer asks a direct question about smoking. The ML bot provides an irrelevant answer where as the ML+BPG bot directly responds to the question.
In some examples, the ML bot had a tendency to produce sentences that were grammatically correct but nonsensical, sensible but grammatically incorrect, or just complete gibberish. We were able to correct such issues via RL. The first three examples of Table 3 present such cases. Occasionally the opposite happened. The last example of Table 3 is one such instance.

This paper extends neural conversational models into the batch reinforcement learning setting. The idea is that you can collect human scoring data for some responses from a dialogue model, however such scores are expensive. Thus, it is natural to use off-policy learning – training a base policy on unsupervised data, deploying that policy to collect human scores, and then learning off-line from those scores.
While the overall contribution is modest (extending off-policy actor-critic to the application of dialogue generation), the approach is well-motivated, and the paper is written clearly and is easy to understand.
My main concern is that the primary dataset used (restaurant recommendations) is very small (6000 conversations). In fact, it is several orders of magnitude smaller than other datasets used in the literature (e.g. Twitter, the Ubuntu Dialogue Corpus) for dialogue generation. It is a bit surprising to me that RNN chatbots (with no additional structure) are able to generate reasonable utterances on such a small dataset. Wen et al. (2016) are able to do this on a similarly small restaurant dataset, but this is mostly because they map directly from dialogue states to surface form, rather than some embedding representation of the context. Thus, it remains to be seen if the approaches in this paper also result in improvements when much more unsupervised data is available.
References:
Wen, Tsung-Hsien, Milica Gasic, Nikola Mrksic, Lina M. Rojas-Barahona, Pei-Hao Su, Stefan Ultes, David Vandyke, and Steve Young. "A Network-based End-to-End Trainable Task-oriented Dialogue System." arXiv preprint arXiv:1604.04562 (2016).