Yet Another Tutorial on Diffusion Models

Introduction

Learning diffusion models can be challenging. To my knowledge, there are at least three theories explaining how diffusion models work:

1. Variational inference (similar to VAE). This results in the famous DDPM and DDIM. It first defines a forward noising process \(q(x_t| x_{t-1})\) and a backward process \(p_\theta(x_{t+1}| x_{t})\). Then it optimizes the evidence lower bound (ELBO) to learn the optimal reverse process, which serves as a generative model.

2. Score-based modeling. This is from solving the stochastic differential equation (SDE): what the denoiser learns is the score function \(\nabla_x \log p_t(x)\) of a stochastic process.

3. Rectified flow. In this view, diffusion models are a way to map samples from two distributions, i.e., between random Gaussians (source) and images (target), and to find smooth non-intersecting paths (i.e., rectify the field) to move samples from source to target.

Diffusion Models from Variational inference

ELBO

Before discussing DDPM, it is helpful to recall the basic idea of the evidence lower bound (ELBO). Suppose we have an observed variable \(x\), a latent variable \(z\), and a probabilistic model \(p_\theta(x|z)\). We want to optimize \(\theta\) to maximize the data likelihood

\[ \log p_{\theta}(x)=\log \int p_{\theta}(x,z)\mathrm{d}z. \]

This integral is hard to compute directly without enumerating \(z\). The trick is to introduce an approximate posterior \(q(z|x)\) that is close to \(p_\theta(z|x)\). This \(q(z|x)\) can be either a fixed distribution (as in most diffusion models) or a learned distribution (as in VAE). Then we can rewrite the likelihood as

\[ \begin{align} \log p_{\theta}(x) &= \log \mathbb{E}_{q(z|x)} \left[ \frac{p_{\theta}(x,z)}{q(z|x)} \right]\\ &\ge \mathbb{E}_{q(z|x)} \left[ \log \frac{p_{\theta}(x,z)}{q(z|x)} \right]. \end{align} \]

The inequality is Jensen's inequality. The right-hand side is called the ELBO:

\[ \mathrm{ELBO}= \mathbb{E}_{z\sim q(z|x)} \left[ \log p_{\theta}(x,z)-\log q(z|x) \right] \le \log p_{\theta}(x). \]

The gap between ELBO and \(p_\theta(x)\) is the KL-divergence:

\[ \begin{align} \log p_{\theta}(x) &= \mathrm{ELBO} +D_{\mathrm{KL}}\left(q(z|x)\,\Vert\,p_{\theta}(z|x)\right). \end{align} \]

In the VAE-style factorization, the negative ELBO can be written in the KL divergence form:

\[ \begin{align} L_{\mathrm{VAE}} &:= -\mathrm{ELBO}\\ &= \mathbb{E}_{q(z|x)} \left[ \log q(z|x)-\log p_{\theta}(x,z) \right]\\ &= \mathbb{E}_{q(z|x)} \left[ \log q(z|x)-\log p(z)-\log p_{\theta}(x|z) \right]\\ &= \mathbb{E}_{q(z|x)} \left[ \log \frac{q(z|x)}{p(z)} \right]- \mathbb{E}_{q(z|x)} \left[ \log p_{\theta}(x|z) \right]\\ &= D_{\mathrm{KL}}\left(q(z|x)\,\Vert\,p(z)\right)- \mathbb{E}_{q(z|x)} \left[ \log p_{\theta}(x|z) \right]. \end{align} \]

This is the form is close to the DDPM expression below.

DDPM

DDPM is the classic formulation that made diffusion models popular in computer vision. They derived the fomulation of the diffusion model with variational inference, similar to VAE. Unfortunately, the fomulation is quite complex since it is one of the early work. Assuming \(x_0\) is a clean image sample, DDPM first defines a forward Markov process that gradually corrupts \(x_0\) with Gaussian noise:

\[ q(x_t | x_{t-1}) = \mathcal{N}\left(\sqrt{\alpha_t}x_{t-1}, (1-\alpha_t)I\right), \qquad \alpha_t = 1-\beta_t. \]

