PaperlayerView exact revision on arXiv

Source-linked AI summary

Stick-Breaking Variational Autoencoders

Eric Nalisnick, Padhraic Smyth

arXiv:1605.06197v3stat.ML

TL;DR

The paper addresses the absence of SGVB inference for Bayesian nonparametric stick-breaking processes. It introduces Kumaraswamy-based inference and uses it to construct SB-VAE models with stochastic latent dimensionality, which experimentally learn more discriminative representations than Gaussian variants. The approach also has limitations from variational truncation and evaluation comparability in one setting.

  • Problem

    SGVB lacked an extension to Bayesian nonparametric processes, limiting its application to models with adaptive latent capacity.

  • Method

    The paper uses the Kumaraswamy distribution to perform SGVB inference for stick-breaking weights and constructs unsupervised and semi-supervised nonparametric VAE variants.

  • Results

    The SB-VAE and semi-supervised variant learn highly discriminative latent representations, with the SB-DGM performing markedly better across almost all semi-supervised experiments.

  • Takeaways & Limitations

    Stick-breaking priors provide differentiable control of model capacity with only linear extra cost in the truncation size.

  • Takeaways & Limitations

    The variational posterior is truncated in the presented approach, and one unsupervised comparison is not strictly comparable because GMVAE uses cluster-level labeling rather than kNN.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper extends SGVB toward Bayesian nonparametric models by enabling inference for stick-breaking weights. It then uses this development to build nonparametric VAE variants with adaptive latent capacity and more discriminative representations.

  • SGVB combines deep neural networks’ expressiveness with probabilistic latent variables’ robustness to uncertainty.
  • The paper addresses SGVB’s conspicuous lack of an extension to Bayesian nonparametric processes.
  • The authors use the Kumaraswamy distribution to overcome the absence of a differentiable non-centered parametrization for Beta approximate posteriors.
  • The resulting VAE and semi-supervised M2 variants use an infinite-capacity hidden layer with as many stick segments as the data requires.
  • For natural-image datasets, stick-breaking priors produce latent representations that better preserve class boundaries and regularize semi-supervised learning.

2 BACKGROUND

The background introduces VAEs, SGVB, and stick-breaking priors as the components combined by the paper. VAEs map between latent and observed variables, SGVB optimizes an ELBO through differentiable sampling, and stick-breaking priors provide infinitely many normalized weights.

  • A VAE uses a density network to map latent variables to observations and an inference network to map observations back to latent variables.
  • The VAE’s prior and approximate posterior have traditionally been modeled as marginally Gaussian.
  • SGVB estimates generative and variational parameters using differentiable Monte Carlo expectations of the evidence lower bound.
  • SGVB requires a differentiable non-centered parametrization that exposes variational parameters while drawing randomness from a fixed distribution.
  • A stick-breaking prior represents a random measure as infinitely many weighted atoms whose weights sum to one.
  • When the break variables follow Beta(1, α0), their infinite weight sequence follows the GEM distribution with concentration parameter α0.

3 SGVB FOR GEM RANDOM VARIABLES

The paper develops SGVB inference for stick-breaking weights by replacing the difficult Beta posterior parametrization with alternatives that support differentiable sampling. The Kumaraswamy distribution is especially useful because it offers a closed-form inverse CDF and an approximable KL divergence from the Beta distribution.

  • The paper focuses on SGVB inference for the sequence of stick-breaking weights, called GEM random variables, while leaving inference for the random measure to future work.
  • The standard Gamma composition of a Beta variable does not provide the differentiable non-centered parametrization required by SGVB because Gamma shape parameters lack one.
  • The asymptotic inverse-CDF approximation for Gamma variables becomes poor as the shape parameter increases, motivating a finite-difference approximation when a ≥ 1.
  • The Kumaraswamy is a two-parameter continuous distribution on the unit interval that can serve as an alternative approximate posterior.
  • Its closed-form inverse CDF supplies the differentiable non-centered parametrization needed for sampling.
  • The Kumaraswamy–Beta KL divergence can be closely approximated in closed form for ELBO computation.
  • A Probit stick-breaking parametrization is made practical by replacing the Gaussian CDF with the closed-form logistic function.

4 STICK-BREAKING VARIATIONAL AUTOENCODERS

The SB-VAE replaces Gaussian latent variables with GEM stick-breaking weights and uses a neural-network inference process to construct the latent representation. Its variational posterior is truncated for computation, while the prior remains infinite-dimensional.

  • The SB-VAE draws latent variables from the GEM distribution, representing the hidden representation as an infinite sequence of stick-breaking weights.
  • The model uses a feedforward architecture to parametrize the infinite-dimensional stick-breaking graphical model.
  • The generative process uses stick-breaking weights π_i with concentration parameter α_0 and a density network likelihood p_θ(x_i|π_i).
  • The inference network samples K fraction variables v_i,k and composes stick segments through a linear-time operation.
  • Inference: The Kth fraction v_i,K is fixed to one so the stick segments sum to one, but this variational truncation does not make the prior finite-dimensional.
  • Inference: Truncation-free posterior experiments slowed optimization without increasing performance, motivating the truncated posterior used here.

