AANN 18/04/2026

A vanilla autoencoder (neural network) is made up of two subnetworks: an encoder, which takes a high-dimensional vector \(x\) (in our example a flattened \(28\times 28\) image in \(\mathbb{R}^{784}\)) and "encodes" it in a lower-dimensional representation \(h\) (in our example \(h\in\mathbb{R}^{15}\); and a decoder, which takes that low-dimensional representation \(h\) and maps it to \(\hat{x}\), which is a reconstruction of the original high-dimensional vector \(x\). In learning how to map data into lower dimensions, and then recover it, the autoencoder learns to perform dimensionality reduction¹. In Figure 1 there is an example of an autoencoder.

In a variational autoencoder (VAE), the lower dimensional representation is a distribution (typically a Gaussian distribution). I.e. the encoder maps \(x\) to \((\boldsymbol{\mu},\boldsymbol{\sigma}^{2})\) (the parameters of a multivariate Gaussian distribution), then the decoder takes a sample from that distribution, \(z\sim\text{N}(\boldsymbol{\mu},\boldsymbol{\sigma}^{2})\), and maps it to a reconstruction \(\hat{x}\)². As we will see below, the loss function for training a VAE tries to keep this latent Gaussian distribution as close to a standard normal as possible. In Figure 2 there is an example of the VAE architecture.

In a conditional VAE we assume that there is some additional conditioning information \(c\) associated with each datum \(x\). In this case, the decoder takes the (concatenated) vector \((z,c)\) and maps it to the reconstruction \(\hat{x}\). In our example \(c\in\mathbb{R}^{10}\) because it is a one-hot-encoding of the digit.

There are two terms in the loss function for training a VAE: the reconstruction loss, which measures the error in the reconstruction \(\hat{x}\) of the truth \(x\), in our example I use the binary cross entropy; and the regularization loss, which measures how far the latent distribution is from a standard normal, in this example I use the Kullback-Leibler divergence (which can be computed in closed form between Gaussian distributions).

There is an implementation of the (C)VAE and the binary cross entropy loss function here.

The MNIST digits data are 28-by-28 pixel images of handwritten digits. In Figure 4, I show some of the instances of the number five. (You can find the script to download and visualise the data here.)

I trained the VAE shown in Figure 2 to sample novel instances of the digit 5. The training process is pretty straight forward so you might be best off just looking at the script implementing it, and the script that samples from the trained model. Some samples from the trained model are shown in Figure 5.

Note that these images are not part of the training data. We have sampled new latent vectors from the standard normal distribution and used these as input to the decoder to get the images. Using the standard normal as our source of random variables is why it is important to have the KL-divergence term in the loss function.

If you were to use one of these trained decoders as a way to sample from your prior distribution, as in PriorVAE and PriorCVAE you'd want to have some confidence the decoder behaves sensibly across the region of parameter space the MCMC is likely to visit. To test this, I sampled random directions and then looked at the images generated as you moved in that direction (i.e. scaled a random vector by a factor of 0, 2 and 4). In Figure 6, examples of the generated images are shown, while the images degrade slightly, the digits are still vaguely five shaped.

I trained the CVAE shown in Figure 3 to sample novel instances of all of the digits. Again, the training process is straight forward so you might be best off just looking at the script implementing it, and the script that samples from the trained model. Some samples from the trained model are shown in Figure 7. I suspect you could get a better image if you optimised the training, but since this is only a proof of principle I won't bother; they are mostly recognisable…

To get an idea for how the CVAE decoder works as a prior distribution within MCMC, we can consider the case where we start with a partial observation of a number and want to estimate what the rest of it looked like (and what number it is!) If you were applying this in the context of spatial epidemiology, this might correspond to estimating the prevalence of disease across a large area when you only have prevalence data from a few sites.

We start by taking one of MNIST digits (from the testing set), and generating a random mask of pixels to observe. In Figure 8 (left) there is a masked number three. We then set up our NumPyro model of this process of masking and observation with noise. A probabilistic program for this might look like the following:

1. \(\theta\sim\text{N}_{25}(0,1)\)
2. \(z\leftarrow\theta[\texttt{1:15}]\)
3. \(c\leftarrow\text{softmax}(\theta[\texttt{16:25}])\)
4. \(\mu_Y\leftarrow\text{Decoder}(z,c)\)
5. \(Y[\texttt{mask}]\sim\text{N}(\mu_Y[\texttt{mask}], {0.1}^2)\)

With this program, we can sample from the posterior distribution and use the mean of the samples to understand what the digit is likely to be. In Figure 8 (right) there is an image of the posterior mean of the digit. The code for this is available here.