Variational Autoencoders (Part 2): From a Bayesian Perspective

Numerical Calculation vs. Sampling Calculation

For readers not very familiar with probability and statistics, two easily confused concepts are numerical calculation and sampling calculation. Some readers expressed similar confusion in "The Three Flavors of Capsules: Matrix Capsules and EM Routing." For example, if the probability density function p(x) is known, the expectation of x is defined as: \mathbb{E}[x] = \int x p(x)dx \tag{1} If we want to perform a numerical calculation (numerical integration), we can select several representative points x_0 < x_1 < x_2 < \dots < x_n and obtain: \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 discuss what "representative" means here, nor methods to improve numerical precision. This is written to contrast with sampling calculation. If we sample several points x_1, x_2, \dots, x_n from p(x), we have: \mathbb{E}[x] \approx \frac{1}{n}\sum_{i=1}^n x_i,\quad x_i \sim p(x) \tag{3} Comparing (2) and (3), the main difference is that (2) includes the calculation of the probability p(x_i), while (3) only involves x_i. This is because in (3), x_i is sampled according to the probability p(x); points with higher probability appear more frequently. Thus, the sampling results already implicitly contain p(x), so there is 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} This is the foundation of Monte Carlo simulation.

KL Divergence and Variation

We usually 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 primary property of KL divergence is non-negativity. If p(x) is fixed, then KL\Big(p(x)\Big\Vert q(x)\Big)=0 \Leftrightarrow p(x)=q(x). If q(x) is fixed, similarly KL\Big(p(x)\Big\Vert q(x)\Big)=0 \Leftrightarrow p(x)=q(x). That is, regardless of which one is fixed, minimizing the KL divergence results in the two being as equal as possible. The rigorous proof of this requires the calculus of variations. In fact, the "V" in VAE (Variational) exists because the derivation of VAE utilizes KL divergence (which in turn involves variational methods).

Of course, KL divergence has a significant problem: if q(x) is zero in a region where p(x) is non-zero, the KL divergence becomes infinite. This is an inherent issue with KL divergence, and we can only try to avoid it. For instance, we use a Gaussian distribution rather than a uniform distribution for the prior of the latent variables for this reason, as mentioned in the previous article.

As a side note, is KL divergence the only way to measure the difference between two distributions? Certainly not. We can refer to the Statistical Distance section on Wikipedia, which introduces many distances. For example, there is a beautiful measure called the 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 not only symmetric but also avoids the infinity problem of KL divergence. However, we still choose KL divergence because we seek not only theoretical beauty but also practical feasibility. KL divergence can be written in the form of an expectation, which allows us to perform sampling calculations. Conversely, the Bhattacharyya distance is not as easy to handle. If one were to try replacing KL divergence with Bhattacharyya distance in the following calculations, they would find it nearly impossible to proceed.

Notation Table

Explaining VAEs inevitably involves many formulas and symbols. Some are listed below:

Directly Facing the Joint Distribution

The starting point remains the same. First, we have a set of data samples \{x_1, \dots, x_n\}, described collectively by x. We hope to describe the distribution \tilde{p}(x) of x using a latent variable z: 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 (standard normal distribution). The goal is for q(x) to approximate \tilde{p}(x). In this way, we (theoretically) describe \tilde{p}(x) and obtain a generative model q(x|z) simultaneously.

Next, we use KL divergence for approximation. What I have never understood is why, starting from the original paper "Auto-Encoding Variational Bayes", VAE tutorials focus on the description of the posterior distribution p(z|x). Perhaps influenced by the EM algorithm, where p(z|x) is difficult to calculate, the authors focused on its derivation.

However, it is actually most straightforward to directly approximate p(x,z). Specifically, define p(x,z)=\tilde{p}(x)p(z|x). Suppose we use a joint probability distribution q(x,z) to approximate p(x,z). We look at their distance using KL divergence: 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} KL divergence is our ultimate goal; we want the two distributions to be as close as possible, so the KL divergence should be minimized. Of course, since p(x,z) also has parameters, it is not just q(x,z) approximating p(x,z); p(x,z) will also actively approximate q(x,z). They approach each other.

Thus 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} By using equation (4) and substituting each x_i, we can perform the calculation. This expression can be further simplified because \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) is the prior distribution of x determined by the samples x_1, x_2, \dots, x_n. Although we might not be able to write its form accurately, it is fixed and exists. Therefore, this term is just a constant. We can write: \mathcal{L}=KL\Big(p(x,z)\Big\Vert q(x,z)\Big) - \text{constant}= \mathbb{E}_{x\sim \tilde{p}(x)} \left[\int p(z|x)\ln \frac{p(z|x)}{q(x,z)} dz\right] \tag{11} Minimizing KL\Big(p(x,z)\Big\Vert q(x,z)\Big) is now equivalent to minimizing \mathcal{L}. Note that the subtracted constant is \mathbb{E}_{x\sim \tilde{p}(x)} \big[\ln \tilde{p}(x)\big], so \mathcal{L} has a lower bound of -\mathbb{E}_{x\sim \tilde{p}(x)} \big[\ln \tilde{p}(x)\big]. Since \tilde{p}(x) is a probability density in the continuous case, it can be greater or less than 1, so -\mathbb{E}_{x\sim \tilde{p}(x)} \big[\ln \tilde{p}(x)\big] is not necessarily non-negative; i.e., the loss can be negative.