5 SEMI-SUPERVISED MODEL

The semi-supervised model extends the stick-breaking VAE with a categorical class-label variable and adapts its objectives according to whether labels are observed. The label and latent-variable inference pathways share parameters and thereby regularize one another.

  • The semi-supervised model introduces a categorical latent variable y_i representing the class label alongside the stick-breaking latent representation.
  • The label posterior q_φ(y_i|x_i) is produced by a nonlinear inference-network function g_y(x_i).
  • Although y and z are modeled as independent in the posterior factorization, they share inference-network parameters and regularize one another.
  • Semi-supervised training uses different objectives depending on whether the label is present or missing, including label log-likelihood for observed labels and entropy terms for missing labels.

6 RELATED WORK

The paper positions its contribution as combining amortized SGVB inference with nonparametric generative priors, distinguishing it from prior work with nonparametric variational distributions or nondifferentiable discrete latent features.

  • The authors state that SGVB and other amortized variational-inference methods had not previously been used with Bayesian nonparametric priors.
  • Prior Variational Gaussian Process work used a nonparametric variational distribution rather than a nonparametric generative model.
  • Adaptive latent-factor models increase latent dimensionality as data require, but the Indian Buffet Process uses discrete latent variables that prevent fully differentiable inference.

7 EXPERIMENTS

The experiments compare stick-breaking and Gaussian latent variables across reconstruction, likelihood, representation structure, and semi-supervised classification. SB-VAE learns more slowly and has worse likelihood, but shows stronger discriminative structure and adaptive latent usage in several evaluations.

  • Experimental setup: Experiments used Frey Faces, MNIST, MNIST+rot, and SVHN, with matched architectures and optimization hyperparameters where Gaussian and stick-breaking models differed primarily in their priors.MNIST+rot was created to test whether rotated digits use more latent variables than non-rotated digits.
  • Density Estimation: SB-VAE optimization proceeds similarly but at a slightly slower pace than the Gaussian VAE across Frey Faces, MNIST, and MNIST+rot.The authors attribute the slower pace to coupled gradients from the recursive latent-variable definition.
  • Density Estimation: The Gaussian VAE achieves better marginal likelihood than all stick-breaking implementations, at approximately 96 versus approximately 98.Among stick-breaking parametrizations, Kumaraswamy outperforms Gamma and Gauss-Logit on both datasets.
  • Adaptive capacity: Rotated MNIST images use 28.7 latent dimensions on average versus 27.4 for non-rotated images, supporting adaptive dimensionality.The model measured the number of breaks required to represent 99% of the stick.
  • Discriminative Qualities: SB-VAE latent representations outperform the Gaussian VAE in kNN classification at every tested k, while t-SNE shows more cohesive and separated digit clusters.The comparison indicates stronger class structure in the stick-breaking latent space despite the Gaussian VAE's better likelihood.
  • Combating Decoder Pruning: The Gaussian VAE prunes decoder weights for unused dimensions, whereas SB-VAE retains decoder weights despite sparse latent representations and shows no apparent component collapsing.The authors conjecture that this increased capacity contributes to better discriminative performance.

8 CONCLUSIONS

The paper extends SGVB to stick-breaking Bayesian nonparametric priors and uses this development to build deep generative models with infinite-dimensional latent variables. These models produce more discriminative latent representations than the popular Gaussian variant, with only linear extra computational cost in the truncation size.

  • The paper extends Stochastic Gradient Variational Bayes to the weights of stick-breaking Bayesian nonparametric priors.
  • The resulting models use infinite-dimensional latent variables and have differentiable control of their capacity.
  • Their latent representations are more discriminative than those of the popular Gaussian variant.
  • The additional computational cost is limited to assembling stick segments, a linear operation in the truncation size.

APPENDIX

The appendix gives the KL divergence between the Kumaraswamy approximate posterior and Beta prior. Its infinite sum arises from a Taylor expansion and is expected to be well approximated by the first few terms.

  • The appendix defines the KL divergence between a Kumaraswamy approximate posterior and a Beta prior.
  • The Kumaraswamy distribution is parameterized by a and b, while the Beta distribution is parameterized by α and β.
  • The KL expression includes Euler’s constant, the Digamma function, the Beta function, and an infinite sum.
  • The infinite sum appears because a Taylor expansion represents Eq[log(1 − v_k)], and the first few terms should approximate it well.

EXPERIMENTS AND OPTIMIZATION

The experiments use released implementations, fixed train-validation-test splits, common optimization settings, and matched architectures across Gaussian and stick-breaking models. The main controlled difference is the latent prior, with the stick-breaking concentration parameter cross-validated.

  • The authors released Theano implementations and ran all experiments on AWS G2.2XL instances.
  • The datasets use specified train-validation-test splits for Frey Faces, MNIST, MNIST+rot, and SVHN.
  • SVHN alone was preprocessed with PCA to 500 dimensions capturing 99.9% of the data variance.
  • All models used minibatches of 100 and AdaM with α = 0.0003, b1 = 0.95, and b2 = 0.999.
  • The latent dimensionality or truncation level was 50 except for Frey Faces, with ReLU activations and no regularization.
Loading 1605.06197v3…