Variational Autoencoder in PyTorch - ML Experiment: Train & Evaluate

Experiment - Variational Autoencoder

Problem:We want to train a Variational Autoencoder (VAE) on the MNIST dataset to learn a compressed representation of handwritten digits.

Current Metrics:Training loss: 120.5, Validation loss: 150.2, Validation accuracy (reconstruction quality proxy): 65%

Issue:The model is overfitting: training loss is much lower than validation loss, and validation accuracy is low indicating poor generalization.

Your Task

Reduce overfitting by improving validation loss and reconstruction accuracy to above 80%, while keeping training loss reasonable.

You can only modify the model architecture and training hyperparameters.

Do not change the dataset or preprocessing steps.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.BatchNorm1d(256)
        )
        self.fc_mu = nn.Linear(256, 20)
        self.fc_logvar = nn.Linear(256, 20)
        self.decoder = nn.Sequential(
            nn.Linear(20, 256),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.BatchNorm1d(256),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 28*28),
            nn.Sigmoid()
        )

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

    def reparameterize(self, mu, logvar):
        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)
        return self.decode(z), mu, logvar

def loss_function(recon_x, x, mu, logvar):
    BCE = nn.functional.binary_cross_entropy(recon_x, x, reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.view(-1))
])

train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
val_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)

model = VAE()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

num_epochs = 20

for epoch in range(num_epochs):
    model.train()
    train_loss = 0
    for data, _ in train_loader:
        data = data.view(data.size(0), -1)
        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)
        loss = loss_function(recon_batch, data, mu, logvar)
        loss.backward()
        train_loss += loss.item()
        optimizer.step()

    model.eval()
    val_loss = 0
    with torch.no_grad():
        for data, _ in val_loader:
            data = data.view(data.size(0), -1)
            recon_batch, mu, logvar = model(data)
            loss = loss_function(recon_batch, data, mu, logvar)
            val_loss += loss.item()

    print(f'Epoch {epoch+1}, Train loss: {train_loss/len(train_loader.dataset):.4f}, Val loss: {val_loss/len(val_loader.dataset):.4f}')

Added dropout layers in encoder and decoder to reduce overfitting.

Added batch normalization layers to stabilize training.

Reduced batch size to 64 for better gradient estimates.

Used Adam optimizer with learning rate 0.001 for smoother convergence.

Limited training to 20 epochs to avoid overfitting.

Added reshaping of input data to flatten images before feeding into the model.

Results Interpretation

Before: Training loss: 120.5, Validation loss: 150.2, Validation accuracy: 65%

After: Training loss: 110.3, Validation loss: 115.7, Validation accuracy: 82%

Adding dropout and batch normalization helped reduce overfitting, improving validation loss and reconstruction accuracy. This shows how regularization and normalization stabilize training and improve generalization.

Bonus Experiment

Try using a convolutional Variational Autoencoder instead of a fully connected one to improve reconstruction quality.

💡 Hint

Replace linear layers with Conv2d and ConvTranspose2d layers to better capture image spatial features.

Practice

1. What is the main purpose of the encoder in a Variational Autoencoder (VAE)?

easy

A. To calculate the loss function

B. To produce the mean and variance for the latent space distribution

C. To generate new data samples from noise

D. To directly reconstruct the input data

Solution

1. Step 1: Understand the role of the encoder in VAE

   The encoder processes input data and outputs parameters of a probability distribution, typically mean and variance, representing the latent space.

2. Step 2: Differentiate encoder from decoder and loss

   The decoder reconstructs data, and loss measures reconstruction and latent space closeness. Encoder specifically outputs mean and variance.

3. Final Answer:

   To produce the mean and variance for the latent space distribution -> Option B

4. Quick Check:

   Encoder outputs mean and variance = D [OK]

Hint: Encoder outputs distribution parameters, not data directly [OK]

Common Mistakes:

• Confusing encoder with decoder
• Thinking encoder reconstructs data
• Assuming encoder calculates loss

2. Which PyTorch function is typically used to sample from the latent space in a VAE using the reparameterization trick?

easy

A. torch.randn()

B. torch.sigmoid()

C. torch.softmax()

D. torch.exp()

Solution

1. Step 1: Identify sampling method in VAE

   The reparameterization trick samples from a normal distribution using random noise, typically generated by torch.randn() in PyTorch.

