Introduction to Diffusion Models

Diffusion Models are currently standing at the forefront of generative modeling. The overarching theme is to "transform" a canonical distribution like $\mathcal{N}(0,I)$ into the distribution we want to sample from. In a previous post, I covered flow matching, in which we model this transformation as an Ordinary Differential Equation (ODE). We train a sufficiently expressive model (nowadays a Transformer neural network or CNN) to learn the gradient vectors defining this ODE, i.e. the flow. Then, at inference time, we sample $z \sim \mathcal{N}(0,I)$ and follow its path in the ODE using the vectors defined by our network. This is typically done via an iterative discrete solver such as Euler steps, though multiple methods exist.

In Diffusion Models, the story sounds very similar, but the "protagonist" objects change a little bit. Instead of ODEs, we work with stochastic differential equations (SDEs), which are like ODEs but also contain an element of random perturbation throughout the dynamics. And instead of the flow, our neural network learns something called a score function.

Brownian Motion and SDEs

First, let's get a feel for what SDEs are. I would recommend the reader have familiarity with ODEs, at least in the context of flow matching models (read my previous post if you want the bare minimum!).

This is a vast and fascinating area of mathematics known as stochastic calculus. It is widely used in physics and finance. For our purposes here we won't need to dive into topics such as Ito integration and Brownian processes (sadly). We will build basic intuition. Warning: do not expect to make sense of all the math here if you're not familiar with stochastic calculus. The more you scratch, the more I'll refer you to a textbook!

Derivatives are hard to formulate in the stochastic world, as random perturbations remove differentiability for the most part. Instead, we will consider ODE dynamics in a discrete, pre-limit fashion.

An ODE is typically defined as $\frac{d}{dt}X_t = u_t(X_t)$. We can write this as:

$$\lim_{h\to 0} \frac{1}{h}(X_{t+h}-X_t) = u_t(X_t)$$

To go to the stochastic setting, simply remove the limit and add noise:

$$ X_{t+h} = X_t + h\,u_t(X_t) + h\,R_t(h) + \underbrace{\sigma_t(W_{t+h}-W_t)}_{\text{Brownian noise}} $$

where $\lim_{h\to 0}\mathbb{E}[\|R_t(h)\|_2^2]^{1/2} \to 0$ captures the error term and $W_t$ is a stochastic process known as Brownian motion. Important properties of $W_t$ include that it is continuous, satisfies $W_0 = 0$, and $W_t - W_s \sim \mathcal{N}(0,(t-s)I_d)$ for all $s < t$. We call $\sigma_t \geq 0$ the diffusion coefficient.

We'll often write an SDE symbolically:

$$dX_t = u_t(X_t)\,dt + \sigma_t\,dW_t,\quad X_0 = x_0$$

If you'd like, you can also think of the following integral equation. It contains a funky kind of integral called the Ito integral, which essentially accumulates the random fluctuations of $W$. Ignoring that, the equation says that $X_t$'s location at time $t$ depends on the velocity field $u$ throughout its previous locations and on the total random fluctuation acting on it since the start of its trajectory.

$$ X_t = x_0 + \int_0^t u_s(X_s)\,ds + \underbrace{\int_0^t \sigma_s\,dW_s}_{\text{Ito Integral}} $$

What we seek when solving an SDE is a stochastic process that satisfies this equation. What does "satisfy" mean here? This is where it gets quite technical and I'd urge you to switch to a textbook. I do want you to appreciate that defining these concepts rigorously is important and difficult — if you ever need to do the math, learn the math first! I'll just say that there are two notions of "solution" to an SDE: a weak solution (in distribution) and a strong solution (in probability).

It turns out that SDEs can be solved (whatever "solving" them means) in many cases, and in fact sometimes they have "unique" solutions (whatever "unique" means). We will live in the happy ML world where such "unique"-ness applies, as our functions are continuously differentiable with bounded derivatives.

