vae-playground 2025
A modular playground for experimenting with different Variational Autoencoder
families in PyTorch, paired with interactive marimo
notebooks. A VAE encodes each input into a distribution over a latent space,
samples a code z from it, and decodes that code back to an image — so the
latent space becomes smooth and generative (shown in the animation above).
The variants
| Model | Key idea |
|---|---|
| Vanilla VAE | the standard ELBO (reconstruction + KL) |
| β-VAE | disentangled latents by weighting the KL term (β > 1) |
| Conditional VAE | class-conditioned generation |
| VQ-VAE | discrete latent codes with a learned codebook |
| WAE-MMD | a Wasserstein objective with an MMD penalty |
What each notebook covers
Every variant has its own marimo notebook with the theory (with LaTeX), interactive hyperparameter controls, live training, original-vs-reconstructed views, a 2D latent-space visualization colored by class, sampling from the prior, and saving. A final comparison notebook loads the saved checkpoints and shows side-by-side reconstructions, latent spaces, loss curves, samples, and an MSE table.
Datasets & hardware
MNIST, Fashion-MNIST, and CIFAR-10, downloaded automatically via torchvision. The trainer auto-detects the best device (Apple Silicon MPS, CUDA, or CPU).