Variational Autoencoders — Learning Latent Representations

A regular autoencoder is like a library where each book has a specific assigned shelf location. The shelves between books are empty — if you reach between two books you get nothing. A VAE is like a library organised by topic, with smooth transitions between subjects — books on cricket shade gradually into books on other sports, then into general fitness. Any point on the shelf has something meaningful. You can navigate by sliding from one location to another and find related content throughout.

The KL divergence term is the librarian enforcing this organisation. Without it, the encoder would cram all books into tiny clusters and leave most of the shelf empty — efficient but not navigable.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np

torch.manual_seed(42)

# ── Full VAE implementation — convolutional ───────────────────────────
class ConvVAE(nn.Module):
    """
    Convolutional VAE for 64×64 RGB images.
    Encoder: image → (μ, log σ²)
    Decoder: z → image
    """
    def __init__(self, latent_dim: int = 128):
        super().__init__()
        self.latent_dim = latent_dim

        # ── Encoder ───────────────────────────────────────────────────
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 32, 4, 2, 1),   # 64 → 32
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, 2, 1),  # 32 → 16
            nn.ReLU(),
            nn.Conv2d(64, 128, 4, 2, 1), # 16 → 8
            nn.ReLU(),
            nn.Conv2d(128, 256, 4, 2, 1),# 8 → 4
            nn.ReLU(),
            nn.Flatten(),                 # 256 × 4 × 4 = 4096
        )
        self.fc_mu      = nn.Linear(4096, latent_dim)
        self.fc_log_var = nn.Linear(4096, latent_dim)

        # ── Decoder ───────────────────────────────────────────────────
        self.fc_decode = nn.Linear(latent_dim, 4096)
        self.decoder   = nn.Sequential(
            nn.ConvTranspose2d(256, 128, 4, 2, 1), # 4 → 8
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 4, 2, 1),  # 8 → 16
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 4, 2, 1),   # 16 → 32
            nn.ReLU(),
            nn.ConvTranspose2d(32, 3, 4, 2, 1),    # 32 → 64
            nn.Sigmoid(),   # pixel values in [0, 1]
        )

    def encode(self, x):
        h       = self.encoder(x)
        mu      = self.fc_mu(h)
        log_var = self.fc_log_var(h)
        return mu, log_var

    def reparameterise(self, mu, log_var):
        if self.training:
            std = torch.exp(0.5 * log_var)
            eps = torch.randn_like(std)
            return mu + std * eps
        else:
            return mu   # at inference: use mean, no sampling noise

    def decode(self, z):
        h = self.fc_decode(z).view(-1, 256, 4, 4)
        return self.decoder(h)

    def forward(self, x):
        mu, log_var = self.encode(x)
        z           = self.reparameterise(mu, log_var)
        x_recon     = self.decode(z)
        return x_recon, mu, log_var

# ── ELBO loss ─────────────────────────────────────────────────────────
def elbo_loss(x_recon, x, mu, log_var, beta: float = 1.0):
    """
    ELBO = Reconstruction loss + β × KL divergence
    beta=1:  standard VAE
    beta>1:  β-VAE — stronger disentanglement, worse reconstruction
    beta<1:  reconstruction focus — sharper outputs, less structured latent
    """
    # Reconstruction: binary cross-entropy per pixel, summed over batch
    recon_loss = F.binary_cross_entropy(x_recon, x, reduction='sum')

    # KL divergence: −0.5 × Σ(1 + log σ² − μ² − σ²)
    # Closed-form for diagonal Gaussian vs standard normal
    kl_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())

    return (recon_loss + beta * kl_loss) / x.size(0)   # normalise by batch size

# ── Shape check ───────────────────────────────────────────────────────
vae = ConvVAE(latent_dim=128)
x   = torch.rand(4, 3, 64, 64)   # batch of 4 random images

x_recon, mu, log_var = vae(x)
loss = elbo_loss(x_recon, x, mu, log_var, beta=1.0)

