Variational Autoencoders (III): Why Does This Work?

One Sample is Enough

Regarding equation (3), we handled it as follows:

1. Notice that \mathbb{E}_{z\sim p(z|x)}\Big[\ln \frac{p(z|x)}{q(z)}\Big] is exactly the divergence KL\Big(p(z|x)\Big\Vert q(z)\Big) between p(z|x) and q(z). Since we assume both are normal distributions, this term can be calculated analytically.

2. For the term \mathbb{E}_{z\sim p(z|x)}\big[-\ln q(x|z)\big], we consider sampling just one point to be representative enough. Thus, this term becomes -\ln q(x|z),\, z\sim p(z|x).

After this treatment, the entire loss can be explicitly written as: \mathcal{L}=\mathbb{E}_{x\sim \tilde{p}(x)} \left[-\ln q(x|z) + KL\Big(p(z|x)\Big\Vert q(z)\Big)\right],\quad z\sim p(z|x) \tag{4}

Wait, some readers might find this objectionable: KL\Big(p(z|x)\Big\Vert q(z)\Big) is calculated beforehand, which is equivalent to sampling infinite points to estimate this term; yet for \mathbb{E}_{z\sim p(z|x)}\big[-\ln q(x|z)\big], we only sample one point. Since both are parts of the loss, is such unequal treatment really okay?

In fact, \mathbb{E}_{z\sim p(z|x)}\Big[\ln \frac{p(z|x)}{q(z)}\Big] can also be calculated by sampling just one point. That is to say, by sampling only one point for everything, equation (3) can be transformed into: \begin{aligned} \mathcal{L} =&\mathbb{E}_{x\sim \tilde{p}(x)} \left[-\ln q(x|z)+\ln \frac{p(z|x)}{q(z)}\right]\\ =&\mathbb{E}_{x\sim \tilde{p}(x)} \Big[-\ln q(x|z)+\ln p(z|x) - \ln q(z)\Big]\,,\quad z\sim p(z|x) \end{aligned} \tag{5} Although this loss differs from the standard VAE, it actually converges to similar results.

Why is One Point Enough?

So, why is sampling one point enough? Under what circumstances is one point sufficient?

First, let me give an example where “one point is not enough”. Let’s look back at equation (2), which can be rewritten as: q(x|z)=\mathop{\text{argmax}}_{q(x|z)} \mathbb{E}_{x\sim\tilde{p}(x)}\Bigg[\ln\Big(\mathbb{E}_{z\sim q(z)}\big[q(x|z)\big]\Big)\Bigg] \tag{6}

If one point were enough—no, let’s be cautious and say k points—then we could write: q(x|z)=\mathop{\text{argmax}}_{q(x|z)} \mathbb{E}_{x\sim\tilde{p}(x)}\Bigg[\ln\left(\frac{1}{k}\sum_{i=1}^k q(x|z_i)\right)\Bigg],\quad z_1,\dots,z_k \sim q(z) \tag{7} And then we could train using gradient descent.

However, such a strategy is unsuccessful. In practice, the number of samples k we can take is generally smaller than the batch size. In this case, maximizing \ln\left(\frac{1}{k}\sum\limits_{i=1}^k q(x|z_i)\right) falls into a “resource war”: During each iteration, various x_i in a batch compete for z_1, z_2, \dots, z_k. Whoever wins the competition gets a large q(x|z) (simply put, whichever x_i finds a z_j exclusive to it—meaning z_j can only generate x_i and nothing else—then q(x_i|z_j) is large). But every sample is equal, and sampling is random; we cannot predict the outcome of each “resource war.” It’s complete chaos! If the dataset is just MNIST, it might be okay because MNIST samples have a clear clustering tendency; if the number of samples k exceeds 10, there’s enough for each x_i to share. But for datasets like faces or ImageNet, which have no obvious clustering tendency and large intra-class variance, the z values are simply not enough to go around. One moment x_i grabs z_j, the next x_{i+1} grabs z_j, and the training fails directly.

Therefore, it is this “too many people, too little food” situation that leads to the failure of training model (7). But why does sampling one point succeed in VAE?

One Point Really is Enough

This requires us to analyze our thoughts on q(x|z) again. We call q(x|z) the generative model part. Generally, we assume it follows a Bernoulli or Gaussian distribution. Given that Bernoulli distributions have limited application scenarios, we assume it is a normal distribution: q(x|z)=\frac{1}{\prod\limits_{k=1}^D \sqrt{2\pi \sigma_{(k)}^2(z)}}\exp\left(-\frac{1}{2}\left\Vert\frac{x-\mu(z)}{\sigma(z)}\right\Vert^2\right) \tag{8} where \mu(z) is the network used to calculate the mean, and \sigma^2(z) is the network used to calculate the variance. Often, we fix the variance, leaving only the mean network.

Note that q(x|z) is just a probability distribution. After sampling z from q(z) and substituting it into q(x|z), we get the specific form of q(x|z). Theoretically, we should sample from q(x|z) again to get x. However, we do not do this; we directly treat the result of the mean network \mu(z) as x. Being able to do this indicates that q(x|z) is a normal distribution with very small variance (if the variance is fixed, it needs to be tuned low before training; if it’s a Bernoulli distribution, this isn’t an issue as it only has one set of parameters). The result of each sampling is almost identical (always the mean \mu(z)). At this point, there is “almost” a one-to-one correspondence between x and z, approaching a deterministic function x=\mu(z).