Your VAE Has Arrived

At this point, we return to our original intention—to obtain a generative model. We write q(x,z) as q(x|z)q(z), which gives: \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} 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} Look, isn’t the term inside the brackets exactly the VAE loss function? We have just changed the notation. We need to find appropriate q(x|z) and q(z) to minimize \mathcal{L}.

Reviewing the whole process, we hardly performed any "unimaginable" formal transformations, yet the VAE emerged. Thus, there is no need to analyze the posterior distribution; by facing the joint distribution directly, we reach the destination faster.

No Splitting!

Given the characteristics of equation (13), we might be tempted to view \mathcal{L} as two separate parts: the expectation of \mathbb{E}_{z\sim p(z|x)}\big[-\ln q(x|z)\big] and the expectation of KL\Big(p(z|x)\Big\Vert q(z)\Big), assuming the problem becomes minimizing two losses separately.

However, this view is inappropriate. KL\Big(p(z|x)\Big\Vert q(z)\Big)=0 implies that z has no discriminative power, so -\ln q(x|z) cannot be small (poor prediction). Conversely, if -\ln q(x|z) is small, then q(x|z) is large (accurate prediction), meaning p(z|x) cannot be too random, so KL\Big(p(z|x)\Big\Vert q(z)\Big) will not be small. These two parts of the loss are actually antagonistic. Therefore, \mathcal{L} cannot be viewed in isolation; it must be viewed as a whole. The smaller the total \mathcal{L}, the closer the model is to convergence. One cannot observe only one part of the loss.

In fact, this is exactly what is sought in GAN models—a total indicator that can signal the training progress of the generative model. VAE models possess this capability naturally, whereas in GANs, such an indicator only appeared with WGAN.

Posterior Distribution Approximation

Currently, q(z), q(x|z), p(z|x) are all unknown, and even their forms are not determined. For the experiment, every term in equation (13) must be explicitly written out.

First, for ease of sampling, we assume z\sim N(0,I), the standard multivariate normal distribution. This solves q(z). What about q(x|z) and p(z|x)? Let’s fit them all with neural networks.

Note: Ideally, if q(x|z) and q(z) were known, the most reasonable estimate for 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 actually the posterior probability estimation step in the EM algorithm; see "From Maximum Likelihood to EM Algorithm: A Consistent Way of Understanding" for details. However, the integral in the denominator is almost impossible to compute, making this approach unfeasible. So, we simply use a general network to approximate it. This might not be optimal, but it is a usable approximation.

Specifically, we assume p(z|x) is also a normal distribution (with independent components), whose mean and variance are determined by x. This "determination" is 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) and \sigma^2(x) are neural networks that take x as input and output the mean and variance, respectively. \mu(x) plays a role similar to an encoder. Since we assumed a Gaussian distribution, the KL divergence term in (13) can be calculated beforehand: 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 provided in the previous article.

Generative Model Approximation

Now only the generative model part q(x|z) remains. What distribution should we choose? The paper "Auto-Encoding Variational Bayes" provides two candidates: Bernoulli distribution or Normal distribution.

What? Normal distribution again? Is that too simplified? However, there is no other way because we need to construct a distribution, not just any function. As a distribution, it must satisfy normalization requirements, and to be easy to calculate, we don’t have many choices.

Sampling Calculation Trick

The previous section did a lot of work just to write equation (13) explicitly. When we assume p(z|x) and q(z) are both Normal distributions, the KL divergence part of (13) is calculated as equation (16). When we assume q(x|z) is a Bernoulli or Gaussian distribution, -\ln q(x|z) is also calculated. What is missing?

Sampling!

The role of p(z|x) is twofold: one is to calculate KL\Big(p(z|x)\Big\Vert q(z)\Big), and the other is to calculate \mathbb{E}_{z\sim p(z|x)}\big[-\ln q(x|z)\big]. The latter means: -\frac{1}{n}\sum_{i=1}^n \ln q(x|z_i),\quad z_i \sim p(z|x) \tag{23} We have assumed p(z|x) is a Normal distribution with parameters calculated by the model. Thus, with the help of the "reparameterization trick", sampling can be performed.

But how many samples are appropriate? VAE is very straightforward: one! So equation (13) becomes very 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} Each term in this equation can be found in equations (16), (19), (21), and (22). Note that for each x in a batch, a z "exclusive" to x must be sampled from p(z|x) to calculate -\ln q(x|z). Because VAE only samples one point for p(z|x), it looks very similar to a standard AE.

The final question is: is one sample enough? In fact, we run multiple epochs, and the latent variables are randomly generated each time. Therefore, when the number of epochs is sufficient, the adequacy of sampling is guaranteed. I have experimented with multiple samples, and the generated samples did not show significant changes.