Variational AutoEncoder

In a lot of generative models courses, the starting point of training a neural network to generate new realistic data is the Variational Autoencoder (VAE). This model has its origins in the AutoEncoder (AE), which serves a different purpose: to reconstruct data:

Figure 2: The standard autoencoder compresses input $\mathbf{x}$ into a fixed latent code $\mathbf{z}$ via the encoder, then reconstructs ${\mathbf{x}}^{\prime }$ via the decoder. $\mathbf{z}$ is a single point with no probabilistic structure.

Formally, AE consists of two parts: an encoder $f_{\varphi }$ that compresses the input $\mathbf{x}$ into a compact latent representation $\mathbf{z}=f_{\varphi }(\mathbf{x})$, and a decoder $g_{\theta }$ that reconstructs the input from that representation $\hat{\mathbf{x}}=g_{\theta }(\mathbf{z})$. The network is trained end-to-end by minimizing a reconstruction loss, typically mean squared error:

$$\mathcal{L}=\Vert \mathbf{x}-g_{\theta }(f_{\varphi }(\mathbf{x})){\Vert }^2$$

The bottleneck forces the encoder to learn a compressed, meaningful representation of the data. Once trained, the latent space can be used for tasks like dimensionality reduction or feature extraction. However, autoencoders have a critical limitation to serve as a generative model: the latent space has no guaranteed structure. Points in latent space are not organized in any principled way, so randomly sampling an arbitrary $\mathbf{z}$ and decoding it often yields garbage. There is no way to smoothly interpolate between examples or generate novel, realistic samples.

The Variational Autoencoder (VAE), introduced by Kingma & Welling (2013), addresses this by imposing a probabilistic structure on the latent space. Instead of mapping $\mathbf{x}$ to a fixed point $\mathbf{z}$, the encoder $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ outputs the parameters of a distribution (usually a Gaussian). A latent vector $\mathbf{z}$ is then sampled from this distribution rather than deterministically computed. The decoder $p_{\theta }(\mathbf{x}\mid \mathbf{z})$ learns to reconstruct $\mathbf{x}$ from these sampled latents. A prior $p(\mathbf{z})$ is placed over the latent space, and the encoder is regularized to stay close to this prior via KL divergence. This shift - from a deterministic bottleneck to a learned posterior - gives the latent space two important properties:

• Continuity: Nearby points in latent space decode to similar outputs. Because the encoder maps each input $\mathbf{x}$ to a distribution over $\mathbf{z}$ rather than a single point, inputs that are similar naturally produce overlapping distributions - and thus neighboring regions in latent space correspond to similar decoded outputs.

• Completeness: Any point sampled from the prior produces a meaningful output. By regularizing the encoder’s posterior $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ to stay close to the prior $\mathcal{N}(0,\mathbf{I})$, the model ensures that the high-probability regions of the latent space are densely covered with meaningful structure, so random samples from the prior reliably decode into coherent outputs.

---

Figure 2: The VAE encoder $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ maps input $\mathbf{x}$ to a Gaussian distribution in latent space, regularized toward the prior $p(\mathbf{z})=\mathcal{N}(0,\mathbf{I})$. A vector $\mathbf{z}$ is then sampled from $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ and decoded by $p_{\theta }(\mathbf{x}\mid \mathbf{z})$ to reconstruct ${\mathbf{x}}^{\prime }$.

1. Constructions of VAE

Suppose we have a dataset of samples drawn i.i.d. from an unknown, complex distribution $p_{data}(\mathbf{x})$. Since the true form of $p_{data}$ is unknown, we cannot generate new sample by drawing from $p_{data}(\mathbf{x})$ directly. The goal of a generative model is to learn a tractable approximation $p_{\theta }(\mathbf{x})$ from this finite dataset by minimizing a divergence ${\mathcal{D}}_f$ between the two distributions.

Figure 3: The goal of a generative model: find parameters $\theta$ that minimize the divergence $\mathcal{D}(p_{data}(\mathbf{x}),p_{\theta }(\mathbf{x}))$ between the true data distribution (blue) and the learned model distribution (yellow). As $\mathcal{D}\to 0$, the model distribution increasingly overlaps with the data distribution.

Once the optimal parameters $\theta$ are found, $p_{\theta }(\mathbf{x})$ can be used to serve as a proxy for $p_{data}(\mathbf{x})$, enabling two key capabilities:

• Generation: Draw new, realistic samples from $p_{data}(\mathbf{x})$ via sampling methods such as Monte Carlo Sampling via $p_{\theta }(\mathbf{x})$.
• Evaluation: Assess how likely a given sample ${\mathbf{x}}^{\prime }$ is under the learned distribution $p_{\theta }(\mathbf{x})$ - for instance, judging whether an image ${\mathbf{x}}^{\prime }$ looks realistic by computing the likelihood $p_{\theta }({\mathbf{x}}^{\prime })$.

In the case of VAEs, ${\mathcal{D}}_f$ is the KL divergence ${\mathcal{D}}_{KL}$:

$${\mathcal{D}}_{KL}(p_{data}(\mathbf{x})\Vert p_{\theta }(\mathbf{x}))$$

KL divergence intuition

$$\begin{array}{rl}{\mathcal{D}}_{KL}(p_{data}(\mathbf{x})\Vert p_{\theta }(\mathbf{x}))& =\int p_{data}(\mathbf{x})\log \frac{p_{data}(\mathbf{x})}{p_{\theta }(\mathbf{x})}d\mathbf{x}\\ & ={\mathbb{E}}_{\mathbf{x}\sim p_{data}(\mathbf{x})}\left[\log \frac{p_{data}(\mathbf{x})}{p_{\theta }(\mathbf{x})}\right]\\ & ={\mathbb{E}}_{\mathbf{x}\sim p_{data}(\mathbf{x})}\left[\log p_{data}(\mathbf{x})-\log p_{\theta }(\mathbf{x})\right]\end{array}$$