For the posterior distribution p(z|x), we also assume it is a normal distribution. Since we previously said z and x are almost one-to-one, this property also applies to the posterior distribution p(z|x). This implies that the posterior distribution will also be a normal distribution with very small variance (readers can verify this from the encoder results of MNIST). This means that the result of each sampling from p(z|x) is almost the same. Since this is the case, there is no difference between sampling once and sampling multiple times, because the results are basically the same every time. Thus, we have explained why we can start from equation (3) and sample only one point to calculate, turning it into equation (4) or (5).

The Prior of the Posterior

Readers might be a bit confused logically: you say q(x|z) and p(z|x) eventually become normal distributions with very small variance, but that’s just the final training result. During modeling, theoretically, we cannot know beforehand how large the variances of q(x|z) and p(z|x) are, so how can we sample just one point first?

I believe this is our prior knowledge of the problem. When we decide to use a certain dataset X for VAE, the dataset itself carries strong constraints. For example, the MNIST dataset has 784 pixels, but in fact, its independent dimensionality is far less than 784. Most obviously, some edge pixels are always 0. Compared to all possible 28 \times 28 images, MNIST is a very small subset. Another example is the poem-writing robot I wrote about a few days ago; the “Tang Poetry” corpus is a very small subset compared to general sentences. Even if we look at the ImageNet dataset with thousands of categories, it is still a tiny subset of all possible infinite images.

Consequently, we think that this dataset X can be projected into a low-dimensional space (latent variable space), and then the latent variables in the low-dimensional space can have a one-to-one correspondence with the original set X. Readers might have noticed: isn’t this just a regular autoencoder? Yes, it essentially means that in the case of a regular autoencoder, we can achieve a one-to-one correspondence between latent variables and the original dataset (perfect correspondence means the variances of p(z|x) and q(x|z) are 0). Then, after introducing the Gaussian prior distribution q(z), roughly speaking, this just applies translation and scaling to the latent variable space, so the variance doesn’t have to be large.

Therefore, we should guess beforehand that the variances of q(x|z) and p(z|x) are small and let the model realize this estimation. Simply put, the operation of “sampling one” is our prior knowledge of the data and the model, a prior for the posterior distribution, and we hope the model can approach this prior knowledge through this assumption.

The whole logic should be:

1. Have an original corpus;
2. Observe the original corpus and infer that it can correspond one-to-one to some latent variable space;
3. Use the “sample one” method to let the model learn this correspondence.

This part is a bit messy—it feels a bit redundant. I hope I haven’t confused the readers. If it feels confusing, just ignore this part.

The Straightforward IWAE

The next example is called the “Importance Weighted Autoencoder” (IWAE), which more bluntly and directly reflects the magic of the posterior distribution. To some extent, it can be seen as an upgraded version of VAE.

The starting point of IWAE is equation (2). It introduces the posterior distribution to rewrite equation (2): \int q(x|z)q(z)dz = \int p(z|x)\frac{q(x|z)q(z)}{p(z|x)}dz=\mathbb{E}_{z\sim p(z|x)}\left[\frac{q(x|z)q(z)}{p(z|x)}\right] \tag{8} In this way, equation (2) changes from sampling from q(z) to sampling from p(z|x). As we discussed, p(z|x) has a small variance, so sampling a few points is enough: \int q(x|z)q(z)dz = \frac{1}{k}\sum_{i=1}^k \frac{q(x|z_i)q(z_i)}{p(z_i|x)},\quad z_1,\dots,z_k\sim p(z|x) \tag{9} Substituting this into equation (2) gives: q(x|z)=\mathop{\text{argmax}}_{q(x|z)} \mathbb{E}_{x\sim\tilde{p}(x)}\Bigg[\ln\left(\frac{1}{k}\sum_{i=1}^k \frac{q(x|z_i)q(z_i)}{p(z_i|x)}\right)\Bigg],\quad z_1,\dots,z_k \sim p(z|x) \tag{10} This is IWAE. To align it with equations (4) and (5), it can be equivalently written as: \begin{aligned} &q(x|z) = \mathop{\text{argmin}}_{q(x|z),p(z|x)} \mathcal{L}_k,\\ \mathcal{L}_k = \mathbb{E}_{x\sim\tilde{p}(x)}\Bigg[&-\ln\left(\frac{1}{k}\sum_{i=1}^k \frac{q(x|z_i)q(z_i)}{p(z_i|x)}\right)\Bigg],\quad z_1,\dots,z_k \sim p(z|x) \end{aligned} \tag{11} When k=1, the above equation is exactly the same as equation (5). So from this perspective, IWAE is an upgrade of VAE.

From the construction process, replacing p(z|x) in equation (8) with any distribution of z is possible; p(z|x) is chosen only because it is focused and easy to sample from. When k is large enough, the specific form of p(z|x) actually becomes unimportant. This indicates that in IWAE, the role of the encoder model p(z|x) is weakened in exchange for an improvement in the generative model q(x|z). Because in VAE, we assume p(z|x) is a normal distribution, which is just an approximation for ease of calculation. The rationality of this approximation also affects the quality of the generative model q(x|z). It can be proven that \mathcal{L}_k can be closer to the lower bound -\mathbb{E}_{x\sim \tilde{p}(x)} \left[\ln \tilde{p}(x)\right] than \mathcal{L}, so the quality of the generative model will be superior.

Intuitively, in IWAE, the degree of approximation of p(z|x) is no longer that important, so a better generative model can be obtained. However, the price is that the quality of the encoding model decreases, because the importance of p(z|x) is reduced, and the model will not focus as much energy on training p(z|x). Therefore, if we hope to obtain a good encoder, IWAE is not advisable.

There is also a work “Tighter Variational Bounds are Not Necessarily Better” which reportedly improves the quality of both the encoder and decoder, though I haven’t fully understood it yet.