Suppose an image is shaped by its digit, whether it is handwritten, and visual details that are not labeled. A variational autoencoder (VAE) can model these relationships with two neural networks: a decoder that generates an image from the labels and a latent variable Z, and an encoder that estimates Z from an image and its labels. This lets a generative model handle a high-dimensional outcome without requiring a hand-written image-generation function. section 5.1.1 section 5.2.2
A 28 × 28 image has 784 pixel values. Z has fewer dimensions and acts as a compressed encoding of image variation. In a causal DAG, the digit and handwriting labels can be treated as observed causes of the image, while Z stands in for unmodeled influences such as stroke thickness or font. The decoder then represents the image’s causal Markov kernel, P(image | Z, digit, is-handwritten). This is a causal interpretation imposed on the model: a conventional VAE does not make Z causal, and its learned encoding is not the same as discovering the actual latent causes. section 5.2.2 section 5.2.6
The two networks have different jobs:
- The decoder takes Z and the observed labels as inputs and produces parameters for generating an image. In the illustrated causal model, the joint distribution factors into P(Z), P(digit), P(is-handwritten), and P(image | Z, digit, is-handwritten). section 5.2.2 section 5.2.6
- The encoder takes an image and its labels and returns parameters of a distribution Q(Z | image, digit, is-handwritten). A guide samples from Q because directly sampling the corresponding conditional distribution under the generative model is impractical. section 5.2.2 section 5.2.3
Training fits both networks through variational inference. Minimizing negative ELBO indirectly reduces reconstruction error and the divergence between the guide’s distribution and the target conditional distribution for Z. Stochastic variational inference performs this optimization on batches rather than the full dataset. In Pyro, the generative model and guide are paired in SVI with Trace_ELBO as the loss; test loss, reconstructed images, and newly generated images provide checks on training progress. section 5.2.3 section 5.2.4
Labeled digit images
Combine MNIST and TMNIST images with digit and handwriting labels, split them into training and test sets, and train the conditional VAE. To inspect reconstruction, encode a test image with its labels, sample Z, then decode it using the same labels. To inspect generation, sample inputs for the generative model and view the resulting image. Better-looking reconstructions and samples during training are useful checks, but they do not establish that Z has recovered real-world causes. section 5.2.1 section 5.2.3 section 5.2.4 section 5.2.2
Listing 5.9 Full VAE class
class VAE(nn.Module):
def __init__(
self,
z_dim=50, #1
hidden_dim=400, #2
use_cuda=USE_CUDA,
):
super().__init__()
self.use_cuda = use_cuda
self.z_dim = z_dim
self.hidden_dim = hidden_dim
self.setup_networks()
def setup_networks(self): #3
self.encoder = Encoder(self.z_dim, self.hidden_dim)
self.decoder = Decoder(self.z_dim, self.hidden_dim)
if self.use_cuda:
self.cuda()
model = model #4
training_model = training_model #4
training_guide = training_guide #4
The full VAE class shows how the encoder, decoder, generative model, and training guide are assembled into one PyTorch module. section 5.2.2
The VAE is one way to fit a high-dimensional kernel within a causal graphical model, not a substitute for specifying and defending that model’s causal assumptions. Variational inference with a VAE can also handle latent variables in a broader causal-inference workflow that estimates an average treatment effect after fitting a causal generative model. section 5.2.6 section 11.6
