12.2.3 Conditional VAEs

Overview

The VAE from the previous two lessons can generate samples, but it can’t generate specific samples. Sampling z ~ N(0, I) and decoding gives you something the model considers plausible — which abstract pattern, which digit, which face, you don’t get to say. For artistic work that’s often fine. For anything that needs targeted output it isn’t.

The Conditional VAE [1] solves this with a structural change that is, mechanically, almost embarrassingly simple: concatenate the desired output label onto the encoder input and onto the decoder input, then train the same way. The encoder learns to compress away the label information (since it’s already provided); the decoder learns to use the label as an explicit control knob. Once trained, you ask for digit 7 by passing [0,0,0,0,0,0,0,1,0,0] to the decoder alongside any latent vector.

This lesson trains a CVAE on the MNIST handwritten-digit dataset [2] — the canonical sanity check for any new generative architecture — and uses it to demonstrate the core property: same latent, different label, different output.

Learning objectives

1. Explain why a vanilla VAE cannot perform targeted generation and what the CVAE adds to fix it.
2. Build the one-hot label concatenation pattern into both encoder and decoder.
3. Train a CVAE on MNIST and verify that the label controls digit identity while the latent controls within-class style.
4. Recognise the “label leakage” failure mode and predict when it shows up.

Quick start — sketch the architecture

Forty lines: encoder + decoder + a generate(label) method. The full training script in Exercise 1 wraps an MNIST DataLoader and an optimiser around the same class.

import torch
import torch.nn as nn

class CVAE(nn.Module):
    def __init__(self, latent_dim=16, num_classes=10):
        super().__init__()
        self.latent_dim = latent_dim
        self.num_classes = num_classes

        # Encoder: image (784) + label (10) -> latent distribution params
        self.encoder = nn.Sequential(
            nn.Linear(784 + num_classes, 256), nn.ReLU(),
            nn.Linear(256, 256), nn.ReLU(),
        )
        self.fc_mu     = nn.Linear(256, latent_dim)
        self.fc_logvar = nn.Linear(256, latent_dim)

        # Decoder: latent (16) + label (10) -> image (784)
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim + num_classes, 256), nn.ReLU(),
            nn.Linear(256, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid(),
        )

    def encode(self, x, y):  return (lambda h: (self.fc_mu(h), self.fc_logvar(h)))(self.encoder(torch.cat([x, y], dim=1)))
    def decode(self, z, y):  return self.decoder(torch.cat([z, y], dim=1))

    def generate(self, labels, n_each=1):
        with torch.no_grad():
            n = len(labels) * n_each
            y = torch.zeros(n, self.num_classes)
            for i, lbl in enumerate(labels):
                y[i*n_each:(i+1)*n_each, lbl] = 1
            z = torch.randn(n, self.latent_dim)
            return self.decode(z, y)

m = CVAE()
print(f"Generated a '7': shape {m.generate([7]).shape}")

Run it with random (untrained) weights and you get noise — the architecture is there but the model knows nothing yet.

Core concepts

Concept 1 — Why vanilla VAEs can’t take requests

In a standard VAE [3], sampling proceeds in two steps: draw z ~ N(0, I), decode x̂ = decoder(z). There is no input slot for “what kind of sample I want.” If you train on a mix of MNIST digits, the encoder learns to map all ten classes into one shared latent space, the decoder learns to handle that whole space, and at generation time every random z collapses to whatever the decoder considers most plausible at that point.

You can cluster the latent space afterwards and use the centroids — that gets you something — but it’s brittle. The clusters drift between training runs, the boundaries between classes are fuzzy, and any new class requires retraining the clustering pipeline. The simpler fix is to tell the model what you want during training, and again during generation. That’s what conditioning is.

The cleanest place to inject the condition is at the input layer. The encoder sees concat(image, label); the decoder sees concat(latent, label). The model can use the label however it wants — including ignoring it — but training pressure pushes it to use the label correctly because the reconstruction loss demands a class-appropriate output for each class-labelled input. Doersch’s 2016 tutorial [4] frames the architectural choice in the broader context of variational inference; Goodfellow, Bengio and Courville [5] treat it as a special case of a general “conditioning on auxiliary information” pattern that extends to text-to-image and beyond.

