Variational Autoencoder

The Variational Autoencoder (VAE) fixes the exact limitation flagged in Autoencoders: a plain autoencoder's latent space isn't guaranteed to be smooth or well-structured, so sampling random points from it rarely produces good new data. The VAE makes the latent space explicitly probabilistic and well-behaved, turning the autoencoder into a genuine generative model.

The Key Structural Change — A Distribution, Not a Point

Instead of the encoder outputting one fixed latent vector \(\mathbf{z}\), a VAE's encoder outputs the parameters of a distribution — a mean \(\boldsymbol\mu\) and a standard deviation \(\boldsymbol\sigma\) — and \(\mathbf{z}\) is then sampled from \(\mathcal{N}(\boldsymbol\mu, \boldsymbol\sigma^2)\).

The Reparameterization Trick

Directly sampling \(\mathbf{z} \sim \mathcal{N}(\boldsymbol\mu,\boldsymbol\sigma^2)\) would break backpropagation — random sampling isn't a differentiable operation with respect to \(\boldsymbol\mu\) and \(\boldsymbol\sigma\). The reparameterization trick sidesteps this: move the randomness into a separate, fixed-distribution noise variable \(\boldsymbol\epsilon\), and compute \(\mathbf{z}\) as a simple, fully differentiable function of \(\boldsymbol\mu\), \(\boldsymbol\sigma\), and \(\boldsymbol\epsilon\) — gradients can now flow cleanly back through \(\boldsymbol\mu\) and \(\boldsymbol\sigma\) via ordinary backpropagation.

The Complete VAE Loss

This is exactly the assembled VAE loss previewed back in Reconstruction Loss — reconstruction loss (from that note) plus KL divergence (from KL Divergence Loss) pulling every input's latent distribution toward a standard normal prior. This KL term is exactly what makes the latent space smooth and well-structured — every region near the origin, in every direction, is encouraged to correspond to some plausible reconstruction, rather than only specific isolated points being meaningful.

Diagram — Why This Enables Generation

Code

import torch
import torch.nn as nn

class VAE(nn.Module):
    def __init__(self, input_dim, latent_dim):
        super().__init__()
        self.fc_mu = nn.Linear(128, latent_dim)
        self.fc_logvar = nn.Linear(128, latent_dim)
        self.encoder_body = nn.Sequential(nn.Linear(input_dim, 128), nn.ReLU())
        self.decoder = nn.Sequential(nn.Linear(latent_dim, 128), nn.ReLU(), nn.Linear(128, input_dim), nn.Sigmoid())

    def forward(self, x):
        h = self.encoder_body(x)
        mu, logvar = self.fc_mu(h), self.fc_logvar(h)
        std = torch.exp(0.5 * logvar)
        epsilon = torch.randn_like(std)
        z = mu + std * epsilon   # the reparameterization trick
        return self.decoder(z), mu, logvar

def vae_loss(x_hat, x, mu, logvar):
    recon_loss = nn.functional.binary_cross_entropy(x_hat, x, reduction='sum')
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + kl_loss

# Generating a NEW sample after training: just sample z from the prior directly
z_new = torch.randn(1, 20)   # sample from N(0,1) -- the SAME distribution the KL term trained toward
# generated_sample = model.decoder(z_new)

Common Mistakes

• Forgetting the reparameterization trick and sampling \(\mathbf{z}\) directly — this breaks gradient flow through \(\boldsymbol\mu\) and \(\boldsymbol\sigma\), since sampling itself isn't a differentiable operation.
• Setting the KL term's weight too high relative to reconstruction — this can cause "posterior collapse," where the encoder essentially ignores the input and the latent space converges to the prior regardless of what's being encoded, producing poor reconstructions.

Interview Relevance

Q: "Why can't you sample from a VAE's latent distribution directly during training without the reparameterization trick?" Sampling is a stochastic, non-differentiable operation — gradients can't flow backward through a random draw with respect to the distribution's parameters (\(\boldsymbol\mu\), \(\boldsymbol\sigma\)). The reparameterization trick rewrites the sampling as a deterministic, differentiable function of \(\boldsymbol\mu\), \(\boldsymbol\sigma\), and a separate fixed-distribution noise variable \(\boldsymbol\epsilon\), letting backpropagation flow through \(\boldsymbol\mu\) and \(\boldsymbol\sigma\) normally.

Practice Question

Why does the KL divergence term in a VAE's loss specifically push the latent distribution toward a standard normal, rather than toward any arbitrary fixed distribution?

Related DL Notes

• Reconstruction Loss
• KL Divergence Loss
• GAN