total_params = sum(p.numel() for p in vae.parameters())
print(f"ConvVAE shapes:")
print(f"  Input:        {tuple(x.shape)}")
print(f"  Reconstructed:{tuple(x_recon.shape)}")
print(f"  μ:            {tuple(mu.shape)}")
print(f"  log σ²:       {tuple(log_var.shape)}")
print(f"  ELBO loss:    {loss.item():.4f}")
print(f"  Parameters:   {total_params:,}")

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import numpy as np
from torch.utils.data import Dataset, DataLoader

torch.manual_seed(42)

# ── Synthetic dataset — fashion product images ────────────────────────
class SyntheticFashionDataset(Dataset):
    """
    Simulates 6 fashion categories with distinct colour signatures.
    In production: torchvision.datasets.ImageFolder with real images.
    """
    CATEGORY_COLOURS = {
        0: [0.8, 0.2, 0.2],  # kurta — warm red
        1: [0.8, 0.6, 0.1],  # saree — gold
        2: [0.2, 0.3, 0.7],  # jeans — denim blue
        3: [0.3, 0.3, 0.3],  # sneakers — grey
        4: [0.6, 0.4, 0.1],  # watch — brown leather
        5: [0.6, 0.1, 0.5],  # handbag — purple
    }

    def __init__(self, n: int = 600):
        self.n = n
        self.labels = np.random.randint(0, 6, n)

    def __len__(self): return self.n

    def __getitem__(self, i):
        label  = self.labels[i]
        colour = self.CATEGORY_COLOURS[label]
        img    = np.random.randn(3, 64, 64) * 0.1
        for c in range(3):
            img[c] += colour[c]
        img = np.clip(img, 0, 1).astype(np.float32)
        return torch.FloatTensor(img), label

dataset    = SyntheticFashionDataset(n=600)
train_size = 480
val_size   = 120
train_ds, val_ds = torch.utils.data.random_split(dataset, [train_size, val_size])
train_ld = DataLoader(train_ds, batch_size=32, shuffle=True)
val_ld   = DataLoader(val_ds,   batch_size=32)

# ── Smaller VAE for demonstration ─────────────────────────────────────
class SmallVAE(nn.Module):
    def __init__(self, latent_dim=32):
        super().__init__()
        self.latent_dim = latent_dim
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 32, 4, 2, 1), nn.ReLU(),   # 32
            nn.Conv2d(32, 64, 4, 2, 1), nn.ReLU(),  # 16
            nn.Conv2d(64, 128, 4, 2, 1), nn.ReLU(), # 8
            nn.Flatten(),
        )
        self.fc_mu      = nn.Linear(128*8*8, latent_dim)
        self.fc_log_var = nn.Linear(128*8*8, latent_dim)
        self.fc_decode  = nn.Linear(latent_dim, 128*8*8)
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(128, 64, 4, 2, 1), nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 4, 2, 1),  nn.ReLU(),
            nn.ConvTranspose2d(32, 3, 4, 2, 1),
            nn.Sigmoid(),
        )

    def encode(self, x):
        h = self.encoder(x)
        return self.fc_mu(h), self.fc_log_var(h)

    def reparameterise(self, mu, lv):
        return mu + torch.exp(0.5 * lv) * torch.randn_like(mu) if self.training else mu

    def decode(self, z):
        return self.decoder(self.fc_decode(z).view(-1, 128, 8, 8))

    def forward(self, x):
        mu, lv = self.encode(x)
        z = self.reparameterise(mu, lv)
        return self.decode(z), mu, lv

vae       = SmallVAE(latent_dim=32)
optimizer = optim.Adam(vae.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)

def kl_anneal(epoch, warmup=10):
    """KL annealing: linearly ramp β from 0 to 1 over warmup epochs."""
    return min(1.0, epoch / warmup)

print(f"Training VAE on fashion dataset:")
print(f"{'Epoch':>6} {'Train loss':>12} {'Recon':>10} {'KL':>8} {'β':>6}")
print("─" * 46)