Here \(\beta_t\) is a small positive number that controls how much noise is added at step \(t\). Equivalently, we can sample \(x_t\) directly with

\[ x_t=\sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_t}\epsilon_t, \qquad \epsilon_t\sim\mathcal{N}(0,I), \qquad \bar{\alpha}_t=\prod_{i=1}^{t}\alpha_i \tag{1} \]

The forward process is a conditional Gaussian distribution to be learned:

\[ p_{\theta}(x_{t-1}| x_t) = \mathcal{N}\left(\mu_{\theta}(x_t,t), \Sigma_t\right). \]

Next, we maximize \(p_\theta(x_0)\) by optimizing the ELBO. Let \(L_{\mathrm{VLB}}\) as the negative ELBO. It can be decomposed into a sum of KL terms similar to \(L_{\mathrm{VAE}}\):

\[ \begin{align} L_{\mathrm{VLB}} &= D_{\mathrm{KL}}\left(q(x_T|x_0)\,\Vert\,p(x_T)\right)\\ &\quad+ \mathbb{E}_{q(x_{1:T}|x_0)} \left[ \sum_{t=2}^{T} D_{\mathrm{KL}}\left(q(x_{t-1}|x_t,x_0)\,\Vert\,p_{\theta}(x_{t-1}|x_t)\right) \right]\\ &\quad- \mathbb{E}_{q(x_1|x_0)}\left[\log p_{\theta}(x_0|x_1)\right]. \end{align} \]

The first term is usually close to zero if \(T\) is large enough, because \(x_T\) is nearly standard Gaussian. The last term is to predict the image \(x_0\) from \(x_1\), which you can ignore for now (it is similar to the middle term in the final formulation). The middle terms are the important ones: they ask the learned reverse transition \(p_{\theta}(x_{t-1}|x_t)\) to match the true posterior \(q(x_{t-1}|x_t,x_0)\).

This reverse posterior \(q(x_{t-1}|x_t,x_0)\) has a closed form. Using \(x_t=\sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_t}\epsilon_t\), we can write it in terms of the noise \(\epsilon_t\):

\[ q(x_{t-1}|x_t,x_0) = \mathcal{N}\left(\tilde{\mu}_t(x_t,\epsilon), \tilde{\beta}_t I\right), \]

where

\[ \begin{align} \tilde{\mu}_t(x_t,x_0)= \frac{1}{\sqrt{\alpha_t}} \left( x_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_t \right), \qquad \tilde{\beta}_t = \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\beta_t. \end{align} \]

When both distributions are Gaussian and the variance is fixed, each KL term has a closed form:

\[ \begin{align} L_{t-1} &= D_{\mathrm{KL}}\left( q(x_{t-1}|x_t,x_0)\,\Vert\,p_{\theta}(x_{t-1}|x_t) \right)\\ &= \frac{1}{2\tilde{\beta}_t} \left\Vert \tilde{\mu}_t(x_t,x_0)-\mu_{\theta}(x_t,t) \right\Vert^2 + C, \end{align} \]

In diffusion model literature, people found predicting noise \(\epsilon_t\) is better in practice. This becomes a weighted mean-squared error between the true noise \(\epsilon_t\) and the predicted noise:

\[ L_{t-1}=\frac{\beta_t^2}{2\tilde{\beta}_t\alpha_t(1-\bar{\alpha}_t)} \left\Vert\epsilon_t-\epsilon_{\theta}(x_t,t)\right\Vert^2 + C. \]

In practice, DDPM usually drops the weighting and uses the simplified objective

\[ L_{\mathrm{simple}}(\theta)= \mathbb{E}_{(t,x_0,x_t)\sim q} \left[ \left\Vert \epsilon_t-\epsilon_{\theta}\left(x_t, t\right) \right\Vert^2 \right]. \]

At inference time, DDPM starts from \(x_T\sim\mathcal{N}(0,I)\) and repeatedly applies the learned reverse transition until \(x_0\). A single reverse step is stochastic:

\[ x_{t-1} = \mu_{\theta}(x_t,t)+\sqrt{\tilde{\beta}_t}\epsilon_t, \qquad \epsilon_t\sim \mathcal{N}(0,I). \]

