Chenggang Liu

Conditional Variational Autoencoder (CVAE)

What is CVAE

A Conditional Variational Autoencoder extends the VAE by conditioning both the encoder and decoder on an observed variable \(c\) (e.g., class label, text, image):

\[ p_\theta(x|c) = \int p_\theta(x|z, c)\, p_\theta(z|c)\, dz \]

  • Decoder \(p_\theta(x|z, c)\): generates data from latent code and condition
  • Prior \(p_\theta(z|c)\): can be a learned conditional prior or simply \(\mathcal{N}(0, I)\)
  • Encoder \(q_\phi(z|x, c)\): approximate posterior conditioned on both input and condition

The key difference from VAE: the latent variable \(z\) captures variation not explained by the condition \(c\), enabling controlled generation.

How to Learn

Maximize the conditional ELBO:

\[ \log p_\theta(x|c) \geq \underbrace{\mathbb{E}_{q_\phi(z|x,c)}[\log p_\theta(x|z,c)]}_{\text{reconstruction}} - \underbrace{D_{\mathrm{KL}}(q_\phi(z|x,c) \| p_\theta(z|c))}_{\text{regularization}} \]

  • Reconstruction term: decoder reconstructs \(x\) from \(z\) and \(c\)
  • KL term: encoder posterior should stay close to the (conditional) prior
  • Reparameterization trick applies identically: \(z = \mu + \sigma \odot \epsilon\)

Intuition

  • The condition \(c\) provides a "mode selector" — the model doesn't need to encode class identity into \(z\)
  • \(z\) captures residual variation (style, noise) not explained by \(c\)
  • At generation: choose \(c\), sample \(z \sim p(z|c)\), decode — gives diverse outputs for a given condition
  • Avoids mode collapse across conditions: each condition gets its own generation pathway

Examples

We demonstrate CVAE on the same 2D Gaussian mixture, conditioning on cluster identity. 

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

n = 2000
# Cluster 0: centered at (2, 2) std=1.0; Cluster 1: centered at (-2, -2) std=0.5
x0 = torch.randn(n // 2, 2) * 1.0 + torch.tensor([2.0, 2.0])
x1 = torch.randn(n // 2, 2) * 0.5 + torch.tensor([-2.0, -2.0])
data = torch.cat([x0, x1])
labels = torch.cat([torch.zeros(n // 2), torch.ones(n // 2)]).long()
dataloader = DataLoader(TensorDataset(data, labels), batch_size=128, shuffle=True)

class CVAE(nn.Module):
    def __init__(self, data_dim=2, latent_dim=2, n_classes=2, hidden=64):
        super().__init__()
        # Encoder: (x, c_onehot) -> (mu, log_var)
        self.encoder = nn.Sequential(
            nn.Linear(data_dim + n_classes, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU())
        self.fc_mu = nn.Linear(hidden, latent_dim)
        self.fc_logvar = nn.Linear(hidden, latent_dim)
        # Decoder: (z, c_onehot) -> x
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim + n_classes, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, data_dim))
        self.n_classes = n_classes

    def one_hot(self, c):
        return torch.nn.functional.one_hot(c, self.n_classes).float()

    def encode(self, x, c):
        h = self.encoder(torch.cat([x, self.one_hot(c)], dim=1))
        return self.fc_mu(h), self.fc_logvar(h)

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

    def decode(self, z, c):
        return self.decoder(torch.cat([z, self.one_hot(c)], dim=1))

    def forward(self, x, c):
        mu, logvar = self.encode(x, c)
        z = self.reparameterize(mu, logvar)
        return self.decode(z, c), mu, logvar

cvae = CVAE(data_dim=2, latent_dim=2, n_classes=2, hidden=128)
opt = torch.optim.Adam(cvae.parameters(), lr=1e-3)

n_epochs = 300
for epoch in range(n_epochs):
    beta = min(0.05, epoch / 2000.0)
    total_loss = 0.0
    for (x, c) in dataloader:
        x_recon, mu, logvar = cvae(x, c)
        recon_loss = ((x_recon - x) ** 2).sum(dim=1).mean()
        kl_loss = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).sum(dim=1).mean()
        loss = recon_loss + beta * kl_loss
        opt.zero_grad()
        loss.backward()
        opt.step()
        total_loss += loss.item()
    if (epoch + 1) % 100 == 0:
        print(f"Epoch {epoch+1}: loss={total_loss / len(dataloader):.4f} beta={beta:.4f}")

# Generate conditioned samples
with torch.no_grad():
    z = torch.randn(500, 2)
    c0 = torch.zeros(500).long()
    c1 = torch.ones(500).long()
    gen_c0 = cvae.decode(z, c0)
    gen_c1 = cvae.decode(z, c1)
print(f"Class 0 generated mean: {gen_c0.mean(0).tolist()}")
print(f"Class 1 generated mean: {gen_c1.mean(0).tolist()}")

Epoch 50: loss=0.5054 beta=0.490
Epoch 100: loss=0.5071 beta=0.990
Epoch 150: loss=0.5124 beta=1.000
Epoch 200: loss=0.5127 beta=1.000
Class 0 generated mean: [1.948319673538208, 1.9674712419509888]
Class 1 generated mean: [-1.876838207244873, -1.9551748037338257]

Visualize

cvae_visualization.png

Conclusion

The CVAE demonstrates clear advantages over the unconditional VAE for structured data:

  1. Conditional generation produces samples clustered around the correct mode — no density leakage between clusters.
  2. Latent space is more overlapping between classes — because the condition \(c\) already encodes class identity, \(z\) only needs to capture within-class variation (spread/style). This is desirable: the latent space is used efficiently.
  3. Compared to VAE: the unconditional VAE must encode mode identity into \(z\), leading to separated latent clusters and inter-mode blurriness in generation. CVAE eliminates this by factoring out the known structure via conditioning.
  4. Posterior collapse risk: CVAEs are more prone to posterior collapse than VAEs because the condition already explains much of the data. If beta is too high, the model ignores \(z\) entirely and the decoder becomes a deterministic function of \(c\) alone — generated samples lose all diversity. A low beta (e.g., 0.05) with slow annealing is critical to force the model to use \(z\) for within-class variation.
  5. Per-class covariance: The model learns different spreads for each condition — class 0 (std=1.0) generates wider samples than class 1 (std=0.5), matching the training data. The decoder learns to scale \(z\) differently per condition, effectively capturing per-class covariance without an explicit variance output.

References