Variational Autoencoder on CelebA

A PyTorch implementation of a Variational Autoencoder (VAE) trained on the CelebA face dataset (202,599 images). The model learns a compact 128-dimensional latent representation of human faces and demonstrates meaningful semantic arithmetic in latent space.

---

📋 Table of Contents

• Architecture
• Key Concepts Implemented
• Latent Space Experiments
• Training Details
• Project Structure
• Setup

---

🏗️ Architecture

Input Image              Latent Space              Reconstructed Image
(3 × 64 × 64)                                        (3 × 64 × 64)
      │                                                     ▲
      ▼                                                     │
 ┌─────────────────────────────────────────────────────────────┐
 │                        ENCODER                              │
 │  Conv2d(3→32, stride=2)   →  (32, 32, 32)                  │
 │  Conv2d(32→64, stride=2)  →  (64, 16, 16)                  │
 │  Conv2d(64→128,stride=2)  →  (128, 8, 8)                   │
 │  Conv2d(128→256,stride=2) →  (256, 4, 4)  → Flatten(4096)  │
 │                          ┌──────┴──────┐                   │
 │                        fc_mu        fc_logvar              │
 │                         μ (128)     log σ² (128)           │
 └─────────────────────────────────────────────────────────────┘
                               │
                  Reparameterization Trick
                  z = μ + ε·σ,  ε ~ N(0, I)
                               │
 ┌─────────────────────────────────────────────────────────────┐
 │                        DECODER                              │
 │  Linear(128 → 4096) → Unflatten → (256, 4, 4)              │
 │  ConvTranspose2d(256→128, stride=2) → (128, 8, 8)          │
 │  ConvTranspose2d(128→64,  stride=2) → (64, 16, 16)         │
 │  ConvTranspose2d(64→32,   stride=2) → (32, 32, 32)         │
 │  ConvTranspose2d(32→3,    stride=2) → (3, 64, 64) + Sigmoid│
 └─────────────────────────────────────────────────────────────┘

• Each convolutional layer is followed by BatchNorm and LeakyReLU (α = 0.2).
• Weights are initialized using Kaiming (He) initialization, suited for LeakyReLU.

---

🧠 Key Concepts Implemented

1. Reparameterization Trick

Enables gradients to flow through the stochastic sampling step:

z = μ + ε · σ     where ε ~ N(0, I)

The randomness is in ε (not learned), so gradients reach μ and σ during backpropagation.

2. ELBO Loss

Loss = Reconstruction Loss + β · KL Divergence

Reconstruction: BCE summed over pixels, averaged over batch
KL:  -0.5 · Σ (1 + log σ² - μ² - σ²)  per latent dimension

The KL term regularizes the latent space toward N(0, I), keeping it dense and traversable.

3. KL Annealing (β-warmup)

β ramps linearly from 0 → 1 over the first 5 epochs. This prevents posterior collapse — without warmup, the KL term dominates early and the encoder learns to ignore the input.

---

🔬 Latent Space Experiments

Face Generation

Sampling z ~ N(0, I) directly and decoding produces novel faces not seen during training.

Smile Arithmetic

The smile attribute is encoded as a consistent direction in latent space:

v_smile = mean(z | smiling) - mean(z | not smiling)

z_new = z_original + α · v_smile

α = -2	α = -1	α = 0	α = +1	α = +2
less smile	slight frown	original	slight smile	full smile

This demonstrates that the VAE has learned a disentangled, attribute-aligned geometry — the smile direction is approximately orthogonal to identity, hair color, and background.

Latent Space Visualization

5,000 test encodings projected to 2D via PCA and colored by the Smiling attribute show partial clustering, confirming the model has organized the latent space by facial attributes.

---

⚙️ Training Details

Hyperparameter	Value
Dataset	CelebA (202,599 face images)
Image size	64 × 64
Latent dimension	128
Batch size	128
Optimizer	Adam (lr = 1e-4)
Epochs	20
KL warmup	5 epochs
Loss (final epoch)	recon ≈ 6418, KL ≈ 114

---

📁 Project Structure

Variational-Autoencoder-CelebA/
├── vae_celeba.ipynb     # full implementation and experiments
└── requirements.txt     # Python dependencies

---

🚀 Setup

pip install torch torchvision matplotlib scikit-learn pandas opendatasets

jupyter notebook vae_celeba.ipynb

📥 The notebook downloads the dataset automatically from Kaggle (requires kaggle.json credentials).

---

Learning the geometry of faces — one latent dimension at a time. 🧬