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

1. Use MSE for reconstruction – L2 loss for continuous data
2. KL annealing – gradually increase beta in VAE for better training
3. Monitor KL divergence – too low = posterior collapse, too high = blurry
4. Latent space regularization – smooth latent spaces enable interpolation
5. Use perceptual loss – for images, match feature representations not pixels
6. 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.