Concept 2 — One-hot labels and the concatenation pattern

The label is a categorical variable, but neural networks want vectors. The translation is the one-hot encoding: label 7 becomes [0,0,0,0,0,0,0,1,0,0] for ten classes. PyTorch’s F.one_hot(labels, num_classes).float() does it in one call.

Once the label is a vector, the conditioning is two lines:

def encode(self, x, labels_onehot):
    combined = torch.cat([x, labels_onehot], dim=1)   # 784 + 10 = 794
    h = self.encoder_layers(combined)
    return self.fc_mu(h), self.fc_logvar(h)

def decode(self, z, labels_onehot):
    combined = torch.cat([z, labels_onehot], dim=1)   # 16 + 10 = 26
    return self.decoder_layers(combined)

The dim=1 matters — it’s the feature dimension, not the batch dimension. Concatenating along dim=0 would silently double your batch size and produce a cryptic shape mismatch a few lines later.

There are alternatives. Embedding lookups replace one-hot vectors with learned dense vectors and scale better when the number of classes is large. FiLM layers [6] feed the condition through a small network that produces multiplicative and additive modulations applied to every layer’s activations — more powerful, more parameters. Cross-attention is what modern text-to-image models like Stable Diffusion use. The concatenation approach in this lesson is the simplest of the family; everything more sophisticated is built on the same conceptual move (give the model the condition; let it learn to use it).

Concept 3 — Training, generation, and what each piece controls

Training is the standard VAE loop with labels threaded through. Same loss (reconstruction + KL), same reparameterisation, same optimiser.

def train_step(model, images, labels, optimizer):
    optimizer.zero_grad()
    y = F.one_hot(labels, model.num_classes).float()

    mu, logvar = model.encode(images, y)
    z = model.reparameterize(mu, logvar)
    recon = model.decode(z, y)

    recon_loss = F.binary_cross_entropy(recon, images, reduction='sum')
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    loss = recon_loss + kl_loss
    loss.backward(); optimizer.step()
    return loss.item()

Generation flips the script. Drop the encoder, sample z ~ N(0, I), decode with whatever label you want:

def generate(model, digit, n=10):
    model.eval()
    with torch.no_grad():
        y = torch.zeros(n, 10); y[:, digit] = 1
        z = torch.randn(n, model.latent_dim)
        return model.decode(z, y)

The interesting demonstration of what conditioning bought you is to fix z and sweep y across all ten classes — same noise, different label. You should get ten visually different digits that nonetheless share something stylistic (stroke thickness, slant angle, vertical position) because the latent vector is encoding those style attributes while the label handles the identity.

Exercises

EXECUTE I.

Train the CVAE on MNIST and inspect every visualisation

conditional_vae.py is the complete pipeline: download MNIST (~11 MB, automatic via torchvision), train for 50 epochs (~5 min CPU, ~2 min GPU), and produce four PNGs covering training history, generated samples, conditional generation, and a latent-space scatter.

conditional_vae.py — full training pipeline

bash Copy

python conditional_vae.py

Reflection

1. Looking at Figure 1, which digits look most realistic, which look most challenging, and what does that tell you about the training data?
2. In Figure 3’s top row, the same z produces ten different digits with the same style. What style attributes do you see preserved?
3. Why do the latent clusters in Figure 5 overlap so heavily — and why is that a good sign for a CVAE rather than a problem?
4. How would a non-conditional VAE’s latent-space scatter look different from Figure 5?

Discussion

1. Digits 1, 0, and 7 typically come out cleanest because they have simple, low-variance shapes in the MNIST training set. Digits 4, 5, 8, and 9 tend to be the hardest — they have more stroke variation in the dataset and more places for the decoder to lose detail.

2. Stroke thickness, slight rotation, and vertical positioning tend to carry across. The latent is encoding the style of the handwriting independently of which digit is being written, which is exactly what disentanglement looks like in practice.

3. Heavy overlap is the structural signature of a working CVAE. In a vanilla VAE, the latent has to separate the classes itself; the clusters are distinct because they have to be. In a CVAE, the label channel already separates the classes, so the latent is free to use the same regions for style features that are common across digit classes. Overlap = the model is using the label.

