Paper:
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.

Review:
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.