VAE & Extensions

1. Variational Inference Framework

1.1 Problem Formulation

Given observed data and latent variables , we want to maximise the marginal likelihood:

This integral is intractable for complex generative models.

1.2 Evidence Lower Bound (ELBO)

We introduce an approximate posterior and derive the ELBO:

ELBO decomposition:

1.3 KL Divergence for Gaussians

When and :

where is the latent dimension.

1.4 Total ELBO

2. Reparameterisation Trick

2.1 The Problem

We cannot backpropagate through stochastic nodes:

The sampling operation is not differentiable.

2.2 The Solution

Reparameterise by sampling from a deterministic distribution:

Now is a deterministic function of , , and , enabling backpropagation.

2.3 Gradient Estimator

The gradient of the ELBO with respect to is:

This is now unbiased and can be estimated with Monte Carlo sampling.

2.4 Log-Variance Parameterisation

To ensure , we parameterise with log-variance:

3. Posterior Collapse

3.1 The Problem

When the decoder is powerful (e.g., autoregressive), the model may learn to ignore the latent variables:

This leads to and the latent code carries no information.

3.2 Theoretical Analysis

The ELBO can be decomposed as:

For a powerful decoder:

• can be maximised by fitting the decoder to directly
• is optimal for this objective
• No incentive to use latent variables

3.3 Mitigation Strategies

KL Annealing (Warm-up):

where is increased from 0 to 1 during training.

Free Bits:

where is a minimum information threshold per dimension.

δ-VAE:

Add a penalty for latent dimensions that collapse:

4. β-VAE and Disentanglement

4.1 β-VAE Objective

For , the model is encouraged to learn disentangled representations where each latent dimension captures a single generative factor.

4.2 Disentanglement Metrics

Mutual Information Gap (MIG):

where is the sorted set of latent dimensions by mutual information with factor .

Disentanglement Metric (DCI):

where is the importance matrix from a classifier.

4.3 β-VAE Theoretical Analysis

The -VAE introduces a trade-off:

For :

• Stronger information bottleneck
• Encourages factorised latent representation
• May reduce reconstruction quality

5. VQ-VAE (Vector Quantised VAE)

5.1 Architecture

VQ-VAE (van den Oord et al., 2017) uses a discrete latent space:

1. Encoder:
2. Quantisation:
3. Decoder:

where is the codebook.

5.2 Loss Function

where is the stop-gradient operator.

5.3 VQ-VAE-2

VQ-VAE-2 (Razavi et al., 2019) uses a hierarchical architecture:

1. Bottom level: Fine-grained features
2. Top level: Global structure

Each level has its own encoder, decoder, and codebook.

5.4 Advantages over Continuous VAE

6. Code Examples

6.1 VAE Implementation

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

class VAE(nn.Module):
    """
    Variational Autoencoder.
    
    Loss: ELBO = E[log p(x|z)] - KL(q(z|x) || p(z))
    """
    
    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20):
        super().__init__()
        
        # Encoder
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU()
        )
        
        # Mean and log-variance
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
        
        # Decoder
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),
            nn.Sigmoid()
        )
    
    def encode(self, x):
        """Encode input to mean and log-variance."""
        h = self.encoder(x)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar
    
    def reparameterise(self, mu, logvar):
        """
        Reparameterisation trick.
        
        z = μ + σ * ε, where ε ~ N(0, I)
        """
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z):
        """Decode latent code to reconstruction."""
        return self.decoder(z)
    
    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterise(mu, logvar)
        x_recon = self.decode(z)
        return x_recon, mu, logvar
    
    def compute_elbo(self, x, x_recon, mu, logvar):
        """
        Compute ELBO loss.
        
        ELBO = reconstruction - KL
        """
        # Reconstruction loss (binary cross-entropy)
        recon_loss = F.binary_cross_entropy(x_recon, x, reduction='sum')
        
        # KL divergence: KL(q(z|x) || p(z))
        # = -0.5 * sum(1 + log(σ²) - μ² - σ²)
        kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        
        elbo = recon_loss + kl_loss
        
        return elbo, recon_loss, kl_loss
    
    def generate(self, num_samples, device='cuda'):
        """Generate samples from prior p(z) = N(0, I)."""
        z = torch.randn(num_samples, self.latent_dim).to(device)
        return self.decode(z)
    
    def interpolate(self, x1, x2, num_steps=10):
        """Interpolate between two points in latent space."""
        mu1, logvar1 = self.encode(x1)
        mu2, logvar2 = self.encode(x2)
        
        z1 = self.reparameterise(mu1, logvar1)
        z2 = self.reparameterise(mu2, logvar2)
        
        interpolations = []
        for alpha in torch.linspace(0, 1, num_steps):
            z = (1 - alpha) * z1 + alpha * z2
            interpolations.append(self.decode(z))
        
        return torch.cat(interpolations)


# Example: Train VAE
vae = VAE(input_dim=784, hidden_dim=400, latent_dim=20)

# Dummy training loop
optimizer = torch.optim.Adam(vae.parameters(), lr=1e-3)

