DDPM from Scratch: MNIST Implementation

A clean, modular implementation of Denoising Diffusion Probabilistic Models (DDPM) applied to the MNIST dataset. This project implements the core mathematics and architectural requirements described in the seminal paper "Denoising Diffusion Probabilistic Models" by Ho et al.

🚀 Overview

This repository provides a complete pipeline for training a generative model that can synthesize handwritten digits from pure Gaussian noise. Unlike standard GANs, DDPMs learn to reverse a gradual multi-step corruption process.

Key Features

• Modular Architecture: Separate components for the UNet, Gaussian Scheduler, and Data Pipeline.
• Optimized UNet: Features Sinusoidal Positional Timestep Embeddings, Group Normalization, and SiLU (Swish) activations for training stability.
• Kaggle Ready: Includes a consolidated script optimized for Kaggle's T4 GPUs.
• Mathematical Precision: Accurate implementation of the Langevin dynamics reverse step.

---

🧠 Mathematical Foundations

1. The Forward Process (Diffusion)

We gradually add Gaussian noise to a clean image $x_0$ over $T$ timesteps ($T=1000$) according to a linear variance schedule $\beta_1, \dots, \beta_T$: $$q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1 - \beta_t} x_{t-1}, \beta_t \mathbf{I})$$

The model is trained to predict the noise $\epsilon$ added to $x_0$ to produce $x_t$ using the closed-form sampling property: $$x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon, \quad \epsilon \sim \mathcal{N}(0, \mathbf{I})$$

2. The Reverse Process (Sampling)

The model $\epsilon_\theta(x_t, t)$ learns to estimate the noise. We sample starting from $x_T \sim \mathcal{N}(0, \mathbf{I})$ and iterate backwards: $$x_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}t}} \epsilon\theta(x_t, t) \right) + \sigma_t z, \quad z \sim \mathcal{N}(0, \mathbf{I})$$

---

🏗️ Architecture: UNet

The backbone is a symmetrical UNet designed for $28 \times 28$ inputs:

• Encoder: Two downsampling stages using MaxPool2d and residual-style Blocks.
• Bottleneck: Dual Block layers to capture high-level latent representations.
• Decoder: ConvTranspose2d for spatial recovery with skip connections to preserve fine-grained details.
• Normalization: GroupNorm (8 groups) is used throughout to stabilize training across various batch sizes.
• Activation: SiLU (Sigmoid Linear Unit) for smooth gradient flow.

---

🛠️ Installation & Usage

Setup

git clone https://github.com/sharksurfauto-byte/ddpm_implementation_on_MNIST.git
cd ddpm_implementation_on_MNIST
pip install torch torchvision tqdm

Training

To train the model on your local machine or a server:

python train.py

The script will automatically download MNIST and save checkpoints every 5 epochs.

Sampling

To generate new images using a trained checkpoint:

python sample.py

Outputs will be saved as ddpm_generated_digits.png.

Kaggle Execution

For immediate training on Kaggle, use the kaggle_ddpm_mnist.py script. It is a standalone file that can be pasted into a single cell.

---

📊 Results

After 20 epochs with a learning rate of 2e-4, the model successfully transitions from static noise to recognizable digit structures, demonstrating the effective learning of the MNIST data manifold.

At T=1000:

At T=0:

📜 References

• Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models.
• Vaswani, A., et al. (2017). Attention is All You Need (for Sinusoidal Embeddings).