Diffusion Models

Destroy, then learn to un-destroy

Imagine you are handed a photograph — a crisp, clean image of a cat. Now imagine you sprinkle a tiny amount of static over it: the cat is still perfectly recognizable, just slightly grainy. Then you add a bit more. And more. After a thousand rounds of adding a little static, the cat is gone — all you have left is pure television snow, indistinguishable from random noise.

That process is easy. Any child can pour sand on a painting. The interesting question is: can you learn to run it backward? Can you look at a noisy picture and infer what it looked like one step earlier, when it was a tiny bit less noisy?

If you can — if you have a neural network that can reliably undo one small step of noise — then you have a generative model. Start with pure random noise, run your denoiser a thousand times, and out comes a photograph of a cat that never existed.

This is the entire idea behind diffusion models. Two processes, running in opposite directions:

• The forward process (fixed, no learning): a Markov chain that gradually adds Gaussian noise over $T$ steps, turning any data sample $x_0$ into pure noise $x_T \sim \mathcal{N}(0, I)$.
• The reverse process (learned): a Markov chain that runs backward in time, iteratively removing noise — step by step, $x_T \to x_{T-1} \to \cdots \to x_0$ — recovering a data sample from nothing but static.

The name "diffusion" comes from the analogy to physical diffusion: ink dropped into water spreads until it's uniformly dissolved. The forward process is the ink spreading; the reverse process is learning to un-dissolve it, concentrating dilute molecules back into a coherent droplet. Thermodynamically, the forward process increases entropy (disorder); the reverse process decreases it (restores structure).

Adding noise, one step at a time

The forward process is a Markov chain. At each step $t$, we take the current state $x_{t-1}$, scale it down slightly, and add a small amount of Gaussian noise:

Here $\beta_t \in (0,1)$ is the variance schedule at timestep $t$, controlling how much noise is added. In code, the step looks like this:

The scale factor $\sqrt{1-\beta_t}$ is not arbitrary — it is chosen so that if $x_{t-1}$ has unit variance, then $x_t$ also has approximately unit variance. This is called the variance-preserving property, and it keeps the signal from blowing up or collapsing as we iterate.

The closed-form shortcut

Here is the most important mathematical fact about the forward process, and the one that makes diffusion models trainable. We can jump directly from $x_0$ to any $x_t$ without iterating through intermediate steps.

Define:

Then by recursively substituting the single-step formula and using the fact that sums of independent Gaussians are Gaussian:

Or in sampling form — and this is the formula you'll use in every training loop:

The quantity $\bar{\alpha}_t$ has a beautiful interpretation: it is the fraction of the original signal that survives at timestep $t$. At $t=0$, $\bar{\alpha}_0 \approx 1$ — the data is nearly intact. At $t=T$, $\bar{\alpha}_T \approx 0$ — the signal is gone, and $x_T$ is effectively pure noise. The signal-to-noise ratio at each step is $\text{SNR}(t) = \bar{\alpha}_t / (1-\bar{\alpha}_t)$.

The forward posterior

One more formula worth understanding — it will become the training target in §6. Given both endpoints $x_0$ and $x_t$, the intermediate step $x_{t-1}$ has a tractable Gaussian posterior:

where:

This posterior is the ground truth for what the reverse process should do. The entire training objective, as we'll see shortly, is about making the learned reverse step match this posterior.

How fast do we destroy?

The noise schedule $\{\beta_1, \beta_2, \ldots, \beta_T\}$ controls how rapidly the data degrades. Two common choices:

Linear schedule

$\beta_t$ increases linearly from $\beta_1 = 10^{-4}$ to $\beta_T = 0.02$ over $T=1000$ steps. This was the original DDPM recipe. Simple but sub-optimal: $\bar{\alpha}_t$ drops too fast in the middle and too slowly at the ends.

Cosine schedule

$\bar{\alpha}_t$ is defined directly as a shifted cosine curve, ensuring a smoother, more gradual transition from signal to noise. Proposed by Nichol & Dhariwal (2021), it gives the model more time in the "interesting" middle region where the denoising task is neither trivial nor hopeless.

The choice of schedule matters more than you might expect. The cosine schedule consistently outperforms the linear one, because it allocates model capacity to the timesteps where the denoising task is actually hard — the middle region where roughly half the signal remains.

Learning to denoise

The reverse process is what the neural network learns. It is parameterized as another Markov chain, running backward in time:

The network takes the noisy sample $x_t$ and the timestep $t$ as input, and produces the mean $\mu_\theta$ of a Gaussian distribution over the slightly-less-noisy $x_{t-1}$. The variance $\sigma_t^2$ is typically fixed (not learned) — set to either $\beta_t$ or $\tilde{\beta}_t$, both of which are known from the noise schedule.

In principle, the network needs to learn a different denoising function for every one of the $T$ timesteps. In practice, a single network handles all timesteps by taking $t$ as an input. The timestep is typically encoded as a sinusoidal positional embedding (the same trick used in transformers for position encoding) and injected into the network's internal layers.

Three ways to think about what the network predicts

The reverse step formula requires predicting the mean $\mu_\theta(x_t, t)$. But there are several equivalent ways to parameterize this — different quantities the network can predict, all of which produce the same mean when plugged back into the formula.

The relationship between them is straightforward. If the network predicts $\hat{\epsilon}$, we can recover $\hat{x}_0$:

And vice versa. The loss and the gradients are identical; only the human interpretation changes.

The beautifully simple loss

The principled training objective for diffusion models derives from the variational lower bound (ELBO), which decomposes into a sum of KL divergences between the learned reverse steps and the true forward posteriors. The derivation is long and involves telescoping products of Gaussians.

But here is the remarkable result: Ho et al. (2020) showed that you can throw away most of the structure and use a dead-simple objective that works even better:

That's it. Sample a training image $x_0$. Pick a random timestep $t$. Sample noise $\epsilon$. Compute the noisy version $x_t$. Ask the network to predict $\epsilon$. Penalize the squared error. Repeat.

Compare this to GAN training, where you must carefully balance a generator and discriminator, manage mode collapse, and tune a fragile adversarial objective. Diffusion training is a single regression loss. It "just works."

The U-Net: an encoder-decoder with skip connections

The denoising network $\epsilon_\theta(x_t, t)$ is almost always a U-Net — a convolutional architecture originally designed for biomedical image segmentation (Ronneberger et al., 2015) that turned out to be perfectly suited for denoising. The key properties:

• Encoder path (downsampling): a series of residual blocks that progressively reduce spatial resolution while increasing channel count. Each level captures features at a different scale.
• Decoder path (upsampling): mirrors the encoder, progressively increasing spatial resolution back to the input size.
• Skip connections: feature maps from the encoder are concatenated with the corresponding decoder level. This preserves fine spatial detail that would otherwise be lost to downsampling — critical for a denoising task where pixel-level precision matters.
• Timestep conditioning: the timestep $t$ is encoded as a sinusoidal embedding and injected into every residual block, typically via adaptive group normalization (AdaGN). This allows the same network to act as a different denoising function at each timestep.
• Self-attention at low resolutions: at the bottleneck (e.g., $16 \times 16$ or $8 \times 8$), self-attention layers capture global, long-range dependencies — "this is a face, not a landscape."

DDPM: the thousand-step walk

Once the network is trained, generation is the reverse process running from $t=T$ down to $t=0$:

  Sample x_T ~ N(0, I)                      ← start from pure noise
  For t = T, T-1, ..., 1:
      z ~ N(0, I) if t > 1, else z = 0      ← stochastic term
      ε̂ = ε_θ(x_t, t)                       ← ask the network
      x_{t-1} = 1/√α_t · (x_t - β_t/√(1-ᾱ_t) · ε̂) + σ_t · z
  Return x_0

Each step calls the neural network once, so with $T=1000$ timesteps, generating a single image requires 1000 forward passes through the U-Net. This is why DDPM sampling is slow. Each step is cheap, but there are a thousand of them.

The stochastic term $z$ adds a small amount of fresh noise at each step, which acts as a corrective jitter that keeps the sample on the data manifold. Setting $z=0$ at the final step ($t=1$) gives a clean output.

DDIM and accelerated sampling

Song et al. (2020) made a clever observation: the DDPM training objective only constrains the marginals $q(x_t \mid x_0)$ at each timestep — not the joint distribution $q(x_{1:T} \mid x_0)$. This means we can define a different reverse process that matches the same marginals but takes bigger steps.

DDIM (Denoising Diffusion Implicit Models) rewrites the update rule with a tunable noise parameter $\sigma_t$:

The key insight: when $\sigma_t = 0$, the process is fully deterministic. Given the same starting noise $x_T$, you always get the same $x_0$. And because the trajectory is deterministic, you can take much larger steps — jumping from $x_T$ to $x_{T-50}$ to $x_{T-100}$ and so on — without the accumulated randomness that would cause DDPM to go off-track.

