Title: STICK-BREAKING VARIATIONAL AUTOENCODERS

Abstract: 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 z i to an observed datapoint x i , and the other acts as an inference model (Salimans & Knowles, 2013) performing the reverse mapping from x i to z i . Together the two form a computational pipeline that resembles an unsupervised autoencoder (Hinton & Salakhutdinov, 2006) . The generative process can be written mathematically as

EQUATION

where p(z) is the prior and p θ (x|z i ) is the density network with parameters θ. The approximate posterior of this generative process, call it q φ (z|x i ), 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) :

EQUATION

for S samples of z i 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 θ (x i |ẑ i,s ) = S s=1 ∂ ∂ẑ i,s log p θ (x i |ẑ i,s ) ∂ẑ 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 ẑ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) (Ishwaran & James, 2001) . The π k s are random weights independent of G 0 , 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:

if it is of the form G(•) = ∞ k=1 π k δ ζ k where δ ζ k is a discrete measure concentrated at ζ k ∼ G 0 , a draw from the base distribution G 0

π k = v 1 if k = 1 v k j<k (1 -v j ) for k > 1 (3)

where v k ∼ Beta(α, β). When v k ∼ 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 (û) ≈ (ûaΓ(a)) 1 a b (4)

for û ∼ 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) = abx a-1 (1 -x a ) 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:

EQUATION

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):

EQUATION

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

EQUATION

where π i is the vector of stick-breaking weights and α 0 is the concentration parameter of the GEM distribution. The likelihood model p θ (x i |π i ) is a density network just as described in Section 2.1.

4.2. INFERENCE:

The inference process-how to draw π i ∼ q φ (π i |z i )-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 v i,k according to one of the parametrizations in Section 3. Next, a linear-time operation composes the stick segments from the sampled fractions:

EQUATION

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 v i,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 y i is introduced that represents a class label. Its distribution is the categorical one: q φ (y i |x i ) = Cat(y|g y (x i )) where g y 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 y i 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 (θ, φ; x i , y i ) = 1 S S s=1 log p θ (x i |π i,s , y i ) -KL(q φ (π i |x i )||p(π i ; α 0 )) + log q φ (y i |x i ) (10)

where log q φ (y i |x i ) is the log-likelihood of the label. And if the label is missing, we optimize

EQUATION

where H[q φ (y i |x i )] 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 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 repository 5 . 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(x i ) ≈ log 1 S s p(x i |ẑ i,s )p(ẑ i,s )/q(ẑ 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- 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.

-log p(x i ) Model MNIST

Figure 3 (a) shows MNIST digits drawn from the SB-VAE by sampling from the prior-i.e. v k ∼ 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. 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 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.

:

During preparation of this draft, the work ofRuiz 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.

Available at http://www.cs.nyu.edu/~roweis/data.html

Available at http://ufldl.stanford.edu/housenumbers/

Available at http://www.iro.umontreal.ca/~lisa/twiki/bin/view.cgi/Public/MnistVariations

Theano implementations available at https://github.com/enalisnick/stick-breaking_ dgms

The 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.

Figures and Tables:

(a) Finite Dimensional (b) Infinite Dimensional (c) The Stick-Breaking Variational Autoencoder.

Figure 1: Subfigures (a) and (b) show the plate diagrams for the relevant latent variable models. Solid lines denote the generative process and dashed lines the inference model. Subfigure (a) shows the finite dimensional case considered in (Kingma & Welling, 2014a), and (b) shows the infinite dimensional case of our concern. Subfigure (c) shows the feedforward architecture of the Stick-Breaking Autoencoder, which is a neural-network-based parametrization of the graphical model in (b).

Figure 2: Subfigure (a) shows test (expected) reconstruction error vs training epoch for the SB-VAE and Gauss VAE on the Frey Faces dataset, subfigure (b) shows the same quantities for the same models on the MNIST dataset, and subfigure (c) shows the same quantities for the same models on the MNIST+rot dataset.

(a) SB-VAE: MNIST Samples Drawn from Prior (b) Gauss VAE: MNIST Samples Drawn from Prior

Figure 3: Subfigure (a) depicts samples from the SB-VAE trained on MNIST. We show the ordered, factored nature of the latent variables by sampling from Dirichlet's of increasing dimensionality. Subfigure (b) depicts samples from the Gauss VAE trained on MNIST.

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.

Marginal likelihood results (estimated)  for Gaussian VAE and the three parametrizations of the Stick-Breaking VAE.

Table content: <table><tr><td/><td/><td>MNIST+rot</td></tr><tr><td>Gauss VAE</td><td>96.80</td><td>108.40</td></tr><tr><td>Kumar-SB VAE</td><td>98.01</td><td>112.33</td></tr><tr><td>Logit-SB VAE</td><td>99.48</td><td>114.09</td></tr><tr><td>Gamma-SB VAE</td><td>100.74</td><td>113.22</td></tr></table>

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 Figure4(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)

Table content: <table><tr><td/><td/><td/><td>t-SNE EmEeddLng of StLcN-BreaNLng VAE'V Latent SSace</td><td>t-SNE EmEeddLng of StLcN-BreaNLng VAE'V Latent SSace</td></tr><tr><td/><td>k=3</td><td>k=5</td><td>k=10</td></tr><tr><td>SB-VAE</td><td>9.34</td><td>8.65</td><td>8.90</td></tr><tr><td>DLGMM</td><td>9.14</td><td>8.38</td><td>8.42</td></tr><tr><td colspan="4">Gauss VAE 28.4 20.96 15.33</td></tr><tr><td colspan="2">Raw Pixels 2.95</td><td>3.12</td><td>3.35</td></tr><tr><td>GMVAE 6</td><td>-</td><td>8.96</td><td>-</td></tr><tr><td colspan="4">(a) MNIST: Test error for kNN on latent</td></tr><tr><td>space</td><td/><td/><td>(b) MNIST SB-VAE</td><td>(c) MNIST Gauss VAE</td></tr></table>

29±.39 7.34±.47 11.78±.39 14.27±.58 27.67±1.39 32.08±4.00 37.07±5.22 61.37±3.60 Gauss-DGM 3.95±.15 4.74±.43 11.55±2.28 21.78±.73 27.72±.69 38.13±.95 36.08±1.49 48.75±1.47 69.58±1.64 kNN 6.13±.13 7.66±.10 15.27±.76 18.41±.01 23.43±.01 37.98±.01 64.81±.34 68.94±.47 76.64±.54 Percent error on three semi-supervised classification tasks with 10%, 5%, and 1% of labels present for training. Our DGM with stick-breaking latent variables (SB-DGM) is compared with a DGM with Gaussian latent variables (Gauss-DGM), and a k-Nearest Neighbors classifier (k=5).

Table content: <table><tr><td/><td colspan="3">MNIST (N=45,000)</td><td colspan="3">MNIST+rot (N=70,000)</td><td colspan="2">SVHN (N=65,000)</td></tr><tr><td/><td>10%</td><td>5%</td><td>1%</td><td>10%</td><td>5%</td><td>1%</td><td>10%</td><td>5%</td><td>1%</td></tr><tr><td>SB-DGM</td><td>4.86±.14 5.</td><td/><td/><td/><td/><td/><td/><td/></tr></table>

Peer review (iclr2017):
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.