Diffusion Models: A Comparative Study

A comprehensive implementation and comparison of three diffusion model approaches on a 2D smiling face distribution, designed for educational and research purposes.

Overview

This project implements and compares three fundamental approaches to generative modeling using diffusion processes:

1. DDPM (Denoising Diffusion Probabilistic Models) - The standard forward/reverse diffusion process
2. SM (Naive Score Matching) - Direct score function learning with inherent limitations
3. NCSM (Noise Conditional Score Matching) - Multi-scale score learning to address SM limitations

All models are trained on a custom 2D "smiling face" probability distribution that features two circular eye modes and a half-moon mouth mode, providing an intuitive visualization of generative model performance.

Project Structure

diffusion/
├── models/                    # Model implementations
│   ├── ddpm.py               # DDPM with noise scheduler and U-Net-style denoising
│   ├── sm.py                 # Naive score matching with multiple loss variants
│   └── ncsm.py               # Noise conditional score matching with annealed sampling
├── dataset.py                # SmilingFaceDataset with analytical PDF and score functions
├── train.py                  # Unified training script for all models
├── sampling.py               # Langevin dynamics and DDPM sampling implementations
├── vis.py                    # Comprehensive visualization and animation system
├── checkpoints/              # Trained model checkpoints
├── figures/                  # Generated visualizations and animations
└── script/                   # Jupyter notebook for experimentation

Key Interfaces

Training Interface (train.py)

Unified Training Command:

python train.py --model <model_type>

Model Types:

• ddpm: Denoising Diffusion Probabilistic Models
• sm: Naive Score Matching
• ncsm: Noise Conditional Score Matching

Key Configuration Parameters:

• num_epochs: Training duration
• batch_size: Batch size for training
• learning_rate: Optimizer learning rate
• hidden_dim: Neural network hidden dimension
• save_every: Checkpoint saving frequency

Installation

1. Install dependencies:

pip install -r requirements.txt
# or with uv
uv pip install -e .

2. Required packages:

• PyTorch >= 2.7.1
• NumPy >= 2.0.2
• Matplotlib >= 3.9.4
• scikit-learn >= 1.6.1
• tqdm >= 4.66.0

Usage Examples

Training Models

# Train DDPM
python train.py --model ddpm

# Train Naive Score Matching
python train.py --model sm

# Train Noise Conditional Score Matching  
python train.py --model ncsm

Using Trained Models

import torch
from models.ddpm import DiffusionModel
from dataset import SmilingFaceDataset
from sampling import create_ddpm_sampler

# Load dataset
dataset = SmilingFaceDataset(device='cuda')

# Load trained model
model = DiffusionModel.load_checkpoint('checkpoints/ddpm/checkpoint_step_40000.pt')

# Create sampler and generate samples
sampler = create_ddpm_sampler(model, device='cuda')
samples, info = sampler.sample(num_samples=1000)

# Visualize results
dataset.plot_comparison(samples)

Key Features

• Unified Training Framework: Single script supports all three model types with consistent configuration
• Analytical Ground Truth: Custom 2D dataset with closed-form PDF and score functions for validation
• Comprehensive Visualization: Automatic generation of loss curves, sample comparisons, score field visualizations, and sampling animations
• Flexible Sampling: Multiple sampling algorithms with trajectory tracking and visualization
• Model Checkpointing: Automatic saving and loading of training checkpoints
• Educational Focus: Clear implementations with extensive documentation and visualization for understanding diffusion concepts

Model Comparisons

The implementation allows direct comparison of:

• Training stability across different approaches
• Sample quality through visual and quantitative metrics
• Score function accuracy compared to analytical ground truth
• Sampling efficiency and convergence properties
• Computational requirements for training and inference

Generated Outputs

The system automatically generates:

• Training loss curves
• Sample comparison plots
• Score field visualizations
• Sampling process animations
• Model performance metrics
• Checkpoint files for model persistence

---

This codebase serves as both a research tool and educational resource for understanding the foundations and practical implementation of diffusion-based generative models.