The practical consequence: you can use a model trained with the DDPM objective and sample from it with DDIM at inference time. No retraining needed. Just change the sampling loop. This is one of the most useful separations in the diffusion toolbox: the training objective and the sampling procedure are decoupled.

Classifier-free guidance: the quality knob

Diffusion models generate diverse samples by default — but sometimes diversity is the enemy. When you prompt "a red sports car on a mountain road," you want the model to commit to that description, not hedge. Classifier-free guidance (CFG) gives you a knob to trade diversity for fidelity to the condition.

The idea: during training, randomly replace the text condition $c$ with a null/empty token $\varnothing$ some fraction of the time (typically 10–20%). This teaches the network to produce both conditional predictions $\epsilon_\theta(x_t, t, c)$ and unconditional ones $\epsilon_\theta(x_t, t, \varnothing)$.

At sampling time, amplify the difference between the two:

When $w=1$, this is standard conditional generation. When $w>1$ (typical values: 3–15), you amplify the effect of the condition — the model pushes harder toward "this is a red car" and away from "this is a generic image." Higher $w$ means more faithful outputs but less diversity; lower $w$ means more creative, surprising, but potentially off-prompt results.

Training a diffusion model on MNIST

Theory is one thing; watching it work is another. Using a simple U-Net trained on MNIST (28×28 grayscale digits), here is what we actually observe when we run the experiments from the companion notebook.

The forward process in action

Apply the closed-form noising formula $x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon$ to a clean digit at timesteps $t \in \{0, 100, 250, 500, 750, 1000\}$. The digit is recognizable up to about $t=250$; by $t=500$ only the coarsest blob remains; by $t=750$ it is indistinguishable from noise to the human eye.

Training dynamics

The loss curve descends smoothly — no mode collapse, no oscillation, no sudden divergence. This is the signature advantage of diffusion over adversarial training. After 15 epochs (~5 minutes on a T4), the model produces recognizable digits. After further training, they become sharp and varied.

DDPM vs DDIM on the same model

Running DDPM sampling (1000 steps) produces high-quality digits — clean edges, good variety. Running DDIM on the same trained model with only 50 steps produces digits of comparable quality in 1/20th the time. Dropping to 20 steps introduces slight blurriness but digits remain identifiable. At 10 steps, quality degrades noticeably. The sweet spot for DDIM on this model is around 50 steps.

What the model predicts at different timesteps

Visualizing the predicted $\hat{x}_0$ (the model's "guess" of the clean image) at various $t$ reveals a fascinating progression. At $t=900$ (very noisy), the model's guess is a blurry blob — "this is probably a digit of some kind." At $t=500$, the shape is roughly right — "this is an 8, not a 3." At $t=100$, the prediction is nearly pixel-perfect — the model is doing fine refinement.

Loss per timestep: where is denoising hardest?

Computing the average loss at each $t$ reveals a curve that peaks in the middle of the schedule — around $t \approx 400\text{–}600$. At very high $t$ (near pure noise), the model can't do much better than guessing the dataset mean; the loss is moderately high but not the worst. At very low $t$ (nearly clean), the task is easy and the loss is small. The peak in the middle is where the model is working hardest — where there's enough signal to attempt reconstruction but enough noise to make it genuinely difficult.

Noise-space interpolation

Fix two random noise vectors $z_1$ and $z_2$, DDIM-sample from each to get two different digit images, then interpolate smoothly between $z_1$ and $z_2$ (using spherical linear interpolation — slerp — to stay on the unit sphere). The intermediate images transition smoothly from one digit to another — a 7 morphs through ambiguous forms into a 2. This demonstrates that the learned mapping from noise space to data space is smooth and continuous, not a fragmented patchwork.

Don't diffuse pixels — diffuse latents

Running diffusion directly on a $512 \times 512 \times 3$ image means the U-Net operates on tensors with nearly 800,000 dimensions at every timestep. Most of these dimensions encode imperceptible high-frequency details and spatial redundancy. The model wastes capacity on information that humans cannot distinguish.

Latent diffusion (Rombach et al., 2022) — commercialized as Stable Diffusion — solves this by separating the problem into two stages:

1. Perceptual compression (autoencoder): Train a VAE to encode images $x$ into a compact latent representation $z = \mathcal{E}(x)$ and decode them back: $\hat{x} = \mathcal{D}(z)$. Typical compression: $8\times$ spatial downsampling with 4 latent channels, giving a 48× reduction in dimensionality.
2. Semantic generation (diffusion): Run the entire diffusion process — forward noise, reverse denoising — in the latent space $z$ rather than pixel space $x$. The U-Net operates on much smaller tensors.

The decomposition is clean: the autoencoder handles perceptual compression (removing imperceptible high-frequency detail), while the diffusion model handles semantic generation (learning the structure and content of images). The U-Net only needs to learn the hard part.

Where the field is going

Flow matching

Flow matching (Lipman et al., 2023) generalizes diffusion by learning a velocity field $v_\theta(x, t)$ that transports samples from a noise distribution to the data distribution along a continuous trajectory. The training objective is simpler than DDPM's:

where $x_t = (1-t)\,x_0 + t\,x_1$ is a simple linear interpolation between noise ($x_0$) and data ($x_1$). The learned trajectories tend to be straighter than diffusion paths, enabling fewer sampling steps. Stable Diffusion 3 and Flux are built on rectified flow matching rather than classical DDPM.

Consistency models

Consistency models (Song et al., 2023) train a function $f_\theta$ that maps any point on a diffusion trajectory directly to the trajectory's endpoint (the clean data), enabling one-step generation. The key constraint is self-consistency: for any two points on the same trajectory, the model must predict the same endpoint. This is enforced during training by requiring $f_\theta(x_t, t) = f_{\theta^-}(\hat{x}_{t'}, t')$ where $\hat{x}_{t'}$ is one ODE solver step away.

