Variational Autoencoders (VIII): Estimating Sample Probability Density

变分自编码器（八）：估计样本概率密度 苏剑林 · 2021-12-09 · Machine Learning probability variational VAE generative model

In the earlier posts of this series, we've come to understand VAEs from several different angles. Generally speaking, people use a VAE to obtain a generative model, or to build a better encoding model — these are the conventional uses of VAE. But beyond these conventional applications, there are also some "niche needs," such as using it to estimate the probability density of $x$, which often comes up in compression tasks.

This post approaches the VAE model from the angle of estimating probability density, and works through its derivation from that perspective.

Two Problems

So-called "estimating probability density" means: given a batch of samples $x_1,x_2,\cdots,x_N\sim \tilde{p}(x)$, we use an undetermined family of probability densities $q_{\theta}(x)$ to fit these samples, where the fitting objective is generally to minimize the negative log-likelihood:

\begin{equation}\mathbb{E}_{x\sim \tilde{p}(x)}[-\log q_{\theta}(x)] = -\frac{1}{N}\sum_{i=1}^N \log q_{\theta}(x_i)\label{eq:mle}\end{equation}more

But this is purely a formal statement — there are still many issues left unresolved, which can broadly be grouped into two big questions:

1. What kind of $q_{\theta}(x)$ should we use for fitting?
2. What method should we use to solve the above objective?

Mixture Models

For the first question, naturally we hope that $q_{\theta}(x)$ has as strong a fitting capacity as possible — ideally, it should be able to fit any probability distribution. Unfortunately, although neural networks theoretically have universal approximation power, that power applies to fitting functions, not to fitting probability distributions. A probability distribution needs to satisfy $q_{\theta}(x)\geq 0$ and $\int q_{\theta}(x) dx=1$, and the latter condition is usually hard to guarantee.

Since we can't do this directly, let's think about it indirectly, by constructing a mixture model:

\begin{equation}q_{\theta}(x) = \int q_{\theta}(x|z)q(z)dz=\mathbb{E}_{z\sim q(z)}[q_{\theta}(x|z)]\label{eq:q}\end{equation}

Here $q(z)$ is usually chosen to be a simple, parameter-free distribution, such as the standard normal distribution; while $q_{\theta}(x|z)$ is a parameterized simple distribution conditioned on $z$ — for example, a standard normal distribution whose mean and variance depend on $z$.

From the perspective of generative modeling, the above model is interpreted as a two-step process: first sample $z$ from $q(z)$, then feed it into $q_{\theta}(x|z)$ to generate $x$. But the focus of this post is estimating probability density, and the reason we choose such a $q_{\theta}(x|z)$ is that it has strong enough capacity to fit complex distributions. The resulting $q_{\theta}(x)$ is expressed as an average over many simple distributions $q_{\theta}(x|z)$. Readers familiar with Gaussian mixture models will know that such a model can achieve very strong fitting power — in principle it can even fit an arbitrary distribution. So the fitting capacity of the distribution is guaranteed.

Importance Sampling

However, equation $\eqref{eq:q}$ cannot be simply integrated out in closed form — or rather, it is precisely because such a distribution cannot be expressed simply and explicitly that it has strong enough fitting power. So to estimate it, we must resort to sampling-based estimation in the form of $\mathbb{E}_{z\sim q(z)}[q_{\theta}(x|z)]$. In practice, though, $z$ and $x$ tend to be high-dimensional, and high-dimensional spaces suffer from the "curse of dimensionality" — meaning that even if we sample millions or tens of millions of points, it's very hard to adequately cover a high-dimensional space, which in turn means it's very hard to accurately estimate $\mathbb{E}_{z\sim q(z)}[q_{\theta}(x|z)]$.

To address this, we need to find a way to shrink the sampling space. First, we typically keep the variance of $q_{\theta}(x|z)$ fairly small. This way, for a given $x$, there won't be too many values of $z$ that make $q_{\theta}(x|z)$ large; for most $z$, the computed $q_{\theta}(x|z)$ will be very close to zero. So we just need to find a way to sample values of $z$ that make $q_{\theta}(x|z)$ large, and we can obtain a good estimate of $\mathbb{E}_{z\sim q(z)}[q_{\theta}(x|z)]$.

Specifically, we introduce a new distribution $p_{\theta}(z|x)$, and assume that the values of $z$ that make $q_{\theta}(x|z)$ large follow this distribution. Then we have

\begin{equation}q_{\theta}(x) = \int q_{\theta}(x|z)q(z)dz=\int q_{\theta}(x|z)\frac{q(z)}{p_{\theta}(z|x)}p_{\theta}(z|x)dz=\mathbb{E}_{z\sim p_{\theta}(z|x)}\left[q_{\theta}(x|z)\frac{q(z)}{p_{\theta}(z|x)}\right]\end{equation}

In this way, we convert "aimless" sampling from $q(z)$ into more targeted sampling from $p_{\theta}(z|x)$. Since the variance of $q_{\theta}(x|z)$ is kept small, the variance of $p_{\theta}(z|x)$ naturally won't be large either, and sampling efficiency improves. Note that from the generative-model viewpoint, $p_{\theta}(z|x)$ is regarded as an approximation to the posterior distribution; but from the viewpoint of estimating probability density, it is really just a pure importance-weighting function — there's no need to give it any special interpretation.

Training Objective

At this point, we've addressed the first problem: what distribution to use, and how to compute it more efficiently. What remains is how to train the model.

Actually, once we have the concept of importance sampling in hand, we no longer need to worry about things like the ELBO — we can just directly work with the objective $\eqref{eq:mle}$. Substituting in the expression for $q_{\theta}(x)$ gives

\begin{equation}\mathbb{E}_{x\sim \tilde{p}(x)}\left[-\log \mathbb{E}_{z\sim p_{\theta}(z|x)}\left[q_{\theta}(x|z)\frac{q(z)}{p_{\theta}(z|x)}\right]\right]\end{equation}

In fact, if at the step $\mathbb{E}_{z\sim p_{\theta}(z|x)}$ we use the reparameterization trick to draw just a single sample $z$, the training objective becomes

\begin{equation}\mathbb{E}_{x\sim \tilde{p}(x)}\left[-\log q_{\theta}(x|z)\frac{q(z)}{p_{\theta}(z|x)}\right],\quad z\sim p_{\theta}(z|x)\end{equation}

This is, in fact, already the standard VAE training objective. If instead we draw $M > 1$ samples, we get

\begin{equation}\mathbb{E}_{x\sim \tilde{p}(x)}\left[-\log \left(\frac{1}{M}\sum_{i=1}^M q_{\theta}(x|z_i)\frac{q(z_i)}{p_{\theta}(z_i|x)}\right)\right],\quad z_1,z_2,\cdots,z_M\sim p_{\theta}(z|x)\end{equation}

This is exactly the "Importance Weighted Autoencoder," from Importance Weighted Autoencoders, which is regarded as a strengthened version of VAE. In summary, by taking the importance sampling perspective, we can bypass the tedious ELBO-based derivations of the conventional VAE, and we can also skip the joint-distribution perspective introduced in Variational Autoencoders (II): From a Bayesian Viewpoint, arriving directly at the VAE model and even its improved variants.

Summary

This post introduced the VAE starting from the goal of estimating the probability density of samples. By combining this with importance sampling, we obtain a quick derivation of VAE that entirely sidesteps the ELBO and its many tedious details.