Autoencoders and Variational Autoencoders
Autoencoders learn compressed representations of data. They encode inputs into a bottleneck and decode them back, learning meaningful features in the process. Variational autoencoders (VAEs) add probabilistic structure, enabling generation of new data points.
Autoencoder Architecture
Autoencoder Fundamentals
An autoencoder has two parts: an encoder that compresses, and a decoder that reconstructs. The bottleneck forces the network to learn the most important features.
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
Basic Autoencoder
def generate_data(n=1000):
"""Generate 2D Swiss roll-like data."""
t = np.linspace(0, 3 * np.pi, n)
x = t * np.cos(t) + np.random.randn(n) * 0.1
y = t * np.sin(t) + np.random.randn(n) * 0.1
return np.column_stack([x, y]).astype(np.float32)
data = generate_data(1000)
class Autoencoder(nn.Module):
def __init__(self, input_dim=2, latent_dim=1):
super().__init__()
self.encoder = nn.Sequential(
nn.Linear(input_dim, 32),
nn.ReLU(),
nn.Linear(32, 16),
nn.ReLU(),
nn.Linear(16, latent_dim)
)
self.decoder = nn.Sequential(
nn.Linear(latent_dim, 16),
nn.ReLU(),
nn.Linear(16, 32),
nn.ReLU(),
nn.Linear(32, input_dim)
)
def forward(self, x):
z = self.encoder(x)
x_recon = self.decoder(z)
return x_recon, z
# Train
ae = Autoencoder(input_dim=2, latent_dim=1).to(device)
optimizer = optim.Adam(ae.parameters(), lr=0.001)
criterion = nn.MSELoss()
dataset = TensorDataset(torch.FloatTensor(data))
loader = DataLoader(dataset, batch_size=64, shuffle=True)
for epoch in range(500):
total_loss = 0
for (batch,) in loader:
batch = batch.to(device)
recon, z = ae(batch)
loss = criterion(recon, batch)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
if (epoch + 1) % 100 == 0:
print(f"Epoch {epoch+1}: Loss={total_loss/len(loader):.4f}")
# Encode all data
with torch.no_grad():
_, latent = ae(torch.FloatTensor(data).to(device))
print(f"Latent representation shape: {latent.shape}")
Convolutional Autoencoder for Images
class ConvAutoencoder(nn.Module):
def __init__(self, latent_dim=32):
super().__init__()
# Encoder
self.encoder = nn.Sequential(
nn.Conv2d(1, 16, 3, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(16, 32, 3, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(32, 64, 7),
nn.ReLU()
)
self.fc_encode = nn.Linear(64 * 4 * 4, latent_dim)
# Decoder
self.fc_decode = nn.Linear(latent_dim, 64 * 4 * 4)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(64, 32, 7),
nn.ReLU(),
nn.ConvTranspose2d(32, 16, 3, stride=2, padding=1, output_padding=1),
nn.ReLU(),
nn.ConvTranspose2d(16, 1, 3, stride=2, padding=1, output_padding=1),
nn.Sigmoid()
)
def encode(self, x):
h = self.encoder(x)
h = h.view(h.size(0), -1)
return self.fc_encode(h)
def decode(self, z):
h = self.fc_decode(z)
h = h.view(h.size(0), 64, 4, 4)
return self.decoder(h)
def forward(self, x):
z = self.encode(x)
return self.decode(z), z
# Generate synthetic MNIST-like data
fake_mnist = torch.randn(1000, 1, 28, 28).abs()
conv_ae = ConvAutoencoder(latent_dim=32).to(device)
optimizer = optim.Adam(conv_ae.parameters(), lr=0.001)
dataset = TensorDataset(fake_mnist)
loader = DataLoader(dataset, batch_size=32, shuffle=True)
for epoch in range(50):
total_loss = 0
for (batch,) in loader:
batch = batch.to(device)
recon, z = conv_ae(batch)
loss = nn.MSELoss()(recon, batch)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
if (epoch + 1) % 10 == 0:
print(f"Epoch {epoch+1}: MSE={total_loss/len(loader):.4f}")
Denoising Autoencoder
Learning to denoise forces the autoencoder to learn robust features.
class DenoisingAutoencoder(nn.Module):
def __init__(self, input_dim=2, latent_dim=1, noise_factor=0.3):
super().__init__()
self.noise_factor = noise_factor
self.encoder = nn.Sequential(
nn.Linear(input_dim, 32),
nn.ReLU(),
nn.Linear(32, 16),
nn.ReLU(),
nn.Linear(16, latent_dim)
)
self.decoder = nn.Sequential(
nn.Linear(latent_dim, 16),
nn.ReLU(),
nn.Linear(16, 32),
nn.ReLU(),
nn.Linear(32, input_dim)
)
def forward(self, x):
# Add noise during training
if self.training:
x_noisy = x + self.noise_factor * torch.randn_like(x)
else:
x_noisy = x
z = self.encoder(x_noisy)
x_recon = self.decoder(z)
return x_recon, z
# Train denoising AE
dae = DenoisingAutoencoder(input_dim=2, latent_dim=1).to(device)
optimizer = optim.Adam(dae.parameters(), lr=0.001)
for epoch in range(500):
dae.train()
total_loss = 0
for (batch,) in loader:
batch = batch.to(device)
recon, z = dae(batch)
loss = nn.MSELoss()(recon, batch)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
if (epoch + 1) % 100 == 0:
print(f"Epoch {epoch+1}: Denoising loss={total_loss/len(loader):.4f}")
Variational Autoencoder (VAE)
VAEs learn a probabilistic latent space, enabling generation of new data.
class VAE(nn.Module):
def __init__(self, input_dim=2, latent_dim=2):
super().__init__()
# Encoder
self.encoder = nn.Sequential(
nn.Linear(input_dim, 32),
nn.ReLU(),
nn.Linear(32, 16),
nn.ReLU()
)
self.fc_mu = nn.Linear(16, latent_dim)
self.fc_logvar = nn.Linear(16, latent_dim)
# Decoder
self.decoder = nn.Sequential(
nn.Linear(latent_dim, 16),
nn.ReLU(),
nn.Linear(16, 32),
nn.ReLU(),
nn.Linear(32, input_dim)
)
def encode(self, x):
h = self.encoder(x)
return self.fc_mu(h), self.fc_logvar(h)
def reparameterize(self, mu, logvar):
"""Sample from N(mu, sigma^2) using reparameterization trick."""
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
return mu + eps * std
def decode(self, z):
return self.decoder(z)
def forward(self, x):
mu, logvar = self.encode(x)
z = self.reparameterize(mu, logvar)
x_recon = self.decode(z)
return x_recon, mu, logvar
def vae_loss(x_recon, x, mu, logvar, beta=1.0):
"""VAE loss = reconstruction + KL divergence."""
recon_loss = nn.MSELoss()(x_recon, x)
kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
return recon_loss + beta * kl_loss, recon_loss, kl_loss
# Train VAE
vae = VAE(input_dim=2, latent_dim=2).to(device)
optimizer = optim.Adam(vae.parameters(), lr=0.001)
for epoch in range(1000):
vae.train()
total_loss = 0
for (batch,) in loader:
batch = batch.to(device)
recon, mu, logvar = vae(batch)
loss, recon_loss, kl_loss = vae_loss(recon, batch, mu, logvar)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
if (epoch + 1) % 200 == 0:
print(f"Epoch {epoch+1}: Total={total_loss/len(loader):.4f}")
# Generate new data
vae.eval()
with torch.no_grad():
z = torch.randn(100, 2).to(device)
generated = vae.decode(z).cpu().numpy()
print(f"Generated {generated.shape[0]} samples")
Latent Space Interpolation
# Interpolate between two real data points
vae.eval()
with torch.no_grad():
z1 = vae.encode(torch.FloatTensor(data[:1]).to(device))[0]
z2 = vae.encode(torch.FloatTensor(data[-1:]).to(device))[0]
# Linear interpolation in latent space
alphas = np.linspace(0, 1, 10)
interpolations = []
for alpha in alphas:
z = (1 - alpha) * z1 + alpha * z2
decoded = vae.decode(z).cpu().numpy()
interpolations.append(decoded[0])
interpolations = np.array(interpolations)
print(f"Interpolation path: {interpolations.shape}")
Beta-VAE for Disentanglement
class BetaVAE(nn.Module):
def __init__(self, input_dim=2, latent_dim=2, beta=4.0):
super().__init__()
self.beta = beta
self.encoder = nn.Sequential(
nn.Linear(input_dim, 32), nn.ReLU(),
nn.Linear(32, 16), nn.ReLU()
)
self.fc_mu = nn.Linear(16, latent_dim)
self.fc_logvar = nn.Linear(16, latent_dim)
self.decoder = nn.Sequential(
nn.Linear(latent_dim, 16), nn.ReLU(),
nn.Linear(16, 32), nn.ReLU(),
nn.Linear(32, input_dim)
)
def reparameterize(self, mu, logvar):
std = torch.exp(0.5 * logvar)
return mu + torch.randn_like(std) * std
def forward(self, x):
h = self.encoder(x)
mu, logvar = self.fc_mu(h), self.fc_logvar(h)
z = self.reparameterize(mu, logvar)
return self.decoder(z), mu, logvar
beta_vae = BetaVAE(beta=4.0).to(device)
print("Beta-VAE with beta=4.0 for stronger disentanglement")
VAE with Categorical Latent (VQ-VAE)
class VQVAE(nn.Module):
def __init__(self, input_dim=2, n_embeddings=10, embedding_dim=2):
super().__init__()
self.encoder = nn.Sequential(
nn.Linear(input_dim, 32), nn.ReLU(), nn.Linear(32, embedding_dim)
)
self.decoder = nn.Sequential(
nn.Linear(embedding_dim, 32), nn.ReLU(), nn.Linear(32, input_dim)
)
# Codebook
self.embeddings = nn.Embedding(n_embeddings, embedding_dim)
self.n_embeddings = n_embeddings
def forward(self, x):
z_e = self.encoder(x)
# Find nearest codebook entry
distances = (z_e.unsqueeze(1) - self.embeddings.weight.unsqueeze(0)).pow(2).sum(-1)
indices = distances.argmin(dim=-1)
z_q = self.embeddings(indices)
# Straight-through estimator
z_q = z_e + (z_q - z_e).detach()
x_recon = self.decoder(z_q)
return x_recon, z_e, z_q
vqvae = VQVAE().to(device)
x = torch.randn(32, 2).to(device)
recon, z_e, z_q = vqvae(x)
print(f"VQ-VAE reconstruction: {recon.shape}, discrete codes: {z_q.shape}")
Best Practices
-
Use MSE for reconstruction – L2 loss for continuous data
-
KL annealing – gradually increase beta in VAE for better training
-
Monitor KL divergence – too low = posterior collapse, too high = blurry
-
Latent space regularization – smooth latent spaces enable interpolation
-
Use perceptual loss – for images, match feature representations not pixels
-
Evaluate with FID – for generative quality
Summary
Autoencoders learn compressed representations; VAEs add probabilistic structure for generation. Master the reparameterization trick, KL divergence, and latent space manipulation to build models that compress, denoise, and generate data.