Neural Variational Inference: Variational Autoencoders and Helmholtz machines

So far we had a little of “neural” in our VI methods. Now it’s time to fix it, as we’re going to consider Variational Autoencoders (VAE), a paper by D. Kingma and M. Welling, which made a lot of buzz in ML community. It has 2 main contributions: a new approach (AEVB) to large-scale inference in non-conjugate models with continuous latent variables, and a probabilistic model of autoencoders as an example of this approach. We then discuss connections to Helmholtz machines — a predecessor of VAEs.

Auto-Encoding Variational Bayes

As noted in the introduction of the post, this approach, called Auto-Encoding Variational Bayes (AEVB) works only for some models with continuous latent variables. Recall from our discussion of Blackbox VI and Stochastic VI, we’re interested in maximizing the ELBO L(Θ,Λ)$\mathcal{L}(\mathrm{\Theta },\mathrm{\Lambda })$:

L(Θ,Λ)=Eq(z∣x,Λ)logp(x,z∣Θ)q(z∣x,Λ)$$\mathcal{L}(\mathrm{\Theta },\mathrm{\Lambda })={\mathbb{E}}_{q(z\mid x,\mathrm{\Lambda })}\log ?\frac{p(x,z\mid \mathrm{\Theta })}{q(z\mid x,\mathrm{\Lambda })}$$

It’s not a problem to compute an estimate of the gradient of the ELBO w.r.t. model parameters Θ$\mathrm{\Theta }$, but estimating the gradient w.r.t. approximation parameters Λ$\mathrm{\Lambda }$ is tricky as these parameters influence the distribution the expectation is taken over, and as we know from the post on Blackbox VI, naive gradient estimator based on score function exhibits high variance.

Turns out, for some distributions we can make change of variables, that is, for some distributions z～q(z∣x,Λ)$z～q(z\mid x,\mathrm{\Lambda })$ can be represented as a (differentiable) transformation g(ε;Λ,x)$g(\epsilon ;\mathrm{\Lambda },x)$ of some auxiliary random variable ε$\epsilon$ whose distribution does not depend on Λ$\mathrm{\Lambda }$. A well-known example of such reparametrization is Gaussian distribution: if z～N(μ,Σ)$z～\mathcal{N}(\mu ,\mathrm{\Sigma })$ then z$z$ can be represented as z=μ+Σ1/2ε$z=\mu +{\mathrm{\Sigma }}^{1/2}\epsilon$ for ε～N(0,I)$\epsilon ～\mathcal{N}(0,I)$. This transformation is called the reparametrization trick. After the reparametrization the ELBO becomes

L(Θ,Λ)=Eε～N(0,I)logp(x,g(ε;Λ,x)∣Θ)q(g(ε;Λ,x)∣Λ,x)≈1L∑l=1Llogp(x,g(ε(l);Λ,x)∣Θ)q(g(ε(l);Λ,x)∣Λ,x)where ε(l)～N(0,I)$$\begin{array}{rl}\mathcal{L}(\mathrm{\Theta },\mathrm{\Lambda })& ={\mathbb{E}}_{\epsilon ～\mathcal{N}(0,I)}\log ?\frac{p(x,g(\epsilon ;\mathrm{\Lambda },x)\mid \mathrm{\Theta })}{q(g(\epsilon ;\mathrm{\Lambda },x)\mid \mathrm{\Lambda },x)}\\ & \approx \frac{1}{L}\sum _{l=1}^L\log ?\frac{p(x,g({\epsilon }^{(l)};\mathrm{\Lambda },x)\mid \mathrm{\Theta })}{q(g({\epsilon }^{(l)};\mathrm{\Lambda },x)\mid \mathrm{\Lambda },x)}\text{where }{\epsilon }^{(l)}～\mathcal{N}(0,I)\end{array}$$

This objective is much better as we don’t need to differentiate w.r.t. expectation’s distribution, essentially putting variational parameters Λ$\mathrm{\Lambda }$ to the same regime as model parameters Θ$\mathrm{\Theta }$. It’s sufficient now to just take gradients of the ELBO’s estimate, and run any optimization algorithm like Adam.

Oh, and if you wonder what Auto-Encoding in Auto-Encoding Variational Bayes means, there’s an interesting interpretation of the ELBO in terms of autoencoding:

L(Θ,Λ)=Eq(z∣x,Λ)logp(x,z∣Θ)q(z∣x,Λ)=Eq(z∣x,Λ)logp(x∣z,Θ)p(z∣Θ)q(z∣x,Λ)=Eq(z∣x,Λ)logp(x∣z,Θ)?DKL(q(z∣Λ,x)∣∣p(z∣Θ))$$\begin{array}{rl}\mathcal{L}(\mathrm{\Theta },\mathrm{\Lambda })& ={\mathbb{E}}_{q(z\mid x,\mathrm{\Lambda })}\log ?\frac{p(x,z\mid \mathrm{\Theta })}{q(z\mid x,\mathrm{\Lambda })}={\mathbb{E}}_{q(z\mid x,\mathrm{\Lambda })}\log ?\frac{p(x\mid z,\mathrm{\Theta })p(z\mid \mathrm{\Theta })}{q(z\mid x,\mathrm{\Lambda })}\\ & ={\mathbb{E}}_{q(z\mid x,\mathrm{\Lambda })}\log ?p(x\mid z,\mathrm{\Theta })?D_{KL}(q(z\mid \mathrm{\Lambda },x)\mid \mid p(z\mid \mathrm{\Theta }))\end{array}$$

Here the first term can be treated as expected reconstruction (x$x$ from the code z$z$) loss, while the second one is just a regularization term.

Variational Autoencoder