As we can see, the KL divergence measures the expected log-likelihood difference between $p_{data}(\mathbf{x})$ and $p_{\theta }(\mathbf{x})$. Therefore, minimizing ${\mathcal{D}}_{KL}(p_{data}(\mathbf{x})\Vert p_{\theta }(\mathbf{x}))$ pushes $p_{\theta }(\mathbf{x})$ to assign high likelihood to real data $\mathbf{x}$ sampled from $p_{data}(\mathbf{x})$.

Now, We can rewrite the KL divergence as follow:

$$\begin{array}{rl}{\mathcal{D}}_{KL}(p_{data}(\mathbf{x})\Vert p_{\theta }(\mathbf{x}))& ={\mathbb{E}}_{\mathbf{x}\sim p_{data}(\mathbf{x})}\left[\log p_{data}(\mathbf{x})-\log p_{\theta }(\mathbf{x})\right]\\ & =-{\mathbb{E}}_{\mathbf{x}\sim p_{data}(\mathbf{x})}\left[\log p_{\theta }(\mathbf{x})\right]+{\mathbb{E}}_{\mathbf{x}\sim p_{data}(\mathbf{x})}\left[\log p_{data}(\mathbf{x})\right]\\ & =-{\mathbb{E}}_{\mathbf{x}\sim p_{data}(\mathbf{x})}\left[\log p_{\theta }(\mathbf{x})\right]+\mathcal{C}\end{array}$$

The constant ${\mathbb{E}}_{\mathbf{x}\sim p_{data}(\mathbf{x})}[\log p_{data}(\mathbf{x})]$ is simply the entropy of $p_{data}(\mathbf{x})$ and is independent of $\theta$. This is very convenient as $p_{data}(\mathbf{x})$ is unknown and minimizing ${\mathcal{D}}_{KL}$ is equivalent to maximizing the expected log-likelihood of the data $\mathbf{x}$ under $p_{\theta }(\mathbf{x})$:

$$\boxed{\underset{\theta }{\mathrm{argmax}}{\mathbb{E}}_{\mathbf{x}\sim p_{data}(\mathbf{x})}[\log p_{\theta }(\mathbf{x})]}$$

Which is precisely the maximum likelihood estimation (MLE) objective. In practice we replace this population expectation ${\mathbb{E}}_{\mathbf{x}\sim p_{data}(\mathbf{x})}$ by its Monte Carlo estimate, yielding the empirical MLE objective now becomes:

$${\hat{\mathcal{L}}}_{\text{MLE}}(\theta ):=-\frac{1}{N}\sum _{i=1}^N\log p_{\theta }({\mathbf{x}}^{(i)})\text{\hspace{1em}(1)}$$

Where $N$ is the number of samples in the dataset. This objective is then optimized via SGD over minibatches.

1.1 Decoder (Generator)

Returning to the autoencoder setting, the goal is to generate a new sample $\mathbf{x}$ from a latent variable $\mathbf{z}$ via a neural network decoder $p_{\theta }(\mathbf{x}\mid \mathbf{z})$. We can express the target distribution $p_{\theta }(\mathbf{x})$ in equation (1) as the marginal distribution:

$$\boxed{p_{\theta }(\mathbf{x})=\int p_{\theta }(\mathbf{x}\mid \mathbf{z})p(\mathbf{z})d\mathbf{z}}$$

Unfortunately, directly optimizing this objective via MLE is intractable: it requires integrating over the entire high-dimensional latent space, and since $p_{\theta }(\mathbf{x}\mid \mathbf{z})$ is a deep, expressive neural network with no closed-form solution, evaluating this integral exactly is computationally infeasible. To make this optimization tractable, we need a way to focus only on latent states $\mathbf{z}$ that are likely to have generated the current input $\mathbf{x}$, rather than integrating over the entire latent space.

1.2 Encoder (Inference Model)

We can reframe the problem: instead of integrating over all possible $\mathbf{z}$, can we identify which latent states $\mathbf{z}$ are most likely to have produced the observed sample $\mathbf{x}$? This leads us to consider the posterior distribution $p_{\theta }(\mathbf{z}\mid \mathbf{x})$, which by Bayes’ rule is:

$$p_{\theta }(\mathbf{z}\mid \mathbf{x})=\frac{p_{\theta }(\mathbf{x}\mid \mathbf{z})p(\mathbf{z})}{p_{\theta }(\mathbf{x})}$$

However, computing this posterior directly is equally intractable, as the denominator $p_{\theta }(\mathbf{x})$ is the same marginal likelihood we started with. This motivates approximating the true posterior with a learned inference model:

$$q_{\varphi }(\mathbf{z}\mid \mathbf{x})\approx p_{\theta }(\mathbf{z}\mid \mathbf{x})$$

And yes, this is exactly the encoder of the VAE! Which can be trained to concentrates probability mass on the $\mathbf{z}$ state that is most relevant to $\mathbf{x}$.

---

2. ELBO (Evidence Lower Bound)

Now that we have a controllable encoder model to generate $\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})$. We can redefine the MLE optimization goal using $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$.