for epoch in range(10):
    x = torch.randn(32, 784)  # Dummy data
    x_recon, mu, logvar = vae(x)
    elbo, recon, kl = vae.compute_elbo(x, x_recon, mu, logvar)
    
    optimizer.zero_grad()
    elbo.backward()
    optimizer.step()
    
    if (epoch + 1) % 5 == 0:
        print(f"Epoch {epoch+1}: ELBO={elbo.item():.2f}, "
              f"Recon={recon.item():.2f}, KL={kl.item():.2f}")

6.2 β-VAE Implementation

class BetaVAE(nn.Module):
    """
    β-VAE for disentangled representation learning.
    
    Loss: E[log p(x|z)] - β * KL(q(z|x) || p(z))
    """
    
    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20, beta=4.0):
        super().__init__()
        self.beta = beta
        self.latent_dim = latent_dim
        
        # Encoder
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU()
        )
        
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
        
        # Decoder
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),
            nn.Sigmoid()
        )
    
    def encode(self, x):
        h = self.encoder(x)
        return self.fc_mu(h), self.fc_logvar(h)
    
    def reparameterise(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z):
        return self.decoder(z)
    
    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterise(mu, logvar)
        return self.decode(z), mu, logvar
    
    def compute_loss(self, x, x_recon, mu, logvar):
        """Compute β-VAE loss."""
        # Reconstruction
        recon_loss = F.binary_cross_entropy(x_recon, x, reduction='sum')
        
        # KL divergence with β weighting
        kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
        kl_loss = self.beta * kl_per_dim.sum()
        
        # Total loss
        total_loss = recon_loss + kl_loss
        
        # Metrics
        kl_per_dim_mean = kl_per_dim.mean(dim=0)
        
        return total_loss, recon_loss, kl_loss, kl_per_dim_mean
    
    def compute_kl_annealed(self, x, x_recon, mu, logvar, current_epoch, warmup_epochs=10):
        """Compute loss with KL annealing."""
        recon_loss = F.binary_cross_entropy(x_recon, x, reduction='sum')
        
        # Anneal β from 0 to self.beta
        beta = min(self.beta, self.beta * (current_epoch / warmup_epochs))
        
        kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
        kl_loss = beta * kl_per_dim.sum()
        
        return recon_loss + kl_loss, recon_loss, kl_loss, beta


class FreeBitsVAE(BetaVAE):
    """
    Free Bits VAE.
    
    Enforces minimum information per latent dimension.
    """
    
    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20, 
                 free_bits=0.5, beta=1.0):
        super().__init__(input_dim, hidden_dim, latent_dim, beta)
        self.free_bits = free_bits
    
    def compute_loss(self, x, x_recon, mu, logvar):
        """Compute Free Bits loss."""
        recon_loss = F.binary_cross_entropy(x_recon, x, reduction='sum')
        
        # KL per dimension
        kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
        
        # Apply free bits: max(free_bits, KL_per_dim)
        kl_per_dim_clamped = torch.clamp(kl_per_dim, min=self.free_bits)
        
        # Sum over dimensions, mean over batch
        kl_loss = self.beta * kl_per_dim_clamped.sum(dim=-1).mean()
        
        total_loss = recon_loss + kl_loss
        
        return total_loss, recon_loss, kl_loss, kl_per_dim.mean(dim=0)


# Example: Compare VAE and β-VAE
vae = VAE(latent_dim=20)
beta_vae = BetaVAE(latent_dim=20, beta=4.0)
free_bits_vae = FreeBitsVAE(latent_dim=20, free_bits=0.5)

x = torch.randn(32, 784)

# VAE
x_recon, mu, logvar = vae(x)
loss_vae, _, kl_vae, _ = vae.compute_elbo(x, x_recon, mu, logvar)

# β-VAE
x_recon, mu, logvar = beta_vae(x)
loss_beta, _, kl_beta, _ = beta_vae.compute_loss(x, x_recon, mu, logvar)

# Free Bits VAE
x_recon, mu, logvar = free_bits_vae(x)
loss_fb, _, kl_fb, _ = free_bits_vae.compute_loss(x, x_recon, mu, logvar)

print(f"VAE: loss={loss_vae.item():.2f}, KL={kl_vae.item():.2f}")
print(f"β-VAE: loss={loss_beta.item():.2f}, KL={kl_beta.item():.2f}")
print(f"Free Bits: loss={loss_fb.item():.2f}, KL={kl_fb.item():.2f}")

6.3 VQ-VAE Implementation

