> ## Documentation Index
> Fetch the complete documentation index at: https://resources.devweekends.com/llms.txt
> Use this file to discover all available pages before exploring further.
# Autoencoders & VAEs
> Master data compression and generative modeling with autoencoders and variational autoencoders
# Autoencoders & Variational Autoencoders
## The Bottleneck Concept
Imagine you need to describe a complex image using only 10 numbers. You'd have to capture the **essential features** and discard the noise. That's exactly what an autoencoder does.
Think of it like the game of Pictionary: you see a detailed photograph and must convey it using only a few quick strokes. Those strokes are your "latent representation" -- they can't capture every pixel, so they encode the most important structural features (shape, pose, dominant colors) and discard the rest (individual pixel noise, fine textures). The better your encoding, the more your partner can reconstruct the original scene from your sketch.
An autoencoder learns to:
1. **Compress** data into a lower-dimensional representation (encoding)
2. **Reconstruct** the original data from this compressed form (decoding)
The magic happens in the **bottleneck** -- a narrow layer that forces the network to learn efficient representations. If the bottleneck is too wide (say, the same dimension as the input), the network can simply memorize every input as-is -- an identity function. If it's too narrow, reconstructions will be blurry or miss important details. Finding the right bottleneck size is the fundamental design decision in autoencoders.
**Why not just use PCA?** PCA (Principal Component Analysis) is a linear autoencoder -- it finds the best linear projection to a lower-dimensional space. Neural network autoencoders generalize this to non-linear compressions, capturing curved manifolds in the data that PCA misses entirely. For complex data like images, the non-linear version recovers dramatically more information at the same compression ratio.
```python theme={null}
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from tqdm import tqdm
# Set random seeds
torch.manual_seed(42)
np.random.seed(42)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Load MNIST
transform = transforms.Compose([
transforms.ToTensor(),
])
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('./data', train=False, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)
print(f"Training samples: {len(train_dataset)}")
print(f"Test samples: {len(test_dataset)}")
```
***
## Standard Autoencoder
The basic autoencoder architecture:
| Component | Function | Shape (MNIST) |
| ---------------- | ------------------------------ | ------------------- |
| **Encoder** | Compress input to latent space | 784 → 128 → 64 → 32 |
| **Latent Space** | Compressed representation | 32 dimensions |
| **Decoder** | Reconstruct from latent space | 32 → 64 → 128 → 784 |
```python theme={null}
class Autoencoder(nn.Module):
"""
Standard Autoencoder with symmetric encoder-decoder.
The encoder and decoder are mirrors of each other -- this symmetry
isn't required, but it's a good starting point and ensures balanced
capacity on both sides of the bottleneck.
"""
def __init__(self, input_dim=784, latent_dim=32):
super().__init__()
self.latent_dim = latent_dim
# Encoder: input -> latent (gradually reduces dimensionality)
# 784 -> 256 -> 128 -> 64 -> 32: each layer compresses by ~2x
self.encoder = nn.Sequential(
nn.Linear(input_dim, 256),
nn.ReLU(), # Non-linearity lets us learn curved manifolds
nn.Linear(256, 128),
nn.ReLU(),
nn.Linear(128, 64),
nn.ReLU(),
nn.Linear(64, latent_dim) # No activation here: latent space is unconstrained
)
# Decoder: latent -> output (mirrors the encoder)
self.decoder = nn.Sequential(
nn.Linear(latent_dim, 64),
nn.ReLU(),
nn.Linear(64, 128),
nn.ReLU(),
nn.Linear(128, 256),
nn.ReLU(),
nn.Linear(256, input_dim),
nn.Sigmoid() # Sigmoid constrains output to [0, 1] to match pixel range
)
def encode(self, x):
"""Compress input to latent representation."""
return self.encoder(x)
def decode(self, z):
"""Reconstruct from latent representation."""
return self.decoder(z)
def forward(self, x):
z = self.encode(x)
return self.decode(z)
# Create autoencoder
autoencoder = Autoencoder(latent_dim=32).to(device)
print(f"Autoencoder parameters: {sum(p.numel() for p in autoencoder.parameters()):,}")
# Test forward pass
test_input = torch.randn(4, 784).to(device)
output = autoencoder(test_input)
print(f"Input shape: {test_input.shape}")
print(f"Output shape: {output.shape}")
```
**Output:**
```
Autoencoder parameters: 380,752
Input shape: torch.Size([4, 784])
Output shape: torch.Size([4, 784])
```
***
## Training the Autoencoder
The autoencoder is trained to minimize **reconstruction loss** -- the difference between input and output. Unlike supervised learning where we have labels, the autoencoder uses the input itself as the target. This is sometimes called "self-supervised" learning.
$$
\mathcal{L}_{recon} = \frac{1}{N} \sum_{i=1}^{N} \|x_i - \hat{x}_i\|^2
$$
Where:
* $x_i$ is the original input
* $\hat{x}_i = \text{Decode}(\text{Encode}(x_i))$ is the reconstructed output
**MSE vs BCE for reconstruction loss:** Use MSE (Mean Squared Error) when outputs are continuous or when the decoder has no activation (or a linear activation). Use BCE (Binary Cross-Entropy) when outputs are in \[0, 1] and the decoder uses a Sigmoid. For MNIST digits (pixel values 0 to 1), both work, but BCE often converges faster because it naturally handles the bounded output range and produces sharper reconstructions.
```python theme={null}
def train_autoencoder(model, train_loader, num_epochs=20, lr=1e-3):
"""
Train autoencoder with reconstruction loss.
"""
optimizer = optim.Adam(model.parameters(), lr=lr)
criterion = nn.MSELoss() # Mean Squared Error for reconstruction
losses = []
for epoch in range(num_epochs):
model.train()
epoch_loss = 0
for batch_idx, (data, _) in enumerate(train_loader):
# Flatten images
data = data.view(data.size(0), -1).to(device)
optimizer.zero_grad()
# Forward pass
reconstructed = model(data)
# Reconstruction loss
loss = criterion(reconstructed, data)
# Backward pass
loss.backward()
optimizer.step()
epoch_loss += loss.item()
avg_loss = epoch_loss / len(train_loader)
losses.append(avg_loss)
if (epoch + 1) % 5 == 0:
print(f"Epoch [{epoch+1}/{num_epochs}] | Loss: {avg_loss:.6f}")
return losses
# Train the autoencoder
print("Training autoencoder...")
losses = train_autoencoder(autoencoder, train_loader, num_epochs=20)
# Plot loss curve
plt.figure(figsize=(10, 4))
plt.plot(losses)
plt.xlabel('Epoch')
plt.ylabel('Reconstruction Loss')
plt.title('Autoencoder Training Loss')
plt.grid(True)
plt.savefig('autoencoder_loss.png')
plt.close()
```
**Output:**
```
Training autoencoder...
Epoch [5/20] | Loss: 0.024512
Epoch [10/20] | Loss: 0.018234
Epoch [15/20] | Loss: 0.015891
Epoch [20/20] | Loss: 0.014523
```
***
## Visualizing Reconstructions
Let's see how well our autoencoder reconstructs images:
```python theme={null}
def visualize_reconstructions(model, test_loader, n_samples=10):
"""
Compare original images with their reconstructions.
"""
model.eval()
# Get a batch of test images
data, labels = next(iter(test_loader))
data = data[:n_samples]
with torch.no_grad():
data_flat = data.view(data.size(0), -1).to(device)
reconstructed = model(data_flat)
reconstructed = reconstructed.view(-1, 1, 28, 28).cpu()
# Plot original vs reconstructed
fig, axes = plt.subplots(2, n_samples, figsize=(15, 3))
for i in range(n_samples):
# Original
axes[0, i].imshow(data[i].squeeze(), cmap='gray')
axes[0, i].axis('off')
if i == 0:
axes[0, i].set_title('Original', fontsize=12)
# Reconstructed
axes[1, i].imshow(reconstructed[i].squeeze(), cmap='gray')
axes[1, i].axis('off')
if i == 0:
axes[1, i].set_title('Reconstructed', fontsize=12)
plt.tight_layout()
plt.savefig('reconstructions.png')
plt.close()
print("Reconstructions saved to 'reconstructions.png'")
visualize_reconstructions(autoencoder, test_loader)
```
***
## Latent Space Visualization
The latent space is where the magic happens. Let's visualize it using t-SNE:
```python theme={null}
from sklearn.manifold import TSNE
def visualize_latent_space(model, test_loader, n_samples=3000):
"""
Visualize latent space using t-SNE.
"""
model.eval()
latent_vectors = []
labels_list = []
with torch.no_grad():
for data, labels in test_loader:
data = data.view(data.size(0), -1).to(device)
z = model.encode(data)
latent_vectors.append(z.cpu())
labels_list.append(labels)
if sum(len(l) for l in latent_vectors) >= n_samples:
break
# Concatenate all latent vectors
latent_vectors = torch.cat(latent_vectors, dim=0)[:n_samples].numpy()
labels_list = torch.cat(labels_list, dim=0)[:n_samples].numpy()
# Apply t-SNE
print("Running t-SNE (this may take a minute)...")
tsne = TSNE(n_components=2, random_state=42, perplexity=30)
latent_2d = tsne.fit_transform(latent_vectors)
# Plot
plt.figure(figsize=(12, 10))
scatter = plt.scatter(latent_2d[:, 0], latent_2d[:, 1],
c=labels_list, cmap='tab10', alpha=0.6, s=5)
plt.colorbar(scatter, label='Digit')
plt.xlabel('t-SNE Dimension 1')
plt.ylabel('t-SNE Dimension 2')
plt.title('Latent Space Visualization (t-SNE)')
plt.savefig('latent_space_tsne.png', dpi=150)
plt.close()
print("Latent space visualization saved!")
visualize_latent_space(autoencoder, test_loader)
```
***
## Convolutional Autoencoder
For images, convolutional autoencoders preserve spatial structure:
```python theme={null}
class ConvAutoencoder(nn.Module):
"""
Convolutional Autoencoder - better for images.
Uses Conv2d for encoding and ConvTranspose2d for decoding.
"""
def __init__(self, latent_dim=64):
super().__init__()
self.latent_dim = latent_dim
# Encoder: 28x28 -> 7x7 -> latent
self.encoder = nn.Sequential(
# 28x28 -> 14x14
nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
# 14x14 -> 7x7
nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
# Flatten
nn.Flatten(),
nn.Linear(64 * 7 * 7, latent_dim)
)
# Decoder: latent -> 7x7 -> 28x28
self.decoder_fc = nn.Linear(latent_dim, 64 * 7 * 7)
self.decoder_conv = nn.Sequential(
# 7x7 -> 14x14
nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
# 14x14 -> 28x28
nn.ConvTranspose2d(32, 1, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.Sigmoid()
)
def encode(self, x):
return self.encoder(x)
def decode(self, z):
x = self.decoder_fc(z)
x = x.view(-1, 64, 7, 7)
return self.decoder_conv(x)
def forward(self, x):
z = self.encode(x)
return self.decode(z)
# Create and test
conv_ae = ConvAutoencoder(latent_dim=64).to(device)
print(f"Conv Autoencoder parameters: {sum(p.numel() for p in conv_ae.parameters()):,}")
# Test with image batch
test_images = torch.randn(4, 1, 28, 28).to(device)
output = conv_ae(test_images)
print(f"Input shape: {test_images.shape}")
print(f"Output shape: {output.shape}")
```
**Output:**
```
Conv Autoencoder parameters: 285,793
Input shape: torch.Size([4, 1, 28, 28])
Output shape: torch.Size([4, 1, 28, 28])
```
***
## Denoising Autoencoder
A denoising autoencoder learns to remove noise from corrupted inputs. The key insight is subtle: by training the network to reconstruct clean data from noisy data, we force the encoder to learn the underlying **structure** of the data rather than memorizing surface-level details. Noise is random and unpredictable, so the only way to reconstruct the clean input is to learn what "normal" data looks like.
This is analogous to how humans learn to read messy handwriting: you don't memorize every possible scrawl, you learn the underlying structure of each letter, which lets you "denoise" any handwriting you encounter.
$$
\mathcal{L}_{denoise} = \|x - D(E(\tilde{x}))\|^2
$$
Where $\tilde{x} = x + \epsilon$ is the noisy input. Note that the loss is computed against the **clean** input $x$, not the noisy version.
```python theme={null}
class DenoisingAutoencoder(nn.Module):
"""
Denoising Autoencoder - learns to remove noise.
"""
def __init__(self, input_dim=784, latent_dim=64, noise_factor=0.3):
super().__init__()
self.noise_factor = noise_factor
self.encoder = nn.Sequential(
nn.Linear(input_dim, 256),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(256, 128),
nn.ReLU(),
nn.Linear(128, latent_dim)
)
self.decoder = nn.Sequential(
nn.Linear(latent_dim, 128),
nn.ReLU(),
nn.Linear(128, 256),
nn.ReLU(),
nn.Linear(256, input_dim),
nn.Sigmoid()
)
def add_noise(self, x):
"""Add Gaussian noise to input."""
noise = torch.randn_like(x) * self.noise_factor
noisy = x + noise
return torch.clamp(noisy, 0, 1) # Keep in valid range
def forward(self, x, add_noise=True):
if add_noise:
x = self.add_noise(x)
z = self.encoder(x)
return self.decoder(z)
def train_denoising_ae(model, train_loader, num_epochs=20):
"""
Train denoising autoencoder.
"""
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.MSELoss()
for epoch in range(num_epochs):
model.train()
epoch_loss = 0
for data, _ in train_loader:
data = data.view(data.size(0), -1).to(device)
optimizer.zero_grad()
# Forward pass with noisy input
noisy_data = model.add_noise(data)
reconstructed = model(noisy_data, add_noise=False)
# Loss against CLEAN data
loss = criterion(reconstructed, data)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
if (epoch + 1) % 5 == 0:
print(f"Epoch [{epoch+1}/{num_epochs}] | Loss: {epoch_loss/len(train_loader):.6f}")
# Train denoising autoencoder
denoising_ae = DenoisingAutoencoder(noise_factor=0.5).to(device)
print("Training Denoising Autoencoder...")
train_denoising_ae(denoising_ae, train_loader, num_epochs=20)
def visualize_denoising(model, test_loader, n_samples=10):
"""
Show: Original -> Noisy -> Denoised
"""
model.eval()
data, _ = next(iter(test_loader))
data = data[:n_samples]
data_flat = data.view(data.size(0), -1).to(device)
with torch.no_grad():
noisy = model.add_noise(data_flat)
denoised = model(noisy, add_noise=False)
fig, axes = plt.subplots(3, n_samples, figsize=(15, 4.5))
for i in range(n_samples):
axes[0, i].imshow(data[i].squeeze(), cmap='gray')
axes[0, i].axis('off')
axes[1, i].imshow(noisy[i].cpu().view(28, 28), cmap='gray')
axes[1, i].axis('off')
axes[2, i].imshow(denoised[i].cpu().view(28, 28), cmap='gray')
axes[2, i].axis('off')
axes[0, 0].set_title('Original', fontsize=12)
axes[1, 0].set_title('Noisy', fontsize=12)
axes[2, 0].set_title('Denoised', fontsize=12)
plt.tight_layout()
plt.savefig('denoising_results.png')
plt.close()
print("Denoising results saved!")
visualize_denoising(denoising_ae, test_loader)
```
***
## Variational Autoencoder (VAE)
VAEs are generative models that learn a **probabilistic** latent space. Instead of encoding to fixed points, VAEs encode to **distributions**.
**Why does this matter?** A standard autoencoder maps each input to a single point in latent space. The problem is that the space *between* those points is undefined -- if you sample a random point in latent space and decode it, you get garbage. A VAE forces the encoder to output a distribution (mean + variance) rather than a point, and the KL divergence term pulls those distributions toward a standard normal. This "fills in" the latent space, making it smooth and continuous -- nearby points decode to similar outputs, and random samples from the prior produce coherent outputs.
### Key Differences from Standard Autoencoders
| Aspect | Standard AE | VAE |
| ------------------------- | ------------------- | ---------------------------------------- |
| **Latent representation** | Deterministic point | Probability distribution |
| **Encoder output** | Single vector $z$ | Mean $\mu$ and variance $\sigma^2$ |
| **Sampling** | Not applicable | Sample from $\mathcal{N}(\mu, \sigma^2)$ |
| **Generation** | Poor | Can generate new samples |
### The VAE Objective: ELBO
The Evidence Lower BOund (ELBO) is the core objective. Think of it as a tug-of-war between two goals:
$$
\mathcal{L}_{VAE} = \underbrace{\mathbb{E}_{q(z|x)}[\log p(x|z)]}_{\text{Reconstruction}} - \underbrace{D_{KL}(q(z|x) \| p(z))}_{\text{KL Divergence}}
$$
| Term | Meaning | Purpose |
| ----------------------- | ----------------------- | ----------------------------------------------------- |
| **Reconstruction term** | Expected log-likelihood | Make outputs similar to inputs |
| **KL Divergence** | Distance from prior | Keep latent distribution close to $\mathcal{N}(0, I)$ |
**The tension:** The reconstruction term wants the encoder to create maximally informative latent codes (spreading them apart to preserve information). The KL term wants all codes to look like a standard normal distribution (pushing them together). The balance between these two forces determines what the latent space looks like -- too much KL pressure and the model ignores the latent code entirely ("posterior collapse"), too little and you can't generate new samples.
The KL divergence for Gaussian has a closed form:
$$
D_{KL}(q(z|x) \| p(z)) = -\frac{1}{2} \sum_{j=1}^{J}(1 + \log(\sigma_j^2) - \mu_j^2 - \sigma_j^2)
$$
```python theme={null}
class VAE(nn.Module):
"""
Variational Autoencoder.
Encoder outputs mu and log_var (log of variance, not variance directly).
Why log_var? Because variance must be positive, and log_var is unconstrained --
the network can output any real number, and exp(log_var) is always positive.
This avoids needing to clip or constrain the network output.
"""
def __init__(self, input_dim=784, hidden_dim=256, latent_dim=20):
super().__init__()
self.latent_dim = latent_dim
# Encoder: x -> hidden
self.encoder_layers = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU()
)
# Latent space parameters
self.fc_mu = nn.Linear(hidden_dim, latent_dim)
self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
# Decoder: z -> x
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 latent distribution parameters.
Returns: mu, log_var
"""
h = self.encoder_layers(x)
mu = self.fc_mu(h)
log_var = self.fc_logvar(h)
return mu, log_var
def reparameterize(self, mu, log_var):
"""
Reparameterization trick: z = mu + std * epsilon
Allows backpropagation through random sampling.
"""
std = torch.exp(0.5 * log_var) # std = exp(log_var / 2) = sqrt(var)
epsilon = torch.randn_like(std) # Sample from N(0,1) -- the randomness source
z = mu + std * epsilon # Shift and scale: now z ~ N(mu, var)
return z
def decode(self, z):
"""Decode latent vector to reconstruction."""
return self.decoder(z)
def forward(self, x):
mu, log_var = self.encode(x)
z = self.reparameterize(mu, log_var)
reconstruction = self.decode(z)
return reconstruction, mu, log_var
def vae_loss(reconstruction, original, mu, log_var, beta=1.0):
"""
VAE Loss = Reconstruction Loss + beta * KL Divergence
Args:
reconstruction: Decoder output
original: Original input
mu: Mean of latent distribution
log_var: Log variance of latent distribution
beta: Weight for KL divergence (beta-VAE). beta=1 is standard VAE.
"""
# Reconstruction loss (binary cross entropy) -- measures how well
# the decoder reproduces the input. Using reduction='sum' (not 'mean')
# ensures the loss scales with image size, preventing the KL term
# from dominating for small images.
recon_loss = F.binary_cross_entropy(reconstruction, original, reduction='sum')
# KL divergence: closed-form solution for two Gaussians (q(z|x) vs N(0,I)).
# Each term has an intuitive meaning:
# log_var: penalizes distributions that are too narrow (overly certain)
# mu^2: penalizes distributions whose mean drifts from the origin
# exp(log_var): penalizes distributions that are too wide
# The constant 1 balances these terms at the optimum (when mu=0, var=1)
kl_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
return recon_loss + beta * kl_loss, recon_loss, kl_loss
# Create VAE
vae = VAE(latent_dim=20).to(device)
print(f"VAE parameters: {sum(p.numel() for p in vae.parameters()):,}")
```
**Output:**
```
VAE parameters: 474,260
```
***
## The Reparameterization Trick
The reparameterization trick is **key** to training VAEs. Here's why:
**Problem:** We need to sample $z \sim \mathcal{N}(\mu, \sigma^2)$, but sampling is not differentiable! If you write `z = torch.normal(mu, sigma)`, PyTorch has no way to compute $\partial z / \partial \mu$ because the sampling operation is stochastic. Backpropagation needs a deterministic computation graph.
**Solution:** Instead of sampling directly, we sample $\epsilon \sim \mathcal{N}(0, I)$ and compute:
$$
z = \mu + \sigma \odot \epsilon
$$
Now the gradient can flow through $\mu$ and $\sigma$ because $\epsilon$ is treated as a constant (it was sampled before the forward pass). The randomness is "externalized" into $\epsilon$, and the rest of the computation is a standard deterministic function that autograd can differentiate. This trick is what made VAEs trainable at all -- without it, the entire probabilistic latent space idea would be a theoretical curiosity.
```python theme={null}
def visualize_reparameterization():
"""
Demonstrate the reparameterization trick.
"""
# Parameters (learned by encoder)
mu = torch.tensor([2.0])
log_var = torch.tensor([0.5]) # variance = exp(0.5) ≈ 1.65
# Standard deviation
std = torch.exp(0.5 * log_var)
# Sample epsilon from N(0, 1)
n_samples = 1000
epsilon = torch.randn(n_samples)
# Reparameterized samples
z_samples = mu + std * epsilon
# Verify distribution
print(f"Target mean: {mu.item():.2f}, Sample mean: {z_samples.mean():.2f}")
print(f"Target std: {std.item():.2f}, Sample std: {z_samples.std():.2f}")
# Plot
plt.figure(figsize=(10, 4))
plt.hist(z_samples.numpy(), bins=50, density=True, alpha=0.7)
plt.axvline(mu.item(), color='r', linestyle='--', label=f'μ = {mu.item():.1f}')
plt.xlabel('z')
plt.ylabel('Density')
plt.title('Samples using Reparameterization Trick')
plt.legend()
plt.savefig('reparameterization_demo.png')
plt.close()
visualize_reparameterization()
```
**Output:**
```
Target mean: 2.00, Sample mean: 2.01
Target std: 1.28, Sample std: 1.29
```
***
## Training the VAE
```python theme={null}
def train_vae(model, train_loader, num_epochs=30, lr=1e-3, beta=1.0):
"""
Train VAE with ELBO loss.
"""
optimizer = optim.Adam(model.parameters(), lr=lr)
train_losses = []
recon_losses = []
kl_losses = []
for epoch in range(num_epochs):
model.train()
epoch_loss = 0
epoch_recon = 0
epoch_kl = 0
for batch_idx, (data, _) in enumerate(train_loader):
data = data.view(data.size(0), -1).to(device)
optimizer.zero_grad()
# Forward pass
reconstruction, mu, log_var = model(data)
# Calculate loss
loss, recon_loss, kl_loss = vae_loss(
reconstruction, data, mu, log_var, beta=beta
)
# Backward pass
loss.backward()
optimizer.step()
epoch_loss += loss.item()
epoch_recon += recon_loss.item()
epoch_kl += kl_loss.item()
# Average losses
n = len(train_loader.dataset)
train_losses.append(epoch_loss / n)
recon_losses.append(epoch_recon / n)
kl_losses.append(epoch_kl / n)
if (epoch + 1) % 5 == 0:
print(f"Epoch [{epoch+1}/{num_epochs}] | "
f"Total: {train_losses[-1]:.4f} | "
f"Recon: {recon_losses[-1]:.4f} | "
f"KL: {kl_losses[-1]:.4f}")
return train_losses, recon_losses, kl_losses
# Train VAE
print("Training VAE...")
train_losses, recon_losses, kl_losses = train_vae(vae, train_loader, num_epochs=30)
# Plot losses
plt.figure(figsize=(12, 4))
plt.subplot(1, 3, 1)
plt.plot(train_losses)
plt.title('Total Loss')
plt.xlabel('Epoch')
plt.subplot(1, 3, 2)
plt.plot(recon_losses)
plt.title('Reconstruction Loss')
plt.xlabel('Epoch')
plt.subplot(1, 3, 3)
plt.plot(kl_losses)
plt.title('KL Divergence')
plt.xlabel('Epoch')
plt.tight_layout()
plt.savefig('vae_training.png')
plt.close()
```
**Output:**
```
Training VAE...
Epoch [5/30] | Total: 163.4521 | Recon: 151.2134 | KL: 12.2387
Epoch [10/30] | Total: 141.8934 | Recon: 127.5621 | KL: 14.3313
Epoch [15/30] | Total: 135.2178 | Recon: 119.8934 | KL: 15.3244
Epoch [20/30] | Total: 131.5623 | Recon: 115.4521 | KL: 16.1102
Epoch [25/30] | Total: 129.3421 | Recon: 112.8934 | KL: 16.4487
Epoch [30/30] | Total: 127.8912 | Recon: 111.2345 | KL: 16.6567
```
***
## Generating New Samples
The true power of VAEs - generating new data by sampling from the latent space!
```python theme={null}
def generate_samples(model, n_samples=64):
"""
Generate new samples by sampling from the prior N(0, I).
"""
model.eval()
with torch.no_grad():
# Sample from standard normal
z = torch.randn(n_samples, model.latent_dim).to(device)
# Decode
samples = model.decode(z)
samples = samples.view(-1, 1, 28, 28).cpu()
return samples
def visualize_generated(model, n_samples=64):
"""
Display grid of generated samples.
"""
samples = generate_samples(model, n_samples)
# Create grid
n_row = int(np.sqrt(n_samples))
fig, axes = plt.subplots(n_row, n_row, figsize=(10, 10))
for i, ax in enumerate(axes.flat):
ax.imshow(samples[i].squeeze(), cmap='gray')
ax.axis('off')
plt.suptitle('VAE Generated Samples', fontsize=16)
plt.tight_layout()
plt.savefig('vae_generated.png')
plt.close()
print(f"Generated {n_samples} new samples!")
visualize_generated(vae, n_samples=64)
```
***
## Latent Space Interpolation
We can smoothly transition between images by interpolating in latent space:
```python theme={null}
def interpolate_latent(model, start_img, end_img, n_steps=10):
"""
Interpolate between two images in latent space.
"""
model.eval()
with torch.no_grad():
# Encode both images
start_flat = start_img.view(1, -1).to(device)
end_flat = end_img.view(1, -1).to(device)
start_mu, start_logvar = model.encode(start_flat)
end_mu, end_logvar = model.encode(end_flat)
# Interpolate between means
interpolations = []
for alpha in np.linspace(0, 1, n_steps):
z = (1 - alpha) * start_mu + alpha * end_mu
decoded = model.decode(z)
interpolations.append(decoded.view(28, 28).cpu())
return interpolations
def visualize_interpolation(model, test_loader):
"""
Show interpolation between two random digits.
"""
# Get two different digits
data, labels = next(iter(test_loader))
# Find a 3 and a 7
idx_3 = (labels == 3).nonzero()[0].item()
idx_7 = (labels == 7).nonzero()[0].item()
img_3 = data[idx_3]
img_7 = data[idx_7]
# Interpolate
interpolations = interpolate_latent(vae, img_3, img_7, n_steps=10)
# Plot
fig, axes = plt.subplots(1, 10, figsize=(15, 1.5))
for i, ax in enumerate(axes):
ax.imshow(interpolations[i], cmap='gray')
ax.axis('off')
ax.set_title(f'{i/(len(axes)-1):.1f}')
plt.suptitle('Latent Space Interpolation: 3 → 7', fontsize=14)
plt.tight_layout()
plt.savefig('interpolation.png')
plt.close()
print("Interpolation saved!")
visualize_interpolation(vae, test_loader)
```
***
## Convolutional VAE
For better image generation, use convolutional layers:
```python theme={null}
class ConvVAE(nn.Module):
"""
Convolutional VAE for better image generation.
"""
def __init__(self, latent_dim=32):
super().__init__()
self.latent_dim = latent_dim
# Encoder
self.encoder_conv = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1), # 28->14
nn.BatchNorm2d(32),
nn.LeakyReLU(0.2),
nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1), # 14->7
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2),
nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), # 7->4
nn.BatchNorm2d(128),
nn.LeakyReLU(0.2),
)
self.flatten_size = 128 * 4 * 4
self.fc_mu = nn.Linear(self.flatten_size, latent_dim)
self.fc_logvar = nn.Linear(self.flatten_size, latent_dim)
# Decoder
self.decoder_fc = nn.Linear(latent_dim, self.flatten_size)
self.decoder_conv = nn.Sequential(
nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=0), # 4->7
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2),
nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1), # 7->14
nn.BatchNorm2d(32),
nn.LeakyReLU(0.2),
nn.ConvTranspose2d(32, 1, kernel_size=3, stride=2, padding=1, output_padding=1), # 14->28
nn.Sigmoid()
)
def encode(self, x):
h = self.encoder_conv(x)
h = h.view(h.size(0), -1)
return self.fc_mu(h), self.fc_logvar(h)
def reparameterize(self, mu, log_var):
std = torch.exp(0.5 * log_var)
eps = torch.randn_like(std)
return mu + std * eps
def decode(self, z):
h = self.decoder_fc(z)
h = h.view(-1, 128, 4, 4)
return self.decoder_conv(h)
def forward(self, x):
mu, log_var = self.encode(x)
z = self.reparameterize(mu, log_var)
return self.decode(z), mu, log_var
# Create Conv VAE
conv_vae = ConvVAE(latent_dim=32).to(device)
print(f"Conv VAE parameters: {sum(p.numel() for p in conv_vae.parameters()):,}")
# Test
test_batch = torch.randn(4, 1, 28, 28).to(device)
output, mu, logvar = conv_vae(test_batch)
print(f"Input: {test_batch.shape}, Output: {output.shape}")
print(f"Latent: mu {mu.shape}, logvar {logvar.shape}")
```
***
## Beta-VAE: Disentangled Representations
Beta-VAE encourages **disentangled** representations by increasing the weight of KL divergence:
$$
\mathcal{L}_{\beta-VAE} = \mathbb{E}[\log p(x|z)] - \beta \cdot D_{KL}(q(z|x) \| p(z))
$$
**What does "disentangled" mean?** In a disentangled representation, each latent dimension controls one independent factor of variation. For faces: dimension 1 might control hair color, dimension 2 controls smile, dimension 3 controls head rotation, etc. Changing one dimension doesn't affect the others. This is powerful because it gives you interpretable, controllable generation.
**How does increasing beta help?** A higher beta forces the posterior closer to the isotropic prior $\mathcal{N}(0, I)$, which has independent dimensions by definition. The encoder must find a way to encode information using statistically independent dimensions, which naturally leads to disentanglement. The cost is reconstruction quality -- the model must discard more information to satisfy the stronger regularization.
| beta Value | Effect |
| ---------- | ------------------------------------------------- |
| beta = 1 | Standard VAE |
| beta > 1 | More disentanglement, less reconstruction quality |
| beta \< 1 | Better reconstruction, more entangled |
**KL annealing tip:** Instead of fixing beta, many practitioners start training with beta=0 (pure autoencoder) and linearly increase it to the target value over the first 10-20% of training. This prevents "posterior collapse" -- a failure mode where the encoder learns to ignore the input and output the prior $\mathcal{N}(0, I)$ for everything, because the KL penalty dominates before the decoder is good enough to use the latent codes.
```python theme={null}
def train_beta_vae(model, train_loader, num_epochs=30, beta=4.0):
"""
Train Beta-VAE for disentangled representations.
"""
optimizer = optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(num_epochs):
model.train()
total_loss = 0
for data, _ in train_loader:
if data.dim() == 4:
data = data.to(device)
else:
data = data.view(-1, 784).to(device)
optimizer.zero_grad()
if hasattr(model, 'encoder_conv'):
# Conv VAE
reconstruction, mu, log_var = model(data)
recon_loss = F.binary_cross_entropy(
reconstruction.view(-1), data.view(-1), reduction='sum'
)
else:
# Linear VAE
reconstruction, mu, log_var = model(data)
recon_loss = F.binary_cross_entropy(
reconstruction, data.view(-1, 784), reduction='sum'
)
kl_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
# Beta-weighted loss
loss = recon_loss + beta * kl_loss
loss.backward()
optimizer.step()
total_loss += loss.item()
if (epoch + 1) % 10 == 0:
print(f"Epoch [{epoch+1}/{num_epochs}] | "
f"Loss: {total_loss/len(train_loader.dataset):.4f}")
return model
# Train Beta-VAE with different beta values
print("\nTraining Beta-VAE (β=4)...")
beta_vae = VAE(latent_dim=10).to(device)
train_beta_vae(beta_vae, train_loader, num_epochs=20, beta=4.0)
```
***
## Exercises
Add an L1 sparsity constraint to encourage sparse representations.
```python theme={null}
# TODO: Modify the autoencoder to add sparsity regularization
# Hint: Add L1 penalty on latent activations
class SparseAutoencoder(nn.Module):
def __init__(self, input_dim=784, latent_dim=64, sparsity_weight=1e-3):
# Your implementation here
pass
```
Solution
```python theme={null}
class SparseAutoencoder(nn.Module):
def __init__(self, input_dim=784, latent_dim=64, sparsity_weight=1e-3):
super().__init__()
self.sparsity_weight = sparsity_weight
self.encoder = nn.Sequential(
nn.Linear(input_dim, 256),
nn.ReLU(),
nn.Linear(256, latent_dim),
nn.ReLU() # ReLU helps sparsity
)
self.decoder = nn.Sequential(
nn.Linear(latent_dim, 256),
nn.ReLU(),
nn.Linear(256, input_dim),
nn.Sigmoid()
)
def forward(self, x):
z = self.encoder(x)
reconstruction = self.decoder(z)
return reconstruction, z
def loss(self, x, reconstruction, z):
# Reconstruction loss
recon_loss = F.mse_loss(reconstruction, x)
# L1 sparsity penalty on latent activations
sparsity_loss = self.sparsity_weight * torch.mean(torch.abs(z))
return recon_loss + sparsity_loss
# Training loop
sparse_ae = SparseAutoencoder().to(device)
optimizer = optim.Adam(sparse_ae.parameters(), lr=1e-3)
for epoch in range(20):
for data, _ in train_loader:
data = data.view(-1, 784).to(device)
optimizer.zero_grad()
reconstruction, z = sparse_ae(data)
loss = sparse_ae.loss(data, reconstruction, z)
loss.backward()
optimizer.step()
if (epoch + 1) % 5 == 0:
# Check sparsity
with torch.no_grad():
_, z = sparse_ae(data)
sparsity = (z < 0.01).float().mean()
print(f"Epoch {epoch+1} | Loss: {loss.item():.4f} | Sparsity: {sparsity:.2%}")
```
Create a VAE that can generate specific digits by conditioning on class labels.
```python theme={null}
# TODO: Implement CVAE that takes class label as input
# The encoder and decoder should both receive the class information
class ConditionalVAE(nn.Module):
def __init__(self, input_dim=784, latent_dim=20, num_classes=10):
# Your implementation here
pass
```
Solution
```python theme={null}
class ConditionalVAE(nn.Module):
def __init__(self, input_dim=784, latent_dim=20, num_classes=10, hidden_dim=256):
super().__init__()
self.latent_dim = latent_dim
# Embedding for class labels
self.label_embedding = nn.Embedding(num_classes, hidden_dim)
# Encoder: x + label -> mu, log_var
self.encoder = nn.Sequential(
nn.Linear(input_dim + hidden_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: z + label -> x
self.decoder = nn.Sequential(
nn.Linear(latent_dim + hidden_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, labels):
label_emb = self.label_embedding(labels)
x_cond = torch.cat([x, label_emb], dim=1)
h = self.encoder(x_cond)
return self.fc_mu(h), self.fc_logvar(h)
def reparameterize(self, mu, log_var):
std = torch.exp(0.5 * log_var)
eps = torch.randn_like(std)
return mu + std * eps
def decode(self, z, labels):
label_emb = self.label_embedding(labels)
z_cond = torch.cat([z, label_emb], dim=1)
return self.decoder(z_cond)
def forward(self, x, labels):
mu, log_var = self.encode(x, labels)
z = self.reparameterize(mu, log_var)
return self.decode(z, labels), mu, log_var
def generate(self, labels, n_per_class=8):
"""Generate samples for given class labels."""
self.eval()
with torch.no_grad():
z = torch.randn(len(labels), self.latent_dim).to(labels.device)
return self.decode(z, labels)
# Usage
cvae = ConditionalVAE().to(device)
# Generate specific digits
labels = torch.arange(10).to(device)
generated = cvae.generate(labels)
print(f"Generated digits 0-9: {generated.shape}")
```
Vector Quantized VAE uses discrete latent codes instead of continuous.
```python theme={null}
# TODO: Implement the vector quantization layer
# Hint: Map continuous vectors to nearest codebook entries
class VectorQuantizer(nn.Module):
def __init__(self, num_embeddings=512, embedding_dim=64):
# Your implementation here
pass
```
Solution
```python theme={null}
class VectorQuantizer(nn.Module):
"""
Vector Quantization layer for VQ-VAE.
Maps continuous vectors to discrete codebook entries.
"""
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):
# z shape: [B, C, H, W] -> [B, H, W, C]
z = z.permute(0, 2, 3, 1).contiguous()
z_flat = z.view(-1, self.embedding_dim)
# Calculate distances to codebook
distances = (
torch.sum(z_flat ** 2, dim=1, keepdim=True)
+ torch.sum(self.embedding.weight ** 2, dim=1)
- 2 * torch.matmul(z_flat, self.embedding.weight.t())
)
# Find nearest codebook entries
encoding_indices = torch.argmin(distances, dim=1)
# Quantize
z_q = self.embedding(encoding_indices).view(z.shape)
# Compute loss
e_latent_loss = F.mse_loss(z_q.detach(), z) # Commitment loss
q_latent_loss = F.mse_loss(z_q, z.detach()) # Codebook loss
loss = q_latent_loss + self.commitment_cost * e_latent_loss
# Straight-through estimator
z_q = z + (z_q - z).detach()
# Back to [B, C, H, W]
z_q = z_q.permute(0, 3, 1, 2).contiguous()
return z_q, loss, encoding_indices
class VQVAE(nn.Module):
def __init__(self, num_embeddings=512, embedding_dim=64):
super().__init__()
self.encoder = nn.Sequential(
nn.Conv2d(1, 32, 4, 2, 1),
nn.ReLU(),
nn.Conv2d(32, embedding_dim, 4, 2, 1),
nn.ReLU()
)
self.vq = VectorQuantizer(num_embeddings, embedding_dim)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(embedding_dim, 32, 4, 2, 1),
nn.ReLU(),
nn.ConvTranspose2d(32, 1, 4, 2, 1),
nn.Sigmoid()
)
def forward(self, x):
z = self.encoder(x)
z_q, vq_loss, _ = self.vq(z)
x_recon = self.decoder(z_q)
return x_recon, vq_loss
# Test
vqvae = VQVAE().to(device)
test_input = torch.randn(4, 1, 28, 28).to(device)
recon, vq_loss = vqvae(test_input)
print(f"VQ-VAE output shape: {recon.shape}")
print(f"VQ loss: {vq_loss.item():.4f}")
```
***
## Key Takeaways
**What You Learned:**
* ✅ **Autoencoders** - Encoder-decoder architecture with bottleneck for compression
* ✅ **Latent Space** - Lower-dimensional representation that captures essential features
* ✅ **Denoising AE** - Learn to remove noise by training with corrupted inputs
* ✅ **VAE Theory** - Probabilistic latent space with ELBO objective
* ✅ **KL Divergence** - Regularizes latent space to match prior distribution
* ✅ **Reparameterization** - Enables backpropagation through sampling
* ✅ **Generation** - Sample from latent space to create new data
* ✅ **Beta-VAE** - Control disentanglement with β hyperparameter
***
## Common Pitfalls
**Autoencoder Mistakes to Avoid:**
1. **Latent dim too large** -- No compression = the network learns the identity function. A good rule of thumb: start with a latent dimension that is 10-50x smaller than the input dimension, then tune based on reconstruction quality vs. downstream task performance.
2. **Latent dim too small** -- Poor reconstructions, lost information. You can diagnose this by plotting reconstruction loss as a function of latent dimension -- the curve will show a sharp elbow where adding more dimensions stops helping.
3. **Ignoring KL collapse (posterior collapse)** -- The VAE's decoder becomes so powerful that it ignores the latent code entirely, and the encoder outputs the prior for every input. Fix with KL annealing (start beta=0, increase linearly), or use free bits (allow a minimum KL per dimension before penalizing).
4. **Wrong reconstruction loss** -- Use BCE for \[0,1] images with Sigmoid output, MSE for continuous data or unbounded outputs. Mismatching the loss and activation leads to poor gradients and blurry results.
5. **Not normalizing inputs** -- Autoencoders work best with normalized data. For images, normalize to \[0,1] for Sigmoid decoders or \[-1,1] for Tanh decoders. Mismatched ranges cause the loss to be dominated by scale differences rather than structural features.
***
## Interview Deep-Dive
**Strong Answer:**
* In a VAE, the encoder outputs parameters of a distribution ($\mu$, $\sigma^2$) rather than a single point. During training, we need to sample $z \sim \mathcal{N}(\mu, \sigma^2)$ and then backpropagate through the entire encoder-decoder pipeline. The problem is that sampling is a stochastic operation -- PyTorch (or any autograd system) cannot compute $\partial z / \partial \mu$ when $z$ was drawn from a random process.
* The reparameterization trick rewrites $z = \mu + \sigma \cdot \epsilon$ where $\epsilon \sim \mathcal{N}(0, I)$ is sampled independently. Now the randomness is in $\epsilon$ (which doesn't depend on any parameters), and $z$ is a deterministic, differentiable function of $\mu$ and $\sigma$. Gradients flow cleanly: $\partial z / \partial \mu = 1$ and $\partial z / \partial \sigma = \epsilon$.
* **Without the trick**, you'd need to use REINFORCE-style gradient estimators (score function estimator), which are unbiased but have extremely high variance. In practice, training becomes so noisy that the model fails to converge for any non-trivial dataset. The reparameterization trick reduces gradient variance by orders of magnitude, making VAE training practical.
* **A senior engineer would note**: the trick only works for distributions where we can express sampling as a deterministic transformation of a fixed base distribution. It works for Gaussians, but not directly for discrete distributions. For discrete latent variables (like VQ-VAE), you need alternatives like the straight-through estimator or Gumbel-Softmax.
**Follow-up: How does the straight-through estimator in VQ-VAE solve a similar problem?**
VQ-VAE uses discrete codebook entries, which have zero gradients everywhere (argmin is piecewise constant). The straight-through estimator "pretends" the quantization step is an identity during the backward pass: gradients from the decoder flow directly to the encoder, bypassing the non-differentiable lookup. It's biased but works remarkably well in practice, and the commitment loss ($\|z_e - \text{sg}[z_q]\|^2$) ensures the encoder stays close to the codebook entries.
**Strong Answer:**
* **Posterior collapse** occurs when the encoder learns to output the prior $\mathcal{N}(0, I)$ for every input, making the latent code uninformative. The decoder compensates by becoming an unconditional generative model (a decoder-only language model, effectively). The KL divergence drops to zero, and the ELBO reduces to just the marginal log-likelihood -- the "variational" part of VAE becomes useless.
* **Why it happens**: the KL penalty encourages the posterior to match the prior. Early in training, the decoder is weak and can't use the latent code effectively. The optimizer finds it easier to minimize KL (by making the posterior equal the prior) than to improve reconstruction (which requires coordinated encoder-decoder learning). Once collapsed, the decoder learns to ignore $z$, and the encoder has no gradient signal to recover.
* **Strategy 1: KL Annealing.** Start with $\beta = 0$ and linearly increase to 1 over the first 10-20% of training. This lets the decoder learn to use the latent code before the KL penalty kicks in. **Trade-off**: adds a hyperparameter (annealing schedule) and doesn't guarantee the model stays out of collapse after annealing completes. Cyclical annealing (repeatedly cycling beta from 0 to 1) can help more.
* **Strategy 2: Free Bits.** Allow each latent dimension a minimum KL of $\lambda$ (typically 0.1-0.5 nats) before penalizing. The modified loss: $\sum_j \max(\lambda, D_{KL}^{(j)})$. This ensures each dimension encodes at least $\lambda$ nats of information. **Trade-off**: the model can still concentrate all information in a few dimensions while others collapse, and the hyperparameter $\lambda$ is sensitive.
* **Strategy 3: Stronger decoder bottleneck.** If the decoder is too powerful (e.g., an autoregressive decoder like PixelCNN), it can model the data without the latent code. Deliberately limiting decoder capacity (fewer layers, smaller hidden dim, removing autoregressive connections) forces it to rely on $z$. **Trade-off**: reconstruction quality degrades, and finding the right balance is empirical.
* **A senior engineer would note**: posterior collapse is fundamentally about the balance of information pathways. If the decoder can "route around" the latent bottleneck, it will. The most robust approach combines KL annealing with a decoder architecture that genuinely needs the latent code (e.g., a simple feedforward decoder with limited capacity).
**Strong Answer:**
* **Standard Autoencoders**: deterministic encoder-decoder with a bottleneck. Best for: dimensionality reduction, feature extraction, denoising, anomaly detection (high reconstruction error = anomaly). Cannot generate new samples because the latent space is unstructured -- points between encoded samples decode to garbage. Use when generation is not needed and you want the simplest, fastest model for compression or representation learning.
* **VAEs**: probabilistic encoder (outputs $\mu$, $\sigma$) with KL regularization against $\mathcal{N}(0, I)$. Best for: generating new samples, learning smooth latent representations, interpolation between data points, disentangled representations (beta-VAE). **Trade-off**: reconstructions are blurrier than standard autoencoders because the KL term trades reconstruction fidelity for latent space regularity. The Gaussian assumption also limits expressiveness -- real data distributions are rarely Gaussian.
* **VQ-VAEs**: discrete latent space using a learned codebook. The encoder maps to continuous vectors, which are then snapped to the nearest codebook entry. Best for: high-fidelity generation (especially when paired with an autoregressive prior over the codebook indices), learning hierarchical discrete representations (VQ-VAE-2 achieves near-photorealistic generation). **Trade-off**: requires more complex training (straight-through estimator, commitment loss, codebook EMA updates), and generation requires a separate prior model (like PixelCNN or a Transformer) trained on the codebook indices.
* **Decision framework**: need compression/anomaly detection? Standard AE. Need smooth generation and interpolation? VAE. Need high-fidelity generation with discrete control? VQ-VAE. Need state-of-the-art generation quality? VQ-VAE-2 with a Transformer prior, or skip autoencoders entirely and use diffusion models.
**Strong Answer:**
* **Core architecture**: treat each user's interaction history as a sparse vector (items rated or interacted with) and train an autoencoder to reconstruct it. The latent representation captures user preferences, and the decoder output for unobserved items becomes the recommendation score. This is the approach behind Variational Autoencoders for Collaborative Filtering (Mult-VAE), which uses a multinomial likelihood and consistently outperforms matrix factorization baselines.
* **Handling missing data**: the input is the user's observed interactions (e.g., a 10,000-dim vector with values only at the 50 items they've interacted with). The loss is computed only over observed entries during training, but at inference time, we decode the full vector and rank the unobserved items by predicted score. The autoencoder learns to "fill in" the missing entries by learning patterns across users.
* **Cold-start problem**: for new users with very few interactions, the encoder has insufficient signal. Strategies: (1) use a hybrid model that incorporates side information (user demographics, item metadata) as additional encoder inputs, (2) use a VAE with a learned prior conditioned on available metadata instead of a standard normal prior, (3) for brand-new users with zero interactions, fall back to popularity-based or content-based recommendations until enough interaction data accumulates.
* **Architecture details**: the encoder uses dropout on the input (dropout rate 0.5) as a form of augmentation -- this is equivalent to a denoising autoencoder and prevents the model from memorizing the training set. Use the multinomial log-likelihood loss rather than MSE, since user interactions are better modeled as counts or implicit feedback, not continuous values.
* **Production considerations**: the latent vectors are compact (128-256 dimensions) and can be precomputed for all users, enabling fast approximate nearest-neighbor retrieval for real-time recommendations. Retrain weekly or use incremental updates with new interaction data.
***
Learn about the cutting-edge generative models behind Stable Diffusion and DALL-E