MNIST VAE
From-scratch Variational Autoencoder in PyTorch with digits generated from pure noise, and the reconstruction-versus-KL tradeoff made visible.
Problem
A plain autoencoder compresses and reconstructs but cannot generate: its latent space is full of dead regions no training image maps to, so decoding a random point produces garbage.
How it was solved
Two changes. The encoder outputs a distribution rather than a point, and a KL term packs every image’s Gaussian toward N(0, I), filling the dead regions and molding the space into a smooth, gap-free shape.
Built from scratch: the two-headed encoder, the reparameterization trick that lets gradients flow back through a sampling step, and the closed-form KL alongside a reconstruction term.
Results
Every digit shown was generated from noise, with no input image.
The samples are slightly blurry, and that blur is the point: it is the reconstruction term losing ground to the KL term. That exact tradeoff is what diffusion models sidestep, which is why their samples are sharper.