Aditya Sridhar

Model · Oct 2026 · runs in your browser

Variational autoencoder

The classifiers before this one map an image to a label. This network maps an image to two numbers and back. Training forces those two numbers to form a smooth map, so every point on it decodes to something that looks like a digit, including points no training image landed on.

Latent dimensions
2
Test ELBO
-137.0 nats
Parameters
1,068,820
Training time
76 s, GPU

Try it

Latent space
  1. 0
  2. 1
  3. 2
  4. 3
  5. 4
  6. 5
  7. 6
  8. 7
  9. 8
  10. 9
Decoded from z = (0.00, 0.00)

Or draw a digit

The encoder places your drawing on the map, and the decoder redraws it from those two numbers.

Each dot is a test digit, placed by the encoder and coloured by its true label; the labels were never shown to the model. Drag between two clusters to see one digit turn into another. The decoded grid view decodes a 12×12 lattice over the same area.

Training run

Reconstruction loss, nats per imagetraintest
14016017131925313740
KL divergence, nats per imagetraintest
5.006.007.0017131925313740

Results

MetricValueSource
Test reconstruction loss (Bernoulli), nats per image129.95training/vae_mnist/log.csv, epoch 40
Test KL to the unit Gaussian, nats per image7.07training/vae_mnist/log.csv, epoch 40
Test evidence lower bound (higher is better)-137.02training/vae_mnist/log.csv, epoch 40
Largest difference, exported ONNX vs PyTorch, over 10,000 test images2.4e-5public/models/vae-mnist/meta.json

The reconstruction term is evaluated at the mean of the encoder's distribution, so the ELBO here is the deterministic bound, not a sampled estimate.

How it works

Encoder and decoder. The encoder is an MLP, 784 → 512 → 256, that outputs a mean and a log-variance for each of the 2 latent dimensions. The decoder mirrors it and outputs a probability for each pixel. The page runs both as ONNX files (2086 KB and 2089 KB).

The loss. Reconstruction alone would give an ordinary autoencoder, which can scatter digits anywhere and leave gaps that decode to noise. The KL term pulls every encoding towards a unit Gaussian, so the codes pack together and the space between them stays meaningful. During training the decoder sees a sample drawn around each mean (the reparameterisation trick), which is what makes the map smooth.

Training. Adam, learning rate 0.002 with cosine decay, batch 128, 40 epochs.

Reproduce

Hardware
NVIDIA GeForce RTX 3050 Laptop GPU, PyTorch 2.6.0+cu124
Commit
0370612
Data
MNIST, 60,000 train / 10,000 test

Environment

pip install torch numpy onnx onnxruntime

Train, export, check the exported files

python training/vae_mnist/train.py