Variational autoencoders learn structured low-dimensional manifolds by optimizing the evidence lower bound

Video. Variational Autoencoders from Scratch! · 1:19:22 · YouTube

I cover the motivation, mathematical derivation, PyTorch implementation, and practical training considerations for variational autoencoders using the CelebA dataset.

The idea

The manifold hypothesis: High-dimensional sensory data naturally lies on a much lower-dimensional continuous manifold where semantic attributes are smoothly organized. Standard autoencoders can compress data, but their latent spaces lack structure, making it difficult to generate valid new samples. Variational autoencoders fix this by forcing the encoder to map inputs to probability distributions over the latent space rather than single deterministic points, encouraging continuous and smooth latent representations. By balancing how well the decoder reconstructs inputs against how closely the latent distributions match a simple Gaussian prior, I build a generative model capable of synthesizing novel data points.

The mechanism

\log p_\theta(x) \ge \mathbb{E}_{q_\phi(z|x)} \left[ \log p_\theta(x|z) \right] - D_{\text{KL}}\left( q_\phi(z|x) \,||\, p(z) \right)

This equation defines the evidence lower bound objective by balancing reconstruction fidelity against latent regularization relative to a prior.

z = \mu_\phi(x) + \sigma_\phi(x) \odot \epsilon, \quad \epsilon \sim \mathcal{N}(0, I)

This reparameterization trick isolates the stochastic sampling step outside the computational graph, allowing backpropagation of gradients through encoder parameters.

D_{\text{KL}}\left(\mathcal{N}(\mu, \operatorname{diag}(\sigma^2)) \,||\, \mathcal{N}(0, I)\right) = -\frac{1}{2} \sum_{j=1}^d \left( 1 + \log(\sigma_j^2) - \mu_j^2 - \sigma_j^2 \right)

This computes the closed-form distance between the predicted diagonal Gaussian posterior and the standard normal prior.

\mathcal{L}_{\text{VAE}} = \text{MSE}(x, \hat{x}) - \beta \cdot D_{\text{KL}}\left( q_\phi(z|x) \,||\, p(z) \right)

This loss objective applies a hyperparameter weight to the KL divergence term to mitigate posterior collapse while preserving reconstruction capability.

Worth knowing

Performing image resizing before converting images to PyTorch tensors prevents severe CPU bottlenecks during data loading.
Setting the KL divergence weight ( $\beta$ ) too high degrades reconstruction fidelity, while setting it to zero collapses the architecture into a deterministic autoencoder.
Direct stochastic sampling blocks gradient flow during backpropagation, requiring explicit extraction of random noise via the reparameterization trick.

Code

The core reparameterization trick, KL divergence calculation, and forward sampling pass are implemented in modeling.py inside the VAE class.

@staticmethod
def sample_z_from_mean_logvar(mu: T.Tensor, log_var: T.Tensor) -> T.Tensor:
    # reparameterization trick
    eps = T.randn(size=mu.shape).to(mu.device)
    z = mu + (T.sqrt(T.exp(log_var)) * eps)
    return z

Use it when / don't use it when

Use it when

You want to learn continuous, semantically meaningful latent representations in a self-supervised manner without labels.
You need a foundational generative building block for downstream tasks or advanced diffusion architectures.

Don't use it when

You require hyper-realistic, photorealistic image output quality without applying hierarchical or vector-quantized extensions.
You only need simple deterministic dimensionality reduction or compression without generating new data points.