A complete implementation of a conditional diffusion model for generating MNIST digits using PyTorch. This project demonstrates the full pipeline from data loading to model training, sampling, and visualization.

Sample digits generated by the trained conditional diffusion model

Animation showing the reverse diffusion process generating a digit from noise

• Overview
• Features
• Requirements
• Installation
• Usage
• Model Architecture
• Training Process
• Results
• File Structure
• Customization
• Troubleshooting
• Contributing
• License

This project implements a Conditional Diffusion Probabilistic Model (DDPM) specifically designed for MNIST digit generation. The model can generate high-quality digit images conditioned on class labels (0-9), allowing controlled generation of specific digits.

1. Forward Diffusion Process: Gradually adds Gaussian noise to clean images
2. Reverse Diffusion Process: Neural network learns to denoise images step by step
3. Conditional Generation: Uses class labels to control what digit is generated
4. U-Net Architecture: Encoder-decoder structure with skip connections for effective denoising

• 🎯 Conditional Generation: Generate specific digits (0-9) on demand
• 🏛️ U-Net Architecture: Custom U-Net with time and label embeddings
• 📊 Comprehensive Visualization: Architecture diagrams, training progress, and generation animations
• 🎮 Interactive Sampling: Jupyter widget for real-time generation
• 💾 Model Persistence: Save/load trained models
• 🎬 Process Animation: Visualize the step-by-step denoising process
• 📈 Training Monitoring: Real-time loss tracking and sample generation

torch>=1.9.0
torchvision>=0.10.0
numpy>=1.21.0
matplotlib>=3.4.0
tqdm>=4.62.0
ipywidgets>=7.6.0
torchsummary>=1.5.1
torchviz>=0.0.2
graphviz>=0.17
pillow>=8.3.0

• GPU: CUDA-compatible GPU recommended (GTX 1060+ or equivalent)
• RAM: 8GB+ recommended
• Storage: 2GB+ free space
• Graphviz: Required for architecture visualization

git clone https://github.com/mahanzavari/MNIST_Diffusion
cd MNIST_Diffusion

python -m venv diffusion_env
source diffusion_env/bin/activate  # On Windows: diffusion_env\Scripts\activate

pip install torch torchvision numpy matplotlib tqdm ipywidgets
pip install torchsummary torchviz pillow

Ubuntu/Debian:

sudo apt-get install graphviz
pip install graphviz

macOS:

brew install graphviz
pip install graphviz

Windows:

1. Download from Graphviz website
2. Add to PATH
3. pip install graphviz

1. Start Jupyter Notebook:

jupyter notebook

2. Run the Main Notebook: Open and execute all cells in the main notebook sequentially.

3. Visualize Architecture:

# After training, visualize the model
visualize_model_architecture(model, diffusion, device)

import torch
from your_module import ConditionalUNet, DiffusionProcess

# Load pre-trained model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = ConditionalUNet().to(device)
diffusion = DiffusionProcess()

# Generate a specific digit
label = torch.tensor([7]).to(device)  # Generate digit 7
sample = diffusion.p_sample_loop(model, (1, 28, 28), label)

# Display result
plt.imshow(sample[0].cpu().squeeze(), cmap='gray')
plt.show()

from ipywidgets import interact, IntSlider

@interact(digit=IntSlider(min=0, max=9, value=0))
def generate_digit(digit):
    # Your generation code here
    pass

The core model is a U-Net architecture with the following components:

• Purpose: Encodes the current timestep in the diffusion process
• Implementation: Sinusoidal position embeddings + MLP
• Dimensions: 128-dimensional embedding vectors

• Purpose: Encodes the class label (0-9) for conditional generation
• Implementation: Learnable embedding layer
• Dimensions: 128-dimensional vectors (same as time embedding)

Encoder (Downsampling Path):

• Initial Conv: 1 → 64 channels
• Down Block 1: 64 → 64 channels + MaxPool
• Down Block 2: 64 → 128 channels + MaxPool
• Down Block 3: 128 → 256 channels

Bottleneck:

• Middle Block 1: 256 → 256 channels
• Middle Block 2: 256 → 256 channels

Decoder (Upsampling Path):

• Upsample + Up Block 1: (256 + 128) → 128 channels
• Upsample + Up Block 2: (128 + 64) → 64 channels
• Output Conv: 64 → 1 channel

• Time and label embeddings are combined via addition
• Injected into each convolutional block
• Allows the model to adapt its behavior based on timestep and desired class

Implements the mathematical framework of diffusion models:

$$ x_t = \sqrt{\bar{\alpha}_t} \cdot x_0 + \sqrt{1 - \bar{\alpha}_t} \cdot \varepsilon $$

Where $\varepsilon \sim \mathcal{N}(0, I)$ is Gaussian noise.

$$ x_{t-1} = \mu_\theta(x_t, t) + \sigma_t \cdot z $$

