Title: Stick-Breaking Variational Autoencoders
Authors: Eric Nalisnick
Padhraic Smyth
Venue: ICLR 2017
Keywords: Deep learning
Unsupervised Learning
Semi-Supervised Learning

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 semi-supervised variant, learn highly discriminative latent representations that often outperform the Gaussian VAE’s.

TL;DR: We define a variational autoencoder variant with stick-breaking latent variables thereby giving it adaptive width.

Review discussion:
The technique is potentially an interesting step towards fusing Black Box Variational Inference and Bayesian Nonparametrics.
However, I'm wondering why the model was not compared with the results in "Semi-supervised Learning with Deep Generative Models" for the semi-supervised learning experiments, they seem to have stronger results than the ones shown Table 2. Why would these results be not comparable ?
Hi, thanks for your question.  Notice that Kingma et al. reports SVHN results only for stacked models (M2+M1, M1+TSVM, etc.).  In order to isolate the effects of using S.B. vs Gauss. latent variables, we chose to experiment with just the (non-stacked) M2 model, for which they report results on just MNIST.
Thus their results are not directly comparable for SVHN.
Yet, notice that with just the M2 model, we get ~37% error at 5% labeled;
whereas with M1+M2, Kingma et al. get 36% at roughly 2% labeled.
Of course their result is better, but they are using two VAEs.
For MNIST, the results are somewhat comparable.  I say 'somewhat' because we use much larger encoders and decoders (3-layer ResNets), which is what allows us to get respectable performance on SVHN with one model.
Still, our MNIST results are roughly indistinguishable for theirs.
For example, on MNIST, Kingma et al. reports 3.92% (plus/minus .63) for 3,000 labeled examples; our Gauss-M2 implementation gets 4.74% (plus/minus .43) for 2,250 labeled examples.  Note there is statistical overlap in the results despite us using less labels.
Even though we have a higher-capacity model, I think we don't get better results because the extra encoder/decoder layers make us more prone to overfitting on MNIST.
1. You mention that Knowles' (18) approximation becomes poor for increasing \alpha. Do you have an analysis about what happens to Kumaraswamy's distribution (wrt the approximation) as you change the parameters?
2. The parameter a_0 is cross-validated in the experiments. I was wondering how sensitive is the overall performance to changes in a_0 (i.e. how crucial is for the cross-validation to be done well)?
3. Any insight about why the marginal likelihood is smaller for SB-VAE? Do you think it might be because of the approximation, or because of the actual probabilistic model having more parameters?
4. Minor question, but I couldn't understand what you mean by  "decay of the associated decoder weights" in the legend of Fig. 2.
Thanks!
Hi, thanks for your questions.
1.  Our use of the Kumaraswamy is an approximation in the variational sense---i.e. we don't know the true posterior so we need to pick a parametric family---but it is NOT a numerical approximation like Knowles' work is.
In fact, even if the Beta had a parametrization amenable to SGVB, there's no reason to think the true posterior is more accurately approximated by a Beta than a Kumaraswamy.  The reason the Beta has, historically, been the default approximation for Stick-Breaking models is because it (usually) gives closed-form expectations, which are not available here in any case.
2.  It's only crucial to set a_0 large enough.  If set too small, the Kumaraswamy parameters can become very big or very small, resulting in numerical instabilities.  The only reason not to set a_0 very large is for training-time considerations.  Runtime aside, I haven't seen a problem with setting it 'too' large.
3.  I think the marginal likelihood is worse than the Gauss-VAE because of the restricted form of stick-breaking latent variables (ie decreasing in expectation, summing to 1).
In theory, since MLPs are universal approximators, these constraints shouldn't be an issue (compared to Gaussian space), but in practice, I think the recursive nature of the stick-breaking weights makes optimization harder, which is supported in the expected reconstruction error plots by the slower learning curve.
4.  By 'decay of the associated decoder weight', I meant that the Gauss-VAE turns off weights in the decoder, setting them to zero.  The SB-VAE doesn't do this and leaves the decoder weights without any systematic sparsity.  'Decay' might be a poor word choice; perhaps 'pruning' is better.
Thanks,
Eric
Can you please elaborate on the form of the distribution p_\theta (x | \pi_i) in equation (8)?
Also, what is the base distribution G_0 in your case? And what is the dimensionality of this distribution?
Thanks
Hi again.
1.  p_\theta: The likelihood's form depends on the data, and thus it will be Gaussian when the data is over the real line and Bernoulli when the data is binary.
In either case, the neural network decoder parametrizes the mean.  The neural network itself consists of one or more hidden layers, and it's input will be the vector of stick-breaking weights.
2.  G_0: There is no base distribution.  The latent variables are drawn only from the GEM component of the Dirichlet Process (DP).  Adding a base distribution will result in a mixture latent representation.  This can be done---we are currently investigating it, in fact (http://bayesiandeeplearning.org/papers/BDL_20.pdf)---but it is impractical as the decoder network needs run K^S times, where K is the number of mixture components (truncation size, in the DP case) and S is the number of stochastic layers.
Best,
Eric
Summary: This is the first work to investigate stick-breaking priors, and corresponding inference methods, for use in VAEs. The background material is explained clearly, as well as the explanation of the priors and posteriors and their DNCP forms. The paper is really well written.
In experiments, they find that stick-breaking priors does not generally improve upon spherically Gaussian priors in the completely unsupervised setting, when measured w.r.t. log-likelihood. The fact that they do report this 'negative' result suggests good scientific taste. In a semi-supervised setting, the results are better.
Comments:
- sec 2.1: There is plenty of previous work with non-Gaussian p(z): DRAW, the generative ResNet paper in the IAF paper, Ladder VAEs, etc.
- sec 2.2: two comma's
- text flow eq 6: please refer to appendix with the closed-form KL divergence
- "The v's are sampled via" => "In the posterior, the v's are sampled via". It's not clear you're talking about the posterior here, instead of the prior.
- The last paragraph of section 4 is great.
- Sec 7.1: "Density estimation" => Technically you're also doing mass estimation.
- Sec 7.1: 100 IS samples is a bit on the low side.
- Figure 3(f). Interesting that k-NN works so well on raw pixels.
Thank you for your review.  Our responses / comments are below.
1. "The fact that they do report this 'negative' result suggests good scientific taste."
We appreciate you recognizing this.
2. "sec 2.1: There is plenty of previous work with non-Gaussian p(z): DRAW, the generative ResNet paper in the IAF paper, Ladder VAEs, etc."
Do you mean non-factorized Gaussians?  In the case of DRAW and Ladder VAEs, we still think they are using factorized Gaussians.  The DRAW paper says, “In our experiments the latent distribution is a diagonal Gaussian” (Section 2.1), and the Ladder VAE paper states, “Conditioned on the stochastic layer below each stochastic layer is specified as a fully factorized gaussian distribution” (Section 2.1).  The Ladder VAE paper says this when describing the regular multi-layer VAE but the same notation is used for their model.
3.  "Figure 3(f). Interesting that k-NN works so well on raw pixels."
Yes, this was a bit surprising to us too.  This blog post reports similar results: http://acgrama.blogspot.com/2012/09/knn-with-euclidean-distance-on-mnist.html
Thanks again,
Eric
When talking about 'the prior latent space', one usually talks about the joint distribution over all the latents: p(z), where 'z' is the concatenation of the individual components of latent space: p(z) = p(z1,z2,z3,...) = p(z1)p(z2|z1)p(z3|z1,z2)...
Regular DRAW indeed has spherical Gaussian priors since the latent space factorizes as independent Gaussians: p(z)=p(z1)p(z2)p(z3)..
However, convolutional DRAW (https://arxiv.org/pdf/1604.08772.pdf), ladder VAE and IAF VAE have non-factorized priors: each mean and variance of each level of the latent hierarchy depends non-linearly on the values their ancestors, such that their joint distribution becomes non-Gaussian.
Thanks, Reviewer #3.  I understand now and agree.  I will edit the draft to make a point of this.
Best,
Eric
The paper attempts to combine Variational Auto-Encoders with the Stick-Breaking process. The motivation is to tackle the component collapsing and have a representation with stochastic dimensionality. To demonstrate the merit of their approach, the authors test this model on MNIST and SVHN in an unsupervised and semi-supervised fashion.
After reading the paper in more detail, I find that the claim that the dimensionality of the latent variable is stochastic does not seem quite correct: all latent variables are "used" (which actually enable backpropagation) but the latent variables are parametrized differently (into $\pi$) and the decoding process is altered as to give the impression of sparsity. The way all these latent variables are used does not involve any marginalization but is very similar to the common soft-gating mechanism already used in LSTM or attentional model.
With respect to the Figure 5b showing the decoder input weights: component collapsing probably does not have the same effect as Gaussian prior. $\pi$ is positive therefore having a very small average value might mean that its value is close to zero most of the time, not requiring any update on the weight. For the standard Gaussian prior, component collapsing means having a very noisy input with no signal involved, which forces the decoder to shut down this channel, i.e. have small incoming weights from this collapsed variable.
Adding a histogram of the latent variables in addition to that might help decide if the associated weights are relatively large because they are actually used or if it's because the inputs are zero anyway.
The semi-supervised results are better than a weaker version of the model used in (Kingma et al., 2014), but as to have a fairer comparison, the results should be compared with the M1+M2 model in that paper, even if that requires also using two VAEs.
Thanks for your comments, Reviewer #2.  Our responses are below.
1. "After reading the paper in more detail, I find that the claim that the dimensionality of the latent variable is stochastic does not seem quite correct: all latent variables are "used" (which actually enable backpropagation) but the latent variables are parametrized differently (into $\pi$) and the decoding process is altered as to give the impression of sparsity. The way all these latent variables are used does not involve any marginalization but is very similar to the common soft-gating mechanism already used in LSTM or attentional model."
We are not sure what you mean by “the latent variables are parametrized differently (into $\pi$) and the decoding process is altered as to give the impression of sparsity”.  Are you referring to how the stick segments are recursively constructed?
While you are correct that the stick segments lie on the simplex and therefore look as if they could be produced by a soft-gating function, this view neglects the underlying probabilistic model and the need for sampling from the simplex (for ELBO optimization).  Generating differentiable samples from the simplex can not be done via existing methods and was an open problem---for example, see the discussion of the NVI Topic Model in my response to Review 1.
And yes, all the latent variables are used during every forward propagation, but this is commonly the case in Bayesian nonparametric models with posterior truncation.
For instance, Blei and Jordan, in their seminal work on VI for Infinite Mixtures (http://www.cs.columbia.edu/~blei/papers/BleiJordan2004.pdf), use a truncated posterior, just as we do, and therefore the number of components is fixed and all are ‘used’ in the sense that they have non-zero weight.  Their generative process is still a formal Bayes nonparametric model, just as ours is.  As we discuss in Section 4, our model does work with an adaptive number of latent variables; it’s just that optimization is slower.
2. "The semi-supervised results are better than a weaker version of the model used in (Kingma et al., 2014), but as to have a fairer comparison, the results should be compared with the M1+M2 model in that paper, even if that requires also using two VAEs."
We agree that the paper can only be improved by making the semi-supervised results more comparable to Kingma et al.’s.  We will try to add these before the discussion period ends.
Yet, we believe the absence of M1+M2 experiments is not a crucial flaw or deficiency.  Kingma et al’s results are not SOTA any longer, and the point of our paper is not semi-supervised learning (as Kingma et al.’s is).
Rather, our aim is to compare SB vs Gauss latent variables in a semi-supervised setting, and we believe our experiment accurately does this (while still having results partially comparable to Kingma et al’s, as explained in a pre-review response).
Thanks again,
Eric
"Are you referring to how the stick segments are recursively constructed? "
Yes.
"While you are correct that the stick segments lie on the simplex and therefore look as if they could be produced by a soft-gating function, this view neglects the underlying probabilistic model and the need for sampling from the simplex (for ELBO optimization).  Generating differentiable samples from the simplex can not be done via existing methods and was an open problem---for example, see the discussion of the NVI Topic Model in my response to Review 1."
Is the simplex in question the Stick-Breaking Weights ? I don't know if that's because it's obvious in hindsight, but I do not see the issue. You are generating differentiable sample v_{k}s from a ]0, 1[ interval and this very problem can be addressed using inverse transform sampling as you did using functions that only have values in ]0, 1[. The KL divergence are between variables that are in ]0, 1[ (with the Beta distribution) and the construction of the stick segment are part of the decoding process and not the inference process.
Using inverse transform sampling to generate differentiable samples for VAEs is not novel: it was mentioned in one of the original VAE papers by Kingma and Welling and was further explored in the Normalizing Flow paper by Rezende et al.
As to sample from a constrained set, it is indeed challenging but the constrained set in question ]0, 1[ can be easily solved.
"Their generative process is still a formal Bayes nonparametric model, just as ours is.  As we discuss in Section 4, our model does work with an adaptive number of latent variables; it’s just that optimization is slower.  "
Does this generative process put zero values on the decoder weights corresponding to \pi_{k}s that were not used during training ? The generative model from (Blei and Jordan, 2004) is not truncated despite using a truncated variational approximation.
I understand that there is a Bayesian nonparametric view behind it but I fail how this view has practical implication on the model if the number of latent variables is not adaptive. Even in the case of adaptive number of variables, I'm not sure I see a strong difference between the adaptation in the number of latent variables and component collapsing. I'll repeat my remark: "With respect to the Figure 5b showing the decoder input weights: component collapsing probably does not have the same effect as Gaussian prior. $\pi$ is positive therefore having a very small average value might mean that its value is close to zero most of the time, not requiring any update on the weight. For the standard Gaussian prior, component collapsing means having a very noisy input with no signal involved, which forces the decoder to shut down this channel, i.e. have small incoming weights from this collapsed variable."
I am not convinced to which extent this paper opens the path to link Bayesian nonparametric with VAEs yet and about the pruning comparison with the standard VAE. I do agree that the simplex representation seems to be more discriminative than a gaussian feature space.
Thanks for your responses and elaborations, Reviewer #2.  Here are our responses and comments...
> Is the simplex in question the Stick-Breaking Weights ?
Yes, the joint space.
> I don't know if that's because it's obvious in hindsight, but I do not see the issue....
As to sample from a constrained set, it is indeed challenging but the constrained set in question ]0, 1[ can be easily solved.
At the risk of being redundant, I’ll attempt to state our position in other words.  Perhaps this is just a point of fundamental disagreement...
There are two approaches to generating differentiable samples on the unit interval, as we see it, and one is easier than the other.  The easier approach is to appropriately transform an unbounded random variable, such as applying the logistic function to a Gaussian sample.  This is our Gauss-Logit baseline, which is also used in the ICLR submission “Neural Variational Inference for Topic Models” (https://openreview.net/forum?id=BybtVK9lg).  If this is all we used to sample the latent variables, your assessment would be fair and accurate.
The harder approach is to use a proper continuous distribution on the unit interval.  This too would be easy if the Beta had a non-centered parameterization---but it doesn’t.
At least two other recent papers have addressed this problem: ”The Generalized Reparametrization Gradient” (https://arxiv.org/abs/1610.02287) (GRG) and “Rejection Sampling Variational Inference” (https://arxiv.org/abs/1610.05683) (RSVI).  Both can produce differentiable Beta samples, but each has a drawback: the former is not a true DNCP (since there is still weak dependence on the variational parameters) and the latter, obviously, might require multiple attempts to produce a valid sample.
Our use of the Kumaraswamy is the only approach we are aware of that [1] has a variational approximation with the natural support (unlike the easier transformation-based approach) and [2] can be sampled via a true non-centered parametrization (unlike GRG and RSVI).
While our proposal is certainly not groundbreaking, having these two qualities makes it a notable contribution supplementary to the GRG and RSVI.
> Does this generative process put zero values on the decoder weights corresponding to \pi_{k}s that were not used during training ? The generative model from (Blei and Jordan, 2004) is not truncated despite using a truncated variational approximation.
Our generative model is not truncated: the decoder weight matrix is theoretically infinitely wide, but the columns at indices higher than the truncation threshold are censored (and hence, don’t need implemented) because they are multiplied by latent variables with \pi_k>T = 0 and thus never influence downstream computation.
> With respect to the Figure 5b showing the decoder input weights: component collapsing probably does not have the same effect as Gaussian prior. $\pi$ is positive therefore having a very small average value might mean that its value is close to zero most of the time, not requiring any update on the weight.  For the standard Gaussian prior, component collapsing means having a very noisy input with no signal involved, which forces the decoder to shut down this channel, i.e. have small incoming weights from this collapsed variable.
We agree that the SB-VAE’s latent variables at higher indices will have small values, and as a result, their corresponding decoder weights will receive small updates.  This is not to say those indices are ‘collapsed’ in the same manner as the Gauss-VAE’s decoder though.  The difference is that, by zeroing-out decoder weights, no data instance can ever use those latent dimensions of the Gaussian VAE.
However, in the SB-VAE, nothing is preventing an instance from using those higher indices if it needs to at any time (even though, in aggregate, those dimensions may look effectively inactive).
Thanks,
Eric
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.
Thanks for your thoughts, Reviewer #1.  Our responses are below, paired with the corresponding segment of your review.
1. "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."
Thanks, this is a good point.  We’ll make different subsections for the generative and inference processes in Section 4 and include and/or elaborate on our answers to your pre-review questions.
2. "Firstly, conditioning the model directly on the stick-breaking weights seems a little odd….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….
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."
Yes, we agree that choosing the simplex for the latent space is one of multiple options.  Other formulations---such as using the remaining length of the stick---are possible.
Ultimately, we made the choice to use the simplex for the following reasons: (1) if the decoder is sufficiently powerful, the simplex should have as much representational capacity as Euclidean space does for the Gaussian VAE, (2) using the weights keeps the model closest to its natural extension to a mixture VAE, and lastly, (3) we found the formulation had some interesting properties, such as being better at classification, that seemed worthy of reporting.  We believe the experiments generating MNIST samples from increasingly higher dimensional latent spaces and examining the number of latent variables used for MNIST+rot provided insight into the model’s nonparametric qualities to some degree.
As for the paper’s high level motivation: our original goal was to connect SGVB and Bayesian Nonparametrics.  Finding a differentiable non-centered parametrization (DNCP) for the Stick-Breaking process was the central challenge to this, and once we found the Kumaraswamy, we formulated the SB-VAE to test if optimization of this somewhat unusual computation graph (i.e. the recurrent dependence across the hidden layer) and this little known distribution was possible.  It turned out that it was, and in addition, we saw some interesting properties that seemed useful for people to know.
Yet, we agree that, while the SB-VAE is a VAE alternative one should try if unhappy with the Gauss-VAE’s performance, the use of the Kumaraswamy as a DNCP for Beta-like or Dirichlet-like random variables is the paper’s most widely applicable lesson.
In fact, in light of recently proposed VAE extensions, the usefulness of the Kumaraswamy is more apparent now than when we began this work.
For instance, current ICLR submissions such as “Deep Unsupervised Clustering with Gaussian Mixture Variational Autoencoders” and “Neural Variational Inference for Topic Models” both workaround the difficulty of priors on the simplex by either, in the case of the former, just not placing priors on the mixture weights, and in the case of the latter, using a Laplace approximation.  The latter work even acknowledges the problem that our paper solves: “But it is difficult to handle the Dirichlet within NVI because it is difficult to develop an effective reparameterization function for the [Reparametrization Trick].”  Moreover, the recent work from the Blei lab still does not describe a true DNCP for the Beta, their solutions being either weakly dependent on the approx. posterior parameters (“The Generalized Reparametrization Gradient, NIPS 2016) or requiring rejection sampling (“Rejection Sampling Variational Inference”, ArXiv).
3. "Perhaps some comparison with a dropout approach or a more thorough discussion related to dropout would make this clearer."
This is a good idea to improve motivation; we can elaborate on the connection to Nested Dropout, perhaps moving it out of the ‘Related Work’ section.
However, for experimental comparisons, we found optimization under Nested Dropout very difficult, even with unit sweeping.  The models were not competitive and thus left out of experiments.  Comparing against the Infinite RBM is another option, but to make comparison fair, we’d have to tie the encoder and decoder weights (since RBMs have only one weight matrix), and it is not totally clear how to do this in a VAE since the networks are asymmetric.
Thanks again,
Eric

Meta-review:
role: PC
data: comment: This paper will make a positive contribution to the conference, especially since it is one of the first to look at stick-breaking as it applies to deep generative models. The paper will make a positive contribution to the conference.

Decision: Accept (Poster)