Aditya Venkatesh's Personal Website

Loss and Reparameterization

VAEs optimize a cleverly designed loss function that balances two goals. First, reconstruction loss measures how well the decoder rebuilds the original input, often using mean squared error for images:$${\mathcal{L}}_{\text{recon}}=\mathrm{\parallel }x-\widehat{x}{\mathrm{\parallel }}^2$$

where $x$ is the input and $\widehat{x}$ the reconstruction.

Second, KL divergence regularizes the latent distribution to stay close to a standard normal prior $\mathcal{N}(0,1)$:$${\mathcal{L}}_{\text{KL}}=D_{\text{KL}}(q(z\mathrm{\mid }x)\mathrm{\parallel }p(z))$$

This prevents the model from memorizing training data and encourages a structured latent space. The full VAE loss is $\mathcal{L}={\mathcal{L}}_{\text{recon}}+\beta {\mathcal{L}}_{\text{KL}}$, where $\beta$ tunes the tradeoff.

Training requires the reparameterization trick to make sampling differentiable. Instead of sampling $z\sim q(z\mathrm{\mid }x)$ directly (which breaks gradients), we compute $z=\mu +\sigma \odot ϵ$ where $ϵ\sim \mathcal{N}(0,1)$. This shifts randomness to a constant, allowing backpropagation through $\mu$ and $\sigma$.

The Magic of Smooth Latent Spaces

VAEs shine in their latent space properties. Because encodings form distributions pulled toward a Gaussian prior, the space becomes continuous and interpolable. Encode two cat images, then linearly interpolate their latent vectors – the decoder spits out convincing morphs between them. This smoothness arises directly from KL regularization, which spaces out representations evenly.

Latent dimensions also capture semantic features. Lower dimensions might encode broad categories like “animal” or “background,” while higher ones handle details like fur texture. This structure makes VAEs ideal for tasks beyond pure reconstruction.