Train a VAE on Digits

Build a variational autoencoder, write its loss with torch.distributions, and train it on handwritten digits.

The data is scikit-learn's 8x8 digits, flattened to 64 pixels and made binary, so each pixel is 0 or 1. A VAE has three parts:

• An encoder maps x to the mean mu and the log variance logvar of a Gaussian over a small latent vector z.
• A reparameterised sample z = mu + exp(0.5 * logvar) * eps, with eps ~ N(0, I). The noise is outside the network, so gradients flow to mu and logvar.
• A decoder maps z to one logit for each pixel.

The loss is the negative ELBO, averaged over the batch:

recon = -sum over pixels of log Bernoulli(x | logits)
kl    =  sum over latent dims of KL( N(mu, sigma) || N(0, 1) )
loss  =  mean over the batch of (recon + kl)

torch.distributions gives you both terms: Bernoulli(logits=...) has a log_prob, and kl_divergence knows the closed form for two Normals.

Then write train_vae(X), which builds a VAE, trains it with mini-batches from a DataLoader, and returns it. This part is scored on the result.

Function signature

class VAE(nn.Module):
    def __init__(self, in_dim=64, latent_dim=8, hidden=128): ...
    def encode(self, x): ...      # -> (mu, logvar), each (n, latent_dim)
    def decode(self, z): ...      # -> logits, (n, in_dim)
    def forward(self, x): ...     # -> (logits, mu, logvar)


def negative_elbo(x, logits, mu, logvar):
def train_vae(X):

Example

Input

x      = torch.tensor([[1.0, 0.0]])
logits = torch.tensor([[0.0, 0.0]])
mu     = torch.tensor([[1.0]])
logvar = torch.tensor([[0.0]])
negative_elbo(x, logits, mu, logvar)

Output

tensor(1.8863)

Reasoning

A logit of 0 gives each pixel probability 0.5, so the reconstruction term is 2 * log 2 = 1.3863. The posterior is N(1, 1), and its KL from N(0, 1) is 0.5 * (1 + 1 - 1 - 0) = 0.5. The sum is 1.8863.

Constraints and notes

• negative_elbo returns a scalar tensor that autograd can differentiate.
• forward must draw a new z with the reparameterisation trick. Do not call .sample(), which blocks the gradient. rsample() is fine.
• train_vae gets a float tensor of shape (n, 64) with values in {0, 1}. It must return a VAE with latent_dim=8.
• On held-out digits, the trained model must reach a negative ELBO below 25 nats, with a mean KL above 1.5 nats. A model that ignores its latent code cannot do this. An untrained model is near 44.
• Keep training under about 10 seconds on a CPU.