Diffusion Transformers (DiT)

Peebles & Xie (2023) replaced U-Nets with vision transformers, showing that scaling laws apply to diffusion architectures: bigger DiTs systematically produce better FID scores, following the same predictable curves as scaling LLMs. DALL-E 3, Stable Diffusion 3, Flux, and Sora all use DiT variants with billions of parameters.

Beyond images

The core framework — corrupt with noise, learn to denoise — is remarkably general. Diffusion models now generate video (Sora, Kling), audio and music (Stable Audio), 3D shapes (point cloud and NeRF diffusion), molecular structures for drug discovery, and even robot action trajectories (diffusion policies). The architecture changes; the principle does not.

Watch denoising happen in 2D

We can't run a real U-Net in the browser, but we can run exact diffusion on a simple 2D distribution. Below, a mixture of five Gaussians serves as the "data." The forward process dissolves the clusters into a uniform cloud of noise. The reverse process — using the analytically exact score function — reassembles them. Drag the slider or hit play to watch.

The score function $\nabla \log q(x_t)$ is computed analytically from the known Gaussian mixture. At each reverse step, particles follow the score toward the nearest cluster center. The result: structure emerges from noise, just as in a real diffusion model — but in two dimensions you can see every particle's journey.

Takeaways & exercises

The six things to remember

1. Diffusion = destroy + un-destroy. The forward process adds noise step by step (fixed, no learning). The reverse process removes it (learned, one Gaussian at a time).
2. The closed-form shortcut $x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon$ makes training embarrassingly parallel across timesteps. You never iterate the forward chain.
3. The training loss is just MSE on noise predictions. No adversarial training, no mode collapse, no careful balancing. This is why diffusion models are so stable.
4. Training and sampling are decoupled. You can train with DDPM's objective and sample with DDIM, DPM-Solver, or any other method — no retraining needed.
5. Latent diffusion is the efficiency insight. Compress images to a 48× smaller latent space first, then run diffusion there. This is what makes Stable Diffusion practical.
6. Classifier-free guidance is the quality knob. Higher $w$ means more faithful to the prompt but less diverse. Typical values: 3–15.

Exercises

Further reading, in order

• Sohl-Dickstein et al. (2015) — Deep Unsupervised Learning using Nonequilibrium Thermodynamics. The original diffusion idea.
• Ho, Jain, Abbeel (2020) — Denoising Diffusion Probabilistic Models. The paper that made it work. Read this one.
• Song et al. (2020) — DDIM. Faster, deterministic sampling.
• Song & Ermon (2019/2021) — Score-Based Generative Modeling / Score SDE. The continuous-time perspective.
• Nichol & Dhariwal (2021) — Improved DDPM. Cosine schedule and learned variance.
• Dhariwal & Nichol (2021) — Diffusion Models Beat GANs. The quality breakthrough.
• Ho & Salimans (2022) — Classifier-Free Diffusion Guidance. The quality knob.
• Rombach et al. (2022) — Latent Diffusion Models. Stable Diffusion.
• Peebles & Xie (2023) — DiT. Transformers replace U-Nets.
• Lipman et al. (2023) — Flow Matching. The generalization beyond diffusion.
• Song et al. (2023) — Consistency Models. One-step generation.
• P. Kulkarni — Module 04 notebook. MNIST diffusion from scratch in PyTorch.