Reverse Diffusion on a 2D Spiral

A minimal from-scratch DDPM, trained on points sampled along a 2D spiral, animated with Manim to show randomly initialized noise being denoised back onto the spiral step by step.

Two sampling variants are compared side by side using the same trained model and the same starting noise:

• Stochastic (DDPM ancestral sampling) — random Gaussian noise is injected at every reverse step, producing visibly jagged, random-walk-like point trajectories.
• Deterministic (DDIM, η=0) — no noise is injected; the same trained model instead follows a smooth, direct path from noise to data.

What's here

• dataset.py — samples points along an Archimedean spiral with a small amount of noise
• model.py — a tiny time-conditioned MLP (~a few hundred parameters) that predicts noise given (x, t)
• diffusion.py — the DDPM forward/reverse process (q_sample, ancestral p_sample_step, and the deterministic DDIM p_sample_step_deterministic)
• train.py — trains the model on the spiral dataset
• sample.py / sample_deterministic.py — run the reverse process from random noise, saving the full per-step trajectory of every point (not just the final result) for animation
• spiral_diffusion.py — the Manim scenes (ReverseDiffusion, ReverseDiffusionDeterministic) that animate those trajectories, with fading comet-tails showing each point's path

Usage

pip install -r requirements.txt

python train.py                    # trains the model -> checkpoints/diffusion.pt
python sample.py                   # stochastic (DDPM) trajectory -> data/trajectory.npy
python sample_deterministic.py     # deterministic (DDIM eta=0) trajectory -> data/trajectory_deterministic.npy

manim -qh spiral_diffusion.py ReverseDiffusion               # renders media video (stochastic)
manim -qh spiral_diffusion.py ReverseDiffusionDeterministic  # renders media video (deterministic)

Training takes well under a minute on a GPU (the model has on the order of tens of thousands of parameters); each render takes a couple of minutes.

Why this exists

Diffusion is usually explained through image models, where the actual denoising trajectory is invisible inside a high-dimensional latent space. Here the data is 2D, so every step of the reverse process can be plotted directly — a random point drifting through noisy or smooth intermediate positions until it lands on the data manifold, which makes the mechanics of the forward/reverse process (and the difference between ancestral and deterministic samplers) directly visible rather than something you have to take on faith.