class VectorQuantiser(nn.Module):
    """Vector Quantisation layer."""
    
    def __init__(self, num_embeddings=512, embedding_dim=64, commitment_cost=0.25):
        super().__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        self.commitment_cost = commitment_cost
        
        # Codebook
        self.embedding = nn.Embedding(num_embeddings, embedding_dim)
        self.embedding.weight.data.uniform_(-1/num_embeddings, 1/num_embeddings)
    
    def forward(self, z_e):
        """
        Quantise encoder output.
        
        Parameters
        ----------
        z_e : Tensor (batch, height, width, embedding_dim)
        
        Returns
        -------
        z_q : Tensor (batch, height, width, embedding_dim)
        loss : scalar
        encoding_indices : Tensor (batch, height, width)
        """
        # Reshape for distance computation
        z_e_flat = z_e.view(-1, self.embedding_dim)
        
        # Compute distances to codebook entries
        distances = (z_e_flat ** 2).sum(dim=1, keepdim=True) + \
                    (self.embedding.weight ** 2).sum(dim=1) - \
                    2 * torch.matmul(z_e_flat, self.embedding.weight.T)
        
        # Find closest codebook entries
        encoding_indices = distances.argmin(dim=1)
        z_q_flat = self.embedding(encoding_indices)
        z_q = z_q_flat.view(z_e.shape)
        
        # Compute losses
        # Codebook loss: update codebook to match encoder outputs
        codebook_loss = F.mse_loss(z_q_flat.detach(), z_e_flat)
        
        # Commitment loss: encourage encoder to commit to codebook
        commitment_loss = F.mse_loss(z_e_flat, z_q_flat.detach())
        
        # Total VQ loss
        vq_loss = codebook_loss + self.commitment_cost * commitment_loss
        
        # Straight-through estimator
        z_q = z_e + (z_q - z_e).detach()
        
        encoding_indices = encoding_indices.view(z_e.shape[:-1])
        
        return z_q, vq_loss, encoding_indices


class VQVAE(nn.Module):
    """
    Vector Quantised VAE.
    
    Uses discrete latent space with learned codebook.
    """
    
    def __init__(self, input_channels=3, hidden_dim=128, embedding_dim=64, 
                 num_embeddings=512):
        super().__init__()
        
        self.latent_dim = embedding_dim
        
        # Encoder
        self.encoder = nn.Sequential(
            nn.Conv2d(input_channels, hidden_dim, 4, 2, 1),
            nn.ReLU(),
            nn.Conv2d(hidden_dim, hidden_dim, 4, 2, 1),
            nn.ReLU(),
            nn.Conv2d(hidden_dim, embedding_dim, 3, 1, 1),
        )
        
        # Vector Quantiser
        self.vq = VectorQuantiser(num_embeddings, embedding_dim)
        
        # Decoder
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(embedding_dim, hidden_dim, 4, 2, 1),
            nn.ReLU(),
            nn.ConvTranspose2d(hidden_dim, hidden_dim, 4, 2, 1),
            nn.ReLU(),
            nn.ConvTranspose2d(hidden_dim, input_channels, 3, 1, 1),
            nn.Tanh()
        )
    
    def forward(self, x):
        # Encode
        z_e = self.encoder(x.permute(0, 2, 3, 1))
        
        # Quantise
        z_q, vq_loss, indices = self.vq(z_e)
        
        # Decode
        x_recon = self.decoder(z_q.permute(0, 3, 1, 2))
        
        return x_recon, vq_loss, indices
    
    def compute_loss(self, x, x_recon, vq_loss):
        """Compute total VQ-VAE loss."""
        # Reconstruction loss
        recon_loss = F.mse_loss(x_recon, x)
        
        # Total loss
        total_loss = recon_loss + vq_loss
        
        return total_loss, recon_loss
    
    def get_codebook_usage(self, indices):
        """Compute codebook usage statistics."""
        unique_codes = indices.unique().numel()
        usage = unique_codes / self.vq.num_embeddings
        return usage


# Example: VQ-VAE
vqvae = VQVAE(input_channels=3, hidden_dim=128, embedding_dim=64, num_embeddings=512)

x = torch.randn(4, 3, 32, 32)
x_recon, vq_loss, indices = vqvae(x)

total_loss, recon_loss = vqvae.compute_loss(x, x_recon, vq_loss)
usage = vqvae.get_codebook_usage(indices)

print(f"Reconstruction loss: {recon_loss.item():.4f}")
print(f"VQ loss: {vq_loss.item():.4f}")
print(f"Total loss: {total_loss.item():.4f}")
print(f"Codebook usage: {usage:.2%}")

7. Summary

1. VAE provides a principled framework for generative modelling via variational inference, optimising the ELBO.

2. Reparameterisation trick enables gradient-based optimisation through stochastic latent variables.

3. Posterior collapse is a common failure mode; KL annealing and free bits help mitigate it.

4. β-VAE encourages disentangled representations by increasing the KL penalty.

5. VQ-VAE uses discrete latent spaces, avoiding posterior collapse and producing sharper samples.

References

• Kingma, D. P., & Welling, M. (2014). Auto-encoding variational Bayes. ICLR.
• Rezende, D. J., Mohamed, S., & Wierstra, D. (2014). Stochastic backpropagation and approximate inference in deep generative models. ICML.
• Higgins, I., et al. (2017). β-VAE: Learning basic visual concepts with a constrained variational framework. ICLR.
• van den Oord, A., et al. (2017). Neural discrete representation learning. NeurIPS.
• Razavi, A., et al. (2019). Generating diverse high-fidelity images with VQ-VAE-2. NeurIPS.
• Bowman, S. R., et al. (2016). Generating sentences from a continuous space. ACL.
• Chen, X., et al. (2017). Variational lossy autoencoder. ICLR.