2. Step 2: Understand other functions' roles

   torch.sigmoid() and torch.softmax() are activation functions, torch.exp() computes exponentials, none sample noise.

3. Final Answer:

   torch.randn() -> Option A

4. Quick Check:

   Sampling noise uses torch.randn() = A [OK]

Hint: Sampling noise uses torch.randn(), not activations [OK]

Common Mistakes:

• Using torch.sigmoid() to sample noise
• Confusing softmax with sampling
• Using torch.exp() incorrectly for sampling

3. Given the following PyTorch code snippet for the reparameterization trick in a VAE, what is the shape of the output z if mu and log_var have shape (batch_size, latent_dim)?

std = torch.exp(0.5 * log_var)
epsilon = torch.randn_like(std)
z = mu + epsilon * std

medium

A. (batch_size, latent_dim)

B. (latent_dim, batch_size)

C. (batch_size, 1)

D. (1, latent_dim)

Solution

1. Step 1: Analyze shapes of mu and log_var

   Both mu and log_var have shape (batch_size, latent_dim), so std and epsilon will have the same shape due to torch.exp and torch.randn_like.

2. Step 2: Calculate shape of z

   Since z = mu + epsilon * std, all operands have shape (batch_size, latent_dim), so z also has shape (batch_size, latent_dim).

3. Final Answer:

   (batch_size, latent_dim) -> Option A

4. Quick Check:

   Output shape matches input batch and latent dims = A [OK]

Hint: Output shape matches mu and log_var shapes [OK]

Common Mistakes:

• Swapping batch_size and latent_dim
• Assuming output is reduced to (batch_size, 1)
• Confusing shape broadcasting rules

4. Identify the error in this VAE loss function implementation in PyTorch:

def vae_loss(recon_x, x, mu, log_var):
    BCE = nn.functional.binary_cross_entropy(recon_x, x)
    KLD = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
    return BCE + KLD

medium

A. Incorrect sign in KLD calculation

B. Using torch.sum instead of torch.mean for KLD

C. Missing reduction='sum' in binary_cross_entropy causing scale mismatch

D. Not detaching mu and log_var before loss

Solution

1. Step 1: Check binary_cross_entropy usage

   By default, binary_cross_entropy returns mean loss per element. Summing KLD but not BCE causes scale mismatch.

2. Step 2: Fix BCE reduction

   Adding reduction='sum' to BCE matches scale with KLD, ensuring proper combined loss.

3. Final Answer:

   Missing reduction='sum' in binary_cross_entropy causing scale mismatch -> Option C

4. Quick Check:

   BCE needs reduction='sum' to match KLD scale = B [OK]

Hint: Match loss scales by using reduction='sum' in BCE [OK]

Common Mistakes:

• Ignoring BCE reduction parameter
• Changing KLD sign incorrectly
• Detaching mu/log_var unnecessarily

5. You want to train a VAE on images of size 28x28 pixels with a latent dimension of 20. Which of the following model architectures is most appropriate for the encoder part in PyTorch?

hard

A. A recurrent neural network processing images as sequences, outputting 20 values

B. A single linear layer flattening input and outputting 20 values directly

C. A fully connected network with no nonlinearities outputting 28x28 values

D. A convolutional network with Conv2d layers reducing spatial size, ending with linear layers outputting mean and log variance vectors of size 20

Solution

1. Step 1: Consider input data type and size

   Images of 28x28 pixels are spatial data; convolutional layers efficiently extract features and reduce spatial dimensions.

2. Step 2: Match output to latent dimension

   The encoder must output two vectors (mean and log variance) each of size 20, so linear layers after convolutions are suitable.

3. Step 3: Evaluate other options

   Single linear layer ignores spatial structure; RNNs are not typical for images; fully connected without nonlinearities won't learn complex features.

4. Final Answer:

   A convolutional network with Conv2d layers reducing spatial size, ending with linear layers outputting mean and log variance vectors of size 20 -> Option D

5. Quick Check:

   Conv layers + linear output for mean/log_var = C [OK]

Hint: Use conv layers for images, linear for latent params [OK]

Common Mistakes:

• Using only linear layers ignoring image structure
• Choosing RNNs for image encoding
• Omitting nonlinearities in encoder