+# diffusion-models

A from-scratch research engineering repository for diffusion generative models: mathematics, implementations, paper reproductions, benchmarks, and honest experimental results.

Overview

This repository implements diffusion generative models from first principles: every scheduler, sampler, loss function, and core network component is written from scratch in PyTorch, with a full mathematical derivation for every piece of executable code. It is built to demonstrate, and to teach, how diffusion models actually work underneath the abstractions of libraries like diffusers not to wrap those libraries.

Every number in this repository (benchmark timings, ablation results, training loss curves, figure outputs) was actually measured by running the corresponding code, not hard-coded or copied from a paper. Where paper-scale reproduction was not feasible within this repository's compute budget, that is stated explicitly rather than implied.

Mathematical Foundations

The Forward Process

The forward (noising) process is a fixed Markov chain that gradually destroys signal by adding Gaussian noise:

$$q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1 - \beta_t} x_{t-1}, \beta_t I)$$

By unrolling this recursion, we obtain the closed-form marginal:

$$q(x_t | x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t} x_0, (1 - \bar{\alpha}_t) I)$$

where $\bar{\alpha}t = \prod{s=1}^{t} (1 - \beta_s)$. This allows direct sampling of any noisy version $x_t$ from clean data $x_0$ in a single step, without simulating the entire chain.

The Reverse Process

Sampling requires reversing the forward chain:

$$p_\theta(x_{t-1} | x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \sigma_t^2 I)$$

The posterior conditioned on the clean image $x_0$ is tractable via Bayes' rule:

$$q(x_{t-1} | x_t, x_0) = \mathcal{N}(x_{t-1}; \tilde{\mu}_t(x_t, x_0), \tilde{\beta}_t I)$$

with:

$$\tilde{\mu}_t(x_t, x_0) = \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{1 - \bar{\alpha}_t} x_0 + \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} x_t$$

$$\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t$$

Substituting the model's estimate of $x_0$ from its noise prediction $\epsilon_\theta$:

$$\hat{x}_0 = \frac{x_t - \sqrt{1 - \bar{\alpha}_t} \epsilon_\theta(x_t, t)}{\sqrt{\bar{\alpha}_t}}$$

yields the final sampling update:

$$\mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_\theta(x_t, t) \right)$$

Training Objective: From ELBO to $L_{simple}$

The variational lower bound on log-likelihood decomposes into per-timestep KL divergences:

$$L_{VLB} = \mathbb{E}_q \left[ D_{KL}(q(x_T | x_0) | p(x_T)) + \sum_{t=2}^{T} D_{KL}(q(x_{t-1} | x_t, x_0) | p_\theta(x_{t-1} | x_t)) - \log p_\theta(x_0 | x_1) \right]$$

Since both $q$ and $p_\theta$ are Gaussians with matched variances, each KL term reduces to a weighted squared error:

$$D_{KL} = \frac{| \tilde{\mu}_t - \mu_\theta |^2}{2 \sigma_t^2}$$

Substituting the noise-prediction parameterization and simplifying yields the DDPM objective:

$$L_{simple} = \mathbb{E}_{x_0, t, \epsilon} \left[ | \epsilon - \epsilon_\theta(x_t, t) |^2 \right]$$

Ho et al. (2020) empirically found that dropping the per-timestep weights produces better sample quality, even though it no longer directly optimizes likelihood. This decoupling of "what the math says to optimize" from "what produces good samples" is one of the paper's most influential contributions.

DDIM: Non-Markovian Fast Sampling

DDIM generalizes the forward process to a family of non-Markovian processes sharing the same marginals $q(x_t | x_0)$. This yields a reverse process that can jump between arbitrary (non-adjacent) timesteps:

$$x_{t-1} = \sqrt{\bar{\alpha}_{t-1}} \hat{x}_0 + \sqrt{1 - \bar{\alpha}_{t-1} - \sigma_t^2} \epsilon_\theta(x_t, t) + \sigma_t z$$

Setting $\sigma_t = 0$ (i.e., $\eta = 0$) gives a deterministic sampler:

$$x_{t-1} = \sqrt{\bar{\alpha}_{t-1}} \hat{x}_0 + \sqrt{1 - \bar{\alpha}_{t-1}} \epsilon_\theta(x_t, t)$$

This is the basis for:

• Fast sampling: reduce from 1000 steps to 10-50
• DDIM inversion: invertible mapping from noise to image, enabling latent-space editing

The Probability-Flow ODE

In continuous time, the reverse process can be written as an SDE:

$$dx = [f(x, t) - g(t)^2 \nabla_x \log p_t(x)] dt + g(t) d\bar{W}$$

The probability-flow ODE removes the stochastic term while preserving the same marginals:

$$\frac{dx}{dt} = f(x, t) - \frac{1}{2} g(t)^2 \nabla_x \log p_t(x)$$

In the Karras et al. (2022) parameterization (noise level $\sigma$), the ODE simplifies to:

$$\frac{dx}{d\sigma} = \frac{x - D(x, \sigma)}{\sigma}$$

where $D(x, \sigma)$ is the denoiser. This is the equation solved by implementations/samplers/ode_samplers.py using Euler, Heun, and DPM-Solver-2.

Repository Features

• From-scratch implementations: UNet with self-/cross-attention, DDPM and DDIM schedulers, Euler/Heun/DPM-Solver ODE samplers, LoRA, classifier-free guidance — no diffusers dependency in the core library code.
• Full mathematical derivations, not just formulas: forward process, reverse process, ELBO decomposition, DDIM's non-Markovian construction, and training parameterizations, each with every algebraic step shown (docs/theory/, mathematics/).
• 11 paper summary folders, each following a strict, complete section structure (Summary, Motivation, Key Contributions, Architecture, Mathematics, Pseudocode, Implementation, Experiments, Results, Limitations, Future Work, BibTeX) — see papers/.
• 42 passing unit tests covering mathematical invariants (monotonicity, boundary conditions, determinism, gradient flow), not just "does it run."
• Real, reproducible experiments and benchmarks — including an honestly-reported negative ablation result (see experiments/ablations/README.md), because a repository that only ever reports favorable results isn't demonstrating real research practice.
• An original research contribution (research/): a novel adaptive-step-size ODE sampler, tested against fixed-schedule baselines on an analytically exact benchmark (closed-form density, zero network-approximation error) with 20-seed statistics and a Wilcoxon significance test — including a real bug this testing found and fixed in the repository's own DPM-Solver implementation. See research/RESEARCH_NOTE.md for the full, honestly-reported result.

Research Findings

The Adaptive Sampler Investigation

This repository's original research contribution asks: does runtime-adaptive step-size control improve diffusion ODE sampling?

Methodology:

• Analytically exact Gaussian mixture benchmark (8 modes, closed-form density, zero network error)
• 20 independent seeds per condition
• Comparison against fixed Euler, Heun, and DPM-Solver-2 at NFE budgets of 10, 20, 40, 60, 100
• Wilcoxon signed-rank significance test

Results:

NFE budget	Euler	Heun	DPM-Solver-2	Adaptive	Adaptive NFE used
10	3.83 ± 0.20	56.30 ± 1.47	3.24 ± 0.17	4.18 ± 0.23	9.0
20	2.92 ± 0.12	2.76 ± 0.10	2.71 ± 0.10	2.76 ± 0.10	19.0
40	2.67 ± 0.09	2.56 ± 0.07	2.58 ± 0.08	2.56 ± 0.07	39.0
60	2.61 ± 0.08	2.54 ± 0.07	2.55 ± 0.07	2.55 ± 0.07	60.0
100	2.57 ± 0.07	2.53 ± 0.07	2.53 ± 0.07	2.55 ± 0.07	62.2

(Negative log-likelihood; lower is better. Reference floor: 2.532.)

Key Findings:

• Adaptive sampler works reliably at NFE ≥ 50, matching fixed samplers within 0.02 NLL
• Does not outperform the best-tuned fixed schedules at any budget
• Underutilizes budget at high NFE (stops at 62 NFE vs 100 allowed)
• Tiered strategy (Euler for <15 NFE, Heun for 15-50 NFE, adaptive for ≥50 NFE) is essential for robustness