How might we simulate an SDE the same way we did an ODE? That's easier — we go back to this equation:

$$X_{t+h} = X_t + h\,u_t(X_t) + h\,R_t(h) + \underbrace{\sigma_t(W_{t+h}-W_t)}_{\text{Brownian noise}}$$

We know that $W_{t+h}-W_t \sim \mathcal{N}(0,h)$, so we can update $X_t$ iteratively as:

$$X_{t+h} = X_t + h\,u_t(X_t) + \sqrt{h}\,\sigma_t\,\varepsilon_t$$

This is simply the Euler–Maruyama iteration, applied with some additional stochasticity.

Diffusion Models

If all you got from the SDE paragraph earlier was the last part on sampling, that's OK! That's really all we'll need. Our perspective will be similar to the flow matching one:

Score Functions and Denoisers

In flow matching, we learn $u_t$ directly. In diffusion models, we leverage a different approach that eventually learns the same thing. We learn an object called the score function. Historically, models were first created to learn the score function as it is quite an intuitive object. As research progressed, we realized that in many cases the flow-matching and score-matching objectives are essentially the same!

As we mentioned earlier, our goal is to learn a continuous transformation taking $p_{\text{init}}$ to $p_{\text{data}}$. As a reminder (refer to the flow matching post), this transformation goes through "intermediate" distributions $p_t$ for $t \in [0,1]$. Think of this as the "forward" process: at $t=0$ we start at $p_{\text{data}}$ and at $t=1$ we arrive at $p_{\text{init}}$. Furthermore, we cannot really "observe" these distributions directly. Instead, we work with their conditional proxies and use marginalization. Let $p_t(x \mid z)$ be the intermediate distribution at time $t$ conditioned on starting at point $z$ on the unknown target space. Then we can express the marginal score function in terms of the conditional score function simply by marginalization:

$$ \nabla \log p_t(x) = \int \nabla\log p_t(x\mid z)\, \frac{p_t(x\mid z)\,p_{\text{data}}(z)}{p_t(x)}\,dz $$

The score function for Gaussian probability paths is a linear function on $x$ and $z$, as is the conditional vector field $u_t(x\mid z) = \bigl(\dot\alpha_t - \tfrac{\dot\beta_t}{\beta_t}\alpha_t\bigr)z + \tfrac{\dot\beta_t}{\beta_t}\,x$. Those two are thus related, via simple algebraic manipulation:

Both the score function and the vector field are in this case examples of denoisers: they are simply linear reparameterizations of the posterior mean $\mathbb{E}_{z\mid x}[z]$, which describes the reconstruction of the original data point, in expectation, given the intermediate point $x$ at time $t$. A conditional denoiser is defined as $D_t(x \mid z) = z$, and via marginalization we get:

$$ D_t(x) = \int z\,\frac{p_t(x\mid z)\,p_{\text{data}}(z)}{p_t(x)}\,dz = \frac{1}{\dot\alpha_t\beta_t - \alpha_t\dot\beta_t} \bigl(\beta_t\,u_t(x) - \dot\beta_t\,x\bigr) $$

The idea is simple: multiple $z$-s lead to the same $x$. The conditional denoiser gives back the $z$ we came from. The marginal one gives back the $z$ we came from in expectation. This is why we have to weigh every $z$ by the probability $p_t(z \mid x)$ that it yielded $x$.

It is often preferable to skip learning the score function and just learn the denoiser instead! These are called Denoising Diffusion Models.

SDE-based Sampling

Recall that our plan for diffusion models was to sample using SDEs:

$$dX_t = u_t(X_t)\,dt + \sigma_t\,dW_t$$

What we want is $X_t \sim p_t$ for $t \in [0,1]$. In flow matching, if you recall, the continuity equation told us that if we use the marginal vector field then we would obtain that guarantee.

Now that we have stochastic dynamics, a different equation yields the guarantee. It's called the Fokker–Planck Equation:

