mini-diffusion

A small educational lab that visualizes forward and reverse diffusion on images. It is meant to build intuition for the math, not to be a full trained generative model like DDPM.

Script: step_by_step_diffusion.py

• Forward (Markov): start from a clean image (x_0) and add noise step by step using a schedule (\beta_t).
• Reverse (oracle): sample from the true posterior (q(x_{t-1} \mid x_t, x_0)) when (x_0) is known — exact Gaussian “denoising” steps, no neural network.
• Closed form: (x_t = \sqrt{\bar{\alpha}_t}, x_0 + \sqrt{1-\bar{\alpha}_t},\varepsilon) with a fixed (\varepsilon) across (t).

In real DDPM / Stable Diffusion, (x_0) is unknown at sampling time; a network (\varepsilon_\theta(x_t, t)) replaces the oracle. This repo does not train that network.

Introductory write-up (Vietnamese): Diffusion Models cơ bản (Viblo).

Simple math (what gets computed)

This repo does not train neural weights (\theta). In full DDPM, (\theta) is learned by minimizing noise-prediction loss; here we only use fixed formulas.

Schedule (same as linear_beta_schedule in code):

1. Choose (T) (--timesteps). Build (\beta_1,\ldots,\beta_T) as evenly spaced numbers from 1e-4 to 0.02.
2. (\alpha_t = 1 - \beta_t).
3. (\bar{\alpha}_t = \alpha_1 \cdot \alpha_2 \cdots \alpha_t) (cumulative product — np.cumprod).

One noisy image at step (t) (closed form, used in the second figure):

[ x_t = \sqrt{\bar{\alpha}_t}, x_0 + \sqrt{1 - \bar{\alpha}_t},\varepsilon,\quad \varepsilon \sim \mathcal{N}(0, I). ]

So: more noise when (t) is large, because (\bar{\alpha}_t) gets smaller (more of (\varepsilon), less of (x_0)).

Tiny numeric toy ((T=4), np.linspace(1e-4, 0.02, 4)):

step (t)	(\beta_t)	(\alpha_t=1-\beta_t)	(\bar{\alpha}_t)
1	0.0001	0.9999	0.9999
2	0.006733	0.993267	0.993167
3	0.013367	0.986633	0.979892
4	0.020000	0.980000	0.960294

(\bar{\alpha}_t) shrinks as (t) increases → the closed-form mix uses more (\varepsilon) and less (x_0).

Forward Markov step (first figure, one step at a time):

[ x_t = \sqrt{\alpha_t}, x_{t-1} + \sqrt{\beta_t},\varepsilon. ]

That is all the arithmetic the demo needs — no matrices, no training loop.

Setup

python -m venv .venv
source .venv/bin/activate   # Windows: .venv\Scripts\activate
pip install -r requirements.txt

torch and torchvision are required because the default data source is MNIST.

Quick start

By default, one MNIST training image is used; PNGs are written to ./out:

python step_by_step_diffusion.py

On first run, MNIST is downloaded into ./data/mnist (override with --mnist-root).

Code flow (diagrams)

Overview of step_by_step_diffusion.py: one script, two figures. GitHub renders the Mermaid blocks below.

Control flow

flowchart TD
  subgraph entry["Entry"]
    A["main()"] --> B["argparse: out, timesteps, source, mnist-*, closed-form-*"]
    B --> C["run(...)"]
  end

  subgraph data["Load x_0"]
    C --> D{"--source"}
    D -->|mnist| E["load_mnist_x0()"]
    D -->|synthetic| F["make_demo_image()"]
    E --> G["Schedule"]
    F --> G
  end

  subgraph sched["Noise schedule"]
    G --> H["linear_beta_schedule()"]
    H --> I["alphas, alphas_cumprod"]
  end

  subgraph chain["Markov + oracle"]
    I --> J["forward_markov_chain"]
    J --> K["reverse_oracle_chain + posterior_coefs"]
  end

  subgraph plots["Figures"]
    J --> M["forward_and_reverse_steps.png"]
    K --> M
    I --> N["Closed-form: linspace t, fixed epsilon"]
    N --> O["closed_form_forward.png"]
  end

Loading

Sequence (call order)

sequenceDiagram
  autonumber
  participant User
  participant main as main
  participant run as run
  participant load as load_x0
  participant fwd as forward_markov_chain
  participant rev as reverse_oracle_chain
  participant plt as matplotlib

  User->>main: CLI args
  main->>run: run(...)
  alt MNIST
    run->>load: load_mnist_x0
  else synthetic
    run->>load: make_demo_image
  end
  load-->>run: x0 in -1..1
  run->>run: linear_beta_schedule, alphas_cumprod
  run->>fwd: x0, schedule, rng
  fwd-->>run: fwd states
  run->>rev: x0, fwd last state, schedule
  rev-->>run: reverse states
  run->>plt: figure 1 PNG
  run->>run: closed-form t grid, fixed eps
  run->>plt: figure 2 PNG

Loading

Data flow (math)

flowchart LR
  subgraph forward["Forward Markov"]
    x0["x_0"] --> x1["x_1"]
    x1 --> xd["..."]
    xd --> xT["x_T"]
  end

  subgraph oracle["Oracle reverse"]
    xT --> xr["q(x_{t-1}|x_t, x_0)"]
    xr --> x0b["x_0"]
  end

  subgraph closed["Closed form"]
    x0c["x_0"] --> xt["x_t"]
    eps["fixed epsilon"] --> xt
    ab["alpha_bar_t"] --> xt
  end

Loading

Outputs

PNG files are written to --out (default ./out). Example run (synthetic source, --timesteps 20):

Top: noise is added each Markov step. Bottom: sampling from (q(x_{t-1}\mid x_t, x_0)) with known (x_0) (not a neural net).

File	Description
forward_and_reverse_steps.png	Top: forward chain; bottom: oracle reverse
closed_form_forward.png	Same (\varepsilon), several (t) values (range is configurable)

Common CLI options

Flag	Default	Notes
--out	./out	Output directory for PNGs
--timesteps	200	Number of noise steps (T) (increase for larger (t) indices)
--source	mnist	synthetic = procedural RGB image
--mnist-root	./data/mnist	MNIST storage / download path
--mnist-index	0	Index into the training set (0 … 59999)
--size	48	Square side length after resize
--closed-form-panels	8	Number of panels in closed_form_forward.png
--closed-form-t-min	0	Smallest (t) in the closed-form figure
--closed-form-t-max	-1	Largest (t) (-1 means (T-1))
--seed	0	RNG seed

Example: focus on large (t) (almost pure noise), with (T=400):

python step_by_step_diffusion.py --timesteps 400 \
  --closed-form-t-min 250 --closed-form-t-max -1

Requirements

• Python 3.10+
• See requirements.txt