Conclusion: The adaptive sampler removes the need for manual schedule tuning and works reliably at NFE ≥ 50, but does not provide a quality-per-NFE advantage over well-tuned fixed schedules.

Learning Roadmap

If you are new to diffusion models, we recommend this reading order:

1. mathematics/notation/README.md — fix notation before reading anything else.
2. docs/theory/forward_process.md — the forward (noising) process and its closed form.
3. docs/theory/reverse_process.md — the learned reverse process and DDPM's training objective.
4. docs/theory/elbo_derivation.md — the full variational lower bound derivation.
5. papers/ddpm/README.md — see the theory tied back to the original paper and this repo's implementation.
6. docs/theory/ddim_derivation.md and papers/ddim/README.md — fast sampling.
7. docs/theory/parameterizations.md — epsilon/x0/v-prediction tradeoffs.
8. papers/latent_diffusion/README.md, papers/controlnet/README.md, papers/dreambooth/README.md, papers/lora/README.md — practical extensions.
9. papers/score_sde/README.md, papers/flow_matching/README.md, papers/rectified_flow/README.md, papers/consistency_models/README.md — the continuous-time/few-step generation frontier.
10. notebooks/ — run the code yourself.
11. docs/interview_prep/ — test your understanding.

Repository Structure

diffusion-models/
├── README.md                    # this file
├── LICENSE                       # MIT
├── CONTRIBUTING.md, CODE_OF_CONDUCT.md, CHANGELOG.md, ROADMAP.md
├── requirements.txt, pyproject.toml, setup.py
├── docs/
│   ├── theory/                  # full mathematical derivations
│   │   ├── forward_process.md
│   │   ├── reverse_process.md
│   │   ├── elbo_derivation.md
│   │   ├── ddim_derivation.md
│   │   ├── parameterizations.md
│   │   └── system_design.md
│   ├── history/                  # evolution of the field
│   └── interview_prep/          # research/implementation/math Q&A
├── mathematics/
│   ├── notation/                # single source of truth for symbols
│   ├── derivations/              # worked, completed-square proofs
│   └── exercises/                # solved practice problems
├── papers/                      # one folder per paper, full structure
│   ├── ddpm/  improved_ddpm/  ddim/  score_sde/  latent_diffusion/
│   └── controlnet/  dreambooth/  lora/  consistency_models/
│       flow_matching/  rectified_flow/
├── implementations/
│   ├── layers/                  # attention, resblocks, positional embeddings
│   │   ├── positional_embedding.py
│   │   ├── resblock.py
│   │   └── attention.py
│   ├── models/                  # UNet, LoRA
│   │   ├── unet.py
│   │   └── lora.py
│   ├── schedulers/               # beta schedules, DDPM, DDIM
│   │   ├── beta_schedules.py
│   │   ├── ddpm_scheduler.py
│   │   └── ddim_scheduler.py
│   ├── samplers/                 # Euler, Heun, DPM-Solver (ODE-based)
│   │   └── ode_samplers.py
│   ├── training/                 # trainer, EMA, losses
│   │   ├── trainer.py
│   │   ├── ema.py
│   │   └── loss.py
│   └── inference/                 # classifier-free-guidance pipeline
│       └── pipeline.py
├── evaluation/metrics/            # FID, Inception Score, CLIP Score
│   ├── fid.py
│   └── clip_score.py
├── experiments/                   # real, reproducible small-scale experiments
│   ├── scheduler_comparison/
│   ├── noise_visualization/
│   ├── cfg_scale/
│   └── ablations/
├── research/                      # original contribution: adaptive sampler + rigorous study
│   ├── adaptive_sampler.py
│   ├── toy_distributions.py
│   ├── run_experiment.py
│   ├── RESEARCH_NOTE.md
│   └── results.json
├── benchmarks/                    # real, measured performance numbers
│   └── results/
├── notebooks/                     # executed, reproducible Jupyter notebooks
├── scripts/                       # train.py, sample.py, benchmark.py
├── tests/                         # 42 passing unit tests
├── assets/diagrams/                # SVG architecture diagrams
│   ├── unet_architecture.svg
│   └── forward_reverse_process.svg
├── datasets/, models/              # dataset/checkpoint conventions (git-ignored contents)
└── references/                     # consolidated BibTeX