DDIM

Score-Based Generative Modeling

Score-based generative modeling describes the same family of ideas from the perspective of score functions. The score of a distribution is the gradient of log density:

\[ s(x) = \nabla_x \log p(x). \]

For diffusion models, we care about the score of noisy data distributions \(p_t(x_t)\) or \(p_{\sigma}(x_{\sigma})\). A score network \(s_{\theta}(x_t,t)\) learns

\[ s_{\theta}(x_t,t)\approx \nabla_{x_t}\log p_t(x_t). \]

For the Gaussian noising process

\[ x_t = \sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_t}\epsilon, \]

the conditional score of \(q(x_t| x_0)\) is

\[ \nabla_{x_t}\log q(x_t| x_0) = -\frac{\epsilon}{\sqrt{1-\bar{\alpha}_t}}. \]

This explains why predicting noise and predicting score are closely related. If a model predicts \(\epsilon_{\theta}(x_t,t)\), then the corresponding score estimate is

\[ s_{\theta}(x_t,t) = -\frac{\epsilon_{\theta}(x_t,t)}{\sqrt{1-\bar{\alpha}_t}}. \]

Score-Based Generative Modeling through SDEs generalizes this view. Instead of a finite sequence of discrete timesteps, the forward process is a stochastic differential equation that gradually turns data into noise. Sampling solves the reverse-time SDE or a related deterministic probability-flow ODE. In practice, DDPM, DDIM, EDM, and score-based models are often different parameterizations of the same underlying denoising problem.

EDM

EDM, short for "Elucidating the Design Space of Diffusion-Based Generative Models", reframes diffusion models around continuous noise magnitude \(\sigma\) instead of discrete timesteps \(t\). The noisy sample is written as

\[ x_{\sigma}=x_0+\sigma\epsilon, \qquad \epsilon\sim\mathcal{N}(0,I). \]

This notation is cleaner because the noise level itself is explicit. The denoiser is trained to recover the clean sample from \(x_{\sigma}\):

\[ D_{\theta}(x_{\sigma},\sigma)\approx x_0. \]

In EDM, a common training objective is a weighted denoising loss

\[ L_{\mathrm{EDM}}(\theta) = \mathbb{E}_{x_0,\sigma,\epsilon} \left[ w(\sigma) \left\Vert D_{\theta}(x_0+\sigma\epsilon,\sigma)-x_0\right\Vert^2 \right]. \]

DMD Distillation

Rectified Flow (Flow Matching)

Diffusion models learn to denoise along a noisy path. Flow matching and rectified flow instead learn a velocity field that transports noise into data through an ordinary differential equation:

\[ \frac{\mathrm{d}x_t}{\mathrm{d}t}=v_{\theta}(x_t,t). \]

A simple path between noise \(x_0\sim p_{\mathrm{noise}}\) and data \(x_1\sim p_{\mathrm{data}}\) is linear interpolation:

\[ x_t = (1-t)x_0 + tx_1. \]

For this path, the target velocity is just

\[ u_t = \frac{\mathrm{d}x_t}{\mathrm{d}t}=x_1-x_0. \]

The flow-matching objective trains a neural network to predict this velocity:

\[ L_{\mathrm{FM}}(\theta) = \mathbb{E}_{t,x_0,x_1} \left[ \left\Vert v_{\theta}(x_t,t) - (x_1-x_0)\right\Vert^2 \right]. \]

At sampling time, we start from noise and integrate the learned ODE from \(t=0\) to \(t=1\). This is conceptually simpler than diffusion sampling because there is no stochastic reverse chain and no need to add noise at each step. The tradeoff is that the choice of path and velocity target matters. Rectified flow tries to learn straighter transport paths, which can reduce the number of ODE solver steps needed for high-quality generation.

The useful mental model is:

1. DDPM learns a stochastic reverse denoising chain.
2. DDIM uses the same denoising model but follows a deterministic sampling path.
3. EDM rewrites the design around continuous noise levels and better practical choices.
4. Score-based modeling describes denoising as learning gradients of noisy data densities.
5. Rectified flow and flow matching replace denoising with velocity-field learning.