2D Generative Models: DDPM, DDIM, GAN, MeanFlow

Four generative model families implemented from scratch in PyTorch and trained on the 2D checkerboard distribution — a deliberately awkward target: multi-modal, with sharp axis-aligned boundaries and large empty regions a model can easily smear across.

Because the target density is known exactly, sample quality is measurable rather than a matter of opinion. Every model is scored with two optimal-transport metrics, energy distance (ED) and 2-Wasserstein distance (WD), which turns the interesting question into a quantitative one: what does each sample actually cost in network evaluations?

Method	Steps (NFE)	ED	WD
GAN	1	0.0127	0.4700
DDPM	1000	0.0025	0.2941
DDIM	1000	0.0025	0.2940
DDIM	500	0.0017	0.2106
DDIM	100	0.0013	0.2129
DDIM	10	0.0596	0.8685
DDIM	1	1.2627	4.0667
MeanFlow	1	0.0037	0.3464

Two results stand out. DDIM at 100 steps is not merely as good as DDPM at 1000 — it is better on both metrics, so 90% of the sampling budget was adding noise rather than signal. And one-step generation is achievable, but not by truncating a diffusion trajectory: 1-step DDIM collapses (ED 1.2627) while MeanFlow reaches ED 0.0037 with a single network evaluation, three times better than the GAN. Getting to 1-NFE required changing the training objective, not shortening the sampler.

Detailed results, ablations over the DDIM η parameter and the time embedding, learning curves and visualisations are in Experiment Results below.

🧩 Setup

To install all required dependencies, simply run:

uv sync

Tip

Ensure that uv is installed beforehand: pip install uv

🏋️ Training

To start training:

# Train DDPM
uv run train.py -m ddpm

# Train GAN
uv run train.py -m gan

# Train MeanFlow
uv run train.py -m meanflow

Training configurations are defined in configs/ directory:

• configs/ddpm.yaml - DDPM configuration (7500 epochs, sinusoidal time embedding)
• configs/gan.yaml - GAN configuration (10000 epochs)
• configs/meanflow.yaml - MeanFlow configuration (7500 epochs, Huber loss)

Note

The default DDPM configuration uses Sinusoidal Positional Embedding, which produces the best model for DDIM sampling. The trained model is saved to checkpoints/diffusion.pth and can be used directly for both DDPM and DDIM inference.

Tip

To experiment with Learned Embedding, modify configs/ddpm.yaml:

model:
  time_embed_type: "learned"  # Change from "sinusoidal" to "learned"

📦 Generating Submission

To generate predictions for submission:

# GAN
bash scripts/run_problem_1.sh   # GAN

# DDPM
bash scripts/run_problem_2a.sh  # DDPM

# DDIM
bash scripts/run_problem_2b.sh  # DDIM

# MeanFlow
bash scripts/run_problem_adv.sh # MeanFlow

Or run directly with Python:

uv run generate_sample.py -m gan
uv run generate_sample.py -m ddpm
uv run generate_sample.py -m ddim
uv run generate_sample.py -m meanflow

Generated files are saved to results/:

• {method}_generated_sample.npy - Generated samples (5000 points)
• {method}_submission.csv - Sample coordinates as CSV
• {method}.png - Visualization comparing generated vs target distribution

🔬 Running DDIM Experiments

To reproduce the DDIM ablation study results (different steps and η values):

uv run run_ddim_experiments.py

This script runs two experiments:

1. Experiment 1: Different Denoising Steps - Tests DDIM with 1000, 500, 100, 10, and 1 steps (η=1.0 fixed)
2. Experiment 2: Different η Values - Tests DDIM with η ∈ {0, 0.25, 0.5, 0.75, 1.0} (50 steps fixed)

Generated comparison plots are saved to:

• results/ddim_steps_comparison.png - ED/WD vs. denoising steps
• results/ddim_eta_comparison.png - ED/WD vs. η value

Note

This script requires a trained DDPM model (checkpoints/diffusion.pth). Run uv run train.py -m ddpm first if the checkpoint doesn't exist.

🧪 Experiment Results

Main Results

Methods	#Step	ED	WD
GAN	1	0.0127	0.4700
DDPM	1000	0.0025	0.2941
DDIM	1000	0.0025	0.2940
DDIM	500	0.0017	0.2106
DDIM	100	0.0013	0.2129
DDIM	10	0.0596	0.8685
DDIM	1	1.2627	4.0667
(MeanFlow)	1	0.0037	0.3464