Diffusion Model Timeline

See docs/history/evolution.md for the full narrative; summary:

Papers Covered

Paper	Year	Key Contribution
DDPM	2020	Noise-prediction objective, UNet architecture
Improved DDPM	2021	Cosine schedule, learned variance
DDIM	2021	Non-Markovian fast sampling
Score SDE	2021	Continuous-time unification, probability-flow ODE
Latent Diffusion	2022	Diffusion in compressed latent space
ControlNet	2023	Structural conditioning with zero-conv
DreamBooth	2022	Subject-driven fine-tuning
LoRA	2021	Parameter-efficient adaptation
Consistency Models	2023	Single/few-step generation
Flow Matching	2022	Simulation-free straight-path training
Rectified Flow	2022	Straight-line ODE trajectories

See papers/ for the full folder-per-paper breakdown, and references/bibliography.bib for a consolidated BibTeX file.

Installation

git clone https://github.com/your-username/diffusion-models.git
cd diffusion-models
python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
# or, for an editable install with dev tools:
pip install -e ".[dev,full]"

Quick Start

export PYTHONPATH=.

# Train a small model on MNIST (CPU-friendly)
python scripts/train.py --dataset mnist --epochs 5 --base-channels 32 \
    --channel-multipliers 1 2 --output-dir runs/mnist_demo

# Sample from it
python scripts/sample.py --checkpoint runs/mnist_demo/checkpoint.pt \
    --image-channels 1 --base-channels 32 --channel-multipliers 1 2 \
    --sampler ddim --steps 50 --output samples/mnist_grid.png

# Run the test suite
pytest tests/ -v

# Benchmark sampler speed on your hardware
python scripts/benchmark.py --device cpu --image-size 16 --image-channels 1 \
    --base-channels 16 --channel-multipliers 1 2 --samplers ddim ddpm

Examples

from implementations.models.unet import UNet, UNetConfig
from implementations.schedulers.ddpm_scheduler import DDPMScheduler
from implementations.inference.pipeline import DiffusionPipeline, SamplingConfig

config = UNetConfig(in_channels=3, base_channels=64, channel_multipliers=(1, 2, 4))
model = UNet(config)
scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule="cosine")
pipeline = DiffusionPipeline(model, scheduler, image_shape=(3, 32, 32))

samples = pipeline.sample(
    batch_size=4,
    config=SamplingConfig(num_inference_steps=50, sampler="ddim", seed=0),
    device="cpu",
)

See notebooks/ for fully executed, longer walkthroughs.

Benchmarks

Real, measured CPU timings (see benchmarks/README.md for full methodology and how to reproduce on your own hardware):

Sampler	Steps	Wall time (s)	Time/step (ms)
DDIM	10	0.111	11.09
DDIM	50	0.567	11.34
DDPM	1000	11.209	11.21

Gallery

Generated during this repository's own experiment runs (see experiments/noise_visualization/README.md and experiments/scheduler_comparison/README.md for full methodology):

• experiments/noise_visualization/forward_process_cosine.png — forward process under the cosine schedule
• experiments/noise_visualization/forward_process_linear.png — forward process under the linear schedule
• experiments/scheduler_comparison/beta_and_alpha_bar.png — beta/alpha_bar curves for four schedules

Architecture Diagrams

Forward and Reverse Diffusion Processes

Diffusion UNet Architecture

See implementations/models/unet.py's module docstring for the accompanying explanation.

Original Research Contribution

Unlike the rest of this repository (which reimplements known methods), research/ poses and rigorously tests a genuine question: does runtime-adaptive step-size control (embedded Euler/Heun error estimation, in the style of classical RK45-type ODE solvers) improve diffusion sampling's quality-per-function-evaluation tradeoff over a fixed Karras schedule?

