Master probabilistic generative modeling with VAEs. Learn autoencoder architecture, variational inference, reparameterization trick, and latent space manipulation for generation and data compression.
Variational Autoencoders (VAEs) are a powerful class of generative models that combine deep learning with probabilistic inference. Unlike GANs which learn through adversarial competition, VAEs take a principled probabilistic approach—learning a smooth, continuous latent space that can generate new data by sampling.
Imagine compressing images into a "code" so compact that you can reconstruct them later, but with a twist: the code is probabilistic, allowing you to generate entirely new images by sampling from the learned distribution. That's the magic of VAEs!
In this lesson, you'll understand:
1. Autoencoder Basics: Encoder compresses data → latent code, Decoder reconstructs data
2. Variational Inference: Adding probabilistic structure to latent space
3. The VAE Objective: Balancing reconstruction loss and KL divergence
4. Reparameterization Trick: How to backpropagate through stochastic sampling
5. Latent Space Exploration: Interpolation, generation, and disentanglement
You've mastered GANs for sharp image generation—now you'll learn VAEs for smooth, controllable generation with principled probabilistic foundations. By the end, you'll build a VAE that generates handwritten digits and explore the continuous latent space.
Ready to combine probability theory with deep learning? Let's dive into Variational Autoencoders! 🎲
Autoencoders are neural networks that learn compressed representations of data. Understanding vanilla autoencoders is the foundation for VAEs.
Core Idea: Learn to compress and decompress data
1import torch
2import torch.nn as nn
3import torch.optim as optim
4import matplotlib.pyplot as plt
5import numpy as np
6from torchvision import datasets, transforms
7from torch.utils.data import DataLoader
8
9print("=" * 80)
10print("AUTOENCODER FUNDAMENTALS: COMPRESSION & RECONSTRUCTION")
11print("=" * 80)
12print()
13
14# ============================================================================
15# PART 1: Vanilla Autoencoder Architecture
16# ============================================================================
17
18class Autoencoder(nn.Module):
19 """
20 Vanilla Autoencoder: Data → Latent Code → Reconstructed Data
21
22 Components:
23 1. Encoder: Compresses input to low-dimensional latent representation
24 2. Latent Space: Bottleneck layer (compressed representation)
25 3. Decoder: Reconstructs input from latent code
26 """
27
28 def __init__(self, input_dim=784, latent_dim=32):
29 super(Autoencoder, self).__init__()
30
31 # Encoder: 784 → 256 → 128 → 32
32 self.encoder = nn.Sequential(
33 nn.Linear(input_dim, 256),
34 nn.ReLU(),
35 nn.Linear(256, 128),
36 nn.ReLU(),
37 nn.Linear(128, latent_dim), # Bottleneck
38 nn.ReLU()
39 )
40
41 # Decoder: 32 → 128 → 256 → 784
42 self.decoder = nn.Sequential(
43 nn.Linear(latent_dim, 128),
44 nn.ReLU(),
45 nn.Linear(128, 256),
46 nn.ReLU(),
47 nn.Linear(256, input_dim),
48 nn.Sigmoid() # Output in [0, 1]
49 )
50
51 def forward(self, x):
52 # Encode
53 latent = self.encoder(x)
54
55 # Decode
56 reconstructed = self.decoder(latent)
57
58 return reconstructed, latent
59
60
61# Initialize autoencoder
62input_dim = 28 * 28 # MNIST flattened
63latent_dim = 32
64
65autoencoder = Autoencoder(input_dim, latent_dim)
66
67print("Vanilla Autoencoder Architecture:")
68print("-" * 80)
69print(f"Input dimension: {input_dim}")
70print(f"Latent dimension: {latent_dim}")
71print(f"Compression ratio: {input_dim / latent_dim:.1f}x")
72print()
73
74# Count parameters
75encoder_params = sum(p.numel() for p in autoencoder.encoder.parameters())
76decoder_params = sum(p.numel() for p in autoencoder.decoder.parameters())
77total_params = encoder_params + decoder_params
78
79print(f"Encoder parameters: {encoder_params:,}")
80print(f"Decoder parameters: {decoder_params:,}")
81print(f"Total parameters: {total_params:,}")
82print()
83
84# Test forward pass
85test_input = torch.randn(4, input_dim)
86test_reconstructed, test_latent = autoencoder(test_input)
87
88print("Forward Pass Test:")
89print(f" Input shape: {test_input.shape}")
90print(f" Latent shape: {test_latent.shape}")
91print(f" Reconstructed shape: {test_reconstructed.shape}")
92print(f" Information bottleneck: {input_dim} → {latent_dim} → {input_dim}")Training Objective: Minimize Reconstruction Error
1def train_autoencoder(model, dataloader, epochs=10, lr=1e-3, device='cpu'):
2 """Train vanilla autoencoder"""
3
4 model.to(device)
5 optimizer = optim.Adam(model.parameters(), lr=lr)
6 criterion = nn.MSELoss() # Mean Squared Error for reconstruction
7
8 history = {'loss': [], 'epoch_samples': []}
9
10 print("=" * 80)
11 print("TRAINING VANILLA AUTOENCODER")
12 print("=" * 80)
13 print()
14
15 for epoch in range(epochs):
16 model.train()
17 epoch_loss = 0
18
19 for batch_idx, (images, _) in enumerate(dataloader):
20 # Flatten images
21 images = images.view(images.size(0), -1).to(device)
22
23 # Forward pass
24 reconstructed, latent = model(images)
25
26 # Reconstruction loss
27 loss = criterion(reconstructed, images)
28
29 # Backward pass
30 optimizer.zero_grad()
31 loss.backward()
32 optimizer.step()
33
34 epoch_loss += loss.item()
35
36 avg_loss = epoch_loss / len(dataloader)
37 history['loss'].append(avg_loss)
38
39 # Save samples for visualization
40 if (epoch + 1) % 5 == 0:
41 with torch.no_grad():
42 model.eval()
43 sample_images = images[:8]
44 sample_reconstructed, _ = model(sample_images)
45 history['epoch_samples'].append(
46 (sample_images.cpu(), sample_reconstructed.cpu(), epoch + 1)
47 )
48
49 print(f"Epoch [{epoch+1}/{epochs}] - Loss: {avg_loss:.6f}")
50
51 print()
52 print("Training Complete!")
53 return history
54
55
56# Example: Loss function explanation
57print("\n" + "=" * 80)
58print("AUTOENCODER LOSS FUNCTION")
59print("=" * 80)
60print()
61
62print("Reconstruction Loss (MSE):")
63print("-" * 80)
64print("Loss = MSE(original, reconstructed)")
65print(" = (1/N) * Σ(x_i - x̂_i)²")
66print()
67print("Where:")
68print(" x_i = original pixel values")
69print(" x̂_i = reconstructed pixel values")
70print(" N = total number of pixels")
71print()
72
73print("Training Process:")
74print("-" * 80)
75print("1. Encoder compresses input → latent code")
76print("2. Decoder reconstructs input from latent code")
77print("3. Loss measures reconstruction quality")
78print("4. Backpropagation updates both encoder and decoder")
79print("5. Model learns to preserve important features in latent code")Problem: Latent Space is Not Structured
1print("\n" + "=" * 80)
2print("VANILLA AUTOENCODER LIMITATIONS")
3print("=" * 80)
4print()
5
6print("❌ Problems:")
7print("-" * 80)
8print("1. DISCONTINUOUS LATENT SPACE")
9print(" • Latent codes are scattered, not continuous")
10print(" • Small changes in latent code → unpredictable changes in output")
11print(" • Cannot interpolate smoothly between points")
12print()
13
14print("2. CANNOT GENERATE NEW DATA")
15print(" • Latent space has 'holes' and irregular structure")
16print(" • Random sampling produces nonsense")
17print(" • Can only reconstruct training data")
18print()
19
20print("3. OVERFITTING TO TRAINING DATA")
21print(" • Model memorizes specific inputs")
22print(" • Poor generalization to new variations")
23print(" • Latent space doesn't capture data distribution")
24print()
25
26print("✅ Solution: VARIATIONAL AUTOENCODERS (VAEs)")
27print("-" * 80)
28print("VAEs add probabilistic structure to latent space:")
29print(" • Force latent space to be continuous (Gaussian prior)")
30print(" • Enable smooth interpolation between points")
31print(" • Allow generation of new data by sampling")
32print(" • Regularize with KL divergence")1print("\n" + "=" * 80)
2print("AUTOENCODER vs VAE: LATENT SPACE STRUCTURE")
3print("=" * 80)
4print()
5
6print("VANILLA AUTOENCODER:")
7print("-" * 40)
8print("Latent Space:")
9print(" • Scattered points (no structure)")
10print(" • Encoder: x → z")
11print(" • Decoder: z → x̂")
12print(" • Loss: Reconstruction only")
13print()
14print("Generation:")
15print(" ✗ Random sample → unpredictable output")
16print(" ✗ Interpolation → discontinuous transitions")
17print()
18
19print("VARIATIONAL AUTOENCODER:")
20print("-" * 40)
21print("Latent Space:")
22print(" • Continuous Gaussian distribution")
23print(" • Encoder: x → μ, σ (parameters of distribution)")
24print(" • Sampling: z ~ N(μ, σ²)")
25print(" • Decoder: z → x̂")
26print(" • Loss: Reconstruction + KL divergence")
27print()
28print("Generation:")
29print(" ✓ Random sample ~ N(0, 1) → valid output")
30print(" ✓ Interpolation → smooth transitions")Autoencoder Strengths:
Autoencoder Weaknesses:
Why VAEs are Better for Generation:
Understanding vanilla autoencoders is crucial—VAEs build on this foundation by adding probabilistic structure and regularization!
VAEs transform autoencoders into generative models through probabilistic modeling and variational inference. Let's understand the mathematical foundation.
Goal: Model the data distribution p(x)
1print("=" * 80)
2print("VARIATIONAL AUTOENCODERS: PROBABILISTIC FRAMEWORK")
3print("=" * 80)
4print()
5
6print("THE GENERATIVE STORY:")
7print("-" * 80)
8print("1. Sample latent code from prior: z ~ p(z) = N(0, I)")
9print("2. Generate data from conditional: x ~ p(x|z)")
10print()
11print("Key Idea:")
12print(" • Latent variable z captures underlying factors")
13print(" • Decoder models p(x|z) (generation process)")
14print(" • Prior p(z) = N(0, I) ensures smooth latent space")
15print()
16
17print("THE INFERENCE PROBLEM:")
18print("-" * 80)
19print("Given data x, infer latent code z:")
20print(" • Want: posterior p(z|x)")
21print(" • Problem: p(z|x) = p(x|z)p(z) / p(x) is intractable")
22print(" • p(x) = ∫ p(x|z)p(z) dz requires integrating over all z")
23print()
24print("Solution: Variational Inference")
25print(" • Approximate p(z|x) with learned distribution q(z|x)")
26print(" • q(z|x) = N(μ(x), σ²(x)) parameterized by encoder")
27print(" • Optimize ELBO (Evidence Lower Bound)")Probabilistic Encoder and Decoder:
1class VAE(nn.Module):
2 """
3 Variational Autoencoder
4
5 Encoder: x → (μ, log_σ²) [outputs distribution parameters]
6 Sampling: z ~ N(μ, σ²) [reparameterization trick]
7 Decoder: z → x̂ [reconstructs from sample]
8 """
9
10 def __init__(self, input_dim=784, latent_dim=32):
11 super(VAE, self).__init__()
12
13 # Encoder: outputs mean and log-variance
14 self.encoder = nn.Sequential(
15 nn.Linear(input_dim, 256),
16 nn.ReLU(),
17 nn.Linear(256, 128),
18 nn.ReLU()
19 )
20
21 # Latent distribution parameters
22 self.fc_mu = nn.Linear(128, latent_dim) # Mean μ
23 self.fc_logvar = nn.Linear(128, latent_dim) # Log-variance log(σ²)
24
25 # Decoder: reconstructs from latent sample
26 self.decoder = nn.Sequential(
27 nn.Linear(latent_dim, 128),
28 nn.ReLU(),
29 nn.Linear(128, 256),
30 nn.ReLU(),
31 nn.Linear(256, input_dim),
32 nn.Sigmoid()
33 )
34
35 def encode(self, x):
36 """Encode input to latent distribution parameters"""
37 h = self.encoder(x)
38 mu = self.fc_mu(h)
39 logvar = self.fc_logvar(h)
40 return mu, logvar
41
42 def reparameterize(self, mu, logvar):
43 """
44 Reparameterization trick: z = μ + σ * ε, where ε ~ N(0, I)
45
46 Why? Allows backpropagation through stochastic sampling
47 """
48 std = torch.exp(0.5 * logvar) # σ = exp(0.5 * log(σ²))
49 eps = torch.randn_like(std) # ε ~ N(0, 1)
50 z = mu + eps * std # z ~ N(μ, σ²)
51 return z
52
53 def decode(self, z):
54 """Decode latent sample to reconstruction"""
55 return self.decoder(z)
56
57 def forward(self, x):
58 # Encode
59 mu, logvar = self.encode(x)
60
61 # Sample latent code
62 z = self.reparameterize(mu, logvar)
63
64 # Decode
65 x_recon = self.decode(z)
66
67 return x_recon, mu, logvar
68
69
70# Initialize VAE
71vae = VAE(input_dim=784, latent_dim=32)
72
73print("\n" + "=" * 80)
74print("VAE ARCHITECTURE")
75print("=" * 80)
76print()
77
78print("Components:")
79print("-" * 80)
80print(f"1. Encoder: {input_dim} → 256 → 128 → (μ, log σ²)")
81print(f" • Outputs: mean and log-variance of latent distribution")
82print(f" • Dimension: {latent_dim}")
83print()
84
85print(f"2. Reparameterization: z = μ + σ * ε")
86print(f" • ε ~ N(0, 1) [random noise]")
87print(f" • z ~ N(μ, σ²) [sampled latent code]")
88print()
89
90print(f"3. Decoder: {latent_dim} → 128 → 256 → {input_dim}")
91print(f" • Reconstructs data from latent sample")
92print(f" • Output: probability distribution over pixels")
93print()
94
95# Test forward pass
96test_input = torch.randn(4, 784)
97test_recon, test_mu, test_logvar = vae(test_input)
98
99print("Forward Pass Test:")
100print(f" Input: {test_input.shape}")
101print(f" μ (mean): {test_mu.shape}")
102print(f" log σ² (log-variance): {test_logvar.shape}")
103print(f" z (sampled): {vae.reparameterize(test_mu, test_logvar).shape}")
104print(f" Reconstruction: {test_recon.shape}")Objective: Maximize Evidence Lower Bound
1def vae_loss(x_recon, x, mu, logvar, beta=1.0):
2 """
3 VAE Loss = Reconstruction Loss + KL Divergence
4
5 ELBO = E[log p(x|z)] - KL(q(z|x) || p(z))
6
7 Maximize ELBO ≡ Minimize: -ELBO
8 """
9
10 # 1. Reconstruction Loss: -E[log p(x|z)]
11 # Measures how well we reconstruct the input
12 recon_loss = nn.functional.binary_cross_entropy(
13 x_recon, x, reduction='sum'
14 )
15
16 # 2. KL Divergence: KL(q(z|x) || p(z))
17 # Regularizes latent space to match prior N(0, I)
18 # For q = N(μ, σ²) and p = N(0, 1):
19 # KL = -0.5 * Σ(1 + log(σ²) - μ² - σ²)
20 kl_divergence = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
21
22 # Total loss
23 total_loss = recon_loss + beta * kl_divergence
24
25 return total_loss, recon_loss, kl_divergence
26
27
28print("\n" + "=" * 80)
29print("VAE LOSS FUNCTION: ELBO")
30print("=" * 80)
31print()
32
33print("Evidence Lower Bound (ELBO):")
34print("-" * 80)
35print("ELBO = E_q(z|x)[log p(x|z)] - KL(q(z|x) || p(z))")
36print()
37print("Component 1: Reconstruction Loss")
38print(" • -E[log p(x|z)]: Negative log-likelihood")
39print(" • Measures: How well decoder reconstructs input")
40print(" • Implementation: Binary cross-entropy or MSE")
41print(" • Goal: Minimize (better reconstruction)")
42print()
43
44print("Component 2: KL Divergence")
45print(" • KL(q(z|x) || p(z)): Regularization term")
46print(" • Measures: Distance between q(z|x) and prior p(z)")
47print(" • q(z|x) = N(μ(x), σ²(x)) [encoder output]")
48print(" • p(z) = N(0, I) [standard normal prior]")
49print(" • Goal: Minimize (match prior distribution)")
50print()
51
52print("Closed-form KL for Gaussians:")
53print(" KL(N(μ, σ²) || N(0, 1)) = -0.5 * Σ(1 + log(σ²) - μ² - σ²)")
54print()
55
56print("Why This Works:")
57print("-" * 80)
58print("1. Reconstruction loss → preserves information")
59print("2. KL divergence → smooths latent space")
60print("3. Balance creates continuous, structured latent space")
61print("4. Can generate by sampling z ~ N(0, 1)")1print("\n" + "=" * 80)
2print("BETA-VAE: DISENTANGLED REPRESENTATIONS")
3print("=" * 80)
4print()
5
6print("Standard VAE Loss:")
7print(" Loss = Reconstruction + KL")
8print()
9
10print("Beta-VAE Loss:")
11print(" Loss = Reconstruction + β * KL")
12print()
13print("Where β controls the trade-off:")
14print()
15
16print("β = 1.0 (Standard VAE):")
17print(" • Balanced reconstruction and regularization")
18print(" • Good generation quality")
19print()
20
21print("β > 1.0 (High Regularization):")
22print(" • Stronger pressure to match prior")
23print(" • More disentangled latent dimensions")
24print(" • Each dimension captures independent factor")
25print(" • Example: β = 4 for faces (one dim = age, another = pose)")
26print(" • Trade-off: Slightly worse reconstruction")
27print()
28
29print("β < 1.0 (Low Regularization):")
30print(" • Better reconstruction quality")
31print(" • Less structured latent space")
32print(" • Harder to generate from prior")VAE Mathematical Framework:
Why KL Divergence Matters:
Without KL term:
With KL term:
Probabilistic vs Deterministic:
| Aspect | Autoencoder | VAE |
|--------|-------------|-----|
| Encoder output | Single point z | Distribution (μ, σ²) |
| Latent space | Scattered | Continuous Gaussian |
| Loss | Reconstruction | Reconstruction + KL |
| Generation | Poor | Good |
VAEs provide a principled probabilistic framework for generative modeling!
The reparameterization trick is the key innovation that makes VAE training possible. It allows gradients to flow through stochastic sampling.
Challenge: Cannot backpropagate through random sampling
1print("=" * 80)
2print("THE REPARAMETERIZATION TRICK")
3print("=" * 80)
4print()
5
6print("THE PROBLEM:")
7print("-" * 80)
8print("Standard sampling: z ~ N(μ, σ²)")
9print()
10print("Cannot compute gradient:")
11print(" ∂z/∂μ = ? (z is random!)")
12print(" ∂z/∂σ = ?")
13print()
14print("Why this matters:")
15print(" • VAE needs to sample z ~ q(z|x)")
16print(" • Must backpropagate through sampling to train encoder")
17print(" • Random sampling blocks gradient flow")
18print()
19
20print("❌ NAIVE APPROACH (DOESN'T WORK):")
21print()
22print("def sample_naive(mu, sigma):")
23print(" z = torch.normal(mu, sigma) # Random sample")
24print(" return z # Cannot compute gradients!")
25print()
26print("Problem: Sampling operation is non-differentiable")Key Insight: Move randomness outside the computation graph
1print("=" * 80)
2print("REPARAMETERIZATION TRICK")
3print("=" * 80)
4print()
5
6print("✅ THE TRICK:")
7print("-" * 80)
8print("Instead of: z ~ N(μ, σ²)")
9print()
10print("Reparameterize:")
11print(" 1. Sample ε ~ N(0, 1) [standard normal]")
12print(" 2. Transform: z = μ + σ * ε [deterministic given ε]")
13print()
14print("Now z ~ N(μ, σ²) AND we can compute gradients:")
15print(" ∂z/∂μ = 1")
16print(" ∂z/∂σ = ε")
17print()
18
19print("Why it works:")
20print(" • ε is random but independent of μ, σ")
21print(" • z is deterministic function of μ, σ, ε")
22print(" • Gradients flow through μ and σ")
23print(" • Randomness (ε) is 'outside' the gradient computation")
24
25
26def reparameterize_detailed(mu, logvar):
27 """
28 Reparameterization trick with detailed explanation
29
30 Args:
31 mu: Mean of latent distribution [batch_size, latent_dim]
32 logvar: Log-variance log(σ²) [batch_size, latent_dim]
33
34 Returns:
35 z: Sampled latent code [batch_size, latent_dim]
36 """
37
38 # Step 1: Compute standard deviation
39 # σ = exp(0.5 * log(σ²)) = exp(log(σ)) = σ
40 std = torch.exp(0.5 * logvar)
41
42 # Step 2: Sample noise from standard normal
43 # ε ~ N(0, 1)
44 eps = torch.randn_like(std) # Same shape as std
45
46 # Step 3: Transform to desired distribution
47 # z = μ + σ * ε ~ N(μ, σ²)
48 z = mu + std * eps
49
50 return z, eps, std
51
52
53print("\n" + "=" * 80)
54print("IMPLEMENTATION DETAILS")
55print("=" * 80)
56print()
57
58# Example
59batch_size, latent_dim = 4, 8
60mu_example = torch.randn(batch_size, latent_dim)
61logvar_example = torch.randn(batch_size, latent_dim)
62
63z_sample, eps_sample, std_sample = reparameterize_detailed(mu_example, logvar_example)
64
65print("Example:")
66print(f" Batch size: {batch_size}")
67print(f" Latent dim: {latent_dim}")
68print()
69
70print(f"Encoder outputs:")
71print(f" μ shape: {mu_example.shape}")
72print(f" log(σ²) shape: {logvar_example.shape}")
73print()
74
75print(f"Reparameterization:")
76print(f" σ = exp(0.5 * log(σ²)): {std_sample.shape}")
77print(f" ε ~ N(0, 1): {eps_sample.shape}")
78print(f" z = μ + σ * ε: {z_sample.shape}")
79print()
80
81print("Verify distribution:")
82print(f" Sample mean: {z_sample.mean():.4f} (should be ≈ μ.mean())")
83print(f" Sample std: {z_sample.std():.4f} (should be ≈ σ.mean())")
84print(f" μ mean: {mu_example.mean():.4f}")
85print(f" σ mean: {std_sample.mean():.4f}")1print("\n" + "=" * 80)
2print("GRADIENT FLOW THROUGH VAE")
3print("=" * 80)
4print()
5
6print("Forward Pass:")
7print("-" * 80)
8print("x → Encoder → (μ, log σ²)")
9print(" ↓")
10print(" ε ~ N(0, 1)")
11print(" ↓")
12print(" z = μ + σ * ε")
13print(" ↓")
14print(" Decoder → x̂")
15print()
16
17print("Backward Pass (Gradients):")
18print("-" * 80)
19print("Loss ← Reconstruction(x, x̂) + KL(μ, σ²)")
20print(" ↓")
21print("∂L/∂x̂ → Decoder")
22print(" ↓")
23print("∂L/∂z → Reparameterization")
24print(" ↓")
25print("∂L/∂μ = ∂L/∂z * ∂z/∂μ = ∂L/∂z * 1")
26print("∂L/∂σ = ∂L/∂z * ∂z/∂σ = ∂L/∂z * ε")
27print(" ↓")
28print("Encoder ← ∂L/∂μ, ∂L/∂σ")
29print()
30
31print("Key Insight:")
32print(" • Gradients flow through μ and σ")
33print(" • ε is treated as constant (no gradient needed)")
34print(" • Both encoder and decoder get trained")1print("\n" + "=" * 80)
2print("WHY LOG-VARIANCE?")
3print("=" * 80)
4print()
5
6print("Encoder outputs log(σ²) instead of σ:")
7print("-" * 80)
8print()
9
10print("Reason 1: NUMERICAL STABILITY")
11print(" • σ² must be positive")
12print(" • log(σ²) ∈ (-∞, ∞) [unconstrained]")
13print(" • Network can output any value")
14print(" • σ² = exp(log(σ²)) always positive")
15print()
16
17print("Reason 2: AVOID EXPONENTIAL BLOWUP")
18print(" • If network outputs σ² directly:")
19print(" → Must constrain with softplus or exp")
20print(" → Can cause numerical instability")
21print(" • With log(σ²):")
22print(" → More stable gradients")
23print(" → Better training dynamics")
24print()
25
26print("Reason 3: KL DIVERGENCE COMPUTATION")
27print(" • KL formula uses log(σ²) directly")
28print(" • No need to compute log of potentially small σ²")
29print(" • Avoids log(0) errors")
30print()
31
32print("Example:")
33logvar_ex = torch.tensor([-2.0, 0.0, 2.0]) # log(σ²)
34sigma_squared = torch.exp(logvar_ex) # σ²
35sigma = torch.exp(0.5 * logvar_ex) # σ
36
37print(f" log(σ²): {logvar_ex}")
38print(f" σ²: {sigma_squared}")
39print(f" σ: {sigma}")1print("\n" + "=" * 80)
2print("SAMPLING STRATEGIES")
3print("=" * 80)
4print()
5
6print("TRAINING: Single Sample per Data Point")
7print("-" * 80)
8print("• Sample once: z = μ + σ * ε")
9print("• Compute loss with this z")
10print("• Efficient: O(1) sampling cost")
11print("• Monte Carlo estimate of expectation")
12print()
13
14print("INFERENCE: Multiple Samples (Optional)")
15print("-" * 80)
16print("• Sample K times: z₁, z₂, ..., zₖ")
17print("• Generate K reconstructions")
18print("• Average for better estimate")
19print("• Reduces variance but increases compute")
20print()
21
22def generate_with_sampling(vae, x, n_samples=10):
23 """Generate multiple reconstructions"""
24 vae.eval()
25
26 with torch.no_grad():
27 mu, logvar = vae.encode(x)
28
29 # Sample multiple times
30 reconstructions = []
31 for _ in range(n_samples):
32 z = vae.reparameterize(mu, logvar)
33 x_recon = vae.decode(z)
34 reconstructions.append(x_recon)
35
36 # Average
37 avg_recon = torch.stack(reconstructions).mean(dim=0)
38
39 return avg_recon, reconstructions
40
41
42print("When to use multiple samples:")
43print(" • Critical applications (high-quality reconstruction)")
44print(" • Uncertainty quantification")
45print(" • Ensemble predictions")
46print()
47print("Usually single sample is enough during training!")The Reparameterization Trick:
Before: z ~ N(μ, σ²) → No gradients
After: z = μ + σ·ε, ε ~ N(0, 1) → Gradients flow!
Why It's Brilliant:
Implementation Tips:
Common Mistakes:
❌ Sampling without reparameterization
❌ Using σ instead of log(σ²)
❌ Forgetting to detach ε from graph
❌ Not checking for NaN in log operations
The reparameterization trick is what makes VAEs trainable—it's a simple yet elegant solution to backpropagating through stochastic operations!
Let's implement a complete VAE for MNIST, train it, and analyze the results. This brings together all the concepts: architecture, ELBO, and reparameterization.
1import torch
2import torch.nn as nn
3import torch.optim as optim
4import torchvision
5from torchvision import transforms
6from torch.utils.data import DataLoader
7import matplotlib.pyplot as plt
8import numpy as np
9
10print("=" * 80)
11print("COMPLETE VAE IMPLEMENTATION FOR MNIST")
12print("=" * 80)
13print()
14
15# ============================================================================
16# PART 1: VAE Model
17# ============================================================================
18
19class MNIST_VAE(nn.Module):
20 """Production-ready VAE for MNIST"""
21
22 def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20):
23 super(MNIST_VAE, self).__init__()
24
25 self.latent_dim = latent_dim
26
27 # Encoder
28 self.encoder = nn.Sequential(
29 nn.Linear(input_dim, hidden_dim),
30 nn.ReLU(),
31 nn.Linear(hidden_dim, hidden_dim),
32 nn.ReLU()
33 )
34
35 # Latent parameters
36 self.fc_mu = nn.Linear(hidden_dim, latent_dim)
37 self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
38
39 # Decoder
40 self.decoder = nn.Sequential(
41 nn.Linear(latent_dim, hidden_dim),
42 nn.ReLU(),
43 nn.Linear(hidden_dim, hidden_dim),
44 nn.ReLU(),
45 nn.Linear(hidden_dim, input_dim),
46 nn.Sigmoid()
47 )
48
49 def encode(self, x):
50 h = self.encoder(x)
51 mu = self.fc_mu(h)
52 logvar = self.fc_logvar(h)
53 return mu, logvar
54
55 def reparameterize(self, mu, logvar):
56 std = torch.exp(0.5 * logvar)
57 eps = torch.randn_like(std)
58 return mu + eps * std
59
60 def decode(self, z):
61 return self.decoder(z)
62
63 def forward(self, x):
64 mu, logvar = self.encode(x)
65 z = self.reparameterize(mu, logvar)
66 x_recon = self.decode(z)
67 return x_recon, mu, logvar
68
69 def generate(self, n_samples=1, device='cpu'):
70 """Generate new samples from prior"""
71 with torch.no_grad():
72 z = torch.randn(n_samples, self.latent_dim).to(device)
73 samples = self.decode(z)
74 return samples
75
76
77# ============================================================================
78# PART 2: Loss Function
79# ============================================================================
80
81def vae_loss_function(x_recon, x, mu, logvar, beta=1.0):
82 """
83 VAE loss: Reconstruction + KL divergence
84
85 Args:
86 x_recon: Reconstructed images
87 x: Original images
88 mu: Latent mean
89 logvar: Latent log-variance
90 beta: Weight for KL term (beta-VAE)
91 """
92 # Reconstruction loss (binary cross-entropy)
93 recon_loss = nn.functional.binary_cross_entropy(
94 x_recon, x, reduction='sum'
95 )
96
97 # KL divergence: KL(N(μ, σ²) || N(0, 1))
98 kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
99
100 # Total loss
101 total_loss = recon_loss + beta * kl_loss
102
103 return total_loss, recon_loss, kl_loss
104
105
106# ============================================================================
107# PART 3: Training Loop
108# ============================================================================
109
110def train_vae(model, train_loader, epochs=10, lr=1e-3, beta=1.0, device='cpu'):
111 """Train VAE with ELBO objective"""
112
113 model.to(device)
114 optimizer = optim.Adam(model.parameters(), lr=lr)
115
116 history = {
117 'loss': [],
118 'recon_loss': [],
119 'kl_loss': [],
120 'generated_samples': []
121 }
122
123 print("=" * 80)
124 print("TRAINING VAE")
125 print("=" * 80)
126 print(f"Epochs: {epochs} | Learning rate: {lr} | Beta: {beta}")
127 print()
128
129 for epoch in range(epochs):
130 model.train()
131 train_loss = 0
132 train_recon = 0
133 train_kl = 0
134
135 for batch_idx, (data, _) in enumerate(train_loader):
136 data = data.view(-1, 784).to(device)
137
138 # Forward pass
139 x_recon, mu, logvar = model(data)
140
141 # Compute loss
142 loss, recon_loss, kl_loss = vae_loss_function(
143 x_recon, data, mu, logvar, beta
144 )
145
146 # Backward pass
147 optimizer.zero_grad()
148 loss.backward()
149 optimizer.step()
150
151 # Accumulate losses
152 train_loss += loss.item()
153 train_recon += recon_loss.item()
154 train_kl += kl_loss.item()
155
156 # Average losses
157 n_samples = len(train_loader.dataset)
158 avg_loss = train_loss / n_samples
159 avg_recon = train_recon / n_samples
160 avg_kl = train_kl / n_samples
161
162 history['loss'].append(avg_loss)
163 history['recon_loss'].append(avg_recon)
164 history['kl_loss'].append(avg_kl)
165
166 # Generate samples
167 if (epoch + 1) % 5 == 0:
168 samples = model.generate(64, device)
169 history['generated_samples'].append(samples.cpu())
170
171 # Print progress
172 print(f"Epoch [{epoch+1}/{epochs}]")
173 print(f" Total Loss: {avg_loss:.4f}")
174 print(f" Recon Loss: {avg_recon:.4f}")
175 print(f" KL Loss: {avg_kl:.4f}")
176 print()
177
178 print("=" * 80)
179 print("TRAINING COMPLETE")
180 print("=" * 80)
181 print()
182
183 return history
184
185
186# ============================================================================
187# PART 4: Data Loading
188# ============================================================================
189
190# Prepare MNIST dataset
191transform = transforms.Compose([
192 transforms.ToTensor(),
193])
194
195train_dataset = torchvision.datasets.MNIST(
196 root='./data',
197 train=True,
198 transform=transform,
199 download=True
200)
201
202train_loader = DataLoader(
203 train_dataset,
204 batch_size=128,
205 shuffle=True,
206 num_workers=2
207)
208
209print("Dataset:")
210print(f" Training samples: {len(train_dataset)}")
211print(f" Batch size: 128")
212print(f" Batches: {len(train_loader)}")
213print()
214
215# ============================================================================
216# PART 5: Initialize and Train
217# ============================================================================
218
219# Create model
220device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
221print(f"Device: {device}")
222print()
223
224vae_model = MNIST_VAE(
225 input_dim=784,
226 hidden_dim=400,
227 latent_dim=20
228)
229
230print("Model Architecture:")
231print(f" Input: 784 (28×28)")
232print(f" Hidden: 400")
233print(f" Latent: 20")
234print(f" Parameters: {sum(p.numel() for p in vae_model.parameters()):,}")
235print()
236
237# Train model
238history = train_vae(
239 vae_model,
240 train_loader,
241 epochs=20,
242 lr=1e-3,
243 beta=1.0,
244 device=device
245)
246
247# ============================================================================
248# PART 6: Visualization
249# ============================================================================
250
251print("Generating visualizations...")
252print()
253
254# Plot 1: Training curves
255fig, axes = plt.subplots(1, 3, figsize=(15, 4))
256
257axes[0].plot(history['loss'], linewidth=2, color='blue')
258axes[0].set_xlabel('Epoch', fontweight='bold')
259axes[0].set_ylabel('Total Loss', fontweight='bold')
260axes[0].set_title('VAE Training Loss', fontweight='bold')
261axes[0].grid(True, alpha=0.3)
262
263axes[1].plot(history['recon_loss'], linewidth=2, color='green')
264axes[1].set_xlabel('Epoch', fontweight='bold')
265axes[1].set_ylabel('Reconstruction Loss', fontweight='bold')
266axes[1].set_title('Reconstruction Quality', fontweight='bold')
267axes[1].grid(True, alpha=0.3)
268
269axes[2].plot(history['kl_loss'], linewidth=2, color='red')
270axes[2].set_xlabel('Epoch', fontweight='bold')
271axes[2].set_ylabel('KL Divergence', fontweight='bold')
272axes[2].set_title('Latent Space Regularization', fontweight='bold')
273axes[2].grid(True, alpha=0.3)
274
275plt.tight_layout()
276plt.show()
277
278# Plot 2: Generated samples
279if len(history['generated_samples']) > 0:
280 final_samples = history['generated_samples'][-1][:64]
281
282 fig, axes = plt.subplots(8, 8, figsize=(12, 12))
283 for i, ax in enumerate(axes.flat):
284 img = final_samples[i].reshape(28, 28)
285 ax.imshow(img, cmap='gray')
286 ax.axis('off')
287
288 plt.suptitle('VAE Generated Digits', fontsize=16, fontweight='bold')
289 plt.tight_layout()
290 plt.show()
291
292# Plot 3: Reconstruction comparison
293vae_model.eval()
294with torch.no_grad():
295 test_images = next(iter(train_loader))[0][:8].view(-1, 784).to(device)
296 recon_images, _, _ = vae_model(test_images)
297
298 test_images = test_images.cpu().reshape(-1, 28, 28)
299 recon_images = recon_images.cpu().reshape(-1, 28, 28)
300
301 fig, axes = plt.subplots(2, 8, figsize=(16, 4))
302
303 for i in range(8):
304 # Original
305 axes[0, i].imshow(test_images[i], cmap='gray')
306 axes[0, i].axis('off')
307 if i == 0:
308 axes[0, i].set_title('Original', fontweight='bold')
309
310 # Reconstructed
311 axes[1, i].imshow(recon_images[i], cmap='gray')
312 axes[1, i].axis('off')
313 if i == 0:
314 axes[1, i].set_title('Reconstructed', fontweight='bold')
315
316 plt.tight_layout()
317 plt.show()1print("=" * 80)
2print("VAE TRAINING ANALYSIS")
3print("=" * 80)
4print()
5
6print("Loss Components:")
7print("-" * 80)
8print(f"Final Total Loss: {history['loss'][-1]:.4f}")
9print(f"Final Recon Loss: {history['recon_loss'][-1]:.4f}")
10print(f"Final KL Loss: {history['kl_loss'][-1]:.4f}")
11print()
12
13print("What the Losses Tell Us:")
14print("-" * 80)
15print("1. TOTAL LOSS (decreasing):")
16print(" ✓ Model is learning")
17print(" ✓ ELBO is being maximized")
18print()
19
20print("2. RECONSTRUCTION LOSS:")
21print(" • Measures pixel-wise accuracy")
22print(" • Lower = better reconstructions")
23print(" • Should decrease steadily")
24print()
25
26print("3. KL DIVERGENCE:")
27print(" • Measures latent space regularization")
28print(" • Pushes q(z|x) toward N(0, 1)")
29print(" • Should stabilize at reasonable value")
30print(" • Too low → posterior collapse (not learning useful representations)")
31print(" • Too high → ignoring latent space")
32print()
33
34print("Healthy Training Indicators:")
35print("-" * 80)
36print("✓ Reconstruction loss decreases")
37print("✓ KL loss stabilizes (not zero, not huge)")
38print("✓ Generated samples look like digits")
39print("✓ Reconstructions preserve digit identity")1print("=" * 80)
2print("COMMON VAE TRAINING ISSUES")
3print("=" * 80)
4print()
5
6print("ISSUE 1: POSTERIOR COLLAPSE")
7print("-" * 40)
8print("Symptom:")
9print(" • KL divergence → 0")
10print(" • Model ignores latent code")
11print(" • Decoder generates same output for all z")
12print()
13print("Solutions:")
14print(" 1. KL annealing: Start with β = 0, gradually increase to 1")
15print(" 2. Reduce decoder capacity")
16print(" 3. Use free bits technique")
17print(" 4. Increase latent dimension")
18print()
19
20print("ISSUE 2: BLURRY RECONSTRUCTIONS")
21print("-" * 40)
22print("Symptom:")
23print(" • Reconstructions look blurry/averaged")
24print(" • Missing fine details")
25print()
26print("Solutions:")
27print(" 1. Use perceptual loss (not just MSE/BCE)")
28print(" 2. Increase latent dimension")
29print(" 3. Try β < 1 (prioritize reconstruction)")
30print(" 4. Use hierarchical VAE")
31print()
32
33print("ISSUE 3: UNSTABLE TRAINING")
34print("-" * 40)
35print("Symptom:")
36print(" • Losses oscillate wildly")
37print(" • NaN values appear")
38print()
39print("Solutions:")
40print(" 1. Lower learning rate")
41print(" 2. Clip gradients")
42print(" 3. Check for log(0) in KL computation")
43print(" 4. Add small epsilon to variance: σ² + 1e-8")
44print()
45
46print("ISSUE 4: MODE COLLAPSE (in generation)")
47print("-" * 40)
48print("Symptom:")
49print(" • Generated samples lack diversity")
50print(" • All samples look similar")
51print()
52print("Solutions:")
53print(" 1. Increase latent dimension")
54print(" 2. Use β > 1 for disentanglement")
55print(" 3. Add diversity loss")
56print(" 4. Check KL is not too high")VAE Training Dynamics:
Monitoring Checklist:
✓ Plot reconstruction loss (should decrease)
✓ Plot KL divergence (should stabilize, not collapse)
✓ Visualize reconstructions (should preserve content)
✓ Visualize generations (should look realistic)
✓ Check latent space (2D plot if latent_dim=2)
Hyperparameter Tuning:
| Parameter | Impact |
|-----------|--------|
| Latent dim | Higher → more capacity, less regularization |
| Beta | Higher → more disentanglement, worse reconstruction |
| Learning rate | Lower → more stable, slower convergence |
| Hidden dim | Higher → more capacity, risk of overfitting |
Training VAEs requires balancing reconstruction quality with latent space structure—the ELBO objective elegantly captures this trade-off!
VAEs learn smooth, continuous latent spaces that enable powerful applications: interpolation, generation, disentanglement, and more. Let's explore the magic of the learned latent space.
Smoothly transition between data points
1import torch
2import matplotlib.pyplot as plt
3import numpy as np
4
5print("=" * 80)
6print("LATENT SPACE INTERPOLATION")
7print("=" * 80)
8print()
9
10def interpolate_latent(vae, x1, x2, n_steps=10, device='cpu'):
11 """
12 Interpolate between two images in latent space
13
14 Process:
15 1. Encode both images to latent space
16 2. Linearly interpolate between latent codes
17 3. Decode interpolated codes
18 """
19 vae.eval()
20
21 with torch.no_grad():
22 # Encode to latent space
23 mu1, _ = vae.encode(x1)
24 mu2, _ = vae.encode(x2)
25
26 # Create interpolation path
27 alphas = torch.linspace(0, 1, n_steps)
28
29 interpolated_images = []
30 for alpha in alphas:
31 # Linear interpolation: z = (1-α)z₁ + αz₂
32 z = (1 - alpha) * mu1 + alpha * mu2
33
34 # Decode
35 img = vae.decode(z)
36 interpolated_images.append(img)
37
38 return torch.cat(interpolated_images, dim=0)
39
40
41# Example: Interpolate between two digits
42print("Interpolating between two MNIST digits:")
43print("-" * 80)
44print()
45
46# Get two sample images
47sample_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)
48x1, x2 = next(iter(sample_loader))[0].view(2, -1).to(device)
49
50# Interpolate
51interpolated = interpolate_latent(vae_model, x1, x2, n_steps=10, device=device)
52
53# Visualize
54fig, axes = plt.subplots(1, 10, figsize=(20, 2))
55for i in range(10):
56 img = interpolated[i].cpu().reshape(28, 28)
57 axes[i].imshow(img, cmap='gray')
58 axes[i].axis('off')
59 axes[i].set_title(f'α={i/9:.1f}', fontsize=9)
60
61plt.suptitle('Smooth Latent Space Interpolation', fontsize=14, fontweight='bold')
62plt.tight_layout()
63plt.show()
64
65print("Key Observation:")
66print(" ✓ Transition is smooth and continuous")
67print(" ✓ Intermediate images are valid digits")
68print(" ✓ Latent space is well-structured")Sample from N(0, 1) to generate new data
1print("\n" + "=" * 80)
2print("GENERATING NEW SAMPLES FROM PRIOR")
3print("=" * 80)
4print()
5
6def generate_from_prior(vae, n_samples=64, device='cpu'):
7 """Generate new samples by sampling from prior p(z) = N(0, I)"""
8 vae.eval()
9
10 with torch.no_grad():
11 # Sample from standard normal
12 z = torch.randn(n_samples, vae.latent_dim).to(device)
13
14 # Decode
15 samples = vae.decode(z)
16
17 return samples
18
19
20# Generate 64 new digits
21print("Generating 64 new digits from prior:")
22print("-" * 80)
23print()
24
25generated = generate_from_prior(vae_model, n_samples=64, device=device)
26
27# Visualize
28fig, axes = plt.subplots(8, 8, figsize=(12, 12))
29for i, ax in enumerate(axes.flat):
30 img = generated[i].cpu().reshape(28, 28)
31 ax.imshow(img, cmap='gray')
32 ax.axis('off')
33
34plt.suptitle('VAE Generated Digits (sampled from N(0, 1))',
35 fontsize=16, fontweight='bold')
36plt.tight_layout()
37plt.show()
38
39print("Why This Works:")
40print(" • Training regularizes latent space to N(0, 1)")
41print(" • Decoder learns p(x|z) for z ~ N(0, 1)")
42print(" • Random z → valid digit")Perform operations in latent space
1print("\n" + "=" * 80)
2print("LATENT SPACE ARITHMETIC")
3print("=" * 80)
4print()
5
6def latent_arithmetic(vae, images, device='cpu'):
7 """
8 Demonstrate latent space arithmetic
9
10 Example: z_a - z_b + z_c = ?
11 """
12 vae.eval()
13
14 with torch.no_grad():
15 # Encode images
16 mus, _ = vae.encode(images.to(device))
17
18 # Latent arithmetic
19 z_a, z_b, z_c = mus[0], mus[1], mus[2]
20
21 # Operation: "remove b, add c"
22 z_result = z_a - z_b + z_c
23
24 # Decode all
25 results = []
26 for z in [z_a, z_b, z_c, z_result]:
27 img = vae.decode(z.unsqueeze(0))
28 results.append(img)
29
30 return results
31
32
33print("Latent Arithmetic Example:")
34print("-" * 80)
35print("Operation: z_A - z_B + z_C = ?")
36print()
37
38# Get three images
39sample_images = next(iter(sample_loader))[0][:3].view(3, -1)
40results = latent_arithmetic(vae_model, sample_images, device)
41
42# Visualize
43fig, axes = plt.subplots(1, 4, figsize=(12, 3))
44titles = ['A', 'B', 'C', 'A - B + C']
45
46for i, (ax, img) in enumerate(zip(axes, results)):
47 img = img.cpu().reshape(28, 28)
48 ax.imshow(img, cmap='gray')
49 ax.axis('off')
50 ax.set_title(titles[i], fontweight='bold', fontsize=12)
51
52plt.tight_layout()
53plt.show()
54
55print("Interpretation:")
56print(" • A - B removes features of B from A")
57print(" • Adding C incorporates C's features")
58print(" • Result combines characteristics")Train VAE with 2D latent space for visualization
1print("\n" + "=" * 80)
2print("2D LATENT SPACE VISUALIZATION")
3print("=" * 80)
4print()
5
6# Train VAE with latent_dim=2
7vae_2d = MNIST_VAE(input_dim=784, hidden_dim=400, latent_dim=2).to(device)
8
9print("Training 2D VAE for visualization...")
10print("(This allows us to see the entire latent space)")
11print()
12
13# Quick training (fewer epochs for demo)
14history_2d = train_vae(vae_2d, train_loader, epochs=10, device=device)
15
16# Encode test set to 2D latent space
17vae_2d.eval()
18latent_codes = []
19labels = []
20
21with torch.no_grad():
22 for images, targets in train_loader:
23 images = images.view(-1, 784).to(device)
24 mu, _ = vae_2d.encode(images)
25 latent_codes.append(mu.cpu())
26 labels.append(targets)
27
28latent_codes = torch.cat(latent_codes, dim=0).numpy()
29labels = torch.cat(labels, dim=0).numpy()
30
31# Plot latent space
32fig, ax = plt.subplots(1, 1, figsize=(10, 10))
33
34scatter = ax.scatter(
35 latent_codes[:5000, 0],
36 latent_codes[:5000, 1],
37 c=labels[:5000],
38 cmap='tab10',
39 alpha=0.5,
40 s=1
41)
42
43ax.set_xlabel('Latent Dimension 1', fontweight='bold', fontsize=12)
44ax.set_ylabel('Latent Dimension 2', fontweight='bold', fontsize=12)
45ax.set_title('2D Latent Space (colored by digit class)',
46 fontweight='bold', fontsize=14)
47plt.colorbar(scatter, label='Digit Class')
48plt.tight_layout()
49plt.show()
50
51print("Observations:")
52print(" ✓ Different digits cluster in latent space")
53print(" ✓ Smooth transitions between clusters")
54print(" ✓ Continuous manifold learned")Explore latent dimensions
1print("\n" + "=" * 80)
2print("LATENT SPACE TRAVERSAL")
3print("=" * 80)
4print()
5
6def traverse_latent_dimension(vae, dim, range_vals=(-3, 3), n_steps=10, device='cpu'):
7 """
8 Traverse a single latent dimension
9
10 Fix all dimensions except one, vary that dimension
11 """
12 vae.eval()
13
14 with torch.no_grad():
15 # Start from zero latent code
16 z_base = torch.zeros(1, vae.latent_dim).to(device)
17
18 # Vary one dimension
19 values = torch.linspace(range_vals[0], range_vals[1], n_steps)
20
21 images = []
22 for val in values:
23 z = z_base.clone()
24 z[0, dim] = val
25
26 img = vae.decode(z)
27 images.append(img)
28
29 return torch.cat(images, dim=0)
30
31
32print("Traversing Latent Dimension 0:")
33print("-" * 80)
34
35# Traverse first dimension
36traversal = traverse_latent_dimension(vae_model, dim=0, device=device)
37
38# Visualize
39fig, axes = plt.subplots(1, 10, figsize=(20, 2))
40for i in range(10):
41 img = traversal[i].cpu().reshape(28, 28)
42 axes[i].imshow(img, cmap='gray')
43 axes[i].axis('off')
44 axes[i].set_title(f'z₀={-3 + i*0.67:.1f}', fontsize=9)
45
46plt.suptitle('Latent Dimension Traversal', fontsize=14, fontweight='bold')
47plt.tight_layout()
48plt.show()
49
50print("Interpretation:")
51print(" • Each latent dimension captures a variation factor")
52print(" • Traversing shows what that dimension controls")
53print(" • In disentangled VAEs, each dim = independent feature")1print("\n" + "=" * 80)
2print("VAE APPLICATIONS")
3print("=" * 80)
4print()
5
6print("1. DATA GENERATION")
7print("-" * 40)
8print("• Generate synthetic training data")
9print("• Augment limited datasets")
10print("• Create diverse examples")
11print()
12
13print("2. ANOMALY DETECTION")
14print("-" * 40)
15print("• Train on normal data")
16print("• High reconstruction error → anomaly")
17print("• Applications: fraud detection, quality control")
18print()
19
20print("3. DIMENSIONALITY REDUCTION")
21print("-" * 40)
22print("• Compress high-dimensional data")
23print("• Alternative to PCA, t-SNE")
24print("• Preserves non-linear structure")
25print()
26
27print("4. DATA DENOISING")
28print("-" * 40)
29print("• Encode noisy input")
30print("• Decode to clean output")
31print("• Works because latent space is regularized")
32print()
33
34print("5. SEMI-SUPERVISED LEARNING")
35print("-" * 40)
36print("• Use unlabeled data for generative modeling")
37print("• Combine with classifier")
38print("• Improves generalization")
39print()
40
41print("6. REPRESENTATION LEARNING")
42print("-" * 40)
43print("• Learn meaningful features")
44print("• Transfer to downstream tasks")
45print("• Disentangled representations")
46print()
47
48print("7. CREATIVE APPLICATIONS")
49print("-" * 40)
50print("• Image editing (change specific attributes)")
51print("• Style transfer")
52print("• Conditional generation")
53print("• Interpolation for animation")Why Latent Space Manipulation Works:
VAE vs GAN Latent Spaces:
| Aspect | VAE | GAN |
|--------|-----|-----|
| Structure | Explicitly regularized (N(0,1)) | Implicitly learned |
| Smoothness | Very smooth (ELBO enforces) | Can have discontinuities |
| Generation quality | Slightly blurry | Sharp, realistic |
| Interpolation | Excellent | Good but less principled |
| Training stability | Stable | Can be unstable |
Best Practices:
✓ Use higher β for disentanglement
✓ Visualize latent space (2D projection)
✓ Check interpolations for smoothness
✓ Test generation quality from prior
✓ Explore different latent dimensions
VAEs provide a principled way to learn structured latent spaces, enabling powerful applications in generation, manipulation, and representation learning!
Implement a complete Variational Autoencoder with reparameterization trick, train with ELBO loss, generate new digits, and explore the continuous latent space through interpolation and traversal.
1# Complete VAE Implementation with Latent Space Exploration
2# Generate new digits and explore smooth latent space
3
4import torch
5import torch.nn as nn
6import torch.optim as optim
7import torchvision
8from torchvision import transforms
9from torch.utils.data import DataLoader
10import matplotlib.pyplot as plt
11import numpy as np
12
13print("=" * 80)
14print("VARIATIONAL AUTOENCODER: GENERATION & LATENT SPACE EXPLORATION")
15print("=" * 80)
16print()
17
18# Set device
19device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
20print(f"Using device: {device}")
21print()
22
23# ============================================================================
24# PART 1: VAE Model with Reparameterization Trick
25# ============================================================================
26
27class VAE(nn.Module):
28 """
29 Variational Autoencoder for MNIST
30
31 Architecture:
32 - Encoder: 784 → 400 → (μ, log σ²) [latent_dim]
33 - Reparameterization: z = μ + σ * ε
34 - Decoder: latent_dim → 400 → 784
35 """
36
37 def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20):
38 super(VAE, self).__init__()
39
40 self.latent_dim = latent_dim
41
42 # Encoder
43 self.encoder = nn.Sequential(
44 nn.Linear(input_dim, hidden_dim),
45 nn.ReLU(),
46 nn.Linear(hidden_dim, hidden_dim),
47 nn.ReLU()
48 )
49
50 # Latent distribution parameters
51 self.fc_mu = nn.Linear(hidden_dim, latent_dim)
52 self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
53
54 # Decoder
55 self.decoder = nn.Sequential(
56 nn.Linear(latent_dim, hidden_dim),
57 nn.ReLU(),
58 nn.Linear(hidden_dim, hidden_dim),
59 nn.ReLU(),
60 nn.Linear(hidden_dim, input_dim),
61 nn.Sigmoid()
62 )
63
64 def encode(self, x):
65 """Encode to latent distribution parameters"""
66 h = self.encoder(x)
67 mu = self.fc_mu(h)
68 logvar = self.fc_logvar(h)
69 return mu, logvar
70
71 def reparameterize(self, mu, logvar):
72 """
73 Reparameterization trick: z = μ + σ * ε
74
75 Allows backpropagation through stochastic sampling
76 """
77 std = torch.exp(0.5 * logvar) # σ = exp(0.5 * log(σ²))
78 eps = torch.randn_like(std) # ε ~ N(0, 1)
79 z = mu + eps * std # z ~ N(μ, σ²)
80 return z
81
82 def decode(self, z):
83 """Decode latent code to reconstruction"""
84 return self.decoder(z)
85
86 def forward(self, x):
87 mu, logvar = self.encode(x)
88 z = self.reparameterize(mu, logvar)
89 x_recon = self.decode(z)
90 return x_recon, mu, logvar
91
92 def generate(self, n_samples=64):
93 """Generate new samples from prior N(0, I)"""
94 with torch.no_grad():
95 z = torch.randn(n_samples, self.latent_dim).to(device)
96 samples = self.decode(z)
97 return samples
98
99
100# Initialize VAE
101latent_dim = 20
102hidden_dim = 400
103
104vae = VAE(input_dim=784, hidden_dim=hidden_dim, latent_dim=latent_dim).to(device)
105
106print("VAE Architecture:")
107print("-" * 80)
108print(f"Input dimension: 784 (28×28)")
109print(f"Hidden dimension: {hidden_dim}")
110print(f"Latent dimension: {latent_dim}")
111print(f"Total parameters: {sum(p.numel() for p in vae.parameters()):,}")
112print()
113
114# ============================================================================
115# PART 2: VAE Loss Function (ELBO)
116# ============================================================================
117
118def vae_loss(x_recon, x, mu, logvar, beta=1.0):
119 """
120 VAE Loss: Reconstruction + β * KL Divergence
121
122 ELBO = E[log p(x|z)] - β * KL(q(z|x) || p(z))
123 """
124
125 # Reconstruction loss (binary cross-entropy)
126 recon_loss = nn.functional.binary_cross_entropy(
127 x_recon, x, reduction='sum'
128 )
129
130 # KL divergence: KL(N(μ, σ²) || N(0, 1))
131 # Closed form: -0.5 * Σ(1 + log(σ²) - μ² - σ²)
132 kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
133
134 # Total loss
135 total_loss = recon_loss + beta * kl_loss
136
137 return total_loss, recon_loss, kl_loss
138
139
140print("VAE Loss Function:")
141print("-" * 80)
142print("ELBO = Reconstruction Loss + β * KL Divergence")
143print()
144print("Component 1: Reconstruction Loss")
145print(" • Binary cross-entropy between input and reconstruction")
146print(" • Encourages accurate reconstruction")
147print()
148print("Component 2: KL Divergence")
149print(" • Distance between q(z|x) and prior p(z) = N(0, 1)")
150print(" • Regularizes latent space to be continuous")
151print(" • β = 1.0 (standard VAE)")
152print()
153
154# ============================================================================
155# PART 3: Load MNIST Dataset
156# ============================================================================
157
158transform = transforms.Compose([
159 transforms.ToTensor(),
160])
161
162train_dataset = torchvision.datasets.MNIST(
163 root='./data',
164 train=True,
165 transform=transform,
166 download=True
167)
168
169train_loader = DataLoader(
170 train_dataset,
171 batch_size=128,
172 shuffle=True,
173 num_workers=2
174)
175
176print("Dataset:")
177print("-" * 80)
178print(f"Training samples: {len(train_dataset)}")
179print(f"Batch size: 128")
180print(f"Batches per epoch: {len(train_loader)}")
181print()
182
183# ============================================================================
184# PART 4: Training Loop
185# ============================================================================
186
187def train_vae(model, dataloader, epochs=20, lr=1e-3, beta=1.0):
188 """Train VAE with ELBO objective"""
189
190 optimizer = optim.Adam(model.parameters(), lr=lr)
191
192 history = {
193 'total_loss': [],
194 'recon_loss': [],
195 'kl_loss': [],
196 'generated_samples': []
197 }
198
199 print("=" * 80)
200 print("TRAINING VAE")
201 print("=" * 80)
202 print()
203
204 for epoch in range(epochs):
205 model.train()
206 epoch_total = 0
207 epoch_recon = 0
208 epoch_kl = 0
209
210 for batch_idx, (images, _) in enumerate(dataloader):
211 images = images.view(-1, 784).to(device)
212
213 # Forward pass
214 x_recon, mu, logvar = model(images)
215
216 # Compute loss
217 total, recon, kl = vae_loss(x_recon, images, mu, logvar, beta)
218
219 # Backward pass
220 optimizer.zero_grad()
221 total.backward()
222 optimizer.step()
223
224 # Accumulate
225 epoch_total += total.item()
226 epoch_recon += recon.item()
227 epoch_kl += kl.item()
228
229 # Average losses
230 n = len(dataloader.dataset)
231 avg_total = epoch_total / n
232 avg_recon = epoch_recon / n
233 avg_kl = epoch_kl / n
234
235 history['total_loss'].append(avg_total)
236 history['recon_loss'].append(avg_recon)
237 history['kl_loss'].append(avg_kl)
238
239 # Generate samples
240 if (epoch + 1) % 5 == 0:
241 samples = model.generate(64)
242 history['generated_samples'].append(samples.cpu())
243
244 # Print progress
245 print(f"Epoch [{epoch+1}/{epochs}]")
246 print(f" Total: {avg_total:.4f} | Recon: {avg_recon:.4f} | KL: {avg_kl:.4f}")
247
248 print()
249 print("=" * 80)
250 print("TRAINING COMPLETE")
251 print("=" * 80)
252 print()
253
254 return history
255
256
257# Train VAE
258epochs = 20
259learning_rate = 1e-3
260beta = 1.0
261
262print(f"Training Configuration:")
263print(f" Epochs: {epochs}")
264print(f" Learning rate: {learning_rate}")
265print(f" Beta: {beta}")
266print()
267
268history = train_vae(vae, train_loader, epochs=epochs, lr=learning_rate, beta=beta)
269
270# ============================================================================
271# PART 5: Visualizations
272# ============================================================================
273
274print("Generating visualizations...")
275print()
276
277# Plot 1: Training curves
278fig, axes = plt.subplots(1, 3, figsize=(15, 4))
279
280axes[0].plot(history['total_loss'], linewidth=2, color='blue')
281axes[0].set_xlabel('Epoch', fontweight='bold')
282axes[0].set_ylabel('Loss', fontweight='bold')
283axes[0].set_title('Total Loss (ELBO)', fontweight='bold')
284axes[0].grid(True, alpha=0.3)
285
286axes[1].plot(history['recon_loss'], linewidth=2, color='green')
287axes[1].set_xlabel('Epoch', fontweight='bold')
288axes[1].set_ylabel('Loss', fontweight='bold')
289axes[1].set_title('Reconstruction Loss', fontweight='bold')
290axes[1].grid(True, alpha=0.3)
291
292axes[2].plot(history['kl_loss'], linewidth=2, color='red')
293axes[2].set_xlabel('Epoch', fontweight='bold')
294axes[2].set_ylabel('KL Divergence', fontweight='bold')
295axes[2].set_title('Latent Space Regularization', fontweight='bold')
296axes[2].grid(True, alpha=0.3)
297
298plt.tight_layout()
299plt.show()
300
301# Plot 2: Generated samples
302if len(history['generated_samples']) > 0:
303 final_samples = history['generated_samples'][-1][:64]
304
305 fig, axes = plt.subplots(8, 8, figsize=(12, 12))
306 for i, ax in enumerate(axes.flat):
307 img = final_samples[i].reshape(28, 28)
308 ax.imshow(img, cmap='gray')
309 ax.axis('off')
310
311 plt.suptitle('VAE Generated Digits (sampled from N(0, 1))',
312 fontsize=16, fontweight='bold')
313 plt.tight_layout()
314 plt.show()
315
316# Plot 3: Reconstruction comparison
317vae.eval()
318with torch.no_grad():
319 test_images = next(iter(train_loader))[0][:8].view(-1, 784).to(device)
320 recon_images, _, _ = vae(test_images)
321
322 test_images = test_images.cpu().reshape(-1, 28, 28)
323 recon_images = recon_images.cpu().reshape(-1, 28, 28)
324
325 fig, axes = plt.subplots(2, 8, figsize=(16, 4))
326
327 for i in range(8):
328 axes[0, i].imshow(test_images[i], cmap='gray')
329 axes[0, i].axis('off')
330 if i == 0:
331 axes[0, i].set_ylabel('Original', fontweight='bold', fontsize=12)
332
333 axes[1, i].imshow(recon_images[i], cmap='gray')
334 axes[1, i].axis('off')
335 if i == 0:
336 axes[1, i].set_ylabel('Reconstructed', fontweight='bold', fontsize=12)
337
338 plt.tight_layout()
339 plt.show()
340
341# ============================================================================
342# PART 6: Latent Space Exploration
343# ============================================================================
344
345print("=" * 80)
346print("LATENT SPACE EXPLORATION")
347print("=" * 80)
348print()
349
350# Interpolation
351print("1. Latent Space Interpolation:")
352print("-" * 80)
353
354vae.eval()
355with torch.no_grad():
356 # Get two images
357 sample_images = next(iter(train_loader))[0][:2].view(2, -1).to(device)
358
359 # Encode
360 mu1, _ = vae.encode(sample_images[0:1])
361 mu2, _ = vae.encode(sample_images[1:2])
362
363 # Interpolate
364 n_steps = 10
365 alphas = torch.linspace(0, 1, n_steps)
366
367 interpolated = []
368 for alpha in alphas:
369 z = (1 - alpha) * mu1 + alpha * mu2
370 img = vae.decode(z)
371 interpolated.append(img)
372
373 interpolated = torch.cat(interpolated, dim=0).cpu()
374
375 # Visualize
376 fig, axes = plt.subplots(1, 10, figsize=(20, 2))
377 for i in range(10):
378 img = interpolated[i].reshape(28, 28)
379 axes[i].imshow(img, cmap='gray')
380 axes[i].axis('off')
381 axes[i].set_title(f'α={alphas[i]:.1f}', fontsize=9)
382
383 plt.suptitle('Smooth Latent Space Interpolation',
384 fontsize=14, fontweight='bold')
385 plt.tight_layout()
386 plt.show()
387
388print(" ✓ Smooth transitions between digits")
389print(" ✓ All intermediate images are valid")
390print()
391
392# Latent dimension traversal
393print("2. Latent Dimension Traversal:")
394print("-" * 80)
395
396with torch.no_grad():
397 # Fix all dims except first, vary first dimension
398 z_base = torch.zeros(1, latent_dim).to(device)
399
400 n_steps = 10
401 values = torch.linspace(-3, 3, n_steps)
402
403 traversal = []
404 for val in values:
405 z = z_base.clone()
406 z[0, 0] = val # Vary first dimension
407 img = vae.decode(z)
408 traversal.append(img)
409
410 traversal = torch.cat(traversal, dim=0).cpu()
411
412 # Visualize
413 fig, axes = plt.subplots(1, 10, figsize=(20, 2))
414 for i in range(10):
415 img = traversal[i].reshape(28, 28)
416 axes[i].imshow(img, cmap='gray')
417 axes[i].axis('off')
418 axes[i].set_title(f'z₀={values[i]:.1f}', fontsize=9)
419
420 plt.suptitle('Latent Dimension 0 Traversal',
421 fontsize=14, fontweight='bold')
422 plt.tight_layout()
423 plt.show()
424
425print(" ✓ Each dimension captures variation")
426print()
427
428# ============================================================================
429# PART 7: Analysis
430# ============================================================================
431
432print("=" * 80)
433print("VAE ANALYSIS")
434print("=" * 80)
435print()
436
437print("Final Metrics:")
438print("-" * 80)
439print(f"Total Loss: {history['total_loss'][-1]:.4f}")
440print(f"Reconstruction Loss: {history['recon_loss'][-1]:.4f}")
441print(f"KL Divergence: {history['kl_loss'][-1]:.4f}")
442print()
443
444print("Key Observations:")
445print("-" * 80)
446print("1️⃣ Reconstruction quality: Digits are recognizable")
447print("2️⃣ Generation quality: Samples from prior look realistic")
448print("3️⃣ Latent space structure: Smooth interpolations")
449print("4️⃣ KL divergence: Balanced (not collapsed, not ignored)")
450print()
451
452print("VAE vs GAN:")
453print("-" * 80)
454print("VAE Advantages:")
455print(" ✓ Stable training (no adversarial dynamics)")
456print(" ✓ Principled probabilistic framework")
457print(" ✓ Smooth latent space (explicit regularization)")
458print(" ✓ Direct optimization of likelihood lower bound")
459print()
460
461print("GAN Advantages:")
462print(" ✓ Sharper images (adversarial training)")
463print(" ✓ Better perceptual quality")
464print()
465
466print("=" * 80)
467print("🎲 VAEs: Probabilistic Generative Models!")
468print("=" * 80)
469print()
470
471print("What you've accomplished:")
472print(" ✓ Built VAE with reparameterization trick")
473print(" ✓ Trained with ELBO (reconstruction + KL divergence)")
474print(" ✓ Generated new digits by sampling from prior")
475print(" ✓ Explored smooth latent space interpolation")
476print(" ✓ Analyzed latent dimensions and their effects")
477print()
478
479print("Next steps:")
480print(" → Try β-VAE for disentanglement (β > 1)")
481print(" → Implement conditional VAE (CVAE)")
482print(" → Explore hierarchical VAEs")
483print(" → Apply to other datasets (faces, images)")
484print()
485
486print("Ready to master Reinforcement Learning? 🤖")This comprehensive VAE implementation demonstrates the complete pipeline for probabilistic generative modeling.
1. VAE Architecture with Reparameterization Trick
Encoder (probabilistic):
Reparameterization:
Decoder (probabilistic):
2. ELBO Loss Function
Two Components:
1ELBO = Reconstruction Loss + KL Divergence
2
3Reconstruction: How well we reconstruct input
4KL: How close q(z|x) is to prior p(z) = N(0, 1)Why Both Matter:
3. Training Dynamics
Loss Curves:
Healthy Training:
4. Latent Space Exploration
Generation from Prior:
Smooth Interpolation:
Latent Traversal:
Why VAE Generation Works:
VAE vs Autoencoder:
| Aspect | Autoencoder | VAE |
|--------|-------------|-----|
| Encoder output | Single point z | Distribution (μ, σ²) |
| Latent space | Irregular | Continuous Gaussian |
| Loss | Reconstruction | Recon + KL |
| Generation | Poor | Good |
| Training | Deterministic | Stochastic |
VAE vs GAN:
| Aspect | VAE | GAN |
|--------|-----|-----|
| Training | Stable | Can be unstable |
| Image quality | Slightly blurry | Sharp |
| Latent space | Explicitly structured | Implicitly learned |
| Framework | Probabilistic | Adversarial |
Posterior Collapse (KL → 0):
Blurry Reconstructions:
Unstable Training:
ELBO Derivation:
1log p(x) = ELBO + KL(q(z|x) || p(z|x))
2 ≥ ELBO (since KL ≥ 0)
3
4ELBO = E_q(z|x)[log p(x|z)] - KL(q(z|x) || p(z))
5 = Reconstruction - KL DivergenceWhy Maximize ELBO:
Closed-Form KL:
For q = N(μ, σ²) and p = N(0, 1):
1KL(q || p) = -0.5 * Σ(1 + log(σ²) - μ² - σ²)This makes VAE training efficient!
Probabilistic Framework:
Applications:
Research Impact:
This VAE implementation shows how probabilistic modeling enables smooth, controllable generation—a different approach than GANs but equally powerful!
Outstanding work! You've mastered Variational Autoencoders—a elegant fusion of deep learning and probabilistic inference. You've learned how VAEs create structured latent spaces through variational inference, implemented the reparameterization trick, and explored smooth latent space manipulation.
Next: In Lesson 18: Reinforcement Learning Fundamentals, we shift gears to learn how agents learn from interaction:
From generative models to decision-making agents! Reinforcement learning powers game AI, robotics, and autonomous systems. Ready to train intelligent agents? Let's dive in! 🤖
Your progress is saved in this browser, on this device only.