One particular application of AEVB framework comes from using neural networks as the model p(x∣z,Θ)$p(x\mid z,\mathrm{\Theta })$ and the approximation q(z∣x,Λ)$q(z\mid x,\mathrm{\Lambda })$. The model has no requirements, and x$x$ can be discrete or continuous (or mixed). z$z$, however, has to be continuous. Moreover, we need to be able to apply the reparametrization trick. Therefore in many practical applications q(z∣x,Λ)$q(z\mid x,\mathrm{\Lambda })$ is set to be Gaussian distribution q(z∣Λ,x)=N(z∣μ(x;Λ),Σ(x;Λ))$q(z\mid \mathrm{\Lambda },x)=\mathcal{N}(z\mid \mu (x;\mathrm{\Lambda }),\mathrm{\Sigma }(x;\mathrm{\Lambda }))$ where μ$\mu$ and Σ$\mathrm{\Sigma }$ are outputs of a neural network taking x$x$ as input, and Λ$\mathrm{\Lambda }$ denotes a set of neural network’s weights — the parameters you optimize the ELBO with respect to (and also Θ$\mathrm{\Theta }$). In order to make reparametrization trick practical, you’d like to be able to compute Σ1/2${\mathrm{\Sigma }}^{1/2}$ quick. You don’t want to actually compute this quantity as it’d be too computationally expensive. Instead you might want to predict Σ1/2${\mathrm{\Sigma }}^{1/2}$ by a neural network in the first place, or consider only diagonal covariance matrices (as it’s done in the paper).

In case of Gaussian approximation q(z∣x,Λ)$q(z\mid x,\mathrm{\Lambda })$ and Gaussian prior p(z∣Θ)$p(z\mid \mathrm{\Theta })$ we can compute KL-divergence DKL(q(z∣Λ,x)∣∣p(z∣Θ))$D_{KL}(q(z\mid \mathrm{\Lambda },x)\mid \mid p(z\mid \mathrm{\Theta }))$ analytically, see the formula at stats.stackexchange. This reduces variance of gradient estimator, though one can still train a VAE estimating KL-divergence using Monte Carlo, just like the other part of the ELBO.

We optimize both the model and the approximation by gradient ascent. This joint optimization pushed both approximation towards the model, and the model towards approximation. This leads not only to efficient inference using the approximation, but also the model is encouraged to learn latent representations z$z$ such that the true posterior p(z∣x,Θ)$p(z\mid x,\mathrm{\Theta })$ is approximately factorial.

This model has generated a lot of buzz because it can be used as a generative model, essentially VAE is an autoencoder with natural sampling procedure: suppose you’ve trained the model, and now want to sample new samples similar to those you used in the training set. To do so you first sample z$z$ from the prior p(z)$p(z)$, and then generate x$x$ using the model p(x∣z,Θ)$p(x\mid z,\mathrm{\Theta })$. Both operations are easy: the first one is a sampling from some standard distribution (like Gaussian, for example), and the second one is just one feed-forward pass followed by another sampling from another standard distribution (Bernoulli, for example, in case x$x$ is a binary image).

If you want to read more on Variational Auto-Encoders, I refer you to a great tutorial by Carl Doersch.

Helmholtz machines

In the end I’d like to add some historical perspective. The idea of two networks, one “encoding” an observation x$x$ to some latent representation (code) z$z$, and another “decoding” it back is definitely not new. In fact, the whole idea is a special case of the Helmholtz Machines introduced by Geoffrey Hinton 20 years ago.

Helmholtz machine can be thought of as a neural network of stochastic hidden layers. Namely, we now have M$M$ stochastic hidden layers (latent variables) h1,…,hM$h_1,\dots ,h_M$ (with deterministic h0=x$h_0=x$) where the layer hk?1$h_{k?1}$ is stochastically produced by the layer hk$h_k$, that is, it is samples from some distribution p(hk?1∣hk)$p(h_{k?1}\mid h_k)$, which as you might have guessed already is parametrized in the same way as in usual VAEs. Actually, VAEs is a special case of a Helmholtz machine with just one stochastic layer (but each stochastic layer contains a neural network of arbitrarily many deterministic layers inside of it).

This image shows an instance of a Helmholtz machine with 2 stochastic layers (blue cloudy nodes), and each stochastic layer having 2 deterministic hidden layers (white rectangles).

The joint model distribution is

p(x,h1,…,hM∣Θ)=p(hM∣Θ)∏m=0M?1p(hm∣hm+1,Θ)$$p(x,h_1,\dots ,h_M\mid \mathrm{\Theta })=p(h_M\mid \mathrm{\Theta })\prod _{m=0}^{M?1}p(h_m\mid h_{m+1},\mathrm{\Theta })$$

And the approximate posterior is the same, but in inverse order:

q(h1,…,hM∣x,Λ)=∏m=1Mp(hm∣hm?1,Θ)$$q(h_1,\dots ,h_M\mid x,\mathrm{\Lambda })=\prod _{m=1}^{M}p(h_m\mid h_{m?1},\mathrm{\Theta })$$

The p(x,h1,…,hM?1∣hM)$p(x,h_1,\dots ,h_{M?1}\mid h_M)$ distribution is usually called a generative network (or model) as it allows one to generate samples from latent representation(s). The approximate posterior q(h1,…,hM∣x,Λ)$q(h_1,\dots ,h_M\mid x,\mathrm{\Lambda })$ in this framework is called a recognition network (or model). Presumably, the name reflects the purpose of the network to recognize the hidden structure of observations.

So, if the VAE is a special case of Helmholtz machines, what’s new then? The standard algorithm for learning Helmholtz machines, the Wake-Sleep algorithm, turns out to be optimizing a different objective. Thus, one of significant contributions of Kingma and Welling is application of the reparametrization trick to make optimization of the ELBO w.r.t. Λ$\mathrm{\Lambda }$ tractable.