11. Generative AI: VAEs, Diffusion and Flows › 11.3 Variational Auto-Encoders

Section 11.3
Variational Auto-Encoders

Illustration of an auto-encoder where the latent representation between the encoder and decoder has been replaced by a sample from a latent distribution resulti
Figure 135: Illustration of an auto-encoder where the latent representation between the encoder and decoder has been replaced by a sample from a latent distribution resulting in a variational autoencoder.

A standard auto-encoder learns a compressed representation of data by mapping an input \(x\) to a latent representation \(z\), and then reconstructing the input as \(\hat{x}\). The architecture consists of and encoder \(E\) and decoder \(D\) with parameters \(\theta\) and \(\phi\), respectively. The encoder computes,

\begin{align} z &= E(x; \theta) \tag{328}\end{align}

and the decoder reconstructs the input as

\begin{align} \hat{x} = D(z; \phi) \tag{329}\end{align}

The model is trained to minimize reconstruction error, e.g.,

\begin{align} L(\theta, \phi) &= \frac{1}{2} \sum_{x \in \cD} \|x - \hat{x}\|^2 = \frac{1}{2} \sum_{x \in \cD} \|x - D(E(x; \theta); \phi)\|^2 \tag{330}\end{align}

over a dataset of examples \(\cD\).

Although auto-encoders can learn compact latent representations, the latent space is often poorly structured. Small perturbations in \(z\) may produce unrealistic outputs, making standard auto-encoders unsuitable for generative modeling.

A variational auto-encoder (VAE) addresses this issue by making the latent representation probabilistic. Instead of mapping an input to a single point \(z\), the encoder predicts a probability distribution, \(q(z \mid x)\), and the decoder then models \(p(x \mid z)\) as shown in Figure 135. Thus, the encoder produces a distribution over latent variables and the decoder generates data conditioned on sampled latent variables.

Typically, the encoder outputs the parameters of a Gaussian distribution,

\begin{align} q(z \mid x) &= \cN(\mu(x), \Sigma(x)) \tag{331}\end{align}

where \(\mu(x)\) is the latent mean and \(\Sigma(x)\) is the latent covariance matrix, both estimated by the encoder and provided as output. In practice \(\Sigma(x)\) is assumed to be diagonal for efficiency.

A discussed in the preliminaries, a latent sample can be obtained using the reparameterization trick,

\begin{align} z = \mu(x) + \Sigma(x)^{1/2} \epsilon \tag{332}\end{align}

where \(\epsilon \sim \cN(0, I)\). This sampled latent vector is then passed through the decoder to generate \(\hat{x}\).

The model is trained by deriving a variational lower bound known as the evidence lower bound (ELBO). Briefly, we introduce an approximate posterior distribution \(q(z \mid x)\) to the true posterior \(p(z \mid x)\) and rewrite the marginal likelihood as

\begin{align} \log p(x) &= \log \left(\int_z p(x, z) \,\text{d}z \right) \tag{333}\\ &= \log \left(\int_z q(z \mid x) \frac{p(x, z)}{q(z \mid x)} \,\text{d}z \right) \tag{334}\\ &\geq E_{q(z \mid x)} \left[ \log \frac{p(x, z)}{q(z \mid x)} \right] & \text{(Jensen's inequality)} \tag{335}\\ &= E_{q(z \mid x)} \left[ \log \frac{p(x \mid z)p(z)}{q(z \mid x)} \right] \tag{336}\\ &= E_{q(z \mid x)} \left[ \log p(x \mid z) \right] - E_{q(z \mid x)} \left[ \log \frac{q(z \mid x)}{p(z)} \right] \tag{337}\end{align}

where we have used Jensen’s inequality on Line 3 to obtain a tractable bound. The second term is a Kullback-Leibler divergence \(D_{\text{KL}} \left(q(z \mid x; \theta) \,\|\, p(z) \right)\). This allows us to write the VAE training objective as

\begin{align} \ell^{\text{ELBO}}(\theta, \phi) &= \underbrace{E_{q(z \mid x)} \left[\log p(x \mid z; \phi)\right]}_{\text{reconstruction loss}} - \underbrace{D_{\text{KL}} \left(q(z \mid x; \theta) \,\|\, p(z) \right)}_{\text{regularization}} \tag{338}\end{align}

where the first term is a reconstruction loss that encourages the decoder to accurately reconstruct the input \(x\) from the latent representation \(z\). For Gaussian decoders this reduces to a squared-error loss. The second term provides regularization so that there is meaningful interpolation between latent points, stable sampling, and generative capability. Without the regularization term the latent space can become fragmented and difficult to sample from. A full training algorithm is shown below.

1:procedure LearnVAE
2:initialize \(\theta, \phi\)▷ initialize encoder and decoder parameters
3:repeat
4:\(\ell = 0\)▷ zero loss
5:for each sample \(x\) in mini-batch do
6:\(\mu, \Sigma = E(x; \theta)\)▷ encoder forward pass
7:\(\epsilon \sim {\cal N}(0, I)\)▷ sample noise
8:\(z = \mu + \Sigma^{1/2} \epsilon\)▷ reparameterization trick
9:\(\hat{x} = D(z; \theta)\)▷ decoder forward pass
10:\(\ell \gets \ell + \frac{1}{2} \| \hat{x} - x \|^2 - \frac{1}{2} \left(\trace{\Sigma} + \|\mu\|^2 - \log \det \Sigma \right)\)▷ update loss
11:end for
12:\(\theta \gets \theta - \eta \nabla_{\theta} \ell\)▷ update encoder parameters
13:\(\phi \gets \phi - \eta \nabla_{\phi} \ell\)▷ update decoder parameters
14:until converged
15:return \(\theta, \phi\)
16:end procedure

Although VAEs are stable and principled generative models, they have several weaknesses. First, the reconstructions are often blurrier than the outputs from other generative models such as GANs or diffusion models because they optimize a likelihood-based objective. Second, they are susceptible to posterior collapse where the decoder ignores the latent variables entirely, \(q(z \mid x) \approx p(z)\). This is especially common when the decoder is too powerful.

Despite these limitations VAEs have found application in a number of areas, and extended to allow the incorporation of a conditioning signal (resulting in CVAEs) and quantized latent spaces (VQ-VAEs), which give sharper reconstructions and mitigate against posterior collapse. VQ-VAEs are widely used in audio generative systems and sequences requiring tokenized latent representations [100].