Variational Autoencoders (Part 2): Starting from a Bayesian Perspective
Origins
A few days ago I wrote the post Variational Autoencoders (Part 1): So That's What's Going On, approaching the variational autoencoder (VAE) from a fairly intuitive angle. In that post's view, a VAE differs little from an ordinary autoencoder — it merely adds noise and imposes a constraint on that noise. However, my original motivation for wanting to understand VAEs was to see exactly how the Bayesian school's probabilistic graphical models combine with deep learning to actually do their work. If all I get is an intuitive understanding, that's clearly not enough.
So I kept thinking about VAEs for a few more days, trying to explain them using more general, probabilistic language. As it turns out, this way of thinking also answers questions that the intuitive understanding can't — such as whether MSE or cross-entropy is better for the reconstruction loss, and how the reconstruction loss and KL loss should be balanced, and so on.
I'd suggest reading this after Variational Autoencoders (Part 1): So That's What's Going On; this post will try not to repeat the content of the previous one.
Preliminaries
Before getting into the description of VAEs, I think it's worth going over some conceptual material first. more
Numerical Computation vs. Sampling-Based Computation
For readers who aren't very familiar with probability and statistics, the two concepts that are easy to confuse are numerical computation and sampling-based computation — a few readers had the same confusion in The Three Flavors of Capsules: Matrix Capsules and EM Routing. For example, given a probability density function $p(x)$, the expectation of $x$ is defined as
$$\mathbb{E}[x] = \int x p(x)dx\tag{1}$$
If we want to compute this numerically (numerical integration), we can choose a number of representative points $x_0 < x_1 < x_2 < \dots < x_n$, giving
$$\mathbb{E}[x] \approx \sum_{i=1}^n x_i p(x_i) \left(x_i - x_{i-1}\right)\tag{2}$$
We won't get into what "representative" means here, nor discuss ways of improving numerical accuracy. The point of writing it this way is to compare it with sampling-based computation. If we instead draw a number of samples $x_1,x_2,\dots,x_n$ from $p(x)$, then we have
$$\mathbb{E}[x] \approx \frac{1}{n}\sum_{i=1}^n x_i,\quad x_i \sim p(x)\tag{3}$$
We can compare $(2)$ with $(3)$: their main difference is that $(2)$ includes an explicit probability computation, while $(3)$ only involves computing $x$. This is because in $(3)$, $x_i$ is drawn from $p(x)$ according to its probability — points $x_i$ with higher probability show up more often — so we can say the sampling process has already baked in $p(x)$, and there's no need to multiply by $p(x_i)$ again.
More generally, we can write
$$\mathbb{E}_{x\sim p(x)}[f(x)] = \int f(x)p(x)dx \approx \frac{1}{n}\sum_{i=1}^n f(x_i),\quad x_i\sim p(x)\tag{4}$$
and this is the basis of Monte Carlo simulation.
KL Divergence and the Calculus of Variations
We typically use KL divergence to measure the difference between two probability distributions $p(x)$ and $q(x)$, defined as
$$KL\Big(p(x)\Big\Vert q(x)\Big) = \int p(x)\ln \frac{p(x)}{q(x)} dx=\mathbb{E}_{x\sim p(x)}\left[\ln \frac{p(x)}{q(x)}\right]\tag{5}$$
The key property of KL divergence is non-negativity: if we fix $p(x)$, then $KL\Big(p(x)\Big\Vert q(x)\Big)=0 \Leftrightarrow p(x)=q(x)$; if we fix $q(x)$, we likewise get $KL\Big(p(x)\Big\Vert q(x)\Big)=0 \Leftrightarrow p(x)=q(x)$. That is, no matter which one is held fixed, minimizing the KL divergence pushes the two towards being as equal as possible. A rigorous proof of this uses the calculus of variations — and in fact the "V" (variational) in VAE comes precisely from the fact that the VAE derivation involves KL divergence (and hence the calculus of variations as well).
Of course, KL divergence has an obvious issue: when $q(x)$ is zero on some region while $p(x)$ is nonzero there, the KL divergence blows up to infinity. This is an inherent problem with KL divergence, and we can only work around it — for instance, this is why we use a Gaussian rather than a uniform distribution as the prior for the latent variable, as mentioned in the previous post Variational Autoencoders (Part 1): So That's What's Going On.
As an aside: is KL divergence the only way to measure the difference between two distributions? Certainly not. Take a look at Wikipedia's section on Statistical Distance, which introduces quite a few distributional distances. One particularly elegant one is the so-called Bhattacharyya distance, defined as
$$D_B\Big(p(x), q(x)\Big)=-\ln\int \sqrt{p(x)q(x)} dx\tag{6}$$
This distance is symmetric and doesn't suffer from KL divergence's infinity problem. And yet we still choose KL divergence, because we need not just theoretical elegance but also practical feasibility — KL divergence can be written as an expectation, which lets us estimate it via sampling. The Bhattacharyya distance, by contrast, is not so amenable to this; if you try to replace the KL divergence in the computations below with the Bhattacharyya distance, you'll quickly find yourself stuck.
Notation Table for This Post
Explaining VAEs inevitably involves a lot of formulas and symbols, so here I list the meaning of some of the expressions in advance:
$$\begin{array}{c|c} \hline x_k, z_k & \text{denote random variable}x,z\text{the}k\text{sample}\\ \hline x_{(k)}, z_{(k)} & \text{denote multivariate variable}x,z\text{the}k\text{component}\\ \hline \mathbb{E}_{x\sim p(x)}[f(x)] & \text{denote pair}f(x)\text{compute expectation, where}x\text{distribution of is}p(x)\\ \hline KL\Big(p(x)\Big\Vert q(x)\Big)& \text{of two distributions}KL\text{divergence}\\ \hline \Vert x\Vert^2& \text{vector}x\text{of}l^2\text{norm, i.e. squared magnitude}\\ \hline \mathcal{L}& \text{notation of loss function in this paper}\\ \hline D,d & D\text{is input}x\text{dimension of,}d\text{is latent variable}z\text{dimension of}\\ \hline \end{array}$$
The Framework
Here we give a concise and direct derivation of the theoretical framework of the VAE by approximating the joint distribution directly.
Confronting the Joint Distribution Head-On
The starting point remains unchanged, so let's restate it. We have a batch of data samples $\{x_1,\dots,x_n\}$, whose overall behavior is described by $x$, and we want to use a latent variable $z$ to describe the distribution $x$ of $\tilde{p}(x)$:
$$q(x)=\int q(x|z)q(z)dz,\quad q(x,z) = q(x|z)q(z)\tag{7}$$
Here $q(z)$ is the prior distribution (a standard normal), and the goal is for $q(x)$ to approximate $\tilde{p}(x)$. In this way, (in theory) we both describe $\tilde{p}(x)$ and obtain a generative model $q(x|z)$ — two birds with one stone.
The next step is to use KL divergence to perform the approximation. But what has always puzzled me is: why, starting from the original paper Auto-Encoding Variational Bayes, do VAE tutorials focus on the description of the posterior distribution $p(z|x)$? Perhaps this is influenced by the EM algorithm — in problems where EM can't be applied, it's precisely because the posterior $p(z|x)$ is intractable, which is why the authors of the VAE paper focused on deriving $p(z|x)$.
But in fact, approximating $p(x,z)$ directly is the most straightforward route. Specifically, define $p(x,z)=\tilde{p}(x)p(z|x)$, and suppose we use a joint probability distribution $q(x,z)$ to approximate $p(x,z)$. Then we can use KL divergence to measure the distance between them:
$$KL\Big(p(x,z)\Big\Vert q(x,z)\Big) = \iint p(x,z)\ln \frac{p(x,z)}{q(x,z)} dzdx\tag{8}$$
The KL divergence is our ultimate objective, since we want the two distributions to be as close as possible, so the smaller the KL divergence, the better. Of course, since $p(x,z)$ also has its own parameters now, it's not simply that $q(x,z)$ approximates $p(x,z)$ — $p(x,z)$ will also actively move to approximate $q(x,z)$; the two approach each other.
So we have
$$\begin{aligned}KL\Big(p(x,z)\Big\Vert q(x,z)\Big) =& \int \tilde{p}(x) \left[\int p(z|x)\ln \frac{\tilde{p}(x)p(z|x)}{q(x,z)} dz\right]dx\\ =& \mathbb{E}_{x\sim \tilde{p}(x)} \left[\int p(z|x)\ln \frac{\tilde{p}(x)p(z|x)}{q(x,z)} dz\right] \end{aligned}\tag{9}$$
This means that using equation $(4)$, once each $x_i$ is substituted in, the computation can proceed. This expression can be simplified further, since $\ln \frac{\tilde{p}(x)p(z|x)}{q(x,z)}=\ln \tilde{p}(x) + \ln \frac{p(z|x)}{q(x,z)}$, and
$$\begin{aligned}\mathbb{E}_{x\sim \tilde{p}(x)} \left[\int p(z|x)\ln \tilde{p}(x)dz\right] =& \mathbb{E}_{x\sim \tilde{p}(x)} \left[\ln \tilde{p}(x)\int p(z|x)dz\right]\\ =&\mathbb{E}_{x\sim \tilde{p}(x)} \big[\ln \tilde{p}(x)\big] \end{aligned}\tag{10}$$
Note that $\tilde{p}(x)$ here is the prior distribution over $x$ determined by the samples $x_1,x_2,\dots,x_n$; although we may not be able to write down its exact form, it is nonetheless fixed and does exist, so this term is just a constant, and we can write
$$\mathcal{L}=KL\Big(p(x,z)\Big\Vert q(x,z)\Big) - \text{const}= \mathbb{E}_{x\sim \tilde{p}(x)} \left[\int p(z|x)\ln \frac{p(z|x)}{q(x,z)} dz\right]\tag{11}$$
At this point, minimizing $KL\Big(p(x,z)\Big\Vert q(x,z)\Big)$ is equivalent to minimizing $\mathcal{L}$. Note that the constant being subtracted is $\mathbb{E}_{x\sim \tilde{p}(x)} \big[\ln \tilde{p}(x)\big]$, so $\mathcal{L}$ has the lower bound $-\mathbb{E}_{x\sim \tilde{p}(x)} \big[\ln \tilde{p}(x)\big]$ — note that $\tilde{p}(x)$ isn't necessarily a probability; in the continuous case $\tilde{p}(x)$ is a probability density, which can be greater than 1 or less than 1, so $-\mathbb{E}_{x\sim \tilde{p}(x)} \big[\ln \tilde{p}(x)\big]$ isn't necessarily non-negative, meaning the loss could actually be negative.
Your VAE Has Arrived
At this point, let's return to our original intent — to obtain a generative model, we write $q(x,z)$ as $q(x|z)q(z)$, giving us
$$\begin{aligned}\mathcal{L} =& \mathbb{E}_{x\sim \tilde{p}(x)} \left[\int p(z|x)\ln \frac{p(z|x)}{q(x|z)q(z)} dz\right]\\ =&\mathbb{E}_{x\sim \tilde{p}(x)} \left[-\int p(z|x)\ln q(x|z)dz+\int p(z|x)\ln \frac{p(z|x)}{q(z)}dz\right]\end{aligned}\tag{12}$$
Or, more concisely,
$$\begin{aligned}\mathcal{L} = &\mathbb{E}_{x\sim \tilde{p}(x)} \left[\mathbb{E}_{z\sim p(z|x)}\big[-\ln q(x|z)\big]+\mathbb{E}_{z\sim p(z|x)}\Big[\ln \frac{p(z|x)}{q(z)}\Big]\right]\\ = &\mathbb{E}_{x\sim \tilde{p}(x)} \Bigg[\mathbb{E}_{z\sim p(z|x)}\big[-\ln q(x|z)\big]+KL\Big(p(z|x)\Big\Vert q(z)\Big)\Bigg] \end{aligned}\tag{13}$$
See — isn't what's in the parentheses exactly the VAE loss function? We've just used different notation. All we need to do now is find suitable $q(x|z)$ and $q(z)$ that minimize $\mathcal{L}$.
Looking back over the whole derivation, we haven't done anything really "hard to think of" in terms of formal manipulation, and yet the VAE has emerged naturally. So there was no need to analyze the posterior distribution at all — by confronting the joint distribution head-on, we reach the destination more directly.
Don't Split It Up!
Given the structure of equation $(13)$, we might be tempted to split $\mathcal{L}$ into two separate pieces: the expectation involving $\mathbb{E}_{z\sim p(z|x)}\big[-\ln q(x|z)\big]$, and the expectation involving $KL\Big(p(z|x)\Big\Vert q(z)\Big)$ — treating the problem as minimizing two losses independently.
However, this view is mistaken. If $KL\Big(p(z|x)\Big\Vert q(z)\Big)=0$ means that $z$ carries no discriminative information at all, then $-\ln q(x|z)$ can't be small (the prediction won't be accurate); conversely, if $-\ln q(x|z)$ is small, then $q(x|z)$ is large, meaning the prediction is accurate, in which case $p(z|x)$ won't be too random, i.e., $KL\Big(p(z|x)\Big\Vert q(z)\Big)$ won't be small either. So the two loss terms are actually in tension with each other. Thus $\mathcal{L}$ shouldn't be viewed as separate pieces but as a whole; it's the overall value of $\mathcal{L}$ getting smaller that indicates the model is converging, rather than watching either loss term in isolation.
In fact, this is exactly what GAN models have long wished for — a single overall metric that tracks the progress of training the generative model. In VAEs, this capability comes naturally built in, whereas in GANs it took until WGAN for such a metric to appear.
Experiments
Up to this point, we've already completed the overall theoretical construction of the VAE. But to actually put it into practice, some more work is needed. In fact the original paper Auto-Encoding Variational Bayes also goes into considerable detail on this part, but unfortunately, many VAE tutorials online only derive things up to equation $(13)$ and stop there without elaborating further.
Approximating the Posterior Distribution
Right now $q(z),q(x|z),p(z|x)$ are all unknown — we haven't even fixed their functional forms — and to run experiments, we need to spell out every term of equation $(13)$ explicitly.
First, for ease of sampling, we assume $z\sim N(0,I)$, i.e., a standard multivariate normal distribution, which settles $q(z)$. What about $q(x|z),p(z|x)$? Just fit it with a neural network, full stop.
Note: if $q(x|z)$ and $q(z)$ were both already known, then the most reasonable estimate of $p(z|x)$ would be:
$$\hat{p}(z|x) = q(z|x) = \frac{q(x|z)q(z)}{q(x)} = \frac{q(x|z)q(z)}{\int q(x|z)q(z)dz}\tag{14}$$
This is exactly the posterior-probability estimation step in the EM algorithm; see
From Maximum Likelihood to the EM Algorithm: A Unified Perspective
for more detail. But in practice, the integral in the denominator is essentially never tractable, so this approach doesn't work. Instead, we simply approximate it with a generic network — this won't necessarily be optimal, but it does give us a usable approximation.
Concretely, we assume that $p(z|x)$ is also a normal distribution (with independent components), whose mean and variance are determined by $x$ — and this "determination" is done by a neural network:
$$p(z|x)=\frac{1}{\prod\limits_{k=1}^d \sqrt{2\pi \sigma_{(k)}^2(x)}}\exp\left(-\frac{1}{2}\left\Vert\frac{z-\mu(x)}{\sigma(x)}\right\Vert^2\right)\tag{15}$$
Here $\mu(x),\sigma^2(x)$ is a neural network that takes $x$ as input and outputs a mean and a variance, and $\mu(x)$ plays a role analogous to the encoder. Now that we've assumed a Gaussian, the KL divergence term in equation $(13)$ can already be computed in closed form:
$$KL\Big(p(z|x)\Big\Vert q(z)\Big)=\frac{1}{2} \sum_{k=1}^d \Big(\mu_{(k)}^2(x) + \sigma_{(k)}^2(x) - \ln \sigma_{(k)}^2(x) - 1\Big)\tag{16}$$
This is what we call the KL loss, which was already given in the previous post.
Approximating the Generative Model
Now all that remains is the generative model part, $q(x|z)$ — what distribution should we choose for it? The paper Auto-Encoding Variational Bayes offers two candidate choices: the Bernoulli distribution, or the normal distribution.
What? A normal distribution again? Isn't that too much of a simplification? But there really isn't much choice — we need to construct an actual distribution, not just any function, so it has to satisfy the normalization requirement, and it also has to be easy to compute. There simply aren't many options that satisfy both.
The Bernoulli Model
Let's start with the Bernoulli distribution, which as everyone knows is just a binary distribution:
$$p(\xi)=\left\{\begin{aligned}&\rho,\, \xi = 1;\\ &1-\rho,\, \xi = 0\end{aligned}\right.\tag{17}$$
So the Bernoulli distribution is only suitable for cases where $x$ is a multi-dimensional binary vector — for example, when $x$ is a binary image (MNIST can be treated this way). In this case, we use a neural network $\rho(z)$ to compute the parameter $\rho$, giving
$$q(x|z)=\prod_{k=1}^D \Big(\rho_{(k)}(z)\Big)^{x_{(k)}} \Big(1 - \rho_{(k)}(z)\Big)^{1 - x_{(k)}}\tag{18}$$
from which we can compute
$$-\ln q(x|z) = \sum_{k=1}^D \Big[- x_{(k)} \ln \rho_{(k)}(z) - (1-x_{(k)}) \ln \Big(1 -\rho_{(k)}(z)\Big)\Big]\tag{19}$$
This tells us that $\rho(z)$ needs to be squashed into the range 0–1 (e.g., with a sigmoid activation), and that cross-entropy should be used as the loss function; here $\rho(z)$ plays a role analogous to the decoder.
The Normal Distribution Model
Next, the normal distribution — this is the same as $p(z|x)$, just with $x,z$ swapped in position:
$$q(x|z)=\frac{1}{\prod\limits_{k=1}^D \sqrt{2\pi \tilde{\sigma}_{(k)}^2(z)}}\exp\left(-\frac{1}{2}\left\Vert\frac{x-\tilde{\mu}(z)}{\tilde{\sigma}(z)}\right\Vert^2\right)\tag{20}$$
Here $\tilde{\mu}(z),\tilde{\sigma}^2(z)$ is a neural network that takes $z$ as input and outputs a mean and a variance, and $\tilde{\mu}(z)$ plays the role of decoder. This gives
$$-\ln q(x|z) = \frac{1}{2}\left\Vert\frac{x-\tilde{\mu}(z)}{\tilde{\sigma}(z)}\right\Vert^2 + \frac{D}{2}\ln 2\pi + \frac{1}{2}\sum_{k=1}^D \ln \tilde{\sigma}_{(k)}^2(z)\tag{21}$$
Very often we fix the variance to be some constant $\tilde{\sigma}^2$, in which case
$$-\ln q(x|z) \sim \frac{1}{2\tilde{\sigma}^2}\Big\Vert x-\tilde{\mu}(z)\Big\Vert^2\tag{22}$$
and this is exactly where the MSE loss function comes from.
So now it's clear: for binary data, we can use a sigmoid activation on the decoder output and cross-entropy as the loss, which corresponds to taking $q(x|z)$ to be a Bernoulli distribution; whereas for general data, we use MSE as the loss, corresponding to taking $q(x|z)$ to be a normal distribution with fixed variance.
Tricks for Sampling-Based Computation
The previous section did all that work simply to be able to spell out equation $(13)$ explicitly. Once we assume $p(z|x)$ and $q(z)$ are both normal distributions, the KL divergence term in equation $(13)$ has already been worked out, giving equation $(16)$; and once we assume $q(x|z)$ is either Bernoulli or Gaussian, $-\ln q(x|z)$ can also be computed. So what's missing now?
Sampling!
$p(z|x)$ serves two purposes: one is for computing $KL\Big(p(z|x)\Big\Vert q(z)\Big)$, and the other for computing $\mathbb{E}_{z\sim p(z|x)}\big[-\ln q(x|z)\big]$ — and $\mathbb{E}_{z\sim p(z|x)}\big[-\ln q(x|z)\big]$ specifically means
$$-\frac{1}{n}\sum_{i=1}^n \ln q(x|z_i),\quad z_i \sim p(z|x)\tag{23}$$
We've already assumed that $p(z|x)$ is normal, with its mean and variance computed by the model, so sampling can be carried out using the "reparameterization trick."
But how many samples should we draw? The VAE takes the most direct approach possible: one! So under this choice, equation $(13)$ becomes extremely simple:
$$\mathcal{L} = \mathbb{E}_{x\sim \tilde{p}(x)} \Bigg[-\ln q(x|z) + KL\Big(p(z|x)\Big\Vert q(z)\Big)\Bigg],\quad z\sim p(z|x)\tag{24}$$
Every term in this expression can be found in equation $(16),(19),(21),(22)$. Note that for each $x$ in a batch, we need to draw its own "personal" $z$ sample from $p(z|x)$ that's specific to $x$ before we can compute $-\ln q(x|z)$. And it's precisely because VAE only draws a single sample at $p(z|x)$ that it ends up looking so much like an ordinary autoencoder.
So the final question is: is drawing just one sample really enough? In practice, we run many epochs, and the latent variable is randomly generated afresh each time — so as long as the number of epochs is large enough, this does in fact guarantee adequate sampling coverage. I've also experimented with drawing multiple samples, and the generated results didn't seem to change noticeably.
In Tribute
This post has laid out the overall pipeline of the VAE from the perspective of Bayesian theory. When examining things from this angle, we need to keep firmly in mind two things: "distribution" and "sampling" — writing down the form of a distribution, and simplifying the process through sampling.
Put simply, since directly describing a complex distribution is generally infeasible, we introduce a latent variable to turn it into a superposition of conditional distributions. At that point, we can make suitable simplifications to both the latent variable's distribution and the conditional distributions (e.g., assuming both are normal), and the parameters of the conditional distribution can be tied to a deep learning model (using deep learning to compute the parameters of the latent distribution). At this point, the shape of a "deep probabilistic graphical model" becomes visible.
Let us pay tribute together to the great Bayes, and to the many other giants who have studied probabilistic graphical models — they are all true and genuine trailblazers.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.