To reconstruct an observed image—or ask how it might look under a different cause—a model needs a way to infer information it cannot observe directly. An encoder maps the observation to a latent representation. In a variational autoencoder (VAE), it works alongside a decoder, which maps the latent representation back toward an image. section 5.2.2 section 9.4.1
For a 28 × 28 digit image, the latent vector Z has fewer dimensions than the 784 image values. The encoder takes the image together with its digit and is-handwritten labels and returns a location and scale for a distribution over Z. A guide function samples Z from that distribution; the decoder then uses Z and the labels to parameterize an image distribution. section 5.2.2 section 5.2.3
Listing 5.7 Implement the encoder
class Encoder(nn.Module): #1
def __init__(self, z_dim, hidden_dim):
super().__init__()
img_dim = 28 * 28 #2
digit_dim = 10 #3
is_handwritten_dim = 1
self.softplus = nn.Softplus() #4
input_dim = img_dim + digit_dim + is_handwritten_dim #5
self.fc1 = nn.Linear(input_dim, hidden_dim) #5
self.fc21 = nn.Linear(hidden_dim, z_dim) #6
self.fc22 = nn.Linear(hidden_dim, z_dim) #6
def forward(self, img, digit, is_handwritten): #7
input = torch.cat([img, digit, is_handwritten], dim=1) #8
hidden = self.softplus(self.fc1(input)) #9
z_loc = self.fc21(hidden) #10
z_scale = torch.exp(self.fc22(hidden)) #10
return z_loc, z_scale
The PyTorch implementation concatenates image values and labels, passes them through a hidden layer, and outputs separate location and positive-scale parameters for Z. section 5.2.2
Training adjusts both networks using variational inference. Reconstruction compares an original image with one produced after encoding and decoding it, but reconstruction error alone is not the training objective: the probabilistic setup uses stochastic variational inference with an ELBO loss. Compare reconstructions and monitor test loss to check progress. section 5.2.3 section 5.2.4
A useful encoding is not automatically a causal one. Z can stand in for unmodeled influences such as stroke thickness or font, but recovering an encoding that preserves image detail does not recover the actual latent causes. Disentangling causal from noncausal factors generally cannot be achieved from an unsupervised Z alone; domain knowledge, interventions, or semi-supervision can help. section 5.2.2 section 5.2.5
Encoding for a shape counterfactual
For a dSprites image, a pretrained encoder first recovers factors including shape, scale, orientation, and position. Another encoder takes the image and its factors to infer an image-noise variable. A decoder uses that inferred variable with an intervened shape factor to generate an image answering: what if the square were a heart? section 9.4.1
Encoders can also impute a latent confounder rather than compress an image. In a game-purchase workflow, an encoder maps observed engagement, won items, and purchases to distribution parameters for Z inside a Pyro guide; stochastic variational inference fits the guide alongside the causal model. section 11.6.5
