VAE CelebA Faces 🧑‍🎨

Results at a glance: 20-dim latent space · trained on a 5K-image CelebA subset · 30 epochs · reconstruction + prior-sampled generation + latent interpolation, all in one notebook

Convolutional Variational Autoencoder trained on CelebA face images. Learns a continuous, regularized latent space that supports three things beyond plain reconstruction: generating brand-new faces by sampling directly from the prior, morphing smoothly between two real faces via latent interpolation, and diagnosing why a vanilla VAE reconstructs the way it does rather than just reporting a loss number. Runs end-to-end in Colab, no external APIs.

Overview

This project implements a VAE that learns a compressed, continuous latent representation of human faces. Unlike a plain autoencoder, a VAE's latent space is regularized to follow a known distribution (a standard normal), which means you can:

• Reconstruct existing faces by encoding and decoding them
• Generate entirely new, plausible faces by sampling directly from the latent prior
• Interpolate smoothly between two real faces by walking through latent space

Architecture: 3-layer convolutional encoder → latent distribution (μ, log σ²) → reparameterization trick → 3-layer transposed-convolutional decoder.

Architecture

3-layer convolutional encoder (32→64→128 channels) → latent distribution (μ, log σ², dim=20) → reparameterization trick → mirrored 3-layer transposed-convolutional decoder (128→64→32 channels) back to a 128×128×3 image.

Results

Trained for 30 epochs on a ~5,000-image subset of CelebA (128×128, latent dim 20).

Training curve

Reconstruction loss drops sharply over the first ~5 epochs, then plateaus almost flat from epoch ~8 onward. Train and test curves track closely throughout, so there's no overfitting — the model simply reaches the capacity limit of a 20-dimensional latent space fairly early. KL divergence stabilizes (noisily) around 65-68 with no posterior collapse, indicating the latent space stayed well-regularized rather than being ignored by the decoder.

Reconstructions

Reconstructions preserve overall face shape, skin tone, and lighting per identity, but are visibly blurry. This is expected behavior for a vanilla VAE with a pixel-wise (BCE) reconstruction loss and a small latent bottleneck: the model is incentivized toward the statistically "average" plausible face rather than a sharp, specific one — not a training bug.

Generated faces (sampled from the prior)

Faces generated by decoding z ~ N(0, I) directly (no input image) are noticeably noisier than the reconstructions, with some background artifacts. This gap between reconstruction and generation quality is a useful diagnostic: it suggests the decoder is somewhat undertrained for regions of latent space that aren't anchored to a specific encoded image, likely due to the small training subset and limited latent capacity.

Latent interpolation

Interpolating between two real faces' latent codes produces a smooth, coherent morph — hair color, hairstyle, and background shift gradually with no jarring jumps between steps. This is the strongest evidence in this project that the latent space is continuous and well-structured, which is the core property a VAE is supposed to deliver.

Known Limitations

• Reconstructions and generations are blurry. A structural property of vanilla VAEs (pixel-wise BCE loss + small latent bottleneck), not a training bug — see Results above.
• Generated (prior-sampled) faces are noisier than reconstructions. Suggests the decoder is undertrained in regions of latent space far from any real encoded image.
• Trained on a subset of CelebA (~5,000 images) rather than the full ~200k, for faster iteration in Colab.
• Loss plateaued by ~epoch 8-10 given the current architecture/data size — further gains would likely need more data or more latent capacity, not just more epochs.
• Natural extensions to address blur: a perceptual (VGG feature) loss term, a larger latent dimension, or a lower KL weight (β-VAE style) — not implemented here, but straightforward additions.

Setup

pip install -r requirements.txt

Data

This notebook needs a folder of face images (CelebA-aligned format). Two options are provided in the notebook:

• Option A — Google Drive: if you already have CelebA on your Drive, mount it and point DATA_DIR at the folder.
• Option B — Automatic fallback download: pulls a subset of CelebA from the nielsr/CelebA-faces dataset on Hugging Face Hub (no login or agreement gate required). This is the recommended path if you're just running the notebook to see it work.

Note on the dataset itself: CelebA is released for non-commercial research use only. This repo does not redistribute the dataset — the fallback option downloads it at runtime from its official Hugging Face mirror.

Usage

Open the notebook in Colab (badge above) or locally, and run cells top to bottom:

1. Config & Reproducibility
2. Data loading (Option A or B)
3. Model / loss / training definitions
4. Train (or resume from a checkpoint)
5. Visualize: loss curves, reconstructions, generated faces, interpolation

Project Structure

vae-faces/
├── VAE_on_CelebA_Faces.ipynb
├── assets/
│   ├── architecture.png 
│   ├── loss_curves.png
│   ├── reconstructions.png
│   ├── generated_faces.png
│   └── interpolation.png            
├── requirements.txt
└── LICENSE
├── .gitignore

License

Code in this repository is released under the MIT License. The CelebA dataset itself is subject to its own non-commercial research license — see the official CelebA page for terms.