$$\begin{array}{rl}{p}_{\theta }(\mathbf{x})& =\int p_{\theta }(\mathbf{z},\mathbf{x})d\mathbf{z}\\ \log p_{\theta }(\mathbf{x})& =\log \int p_{\theta }(\mathbf{z},\mathbf{x})d\mathbf{z}\\ & =\log \int q_{\varphi }(\mathbf{z}\mid \mathbf{x})\frac{p_{\theta }(\mathbf{z},\mathbf{x})}{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}d\mathbf{z}\\ & =\log {\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\frac{p_{\theta }(\mathbf{z},\mathbf{x})}{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\right]\end{array}$$

The learning objective is now tractable. Now according to Jensen’s inequality. We have the the evidence lower bound ${\mathcal{L}}_{ELBO}$ where:

$$\begin{array}{r}\log {\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\frac{p_{\theta }(\mathbf{z},\mathbf{x})}{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\right]\ge {\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\log \frac{p_{\theta }(\mathbf{z},\mathbf{x})}{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\right]={\mathcal{L}}_{ELBO}\end{array}\text{\hspace{1em}(2)}$$

Deriving further, we can see that ${\mathcal{L}}_{ELBO}$ consists of 2 terms:

$$\begin{array}{rl}{\mathcal{L}}_{ELBO}& ={\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\log \frac{p_{\theta }(\mathbf{z},\mathbf{x})}{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\right]\\ & ={\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\log \frac{p_{\theta }(\mathbf{x}\mid \mathbf{z})p(\mathbf{z})}{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\right]\\ & ={\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\log p_{\theta }(\mathbf{x}\mid \mathbf{z})\right]-{\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\log \frac{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}{p(\mathbf{z})}\right]\\ & =\boxed{\underset{\text{reconstruction error}}{\underbrace{{\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\log p_{\theta }(\mathbf{x}\mid \mathbf{z})\right]}}-\underset{\text{regularizing term}}{\underbrace{{\mathcal{D}}_{KL}(q_{\varphi }(\mathbf{z}\mid \mathbf{x})\Vert p(\mathbf{z}))}}}\end{array}\text{\hspace{1em}(3)}$$

• Reconstruction term - ${\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\log p_{\theta }(\mathbf{x}\mid \mathbf{z})\right]$: This is the reconstruction objective from the standard AE, but now evaluated only over $\mathbf{z}$ sampled from the encoder, making it tractable.

• Regularization term - ${\mathcal{D}}_{KL}(q_{\varphi }(\mathbf{z}\mid \mathbf{x})\Vert p(\mathbf{z}))$: This penalizes the encoder’s posterior $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ for deviating from the prior $p(\mathbf{z})=\mathcal{N}(0,\mathbf{I})$, enforcing the latent space structure needed for generation.

Why is the learning objective tractable now?

The original MLE objective $\log p_{\theta }(\mathbf{x})=\log \int p_{\theta }(\mathbf{x}\mid \mathbf{z})p(\mathbf{z})d\mathbf{z}$ is intractable because it requires integrating $p_{\theta }(\mathbf{x}\mid \mathbf{z})$ - a neural network - over the entire latent space. The ELBO resolves this in two key ways:

1. Replacing the integral with a tractable expectation: Instead of integrating over all $\mathbf{z}$, the reconstruction term ${\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}[\log p_{\theta }(\mathbf{x}\mid \mathbf{z})]$ only requires sampling $\mathbf{z}$ from the encoder $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$, which concentrates mass on the latent regions most relevant to $\mathbf{x}$.

2. A closed-form KL term: $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ is usually modeled as a simple distribution, usually a gaussian and the KL divergence between two Gaussians has a closed-form solution - no integration is needed at all. And it can also be easily trainable via the reparameterization trick

Together, the two terms create a natural tension: maximizing ${\mathcal{L}}_{ELBO}$ encourages the decoder to recover the original input $\mathbf{x}$ as accurately as possible from latent samples $\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ (reconstruction term), while the regularization term pulls the encoder’s posterior $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ back toward the prior $p(\mathbf{z})$. The VAE learns by striking a balance between these two competing objectives.

ELBO as a Divergence Bound

So what is the relationship between ELBO and the true MLE goal $p_{\theta }(\mathbf{x})$? Recall that maximum likelihood training amounts to minimizing the KL divergence between $p_{data}(\mathbf{x})$ and the learned distribution $p_{\theta }(\mathbf{x})$:

$${\mathcal{D}}_{KL}(p_{data}(\mathbf{x})\Vert p_{\theta }(\mathbf{x}))$$

Since this term is intractable in general, the variational framework of VAE introduces a joint comparison $\mathbf{z}$. Specifically, consider two join distributions

• Generative Join (Decoder): $p_{\theta }(\mathbf{z},\mathbf{x})$
• Inference Join (Encoder): $q_{\varphi }(\mathbf{z},\mathbf{x})$

The total error bound is to match these join together is:

$$\begin{array}{rl}{\mathcal{D}}_{KL}(q_{\varphi }(\mathbf{x},\mathbf{z})\Vert p_{\theta }(\mathbf{x},\mathbf{z}))& =\iint q_{\varphi }(\mathbf{x},\mathbf{z})\log \frac{q_{\varphi }(\mathbf{x},\mathbf{z})}{p_{\theta }(\mathbf{x},\mathbf{z})}d\mathbf{x}d\mathbf{z}\\ & =\iint p_{data}(\mathbf{x})q_{\varphi }(\mathbf{z}\mid \mathbf{x})\log (\frac{p_{data}(\mathbf{x})q_{\varphi }(\mathbf{z}\mid \mathbf{x})}{p_{\theta }(\mathbf{x})p_{\theta }(\mathbf{z}\mid \mathbf{x})})d\mathbf{z}d\mathbf{x}\\ & =\int p_{data}(\mathbf{x})\log (\frac{p_{data}(\mathbf{x})}{p_{\theta }(\mathbf{x})})d\mathbf{x}\\ & +\iint p_{data}(\mathbf{x})q_{\varphi }(\mathbf{z}\mid \mathbf{x})\log (\frac{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}{p_{\theta }(\mathbf{z}\mid \mathbf{x})})d\mathbf{z}d\mathbf{x}\\ & =\underset{\text{True Modeling Error}}{\underbrace{{\mathcal{D}}_{KL}(p_{data}(\mathbf{x})\Vert p_{\theta }(\mathbf{x}))}}\\ & +\underset{\text{Inference Error}}{\underbrace{{\mathbb{E}}_{\mathbf{x}\sim p_{data}(\mathbf{x})}\left[{\mathcal{D}}_{KL}(q_{\varphi }(\mathbf{z}\mid \mathbf{x})\Vert p_{\theta }(\mathbf{z}\mid \mathbf{x}))\right]}}\end{array}$$

Thus we have

$${\mathcal{D}}_{KL}(q_{\varphi }(\mathbf{x},\mathbf{z})\Vert p_{\theta }(\mathbf{x},\mathbf{z}))\ge {\mathcal{D}}_{KL}(p_{data}(\mathbf{x})\Vert p_{\theta }(\mathbf{x}))$$

Where equality happens when inference error is zeros, which also means the encoder $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ perfectly model the unknow posterior distribution $p_{\theta }(\mathbf{z}\mid \mathbf{x})$.

Note that ${\mathcal{L}}_{ELBO}$ can also be rewritten as :

$$\begin{array}{r}{\mathcal{L}}_{ELBO}=\log p_{\theta }(\mathbf{x})-{\mathcal{D}}_{KL}(q_{\varphi }(\mathbf{z}\mid \mathbf{x})\Vert p_{\theta }(\mathbf{z}\mid \mathbf{x}))\\ \to \log p_{\theta }(\mathbf{x})-{\mathcal{L}}_{ELBO}={\mathcal{D}}_{KL}(q_{\varphi }(\mathbf{z}\mid \mathbf{x})\Vert p_{\theta }(\mathbf{z}\mid \mathbf{x}))\end{array}$$

We can see that the gap between the true log-likelihood $\log p_{\theta }(\mathbf{x})$ and the ELBO is precisely the inference error of the current sample $\mathbf{x}$. Maximizing the ELBO therefore directly reduces this gap. Specifically, optimizing the encoder $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ tightens the bound by bringing the approximate posterior closer to the true one $p_{\theta }(\mathbf{z}\mid \mathbf{x})$, while optimizing the decoder $p_{\theta }(\mathbf{x}\mid \mathbf{z})$ pushes the $p_{\theta }(\mathbf{x})$ itself upward - lifting the entire lower bound and improving the overall log-likelihood.

---

3. Gaussian VAEs

The most common instantiation of the VAE framework is the Gaussian VAE, where the encoder, decoder and prior are modeled as Gaussians.

Figure 4: Overview of the Gaussian VAE. Each input $\mathbf{x}$ is encoded into a class-conditional Gaussian $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ (colored clusters). The aggregate posterior $q_{\varphi }(\mathbf{z})={\mathbb{E}}_{\mathbf{x}\sim p_{\text{data}}(\mathbf{x})}[q_{\varphi }(\mathbf{z}\mid \mathbf{x})]$ is matched to the isotropic prior $p(\mathbf{z})$ via the KL term in the ELBO. Samples from $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ are decoded by $p_{\theta }(\mathbf{x}\mid \mathbf{z})$ to produce reconstructions ${\mathbf{x}}^{\prime }$, whose marginal $p_{\theta }(\mathbf{x})$ approximates the data distribution.

3.1 The encoder part

For each input $\mathbf{x}$, the encoder produces a Gaussian distribution centered at ${\mathbf{\mu }}_{\varphi }(\mathbf{x})$ with variance ${\mathbf{\sigma }}_{\varphi }^2(\mathbf{x})$, so that similar inputs yield overlapping distributions in the latent space:

$$\begin{array}{rl}\mathbf{z}& ={\mathbf{\mu }}_{\varphi }(\mathbf{x})+{\mathbf{\sigma }}_{\varphi }(\mathbf{x})\odot \mathbf{\epsilon },\mathbf{\epsilon }\sim \mathcal{N}(0,\mathbf{I})\\ & \Rightarrow q_{\varphi }(\mathbf{z}\mid \mathbf{x})=\mathcal{N}(\mathbf{z};{\mathbf{\mu }}_{\varphi }(\mathbf{x}),\text{diag}({\mathbf{\sigma }}_{\varphi }^2(\mathbf{x})))\end{array}$$

This is the reparameterization trick: by expressing $\mathbf{z}$ as a deterministic function of $\varphi$ and a fixed noise variable $\mathbf{\epsilon }$, the stochasticity is separated from the parameters, making the sampling step differentiable and allowing gradients to flow back through $\mathbf{z}$ to the encoder.

Since the prior $p(\mathbf{z})=\mathcal{N}(0,\mathbf{I})$ is also Gaussian, the KL divergence between two Gaussians admits a closed-form solution - no numerical integration required:

$$\begin{array}{r}{\mathcal{L}}_{KL}={\mathcal{D}}_{KL}(\mathcal{N}({\mathbf{\mu }}_{\varphi },\text{diag}({\mathbf{\sigma }}_{\varphi }^2))\Vert \mathcal{N}(0,\mathbf{I}))\\ \boxed{{\mathcal{L}}_{KL}=-\frac{1}{2}\sum _{j=1}^d\left(1+\log {\sigma }_j^2-{\mu }_j^2-{\sigma }_j^2\right)}\end{array}$$

With $d$ is the number of dimension of the latent space.

Derivation of closed-form KL loss

Since both $q_{\varphi }(\mathbf{z}\mid \mathbf{x})=\mathcal{N}({\mathbf{\mu }}_{\varphi },\text{diag}({\mathbf{\sigma }}_{\varphi }^2))$ and $p(\mathbf{z})=\mathcal{N}(0,\mathbf{I})$ are diagonal, the KL factorizes over dimensions. It suffices to derive for a single scalar dimension $z\sim \mathcal{N}(\mu ,{\sigma }^2)$ vs $z\sim \mathcal{N}(0,1)$:

$$\begin{array}{rl}{\mathcal{D}}_{KL}(\mathcal{N}(\mu ,{\sigma }^2)\Vert \mathcal{N}(0,1))& =\int \mathcal{N}(z;\mu ,{\sigma }^2)\log \frac{\mathcal{N}(z;\mu ,{\sigma }^2)}{\mathcal{N}(z;0,1)}dz\\ & ={\mathbb{E}}_q\left[\log \mathcal{N}(z;\mu ,{\sigma }^2)-\log \mathcal{N}(z;0,1)\right]\\ & ={\mathbb{E}}_q\left[\left(-\frac{1}{2}\log (2\pi {\sigma }^2)-\frac{(z-\mu {)}^2}{2{\sigma }^2}\right)-\left(-\frac{1}{2}\log (2\pi )-\frac{z^2}{2}\right)\right]\\ & ={\mathbb{E}}_q\left[-\frac{1}{2}\log {\sigma }^2-\frac{(z-\mu {)}^2}{2{\sigma }^2}+\frac{z^2}{2}\right]\\ & =-\frac{1}{2}\log {\sigma }^2-\frac{1}{2{\sigma }^2}\underset{={\sigma }^2}{\underbrace{{\mathbb{E}}_q[(z-\mu {)}^2]}}+\frac{1}{2}\underset{={\sigma }^2+{\mu }^2}{\underbrace{{\mathbb{E}}_q[z^2]}}\\ & =-\frac{1}{2}\log {\sigma }^2-\frac{1}{2}+\frac{{\sigma }^2+{\mu }^2}{2}\\ & =-\frac{1}{2}\left(1+\log {\sigma }^2-{\mu }^2-{\sigma }^2\right)\end{array}$$

Summing over all $d$ independent dimensions:

$$\boxed{{\mathcal{D}}_{KL}(q_{\varphi }(\mathbf{z}\mid \mathbf{x})\Vert p(\mathbf{z}))=-\frac{1}{2}\sum _{j=1}^d\left(1+\log {\sigma }_j^2-{\mu }_j^2-{\sigma }_j^2\right)}$$

Taking the gradient of the KL term with respect to ${\mu }_j$ and ${\sigma }_j^2$ we have:

$$\begin{array}{rl}\frac{\partial }{\partial {\mu }_j}{\mathcal{D}}_{KL}& =\frac{\partial }{\partial {\mu }_j}\left[-\frac{1}{2}(1+\log {\sigma }_j^2-{\mu }_j^2-{\sigma }_j^2)\right]={\mu }_j\\ \frac{\partial }{\partial {\sigma }_j^2}{\mathcal{D}}_{KL}& =-\frac{1}{2}\left(\frac{1}{{\sigma }_j^2}-1\right)\end{array}$$

Setting these to zero $\Rightarrow {\hat{\mu }}_j=0,{\hat{\sigma }}_j^2=1$. Therefore minimizing the KL term alone pushes the encoder toward:

$$q_{\varphi }(\mathbf{z}\mid \mathbf{x})⟶\mathcal{N}(0,\mathbf{I})=p(\mathbf{z})$$

This is why the reconstruction term is essential: it pulls ${\mu }_j$ away from zero and ${\sigma }_j^2$ toward smaller values to make $\mathbf{z}$ informative about $\mathbf{x}$.

3.2 The Decoder part

To counteract collapse from the regularization term, the reconstruction term enforces that $\mathbf{z}$ remains informative about $\mathbf{x}$. Specifically, the decoder is trained to output a sample ${\mathbf{x}}^{\prime }$ that resembles the original input as closely as possible, given a latent vector $\mathbf{z}$ drawn from the encoder’s posterior $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$. Note that ${\mathbf{x}}^{\prime }$ need not be identical to $\mathbf{x}$:

$$\begin{array}{r}\mathbf{x}={\mathbf{\mu }}_{\theta }(\mathbf{z})+\sigma \odot \mathbf{\epsilon }\sim \mathcal{N}(0,\mathbf{I})\\ \Rightarrow p_{\theta }(\mathbf{x}\mid \mathbf{z})=\mathcal{N}(\mathbf{x};{\mathbf{\mu }}_{\theta }(\mathbf{z}),\mathbf{I}\sigma )\end{array}$$

Here ${\mathbf{\mu }}_{\theta }(\mathbf{z})$ is the output of a neural network decoder, and $\sigma$ is a fixed hyperparameter controlling the spread of the output distribution - large $\sigma$ allows more deviation from the input, while small $\sigma$ forces the reconstruction to stay close to the input $\mathbf{x}$. The reconstruction loss can now be rewritten as:

$$\begin{array}{rl}{\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\log p_{\theta }(\mathbf{x}\mid \mathbf{z})\right]& ={\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[-\frac{1}{2{\sigma }^2}\Vert \mathbf{x}-{\mathbf{\mu }}_{\theta }(\mathbf{z}){\Vert }^2\right]+\log \left(\frac{1}{\sqrt{2\pi {\sigma }^2}}\right)\\ & \propto -{\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\Vert \mathbf{x}-{\mathbf{\mu }}_{\theta }(\mathbf{z}){\Vert }^2\right]\end{array}$$ $$\boxed{{\mathcal{L}}_{recon}={\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\Vert \mathbf{x}-{\mathbf{\mu }}_{\theta }(\mathbf{z}){\Vert }^2\right]}$$

This is equivalent to minimizing the expected MSE between the input $\mathbf{x}$ and the decoder output ${\mathbf{\mu }}_{\theta }(\mathbf{z})$ - which is similar to the original AE loss.

3.3 Overall Training Procedure

With both the encoder and decoder defined, the full training procedure follows directly from maximizing the ELBO. Each training step processes a minibatch of inputs:

$$\begin{array}{l}\text{Algorithm: Training Gaussian VAE with ELBO}\\ \text{Input: }\text{Dataset }\mathcal{X}=\{{\mathbf{x}}^{(1)},\dots ,{\mathbf{x}}^{(N)}\},\text{ batch size }B,\text{ learning rate }\eta \\ \text{Output: }\text{Encoder parameters }\varphi ,\text{ Decoder parameters }\theta \\ 1.\text{Initialize }\varphi ,\theta \text{ randomly}\\ 2.\text{repeat}\\ 3.\text{Sample minibatch }\{{\mathbf{x}}^{(1)},\dots ,{\mathbf{x}}^{(B)}\}\sim \mathcal{X}\\ 4.\text{// Encoder forward pass}\\ 5.({\mathbf{\mu }}^{(i)},{\mathbf{\sigma }}^{(i)})\leftarrow {\text{Encoder}}_{\varphi }({\mathbf{x}}^{(i)})\forall i\\ 6.\text{// Reparameterization trick}\\ 7.{\mathbf{\epsilon }}^{(i)}\sim \mathcal{N}(0,\mathbf{I})\\ 8.{\mathbf{z}}^{(i)}\leftarrow {\mathbf{\mu }}^{(i)}+{\mathbf{\sigma }}^{(i)}\odot {\mathbf{\epsilon }}^{(i)}\\ 9.\text{// Decoder forward pass}\\ 10.{\hat{\mathbf{x}}}^{(i)}\leftarrow {\text{Decoder}}_{\theta }({\mathbf{z}}^{(i)})\forall i\\ 11.\text{// Compute ELBO}\\ 12.{\mathcal{L}}_{\text{recon}}\leftarrow \frac{1}{B}{\sum }_i\log p_{\theta }({\mathbf{x}}^{(i)}\mid {\mathbf{z}}^{(i)})\\ 13.{\mathcal{L}}_{\text{KL}}\leftarrow \frac{1}{B}{\sum }_i{\mathcal{D}}_{KL}(q_{\varphi }(\mathbf{z}\mid {\mathbf{x}}^{(i)})\Vert p(\mathbf{z}))\\ 14.{\mathcal{L}}_{\text{ELBO}}\leftarrow {\mathcal{L}}_{\text{recon}}-{\mathcal{L}}_{\text{KL}}\\ 15.\text{// Update parameters}\\ 16.(\varphi ,\theta )\leftarrow (\varphi ,\theta )+\eta \cdot {\nabla }_{\varphi ,\theta }{\mathcal{L}}_{\text{ELBO}}\\ 17.\text{until }\text{convergence}\end{array}$$

---

4. Drawbacks of Gaussian VAEs

Despite its elegance, the Gaussian VAE has several well-known limitations:

4.1 Blurry reconstructions

Modeling $p_{\theta }(\mathbf{x}\mid \mathbf{z})$ as a Gaussian with a fixed variance corresponds to minimizing MSE, which tends to average over multiple plausible reconstructions using the sampled code $\mathbf{z}$.

Proof of blurry reconstructions in VAEs

Recall the per-sample reconstruction loss:

$${\mathcal{L}}_{recon}={\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\Vert \mathbf{x}-{\mathbf{\mu }}_{\theta }(\mathbf{z}){\Vert }^2\right]$$

When training over the full dataset, we optimize its expectation over all $\mathbf{x}\sim p_{data}$:

$$\begin{array}{rl}& {\mathbb{E}}_{\mathbf{x}\sim p_{data},\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\Vert \mathbf{x}-{\mathbf{\mu }}_{\theta }(\mathbf{z}){\Vert }^2\right]\\ & =\iint p_{data}(\mathbf{x})q_{\varphi }(\mathbf{z}|\mathbf{x})\Vert \mathbf{x}-{\mathbf{\mu }}_{\theta }(\mathbf{z}){\Vert }^{2}d\mathbf{x}d\mathbf{z}\\ & =\iint q_{\varphi }(\mathbf{z})q_{\varphi }(\mathbf{x}|\mathbf{z})\Vert \mathbf{x}-{\mathbf{\mu }}_{\theta }(\mathbf{z}){\Vert }^{2}d\mathbf{x}d\mathbf{z}\\ & =\int q_{\varphi }(\mathbf{z})\left[\int q_{\varphi }(\mathbf{x}|\mathbf{z})\Vert \mathbf{x}-{\mathbf{\mu }}_{\theta }(\mathbf{z}){\Vert }^{2}d\mathbf{x}\right]d\mathbf{z}\\ & ={\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z})}\left[{\mathbb{E}}_{\mathbf{x}\sim q_{\varphi }(\mathbf{x}\mid \mathbf{z})}\left[\Vert \mathbf{x}-{\mathbf{\mu }}_{\theta }(\mathbf{z}){\Vert }^2\right]\right]\end{array}$$

Since $\theta$ only appears in the inner expectation and has no effect on the aggregate posterior $q_{\varphi }(\mathbf{z})$, the outer expectation acts as a constant weight. It suffices to minimize the inner term with respect to ${\mathbf{\mu }}_{\theta }(\mathbf{z})$ for each fixed $\mathbf{z}$. Taking the gradient and setting it to zero:

$$\begin{array}{rl}\frac{\partial }{\partial {\mathbf{\mu }}_{\theta }}{\mathbb{E}}_{\mathbf{x}\sim q_{\varphi }(\mathbf{x}\mid \mathbf{z})}\left[\Vert \mathbf{x}-{\mathbf{\mu }}_{\theta }(\mathbf{z}){\Vert }^2\right]& ={\mathbb{E}}_{\mathbf{x}\sim q_{\varphi }(\mathbf{x}\mid \mathbf{z})}\left[-2\left(\mathbf{x}-{\mathbf{\mu }}_{\theta }(\mathbf{z})\right)\right]=0\\ & \Rightarrow \boxed{{\mathbf{\mu }}_{\theta }^{\ast }(\mathbf{z})={\mathbb{E}}_{\mathbf{x}\sim q_{\varphi }(\mathbf{x}\mid \mathbf{z})}\left[\mathbf{x}\right]}\end{array}$$

The optimal decoder output is the conditional mean of $\mathbf{x}$ given $\mathbf{z}$ under the encoder’s inverse distribution $q_{\varphi }(\mathbf{x}\mid \mathbf{z})$. When multiple distinct images $\mathbf{x}$ map to similar latent codes $\mathbf{z}$, the MSE loss forces the decoder to output their average - producing blurry reconstructions.

For image data, this produces blurry outputs rather than sharp, realistic samples.

So what could be done to improve this? - One natural idea is to make the encoder or decoder more expressive by adding more layers to capture richer, more complex latent structure. However, increasing depth alone does not solve the Limited posterior expressiveness problem of the encoder. Moreover, when the decoder grows too powerful, it can reconstruct $\mathbf{x}$ without relying on $\mathbf{z}$ at all, triggering posterior collapse.

4.2 Limited posterior expressiveness

The diagonal Gaussian assumption for $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ restricts the variational family distribution to an axis-aligned ellipsoid (i.e., zero off-diagonal covariance). If the true posterior $p_{\theta }(\mathbf{z}\mid \mathbf{x})$ has complex, multimodal, or highly correlated structure, a single Gaussian cannot capture it - leading to a persistently loose ELBO bound regardless of the encoder capacity.

4.3 Posterior Collapse

Why does this happen? One key reason is that when the decoder becomes too expressive (a sufficiently deep neural network), at some point during training, for some $\mathbf{x}\sim p_{data}(\mathbf{x})$, the decoder finds it easier to simply approximate the data distribution directly:

$$p_{\theta }(\mathbf{x}\mid \mathbf{z})\approx p_{data}(\mathbf{x})$$

This sounds like a win for reconstruction quality, but it also breaks the encoder. In other words, the decoder has learned to ignore $\mathbf{z}$, the VAE’s output now barely changes no matter what latent code it receives. Recall the ELBO:

$${\mathcal{L}}_{ELBO}=\underset{\text{reconstruction}}{\underbrace{{\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}\mid \mathbf{x})}[\log p_{\theta }(\mathbf{x}\mid \mathbf{z})]}}-\underset{\text{regularization}}{\underbrace{{\mathcal{D}}_{KL}(q_{\varphi }(\mathbf{z}\mid \mathbf{x})\Vert p(\mathbf{z}))}}$$

When the decoder is powerful enough to reconstruct $\mathbf{x}$ without using $\mathbf{z}$, by exploiting its own internal structure, the reconstruction term becomes approximately constant with respect to $\mathbf{z}$. The gradient signal that would normally force the encoder to encode information into $\mathbf{z}$ disappears. The optimizer then finds the path of least resistance: collapse the KL term to zero by driving $q_{\varphi }(\mathbf{z}\mid \mathbf{x})\to p(\mathbf{z})$. At this point, $\mathbf{z}$ and $\mathbf{x}$ become statistically independent, the latent code carries no information about the input, and the decoder can no longer be used to control the output. The VAE reduces to a decoder-only model. so the ELBO degenerates to:

$${\mathcal{L}}_{ELBO}\approx \log p_{\theta }(\mathbf{x})$$

Proof for the independence of $\mathbf{x}$ and $\mathbf{z}$ after postierior corruption

We can rewrite the regularization term, averaged over all $\mathbf{x}\sim p_{data}(\mathbf{x})$, as:

$$\begin{array}{rl}{\mathbb{E}}_{\mathbf{x}\sim p_{data}(\mathbf{x})}\left[D_{KL}(q_{\varphi }(\mathbf{z}\mid \mathbf{x})\mid p(\mathbf{z}))\right]& =\iint p_{data}(\mathbf{x})q_{\varphi }(\mathbf{z}\mid \mathbf{x})\log \frac{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}{p(\mathbf{z})}d\mathbf{z}d\mathbf{x}\\ & =\iint q_{\varphi }(\mathbf{z},\mathbf{x})\log (\frac{q_{\varphi }(\mathbf{x}\mid \mathbf{z})q_{\varphi }(\mathbf{z})}{p(\mathbf{z})p_{data}(\mathbf{x})})d\mathbf{z}d\mathbf{x}\\ & =\iint q_{\varphi }(\mathbf{z},\mathbf{x})\log (\frac{p_{\varphi }(\mathbf{x}\mid \mathbf{z})}{p_{data}(\mathbf{x})})d\mathbf{z}d\mathbf{x}+\iint q_{\varphi }(\mathbf{z},\mathbf{x})\log (\frac{q_{\varphi }(\mathbf{z})}{p(\mathbf{z})})d\mathbf{z}d\mathbf{x}\\ & =\iint q_{\varphi }(\mathbf{z},\mathbf{x})\log (\frac{q_{\varphi }(\mathbf{x},\mathbf{z})}{p_{data}(\mathbf{x})q_{\varphi }(\mathbf{z})})d\mathbf{z}d\mathbf{x}+\int q_{\varphi }(\mathbf{z})\log (\frac{q_{\varphi }(\mathbf{z})}{p(\mathbf{z})})d\mathbf{z}\\ & =\boxed{\underset{\text{mutual information}}{\underbrace{\mathcal{I}(\mathbf{x},\mathbf{z})}}+D_{KL}(q_{\varphi }(\mathbf{z})\Vert p(\mathbf{z}))}\end{array}$$

This decomposition reveals what posterior collapse actually destroys. When the regularization term is driven to zero, both components must vanish simultaneously:

• $\mathcal{I}(\mathbf{x},\mathbf{z})\to 0$ - The latent code $\mathbf{z}$ becomes statistically independent of $\mathbf{x}$ (recall independent property), i.e. $q_{\varphi }(\mathbf{x},\mathbf{z})=q_{\varphi }(\mathbf{z})p_{data}(\mathbf{x})$.
• $D_{KL}(q_{\varphi }(\mathbf{z})\Vert p(\mathbf{z}))\to 0$ - The aggregate posterior $q_{\varphi }(\mathbf{z})$ collapses to the isotropic Gaussian $p(\mathbf{z})=\mathcal{N}(0,\mathbf{I})$. The latent space loses all input class-specific structure, the per-class mixture components that distinguish different types of $\mathbf{x}$ are wiped out.

4.4 Mismatch between aggregate posterior and prior

Another weakness of standard VAEs is that even if each individual posterior $q_{\varphi }(\mathbf{z}\mid \mathbf{x})$ is close to the prior, the aggregated posterior $q_{\varphi }(\mathbf{z})$ may not match $p(\mathbf{z})=\mathcal{N}(0,\mathbf{I})$ at some regions. This mismatch creates “holes” in the latent space - regions with high prior probability but low posterior density - causing poor sample quality at generation time.

---

5. Conclusion

The Variational Autoencoder is a foundational generative model that elegantly combines probabilistic inference with deep learning. By replacing the deterministic bottleneck of a standard autoencoder with a learned posterior distribution, VAEs endow the latent space with a structured, continuous geometry that supports both generation and interpolation. The ELBO provides a tractable training objective that simultaneously encourages faithful reconstruction and regularizes the latent space toward a simple prior - a tension that lies at the heart of all latent-variable generative models.

That said, the Gaussian VAE is far from perfect. The four drawbacks discussed above - blurry reconstructions from the MSE objective, posterior collapse, limited posterior expressiveness from the diagonal Gaussian assumption, and aggregate posterior mismatch - are not merely implementation details; they are fundamental limitations that arise from the design choices made to keep the ELBO tractable.

What comes next? Two important lines of work build directly on these observations to solve standard VAE’s limitations:

• Hierarchical VAEs (HVAEs) address the expressiveness problem by stacking multiple layers of stochastic latent variables. Rather than compressing $\mathbf{x}$ into a single $\mathbf{z}$, HVAEs learn a hierarchy ${\mathbf{z}}_1,{\mathbf{z}}_2,\dots ,{\mathbf{z}}_L$ where each layer captures structure at a different level of abstraction. This allows the model to represent far richer posteriors, and the ELBO generalizes naturally to the hierarchical setting.

• Denoising Diffusion Probabilistic Models (DDPMs) take a different philosophical path. Instead of learning a compact latent code, diffusion models define a fixed forward process that gradually corrupts data with Gaussian noise over $T$ steps, then learn to reverse this process step by step. Remarkably, this can be seen as a special case of a hierarchical latent-variable model where the encoder is fixed (the forward noising process) and only the decoder (the denoising network) is learned. This design sidesteps the blurry-reconstruction and posterior-collapse problems entirely - the fixed encoder cannot collapse, and the step-by-step denoising objective enforces sharp, high-frequency detail at each scale.

---

References

[1]— Kingma, D. P., & Welling, M. (2013). Auto-encoding variational Bayes.

[2]—The Principles of Diffusion Models. (n.d.). https://the-principles-of-diffusion-models.github.io/

[3]—Wikipedia contributors. (n.d.). Jensen’s inequality. Wikipedia. https://en.wikipedia.org/wiki/Jensen%27s_inequality

[4]—Wikipedia contributors. (n.d.). Monte Carlo method. Wikipedia. https://en.wikipedia.org/wiki/Monte_Carlo_method

[5]—Wikipedia contributors. (n.d.). Reparameterization trick. Wikipedia. https://en.wikipedia.org/wiki/Reparameterization_trick

Notations

See the notation reference for a summary of symbols used across all notes.

---

AI Acknowledgement

And yes, I had AI helped with wording and structure (Claude by Anthropic) (•‿•)