4. A vanilla VAE’s scatter would show ten distinct, well-separated clusters with sharp boundaries (e.g., as in the Lesson 12.2.1 plot). Each cluster covers exactly one class because the latent space is doing all the work of identity separation.

LATENT_DIM = 2      # tiny — limited style capacity
LATENT_DIM = 16     # default
LATENT_DIM = 64     # large — risks posterior collapse

loss = recon_loss + 0.0 * kl_loss   # KL disabled

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleCVAE(nn.Module):
    def __init__(self, latent_dim=16, num_classes=10):
        super().__init__()
        self.latent_dim = latent_dim
        self.num_classes = num_classes

        # TODO: 784 (image) + 10 (label) = 794 input features
        self.encoder_fc1 = nn.Linear(???, 256)
        self.encoder_fc2 = nn.Linear(256, 128)
        self.fc_mu      = nn.Linear(128, latent_dim)
        self.fc_logvar  = nn.Linear(128, latent_dim)

        # TODO: latent_dim + 10 (label) input features
        self.decoder_fc1 = nn.Linear(???, 128)
        self.decoder_fc2 = nn.Linear(128, 256)
        self.decoder_fc3 = nn.Linear(256, 784)

    def encode(self, x, y):
        combined = ???                  # TODO: concat along dim=1
        h = F.relu(self.encoder_fc1(combined))
        h = F.relu(self.encoder_fc2(h))
        return self.fc_mu(h), self.fc_logvar(h)

    def reparameterize(self, mu, logvar):
        # TODO: std = exp(0.5*logvar); z = mu + std * eps
        pass

    def decode(self, z, y):
        combined = ???                  # TODO: concat along dim=1
        h = F.relu(self.decoder_fc1(combined))
        h = F.relu(self.decoder_fc2(h))
        return torch.sigmoid(self.decoder_fc3(h))

def cvae_loss(recon, x, mu, logvar):
    # TODO: BCE reconstruction + KL divergence
    pass

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleCVAE(nn.Module):
    def __init__(self, latent_dim=16, num_classes=10):
        super().__init__()
        self.latent_dim = latent_dim
        self.num_classes = num_classes

        self.encoder_fc1 = nn.Linear(784 + 10, 256)
        self.encoder_fc2 = nn.Linear(256, 128)
        self.fc_mu      = nn.Linear(128, latent_dim)
        self.fc_logvar  = nn.Linear(128, latent_dim)

        self.decoder_fc1 = nn.Linear(latent_dim + 10, 128)
        self.decoder_fc2 = nn.Linear(128, 256)
        self.decoder_fc3 = nn.Linear(256, 784)

    def encode(self, x, y):
        combined = torch.cat([x, y], dim=1)
        h = F.relu(self.encoder_fc1(combined))
        h = F.relu(self.encoder_fc2(h))
        return self.fc_mu(h), self.fc_logvar(h)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        return mu + std * torch.randn_like(std)

    def decode(self, z, y):
        combined = torch.cat([z, y], dim=1)
        h = F.relu(self.decoder_fc1(combined))
        h = F.relu(self.decoder_fc2(h))
        return torch.sigmoid(self.decoder_fc3(h))

def cvae_loss(recon, x, mu, logvar):
    recon_loss = F.binary_cross_entropy(recon, x, reduction='sum')
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + kl_loss

Summary

Common pitfalls

• Concatenating along dim=0 (batch) instead of dim=1 (features) — silently doubles the batch and breaks shape-matching downstream.
• Posterior collapse — KL weight too high, model stops using the latent space and produces the average digit for every label.
• Label leakage — KL weight too low, encoder hides digit identity in the latent vector, label conditioning becomes ineffective. The top-row test in Figure 3 is the diagnostic.
• Mismatched one-hot encoding — using labels.float() instead of F.one_hot(labels, 10).float() feeds a scalar where the decoder expects a 10-vector; the resulting weight matrix dimension errors are sometimes hard to read.
• Forgetting model.eval() before generation — leaves BatchNorm and Dropout in training mode and produces noisier outputs than the model is actually capable of.

References