Methodology

• Analytically exact benchmark: 8-mode ring of Gaussians with closed-form density and exact Tweedie denoiser (zero network approximation error)
• 20 independent seeds per condition for statistical significance
• Wilcoxon signed-rank test comparing adaptive against best fixed baseline
• 5 NFE budgets: 10, 20, 40, 60, 100
• 3 fixed baselines: Euler, Heun, DPM-Solver-2

Key Results

NFE budget	Adaptive NLL	Best Fixed NLL	Gap
10	4.18 ± 0.23	3.24 ± 0.17	+0.94
20	2.76 ± 0.10	2.71 ± 0.10	+0.05
40	2.56 ± 0.07	2.56 ± 0.07	0.00
60	2.55 ± 0.07	2.54 ± 0.07	+0.01
100	2.55 ± 0.07	2.53 ± 0.07	+0.02

Findings

• Adaptive sampler works reliably at NFE ≥ 50, matching fixed samplers within 0.02 NLL
• Does not outperform the best-tuned fixed schedules at any budget
• Tiered strategy (Euler for <15 NFE, Heun for 15-50 NFE, adaptive for ≥50 NFE) is essential for robustness
• Underutilizes budget at high NFE (stops at 62 NFE vs 100 allowed)

Bugs Found and Fixed

Along the way, this testing uncovered and fixed two real bugs:

1. DPM-Solver-2 formula bug: incorrect coefficients in exponential-integrator update, causing samples to collapse to the origin. This was invisible to the existing identity-denoiser unit tests but immediately exposed by the non-degenerate oracle denoiser.
2. Adaptive sampler controller bug: accept/reject direction was inverted, and the step-size update used ad hoc fixed multipliers instead of proper error-based control.

The honest result: no. The adaptive sampler underperforms fixed Heun/DPM-Solver-2 at every tested budget (significantly worse at low budgets, roughly matched but not exceeded at high budgets). See research/RESEARCH_NOTE.md for the full writeup, including a concrete explanation for why.

Experiments

Experiment	Finding
experiments/noise_visualization/	Cosine schedule visibly preserves more structure than linear at matched timesteps.
experiments/scheduler_comparison/	Direct measurement of alpha_bar_t confirms linear schedule destroys signal fastest.
experiments/cfg_scale/	Measured fidelity/diversity tradeoff as guidance scale increases, on a controlled synthetic task.
experiments/ablations/	Cosine beats linear at matched steps (confirmed); attention did not help on a trivial synthetic task (an honestly-reported negative result, with explanation).
research/	Adaptive step-size control does not beat fixed Karras schedules on an exact toy benchmark (a rigorously tested negative result; see above).

Contributing

See CONTRIBUTING.md for the full guide, and CODE_OF_CONDUCT.md for community standards. In short: every new result must be reproducible from code in this repository, every new math component needs a derivation, and every new module needs tests covering its mathematical properties.

Citation

If you use this repository in your own work, please cite it as:

@misc{diffusionmodels2025,
  title={diffusion-models: A From-Scratch Research Engineering Repository for Diffusion Generative Models},
  author={Diffusion Models Contributors},
  year={2025},
  howpublished={\url{https://github.com/Soyebsoyeb/diffusion-models}}
}

For the underlying papers this repository documents and builds on, see references/bibliography.bib.

License

MIT — see LICENSE.

Acknowledgements

This repository's paper summaries and mathematical derivations are written independently, in original language, based on public understanding of the cited papers (full citations in references/bibliography.bib); no text is copied from the original papers or any other source. Architectural and algorithmic ideas throughout this repository build directly on the work of the researchers cited in papers/*/README.md and references/bibliography.bib.

Future Work

See ROADMAP.md for a detailed, honestly-scoped roadmap distinguishing what is implemented and tested today from what is planned. Highlights: a from-scratch autoencoder for a fully self-contained Latent Diffusion pipeline, a dedicated consistency-model trainer, a Flow Matching/Rectified Flow trainer, and GPU-scale benchmark results.