Importance Sampling in Disguise: Inside the Objectives of VAEs and Diffusion Models

Variational Autoencoders (VAEs) and Denoising Diffusion Probabilistic Models (DDPMs) are often presented as different generative models. However, while studying their training objectives, I noticed that both rely on the same underlying mathematical idea: estimating otherwise intractable quantities using random samples. This idea comes from Monte Carlo sampling, a general technique for approximating expectations that cannot be computed exactly. In this article, I will show you how this idea appears naturally in the loss functions used to train both VAEs and diffusion models. In fact, diffusion models can be interpreted as a particular structured form of hierarchical VAE, which helps clarify their underlying relationship to variational inference.

To understand where this sampling method come from, I will first revisit the basics of Bayesian inference and see how it leads naturally to the learning objectives used in modern generative models. I am going to explain various components of the Bayesian inference step-by-step while focusing on VAE, rather than DDPM. If you need to learn more about Bayes’ rule and how to derive the cost function, you can refer to my previous articles: “*What is Bayesian belief?”, “[Finally! A Clear Derivation of the VAE KL Loss](https://medium.com/@jpark7/finally-a-clear-derivation-of-the-vae-kl-loss-4cb38d2e47b3)”, and “[From Noise to Clarity: The Math Behind Denoising Diffusion Models](https://medium.com/@jpark7/from-noise-to-clarity-the-math-behind-denoising-diffusion-models-cf72900905c7)”*.

From Bayesian Theory to VAE

Many machine learning textbooks explain that, in Bayesian inference, the posterior distribution p(z|x), can be expressed as the likelihood p(x|z) multiplied by the prior p(z), divided by the marginal likelihood p(x) (aka, total probability) as illustrated below.

The posterior distribution, p(z|x) tells us how probable different values z are, given the sample data x. In this framework, this conditional distribution p(z|x) represents the neural network’s inference, its “best guess” regarding which underlying features most likely produced the specific input x. This probabilistic belief is mathematically encoded within the model’s learned parameters.

To me, the textbook explanation remains too abstract to fully understand the VAE’s mechanics. Here is my take on this idea.

To make this concrete, consider images as sample data x. A 64 x 64 color image can be thought of as a single point in a 12,288 high-dimensional space (64 x 64 x 3). If you picked a point in this space at random, you would not see a face or a dog; you would see meaningless static noise. The VAE’s job is to find the data manifold, a small organized “neighborhood” within this high-dimensional space where realistic (coherent) images actually live. The latent variable z provides a lower-dimensional representation that explains this region. The encoder network learns to map an image x to a latent distribution over z, locating where that image lives in the latent space.

Figure 1: High-Dimensional Image Space, Latent Space, and the Data Manifold

The latent space is a lower-dimensional abstract space z in which a model represents the hidden characteristics of images, such as object identity, pose, viewpoint, color, shape, size, background and so on. However, the exact meaning of each dimension of the latent space is generally not known or directly interpretable. The latent space is therefore an abstract representation, not the actual set of images. During the decoding phase, the decoder network maps points from the latent space to the data manifold in the image space, generating images corresponding to those latent variables. In this way, the data manifold lives in the image space, while the latent space is an internal representation learned by the model, and the decoder learns the mapping between them.

The latent distribution over z is the key to the generative process. Instead of mapping an image to a single point in the latent space, the encoder network describes a region by outputting two values:

1. The Mean: The center point of where the image lives in the latent space.
2. The Variance: The “fuzziness” or spread around that center.

By outputting these Gaussian parameters, the encoder defines a probability distribution, q(z|x), which serves as an approximation of the true (but intractable) posterior p(z|x). I will explain why the true posterior is intractable later. A specific latent vector z can then be sampled from this distribution. This means that the latent space is not just a collection of isolated points, but a continuous landscape where similar images tend to be located near one another in the latent space. The decoder p(x|z) then maps this sampled latent vector to a point in the image space, generating a new image corresponding to z.

This probabilistic representation illustrates the fundamental generative process of a VAE. As you may have noticed, the encoder outputs a list of mean and variance pairs that define a Gaussian distribution for each latent dimension, from which the latent vector z is sampled. The main reason for using Gaussian distributions is that it is easier to sample from them.

Figure 2: VAE Architecture

Note that if z were represented as a single deterministic point, rather than a distribution, the model would instead be a standard autoencoder (AE) which can be thought of as a mechanism that compresses an image into a latent representation and then, reconstruct the original image from that representation.

Understanding the Learning Objective Through Bayes’ Rule

Both encoder q(z|x) and decoder p(x|z) are functional forms of probability distributions whose parameters are modeled by neural networks. As an analogy, consider the heights of middle school students. If you collect height data from students aged 10 to 13, you can fit a Gaussian distribution by estimating its parameters, mean and variance. You can imagine the functional form of this distribution as:

Once you calculate the mean and variance from the data (determining parameters of the model from data), you can sample from that distribution to generate realistic heights that are most probable for that age group.

Similarly, in a VAE, neural networks learn the parameters of more complex probability distributions that describe the latent variables and the generated data. Unlike the simple Gaussian example above, there are no closed form solutions for these distributions. Instead, the parameters are learned during the training process.

Imagine a subset of all possible images, such as 10,000 sampled images of faces. Each image x could have been generated by many different latent variables z, for example, combinations of pose, lighting, and expression. Bayes’ rule tells me the true relationship between these variables:

Both sides represent the same joint probability p(x, z), but they express it from two different perspectives. From a generative perspective, the model samples a latent variable and generates an image according to the likelihood p(x|z), which is modeled by the decoder. From an inference perspective, I observe an image x and want to infer its latent cause z, corresponding to the posterior distribution p(z|x), which is approximated by the encoder q(z|x).

As stated previously, the relationship between the latent space and the image space can be viewed as a mapping between them. For example, a latent variable z1 may be mapped to a happy face, z2 to a sad face and z3 to a frowning face. The joint probability describes how strong each latent point is associated with each image, and high probability connections form a structured relationship between the two spaces. During training, the model learns to increase the probability that these relationships hold. This forms the learning objective of a VAE.

Naturally, the learning objective might appear to be maximizing the joint probability p(x, z). However, the problem is that the latent variable z is unknown for each data point x. In the training dataset, I can only observe data points, x1, x2, …, xn, but not their corresponding latent variables z. Therefore, I need to consider all possible latent variables. Instead of assuming a single z, the model evaluates how probable an image is under any possible latent explanation. In this way the model becomes agnostic to the unknown latent cause.

Suppose possible latent causes are z1: happy face + gray color, z2: happy face + head tilted, z3: happy face + bright light. The joint probabilities are p(x, z1), p(x, z2) and p(x, z3). Since I don’t know the true latent causes, I want the model to assign a high probability of the image x if any explanation makes the image likely. Accounting for all possible latent causes means summing their joint probabilities, which gives the marginal likelihood: p(x) = p(x, z1) + p(x, z2) + p(x, z3).

The key insight is that p(x, z) represents the probability of a specific latent explanation, whereas p(x) represents the probability of the data under all possible explanations. Since the true latent cause is unknown, I need to maximize p(x), which measures how well the model explains the observed data overall.

Consequently, the VAE training objective can be formulated as maximizing the marginal likelihood (also called the model evidence) of the data p(x). This quantity is obtained by considering all possible latent variables z, evaluating how likely the data point x is given each z, and integrating over all such possibilities.

Why Do I Need Sampling?

Although maximizing the marginal likelihood p(x) provides a learning objective, computing it exactly is generally intractable because, in the continuous case, the marginal likelihood requires integrating over all possible latent variables.

This integral is difficult to compute because the latent variable z can be high-dimensional and the likelihood p(x|z) is modeled by a neural network, which has no closed form solution. As a result, evaluating this integral would require infinitely many possible latent variables.

Instead of evaluating all possible latent explanations, the model approximates the integral by sampling a finite number of latent variables. In other words, the model evaluates the likelihood of the image under several sampled latent variables and uses these samples to approximate the contribution of all possible latent causes.

In mathematics, sampling can be expressed in terms of the expectation of a continuous random variable. This connection arises from the law of large numbers and the idea of Monte Carlo integration. The key idea is that an integral can be interpreted as the limit of a sum. Instead of computing the entire area under a probability density function exactly, I can approximate the integral by drawing a finite number of independent and identically distributed (iid) samples from the distribution.

More formally, many integrals in probability can be written as expectations. Suppose I want to compute the expectation of a function f(z) under a probability distribution p(z):

When this integral cannot be computed analytically, it can be approximated using Monte Carlo sampling. If I draw N independent samples z1, z2, …, zn from the distribution p(z), the expectation can be approximated by the sample average:

In the context of VAEs, this idea allows the model to approximate intractable integrals over the latent space. Instead of evaluating every possible latent variable z, the model samples a finite number of latent variables and evaluate how well each z explains the observed data. By averaging over these sampled latent variables, the model can approximate the contribution of all possible latent causes.

As a reminder, the marginal likelihood of the data is expressed as follow:

Directly sampling from the prior p(z) may not be efficient, since many sampled latent variables may not explain the observed data. A common way to address this is importance sampling, which introduces a proposal distribution q(z|x) that focuses sampling on regions of the latent space that are more likely to explain the data.

In most VAEs, the prior distribution p(z) is chosen to be a standard Gaussian N(0, 1). The encoder then learns an approximate posterior q(z|x), which is also modeled as a Gaussian distribution whose mean and variance are predicted by a neural network. Basically, p(z) describes where latent variables should live before seeing data and q(z|x) describes where a specific image likely lives in latent space.

Importance Sampling in VAE Loss Function

Naturally, I need to approximate p(z|x) with q(z|x) while keeping q(z|x) close to the prior p(z). This is done by multiplying and dividing the integrand by q(z|x). Nothing changes here.

Rearranging the term gives

which can be interpreted as an expectation with respect to q(z|x):

In this formulation, samples are drawn from the proposal distribution q(z|x), and each sample is weighted by the importance ratio:

This approach concentrates samples in regions of the latent space that are more relevant for explaining the observed data, allowing the model to approximate the marginal likelihood more efficiently.

Note that the term,

measures how different the proposal distribution q(z|x) is from the prior p(z). When these distributions are similar, the ratio is close to one (1), resulting in more stable estimates. Basically, I need the encoder to produce latent representations that are well aligned with the prior. To achieve this, the VAE architecture includes a Kullback-Leibler (KL) divergence term in the loss function, which explicitly measures the difference between two distributions, q(z|x) and p(z) and encourage the latent representations of all sampled images to be organized in a way that matches the prior.

In the VAE paper, the learning objective is to maximize the log likelihood log p(x) instead of p(x). There are a few reasons for this approach, but in the context of this article, I am going to focus on Jensen’s inequality which gives a tractable lower bound known as the evidence lower bound (ELBO). Maximizing this bound inherently leads to the maximization of the log likelihood. This approach allows the model to approximate the intractable marginal likelihood p(x) while learning both the encoder and decoder networks.

Let’s derive the VAE learning objective. First, I will apply log to the expectation.

If I use Jensen’s inequality,

the lower bound of log p(x) is

After some algebraic manipulation, I obtain three (3) expectation terms.

I can identify both the cross entropy and entropy terms, and their difference corresponds to the relative entropy, also known as the KL divergence.

Cross Entropy

Entropy

The importance ratio inside the log decomposes into two terms: a likelihood term and KL divergence term.

The first term measures how well the decoder reconstructs the original data, and corresponds to the reconstruction log likelihood. If I choose p(x|z) to be a Gaussian, this term reduces to the mean squared error (MSE). Since I did not cover this in my previous article, “Finally! A Clear Derivation of the VAE KL Loss”, I will briefly derive it here.

So maximizing p(x|z) is equivalent to minimizing MSE.

The second term is the KL divergence term which measures the difference between two distributions, q(z|x) and p(z) as I stated earlier.

The effectiveness of importance sampling depends on aligning the proposal distribution with the true posterior. VAEs apply this concept by using an encoder to learn the approximate posterior q(z|x), regularizing it toward the prior p(z) via KL divergence. Consequently, the VAE approach acts as a learnable form of importance sampling to efficiently approximate the marginal likelihood.

I hope this article is helpful!