Auto-Encoding Variational Bayes

The problem

Diederik P. Kingma and Max Welling posted "Auto-Encoding Variational Bayes" to arXiv on 20 December 2013. The paper took on efficient inference and learning in directed probabilistic models that have continuous latent variables, intractable posterior distributions and large datasets.^[1]

A directed probabilistic model describes data through hidden quantities — the latent variables — plus a graph of conditional dependencies linking them to the observations. Learning such a model requires the posterior: the distribution over the latent variables given the data. For continuous latents that posterior usually has no closed form, and the integral that defines it cannot be evaluated exactly. Classical answers either restrict the model family so that the posterior stays tractable, or run a sampling procedure per data point, which becomes prohibitively expensive when the dataset is large. The paper's setting — continuous latents, an intractable posterior, and data at scale — is exactly where those answers break down.^[1]

The method

The authors made two moves. The first is a reparameterisation of the variational lower bound. Variational methods replace the true posterior with a simpler approximate distribution and then push that approximation as close to the posterior as possible by maximising a lower bound on the data likelihood. The obstacle is that the bound contains an expectation over samples drawn from the approximation, and a sampling step sitting in the middle of a computation blocks gradients. The reparameterisation rewrites a sample from the approximate posterior as a deterministic, differentiable transformation of the distribution's parameters and an independent noise variable. Sampling then looks like any other differentiable operation, and gradients can pass through it to the parameters that generated the sample.^[1]

With that rewrite, the authors obtained a lower bound estimator that can be optimised directly with standard stochastic gradient methods — the same machinery used to train neural networks — under mild differentiability conditions, and they showed that it still works in the intractable case.^[1] The second move concerns inference itself. For i.i.d. datasets with continuous latent variables per data point, the authors showed that posterior inference can be made efficient by fitting an approximate inference model, which they call a recognition model, to the intractable posterior. Instead of solving an optimisation problem afresh for every data point, one trains a single model that maps an observation to the parameters of its approximate posterior.^[1]

What the paper showed

The results reported in the paper are properties of the estimator and of the inference scheme rather than benchmark scores. The reparameterised lower bound yields an estimator that is straightforward to optimise with standard stochastic gradient methods, and that estimator scales to large datasets.^[1] Posterior inference, in turn, becomes efficient through the recognition model for i.i.d. data with continuous latents.^[1] Together these two claims cover the whole loop: a generative model and an approximate inference model can both be trained by gradient descent on the same objective.

Why it matters

The significance lies in the connection the paper draws between two normally separate worlds. Approximate Bayesian inference had its own algorithms and its own computational habits; neural-network training had stochastic gradient descent. By expressing the variational objective in a form whose only requirements are differentiability and the ability to sample, the paper turned inference into an optimisation problem of the ordinary kind, one that can be run on data in small random batches and scaled up with the dataset.^[1] The recognition model does the same for the per-datapoint cost: amortising inference across the dataset is what makes the approach usable when there are many observations rather than one.^[1] The paper's own framing of the problem — a lower bound estimator plus a learned approximation to an intractable posterior — is the contribution it claims, and that framing is what makes latent-variable models with continuous hidden variables trainable at large-dataset scale.