$$ \partial_t p_t(x) = -\operatorname{div}(p_t u_t)(x) + \frac{\sigma_t^2}{2}\,\Delta p_t(x) $$

This is the continuity equation we know, but with an additional Laplacian term $\Delta p_t(x) = \sum_{i=1}^d \tfrac{\partial^2}{\partial x_i^2} p_t(x)$. We're now entering the realm of Langevin Dynamics! Diving deep here would be very fun, but for now we can think in terms of physical analogs: the Laplacian term expresses the "flattening" of the concentration of the probability mass towards neighboring points as time increases. Remember, our SDE is driven by a deterministic term, but also by random noise through the Brownian motion. Even if the deterministic drift dissipates, the variance of the noise will cause the trajectory to diffuse.

Mathematically, we can use the Fokker–Planck equation to figure out what the deterministic vector field should be to obtain $X_t \sim p_t$. I'll just give the answer and we can verify it:

In other words, we need the vector field $u_t(X_t)$ to be "pushed" additionally towards the direction of steepest ascent on the distribution we desire. Without this, had we just kept $u_t(X_t)$ as it was for the ODE, the Brownian motion would drown out our deterministic drift.

Think about it a little! With this simple deterministic intervention we are able to guarantee a well-defined behavior even in the presence of noise!

Verification is simple math: plug the equation above into Fokker–Planck. The score function comes from the Laplacian term via the identity $\operatorname{div}(\nabla p_t) = \Delta p_t$.

Score Matching

To snap back to reality, sampling requires learning the score function! Let us parameterize it using a neural network $s_t^\theta : \mathbb{R}^d \times [0,1] \to \mathbb{R}^d$. In the same way as we did for the flow-matching vector fields, we learn the conditional score matching loss, averaged over time, $z$, and forward noising processes. That loss equals the marginal loss up to a constant, which gets wiped during gradient-based optimization.

• Score Matching Loss: $L_{SM}(\theta) = \mathbb{E}_{t,\,z,\,x\sim p_t(\cdot\mid z)} \bigl[\|s_t^\theta(x) - \nabla\log p_t(x)\|_2^2\bigr]$
• Conditional Score Matching Loss: $L_{CSM}(\theta) = \mathbb{E}_{t,\,z,\,x\sim p_t(\cdot\mid z)} \bigl[\|s_t^\theta(x) - \nabla\log p_t(x\mid z)\|_2^2\bigr]$

Proving this is pretty much identical to the flow matching loss theorem.

Outro

That's pretty much it! We went from flow matching to diffusion by introducing stochasticity in the continuous transformation of the known distribution to the unknown one. This introduces a diffusion Laplacian term via the Fokker–Planck equation, which we account for by learning the score function $\nabla\log p_t(x)$. It turns out that, under Gaussian paths, we just need to learn that function — and that just transforms to learning to denoise: we simply train a network to predict the noise we used to corrupt samples.

Till next time!

Appendix: The Classical DDPM Derivation

The perspective we took above is the modern, continuous-time one. Historically, DDPM (Ho et al., 2020) was derived using a discrete Markov chain and a variational lower bound on the log-likelihood. Because this is often the way people present and talk about diffusion models, covering it in its mathematical entirety would be useful, so I'll do it here. Note: I don't think it's very necessary to handle these details in practice, because the SDE approach just makes things so much easier.

The Discrete Markov Chain Setup

Fix a data point $x_0 \sim p_{\text{data}}$ and a sequence of $T$ timesteps. The forward (encoder) process progressively corrupts $x_0$ into pure noise:

$$ q(x_{1:T} \mid x_0) \;=\; \prod_{t=1}^{T} q(x_t \mid x_{t-1}), \qquad q(x_t \mid x_{t-1}) = \mathcal{N}\!\left(x_t;\,\sqrt{1-\beta_t}\,x_{t-1},\;\beta_t I\right) $$