Observation / Insight:

1. Diffusion models outperform GAN: DDPM achieves significantly lower ED (0.0025) compared to GAN (0.0127), demonstrating better distribution matching.
2. DDIM achieves optimal performance at 100-500 steps: The best ED (0.0013) is achieved at 100 steps, suggesting diminishing returns beyond this point.
3. Quality degrades rapidly below 10 steps: DDIM with 1 step shows poor performance (ED=1.2627), indicating insufficient denoising.
4. MeanFlow achieves competitive 1-step generation: With ED=0.0037, MeanFlow significantly outperforms both GAN (0.0127) and 1-step DDIM (1.2627) for single-step generation.

Ablation Study

DDPM

#Step: 1000

Time Embedding	ED	WD
Learned Embedding	0.0052	0.2980
Sinusoidal Positional Embedding	0.0025	0.2941

Observation / Insight: Sinusoidal Positional Embedding outperforms Learned Embedding on both metrics (ED: 0.0025 vs 0.0052, WD: 0.2941 vs 0.2980). This suggests that the fixed frequency-based encoding provides better inductive bias for capturing temporal information in the diffusion process, while learned embeddings may require more training data or epochs to converge to similar performance.

DDIM

#Step: 50

eta	ED	WD
0	0.0112	0.4737
0.25	0.0067	0.3559
0.5	0.0068	0.3559
0.75	0.0045	0.3733
1	0.0046	0.3226

Observation / Insight:

1. Stochasticity improves multi-modal distribution learning: Higher η values (0.75-1.0) achieve better ED (0.0045-0.0046) compared to deterministic sampling η=0 (0.0112).
2. η=1 achieves the best Wasserstein Distance: WD=0.3226 at η=1 outperforms all other η values, suggesting that DDPM-like stochastic sampling better captures the checkerboard's sharp boundaries.
3. Deterministic DDIM (η=0) struggles with multi-modal distributions: The higher ED at η=0 indicates that deterministic sampling may collapse modes in the checkerboard pattern.

Learning Curve Analysis

Training steps vs Training Loss

• GAN (Generator Loss):

• GAN (Discriminator Loss):

• DDPM:

• MeanFlow:

Observation / Insight:

1. GAN shows adversarial dynamics: Generator and Discriminator losses oscillate as they compete, with the Discriminator typically maintaining lower loss as it learns to distinguish real from fake samples.
2. DDPM exhibits stable convergence: The MSE loss decreases smoothly and stabilizes, indicating consistent learning of the noise prediction task without the instability issues common in GANs.
3. MeanFlow converges similarly to DDPM: Using Huber loss with JVP-based training, MeanFlow shows stable convergence comparable to DDPM.

Per 100 epochs vs Energy Distance

Observation / Insight: DDPM and MeanFlow achieve lower Energy Distance more quickly than GAN. DDPM converges to the lowest ED, while MeanFlow achieves competitive performance with its 1-NFE architecture.

Per 100 epochs vs 2-Wasserstein Distance

Observation / Insight: Similar to Energy Distance, diffusion-based methods (DDPM, MeanFlow) outperform GAN in Wasserstein Distance. The gap between GAN and diffusion models highlights the advantage of iterative denoising for capturing complex multi-modal distributions.

Visualization

GIF of training progression (learned distribution over epochs)

• GAN:

• DDPM:

• MeanFlow:

Observation / Insight:

1. GAN learns the general structure quickly but may exhibit mode oscillation during training, with some regions of the checkerboard being captured better than others at different epochs.
2. DDPM shows gradual refinement: Starting from random noise, DDPM progressively learns to generate cleaner checkerboard patterns, with boundaries becoming sharper over epochs.
3. MeanFlow achieves rapid convergence: Leveraging flow matching, MeanFlow quickly learns to generate high-quality samples with its 1-NFE architecture.

Figure of final learned distribution

• GAN:

• DDPM:

• DDIM:

• MeanFlow:

Observation / Insight:

1. All models successfully learn the checkerboard pattern, demonstrating the effectiveness of generative modeling on 2D synthetic data.
2. DDPM/DDIM produce sharper boundaries compared to GAN, which may show slightly more scattered points near the edges.
3. DDIM with accelerated sampling maintains comparable quality to DDPM while using fewer denoising steps.
4. MeanFlow achieves competitive quality with 1-step generation, making it highly efficient for real-time applications.