The challenge
A plain autoencoder learns to compress and reconstruct data, but does not give a latent space with a well-defined probability distribution. The goal was to implement a variational autoencoder whose latent representation is a distribution, and to make the training objective (reconstruction plus regularisation) flexible and testable.
The solution
The encoder outputs the mean and log-variance of a Gaussian latent distribution, and latent vectors are drawn with the reparameterisation trick so the model can be trained with gradients. The objective combines a reconstruction loss with the KL divergence to a standard normal prior, with an adjustable β weight.
How it was built
I implemented the model in PyTorch with an MLP encoder and decoder with separate mean and log-variance heads, Gaussian reparameterisation, KL divergence from the diagonal Gaussian posterior to a standard normal prior and a choice of reconstruction loss (BCE, MSE or L1), combined in a β-weighted ELBO. Latent dimension, hidden dimension and β are configurable; the default setup targets 28×28 images with 16 latent dimensions. The implementation was written against tests that check output shapes, that the loss is finite and that the β weighting works.
Result
This project documents an implementation, not a measured model performance. The architecture, loss functions and β-weighted ELBO are implemented as described, but I have not preserved quantitative results for reconstruction quality, training curves or generated samples, so the page gives no such figures.
Context
The page makes no claims about how well the model reconstructs or generates images. The project shows an implementation of latent-variable modelling (reparameterisation, KL regularisation and a composite loss), not a finished trained model.