where $\{\beta_t\}_{t=1}^T$ is a fixed noise schedule with $\beta_t \in (0,1)$. Each step adds a small amount of Gaussian noise, so the chain is Markovian: $x_t$ depends on the past only through $x_{t-1}$. The encoder is fixed and requires no learning.

The reverse (decoder) process is a learned Markov chain running backwards, starting from $x_T \sim \mathcal{N}(0, I)$:

$$ p_\theta(x_{0:T}) \;=\; p(x_T)\prod_{t=1}^{T} p_\theta(x_{t-1} \mid x_t), \qquad p_\theta(x_{t-1} \mid x_t) = \mathcal{N}\!\left(x_{t-1};\;\mu_\theta(x_t, t),\;\Sigma_\theta(x_t, t)\right) $$

Our goal is to train $\theta$ so that $p_\theta(x_0) \approx p_{\text{data}}(x_0)$. We can't maximize $\log p_\theta(x_0)$ directly because it requires marginalizing over all latents $x_{1:T}$, which is intractable. Instead, we maximize a lower bound.

The ELBO via Jensen's Inequality

The standard ELBO argument from VAEs gives:

$$ \log p_\theta(x_0) = \log \int p_\theta(x_{0:T})\,dx_{1:T} = \log \mathbb{E}_{q(x_{1:T}\mid x_0)}\!\left[\frac{p_\theta(x_{0:T})}{q(x_{1:T}\mid x_0)}\right] $$

Since $\log$ is concave, Jensen's inequality gives $\log \mathbb{E}[\cdot] \geq \mathbb{E}[\log(\cdot)]$, so:

$$ \log p_\theta(x_0) \;\geq\; \underbrace{\mathbb{E}_{q(x_{1:T}\mid x_0)}\!\left[\log p_\theta(x_{0:T}) - \log q(x_{1:T}\mid x_0)\right]}_{=:\;\mathcal{L}_{\text{ELBO}}(\theta)} $$

We maximize $\mathcal{L}_{\text{ELBO}}(\theta)$ over $\theta$. The gap between $\log p_\theta(x_0)$ and the ELBO is exactly $D_{\text{KL}}(q(x_{1:T}|x_0)\,\|\,p_\theta(x_{1:T}|x_0)) \geq 0$.

Expanding via the Markov Factorizations

Now we expand both the numerator and denominator using their respective Markov factorizations:

$$ \log p_\theta(x_{0:T}) - \log q(x_{1:T}\mid x_0) = \log p(x_T) + \sum_{t=1}^{T}\log p_\theta(x_{t-1}\mid x_t) - \sum_{t=1}^{T}\log q(x_t\mid x_{t-1}) $$

The key insight is to use the reverse-conditional of the forward process, $q(x_{t-1}|x_t, x_0)$, which is tractable whenever we condition on $x_0$:

$$ q(x_t \mid x_{t-1}) = q(x_{t-1}\mid x_t, x_0)\;\frac{q(x_t\mid x_0)}{q(x_{t-1}\mid x_0)} $$

Substituting this and telescoping the $q(x_t | x_0)$ terms (they cancel like a telescoping product), the ELBO decomposes into three interpretable pieces:

$$ \mathcal{L}_{\text{ELBO}} = \underbrace{\mathbb{E}_q\!\left[\log p_\theta(x_0 \mid x_1)\right]}_{\text{reconstruction}} - \underbrace{D_{\text{KL}}\!\left(q(x_T \mid x_0)\,\|\,p(x_T)\right)}_{\text{prior matching}} - \sum_{t=2}^{T}\underbrace{\mathbb{E}_q\!\left[D_{\text{KL}}\!\left(q(x_{t-1}\mid x_t, x_0)\,\|\,p_\theta(x_{t-1}\mid x_t)\right)\right]}_{\text{denoising terms } L_{t-1}} $$