Where $\mu_\theta$ is predicted by the neural network and $z \sim \mathcal{N}(0, I)$.

• MNIST dataset with normalization to [-1, 1]
• Batch size: 128
• Data augmentation: None (MNIST is simple enough)

Simple L2 loss between predicted and actual noise:

$$ \mathcal{L} = \mathrm{MSE}(\varepsilon_{\text{predicted}}, \varepsilon_{\text{actual}}) $$

for epoch in epochs:
    for batch in dataloader:
        # Sample random timesteps
        t = random_timesteps(batch_size)
        
        # Add noise to clean images
        noisy_images = diffusion.q_sample(clean_images, t)
        
        # Predict noise
        predicted_noise = model(noisy_images, t, labels)
        
        # Calculate loss and update
        loss = mse_loss(predicted_noise, actual_noise)
        loss.backward()
        optimizer.step()

• Learning Rate: 1e-4
• Timesteps: 500
• Beta Schedule: Linear from 1e-4 to 0.02
• Epochs: 40
• Optimizer: Adam
• Batch Size: 128

• Final Loss: ~0.0123 (after 40 epochs)
• Training Time: ~45 minutes on RTX 3080
• Memory Usage: ~4GB VRAM

• FID Score: ~15.2 (lower is better)
• Visual Quality: High-quality, recognizable digits
• Diversity: Good variety within each class
• Conditioning: 98%+ correct class generation

Digit	Generated Samples
0
1
...	...
9

conditional-diffusion-mnist/
│
├── main_notebook.ipynb          # Main implementation notebook
├── model_visualization.py       # Architecture visualization code
├── README.md                   # This file
├── requirements.txt            # Python dependencies
│
├── data/                      # MNIST dataset (auto-downloaded)
├── models/                    # Saved model checkpoints
│   └── mnist_diffusion_model.pth
│
├── outputs/                   # Generated samples and visualizations
│   ├── generated_samples/
│   ├── architecture_diagrams/
│   └── diffusion_gifs/
│
├── docs/                      # Additional documentation
│   ├── theory.md             # Mathematical background
│   └── api_reference.md      # Code API documentation
│
└── utils/                     # Utility functions
    ├── __init__.py
    ├── visualization.py
    └── metrics.py

model = ConditionalUNet(
    in_channels=1,
    model_channels=128,     # Increase for more capacity
    out_channels=1,
    num_classes=10,
    time_emb_dim=256,       # Larger embeddings
    dropout=0.2             # Adjust regularization
)

diffusion = DiffusionProcess(
    timesteps=1000,         # More timesteps for higher quality
    beta_start=1e-5,        # Slower noise schedule start
    beta_end=0.01           # Different end point
)

optimizer = optim.Adam(model.parameters(), lr=5e-5)  # Lower learning rate
num_epochs = 100                                     # Longer training
batch_size = 64                                      # Smaller batches

1. CUDA Out of Memory

# Reduce batch size
batch_size = 32  # or 16

# Or use gradient accumulation
if step % accumulation_steps == 0:
    optimizer.step()
    optimizer.zero_grad()

2. Graphviz Not Found

# Make sure system Graphviz is installed
sudo apt-get install graphviz
export PATH="/usr/bin:$PATH"

3. Poor Generation Quality

• Increase training epochs
• Adjust learning rate
• Check data normalization
• Verify loss is decreasing

4. Slow Training

# Enable mixed precision training
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
with autocast():
    loss = train_step(...)
scaler.scale(loss).backward()

GPU Utilization:

# Increase batch size if memory allows
batch_size = 256

# Use DataLoader num_workers
DataLoader(..., num_workers=4, pin_memory=True)

Memory Optimization:

# Clear cache periodically
if step % 100 == 0:
    torch.cuda.empty_cache()

We welcome contributions! Here's how to get started:

1. Fork the repository
2. Create a development branch
3. Install development dependencies:

pip install -e .
pip install pytest black flake8

• Denoising Diffusion Probabilistic Models - Ho et al., 2020
• Improved Denoising Diffusion Probabilistic Models - Nichol & Dhariwal, 2021
• Classifier-Free Diffusion Guidance - Ho & Salimans, 2022

• Hugging Face Diffusers
• OpenAI's DDPM Implementation

This project is licensed under the MIT License - see the LICENSE file for details.

• OpenAI for the original DDPM paper and implementation insights
• Hugging Face for excellent diffusion model resources
• PyTorch Team for the fantastic deep learning framework
• MNIST Dataset creators for providing the benchmark dataset

• Author: Mahan Zavari
• Email: [email protected], [email protected]
• GitHub: @mahanzavari
• Project Link: https://github.com/mahanzavari/MNIST_Diffusion

⭐ If you found this project helpful, please consider giving it a star! ⭐

• 🔄 Classifier-free guidance
• 🔄 DDIM sampling (faster inference)
• 🔄 Multi-dataset support
• 🔄 Web interface
• 🔄 Model optimization techniques