for epoch in range(1, 21):
    beta = kl_anneal(epoch, warmup=10)
    vae.train()
    total_loss = recon_total = kl_total = 0

    for imgs, _ in train_ld:
        optimizer.zero_grad()
        x_recon, mu, lv = vae(imgs)

        recon = F.binary_cross_entropy(x_recon, imgs, reduction='sum') / imgs.size(0)
        kl    = -0.5 * torch.sum(1 + lv - mu.pow(2) - lv.exp()) / imgs.size(0)
        loss  = recon + beta * kl

        loss.backward()
        nn.utils.clip_grad_norm_(vae.parameters(), 1.0)
        optimizer.step()

        total_loss   += loss.item()
        recon_total  += recon.item()
        kl_total     += kl.item()

    scheduler.step()
    n = len(train_ld)
    if epoch % 4 == 0:
        print(f"  {epoch:>4}  {total_loss/n:>12.2f}  {recon_total/n:>10.2f}  "
              f"{kl_total/n:>8.2f}  {beta:>6.2f}")

print(f"
KL annealing: β starts at 0 → 1 over 10 epochs")
print(f"  Without annealing: KL collapses all encodings to N(0,I) immediately")
print(f"  With annealing: encoder first learns reconstruction, then latent structure")

import torch
import torch.nn as nn
import numpy as np

# Assume vae is trained from previous section
vae.eval()

# ── Application 1: Interpolation between two fashion items ───────────
print("=" * 50)
print("1. LATENT SPACE INTERPOLATION")
print("=" * 50)

# Two images from different categories
img_kurta  = torch.rand(1, 3, 64, 64) * 0.6 + 0.2   # warm colours
img_jeans  = torch.rand(1, 3, 64, 64) * 0.3 + 0.1   # cool dark colours

with torch.no_grad():
    mu_kurta, _  = vae.encode(img_kurta)
    mu_jeans, _  = vae.encode(img_jeans)

    print("Interpolating kurta → jeans:")
    for t in [0.0, 0.25, 0.5, 0.75, 1.0]:
        z_interp  = (1 - t) * mu_kurta + t * mu_jeans
        img_interp = vae.decode(z_interp)
        print(f"  t={t:.2f}: mean={img_interp.mean():.4f}  "
              f"R={img_interp[0,0].mean():.3f}  "
              f"B={img_interp[0,2].mean():.3f}")

# ── Application 2: Random generation from prior ───────────────────────
print("
" + "=" * 50)
print("2. GENERATING NEW FASHION ITEMS")
print("=" * 50)

with torch.no_grad():
    # Sample from standard normal — the prior p(z)
    z_samples  = torch.randn(8, vae.latent_dim)
    generated  = vae.decode(z_samples)
    print(f"Generated {len(generated)} new images from N(0,I)")
    for i, img in enumerate(generated):
        print(f"  Item {i+1}: mean={img.mean():.4f}  std={img.std():.4f}")

# ── Application 3: Anomaly detection ─────────────────────────────────
print("
" + "=" * 50)
print("3. ANOMALY DETECTION")
print("=" * 50)
print("""
# Core idea: normal items reconstruct well, anomalies reconstruct poorly.
# Reconstruction error = how "normal" an item is.

# Train VAE only on normal (non-defective) product images.
# At inference: compute reconstruction error for each item.
# High error → anomaly (defect, wrong item, etc.)

def anomaly_score(vae, image, n_samples=10):
    vae.eval()
    with torch.no_grad():
        mu, log_var = vae.encode(image)
        # Average reconstruction error over multiple samples
        # (reduces noise from stochastic sampling)
        errors = []
        for _ in range(n_samples):
            z     = vae.reparameterise(mu, log_var)
            recon = vae.decode(z)
            error = F.mse_loss(recon, image, reduction='mean').item()
            errors.append(error)
        return np.mean(errors)

# Normal item: low reconstruction error
# Defective item: high reconstruction error
# Threshold at 95th percentile of validation set errors

# Used by: quality control at garment factories,
# fraud detection (unusual transaction patterns),
# medical imaging (lesion detection)
""")

# ── Application 4: Attribute manipulation ─────────────────────────────
print("4. LATENT ATTRIBUTE MANIPULATION")
print("""
# Learn direction vectors for specific attributes:
# colour_direction = mean(blue_items_latents) - mean(red_items_latents)

# Then at inference:
# z_blue_version = z_original + alpha * colour_direction
# decoded = vae.decode(z_blue_version)
# → same product, different colour, without re-photographing it

# Used by Myntra for product colour variations.
# Used by DoorDash to generate food images in different presentation styles.
""")

import torch
import numpy as np

# ── β-VAE: just change beta in the loss ──────────────────────────────
def beta_vae_loss(x_recon, x, mu, log_var, beta: float):
    """
    β-VAE loss with controllable disentanglement.
    beta=1:   standard VAE — best reconstruction
    beta=4:   moderate disentanglement
    beta=10+: strong disentanglement, blurrier reconstruction
    """
    recon = torch.nn.functional.binary_cross_entropy(
        x_recon, x, reduction='sum') / x.size(0)
    kl    = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp()) / x.size(0)
    return recon + beta * kl, recon.item(), kl.item()

# ── Demonstrate effect of β on KL and reconstruction ─────────────────
print("Effect of β on training dynamics:")
print(f"{'β':>6} {'Effect':>50}")
print("─" * 60)
betas_guide = [
    (0.0,  'Pure reconstruction — AE mode, no latent regularisation'),
    (0.5,  'Light regularisation — sharp images, some structure'),
    (1.0,  'Standard VAE — balanced reconstruction and structure'),
    (4.0,  'Moderate β-VAE — some disentanglement, slightly blurry'),
    (10.0, 'Strong β-VAE — clear disentanglement, blurry outputs'),
    (100., 'Over-regularised — all encodings collapse to N(0,I)'),
]
for b, effect in betas_guide:
    print(f"  {b:>4.1f}  {effect}")

# ── Measuring disentanglement — latent traversal ──────────────────────
print("""
# Latent traversal: fix all dimensions, vary one → visualise effect
# Well-disentangled: varying dim 3 changes ONLY colour
# Poorly-entangled: varying dim 3 changes colour AND shape AND texture

def latent_traversal(vae, image, dim: int, values: list):
    vae.eval()
    with torch.no_grad():
        mu, _ = vae.encode(image)
        images = []
        for v in values:
            z        = mu.clone()
            z[0, dim] = v         # vary only this dimension
            images.append(vae.decode(z))
        return images

# Traversal values — sweep from -3 to +3 (3 standard deviations)
# traversal_values = torch.linspace(-3, 3, 10)
# for dim in range(vae.latent_dim):
#     imgs = latent_traversal(vae, test_image, dim, traversal_values)
#     # If imgs show monotonic smooth change in one attribute → disentangled
""")

# ── Choosing β for production ─────────────────────────────────────────
print("Practical β selection guide:")
use_cases = [
    ('Image reconstruction / compression',    1.0,  'Prioritise fidelity'),
    ('Data augmentation',                      1.0,  'Good quality variations'),
    ('Anomaly detection',                      1.0,  'Sharp recon error signal'),
    ('Attribute manipulation',                 4.0,  'Partial disentanglement'),
    ('Drug discovery (molecular generation)',  4.0,  'Structured chemical space'),
    ('Interpretability research',             10.0,  'Maximum disentanglement'),
]
print(f"  {'Use case':<40} {'β':>5}  Reason")
print("  " + "─" * 60)
for use, b, reason in use_cases:
    print(f"  {use:<40} {b:>5.1f}  {reason}")

Sampling from a distribution stays a fundamentally stochastic operation — no trick turns randomness itself into something differentiable. What the reparameterisation trick actually does is move the randomness out of the path that needs a gradient: instead of sampling z directly from N(mu, sigma squared), it samples an independent epsilon from N(0, I) and computes z as a deterministic function of mu, sigma, and epsilon. Gradients flow cleanly through that deterministic function with respect to mu and sigma; the random part is sidestepped, not made differentiable.

Blurriness is a predictable mathematical consequence of the objective, not evidence of a weaker architecture. Pixel-wise reconstruction loss (binary cross-entropy or MSE) is minimised by outputting the average of plausible values whenever the decoder is uncertain between them, and an average of several sharp options looks smeared. This is fixable without abandoning the VAE framework — adding a perceptual loss, adding an adversarial loss on top of the ELBO (as in VQ-GAN-style hybrids), or moving to a discrete latent space (VQ-VAE) all address the same root cause directly.

The KL term has nothing to do with reconstruction quality — its entire job is regularising the shape of the encoder's output distribution to stay close to a standard normal, which is what makes the latent space continuous and sample-able in the first place. It actually competes with reconstruction rather than assisting it, which is the whole tension the ELBO is built to balance. Turning the KL weight up, as in beta-VAE, deliberately trades away some reconstruction fidelity in exchange for a more organised, disentangled latent space — proof that the two terms pull in different directions, not the same one.

Posterior collapse is a specific, diagnosable failure — the encoder gives up and outputs close to N(0, I) for every input regardless of content, and the decoder compensates by learning to generate a generic average output that ignores z entirely. The overall loss can look like it is decreasing the whole time, which is exactly what makes it deceptive to spot from the training curve alone. It is a symptom of the latent code being ignored, not the network failing to learn anything, and it is fixed with known techniques — KL annealing, free bits, reducing decoder capacity — rather than being a sign the VAE approach itself is wrong for the data.

As the production examples above show, a VAE is frequently a component living inside a diffusion pipeline rather than a competing headline architecture — latent diffusion models like Stable Diffusion use a VAE as their compression stage. Hybrids combining a VAE encoder with a GAN-style discriminator loss on top of the ELBO exist too. Treating these three as three separate competing tools you must pick exactly one of obscures how often real production systems combine them instead.

The ELBO is a lower bound on the log likelihood of the data, and maximising it splits into two terms. The reconstruction term, E of log p(x given z), measures how well the decoder reproduces the input from a sampled latent code — maximising it pushes the decoder toward accurate reconstructions. The KL divergence term measures how far the encoder's distribution q(z given x) sits from the prior p(z), a standard normal — minimising it keeps the latent space smooth and centred instead of letting the encoder scatter arbitrary codes anywhere. These two terms are in direct tension, and that tension is what produces a latent space that is both informative and navigable.

Without it, z is sampled directly from a distribution parameterised by mu and sigma, and sampling is not a differentiable operation with respect to those parameters — gradients cannot flow backward through a random draw, so the encoder could never be trained end to end with standard backpropagation. The reparameterisation trick rewrites the sample as a deterministic function of mu, sigma, and an independent noise term epsilon drawn from N(0, I). Since epsilon carries no learnable parameters, gradients flow cleanly through the mu and sigma paths, and the whole network trains with ordinary backprop.

The pixel-wise reconstruction loss is minimised by predicting the average of plausible pixel values whenever the model is uncertain, and an average of several sharp possibilities reads as blur. A GAN sidesteps this because its loss comes from a discriminator judging overall realism rather than penalising each pixel independently, so it has no incentive to average. Diffusion sidesteps it differently — by refining the image gradually over many steps rather than committing to a single-shot reconstruction, it never has to resolve all its uncertainty in one averaged guess.

Posterior collapse is when the encoder stops encoding useful information and outputs close to the prior, N(0, I), for every input, while the decoder learns to reconstruct a generic average output without actually using z. It typically happens when the KL term dominates too early in training — before the reconstruction signal has had a chance to teach the encoder anything worth preserving — or when the decoder is powerful enough to reconstruct reasonably well without relying on the latent code at all. The standard fix is KL annealing: start training with the KL weight at zero and ramp it up over the first several epochs, so the encoder learns to encode meaningful structure before the regularisation pressure kicks in.

A plain autoencoder's latent space is never regularised to be continuous — the model is only ever trained to correctly decode the exact codes it assigned to training examples, so the empty space between two encoded points can decode to something meaningless. A VAE's KL term forces every encoded distribution to overlap with the same standard normal prior, which packs the whole space with meaningfully decodable regions rather than leaving gaps. That is precisely why walking a straight line between two encoded means and decoding along the way produces a smooth, semantically sensible transition instead of noise.