• Reconstruction: how well can we recover $x_0$ from the slightly-noised $x_1$? This is just a log-likelihood term and is often treated as a single Gaussian regression step.
• Prior matching: how close is the fully-noised $x_T$ to our prior $\mathcal{N}(0,I)$? With a well-chosen schedule this term is negligible (effectively zero for large $T$).
• Denoising terms $L_{t-1}$: these are the heart of the loss. Each one asks: how close is the learned reverse step $p_\theta(x_{t-1}|x_t)$ to the true reverse conditional $q(x_{t-1}|x_t, x_0)$?

The Training Objective

Since both $q(x_{t-1}|x_t, x_0)$ and $p_\theta(x_{t-1}|x_t)$ are Gaussian, each KL term $L_{t-1}$ has a closed form. Setting $\Sigma_\theta = \tilde\beta_t I$ (fixed), minimizing $L_{t-1}$ over $\theta$ is equivalent to minimizing:

$$ L_{t-1} \;\propto\; \mathbb{E}_{x_0,\,\varepsilon}\!\left[ \left\|\tilde\mu_t\!\left(x_t, x_0\right) - \mu_\theta(x_t, t)\right\|^2 \right] $$

Substituting $x_0 = (x_t - \sqrt{1-\bar\alpha_t}\,\varepsilon)/\sqrt{\bar\alpha_t}$ and reparameterizing the mean prediction to predict noise instead, we recover the famous DDPM loss:

Connection to the SDE / Score-Matching Framework

At first glance the discrete DDPM derivation looks completely different from the SDE story. They are, in fact, the same theory at different scales of discretization. Here is the bridge:

1. The continuous-time limit. Take $T \to \infty$ and $\beta_t \to 0$ with $\beta_t \approx \bar\beta / T$ for a fixed schedule $\bar\beta$. In this limit the discrete Markov chain $q(x_t|x_{t-1})$ converges to the SDE:

$$ dX_t = -\tfrac{1}{2}\beta(t)\,X_t\,dt + \sqrt{\beta(t)}\,dW_t $$

where $\beta(t)$ is the continuous-time noise schedule. This is precisely the Variance-Preserving (VP) SDE of Song et al. (2021), and it is an instance of the Ornstein–Uhlenbeck process we saw earlier. Its marginals are Gaussian with mean $\sqrt{\bar\alpha(t)}\,x_0$ and variance $(1-\bar\alpha(t))I$, recovering our Gaussian probability path.

2. The ELBO denoising term becomes score matching. In the continuous limit, each $L_{t-1}$ becomes an infinitesimal contribution. Summing them and taking $T\to\infty$ gives:

$$ \sum_{t=2}^{T} L_{t-1} \;\longrightarrow\; \mathbb{E}_{t,\,x_0,\,x_t}\!\left[\lambda(t)\,\left\|s_\theta(x_t,t) - \nabla\log q(x_t\mid x_0)\right\|^2\right] $$

where $s_\theta$ is the score network and $\lambda(t) = (1-\bar\alpha(t))$ is a weighting coming from the KL calculation. This is exactly the conditional score matching loss $L_{CSM}$ from the main article! Since $\nabla\log q(x_t|x_0) = -(x_t - \sqrt{\bar\alpha_t}\,x_0)/(1-\bar\alpha_t)$ is simply the scaled noise $-\varepsilon/\sqrt{1-\bar\alpha_t}$, predicting the score is the same as predicting the noise — which is what the DDPM loss does.

3. Reverse process = score-guided SDE. Just as the continuous forward SDE above has a time-reversal (Anderson, 1982), the discrete Markov chain has a learned reverse chain. In the continuous limit, the reverse SDE is:

$$ dX_t = \left[-\tfrac{1}{2}\beta(t)\,X_t - \beta(t)\,\nabla\log p_t(X_t)\right]dt + \sqrt{\beta(t)}\,dW_t $$

which is precisely the SDE-based sampling formula from the main article, with $u_t(x) = -\frac{1}{2}\beta(t) x$ (the deterministic drift of the OU process) and $\sigma_t = \sqrt{\beta(t)}$.