34.2 Score-Based Diffusion and Flow Matching

1. The Score Function

Every image synthesized by Stable Diffusion, every protein backbone proposed by RFDiffusion, and every MRI scan reconstructed from undersampled measurements relies on a single mathematical object: the score function. Without it, these systems would have no way to navigate the space of possible outputs toward realistic ones.

Suppose you need to sample from a probability distribution so complex that computing the probability of even a single point is intractable, yet at every point you can compute which direction probability increases fastest. That directional signal alone turns out to be enough to generate perfect samples, and it is called the score function: the gradient of the log-density with respect to the data.

For a probability distribution \(p(x)\) over \(\mathbb{R}^d\), the score function is defined as:

The score function is a vector field over the entire data space. At every point \(x\), it returns a vector pointing toward steepest probability increase, with magnitude proportional to the slope. This matters because the score encodes a distribution's shape without requiring the normalizing constant (the denominator in Bayes' rule), which is typically intractable for high-dimensional data. Taking the gradient of \(\log p(x)\) cancels the normalizing constant entirely, since \(\nabla_x \log \frac{\tilde{p}(x)}{Z} = \nabla_x \log \tilde{p}(x)\), leaving only the unnormalized density. Use score-based methods over VAEs or normalizing flows when you need high-fidelity generation and can afford iterative sampling. Prefer VAEs when single-pass speed is critical or when downstream tasks require an explicit latent space. Figure 34.2.1 illustrates the forward and reverse diffusion process alongside the probability flow ODE and flow matching comparison.

The score has a beautiful geometric interpretation: at any point \(x\), it points in the direction of steepest increase of the log-density. Following the score via Langevin dynamics (an iterative sampling algorithm that generates samples from a distribution by taking small steps along the score direction plus random noise, named after the physicist Paul Langevin) produces samples from \(p(x)\):

where \(\eta\) is a step size. As \(\eta \to 0\) and the number of steps grows, the iterates converge in distribution to samples from \(p(x)\). The noise term \(\sqrt{\eta} \, \epsilon_t\) is essential: without it, the dynamics would converge to a mode of \(p(x)\) rather than sampling from it. This connection between score estimation and sampling is why score-based models work.

The challenge is estimating \(\nabla_x \log p(x)\) without knowing \(p(x)\). Directly regressing a neural network \(s_\theta(x) \approx \nabla_x \log p(x)\) requires the unknown true score as a target. The breakthrough of denoising score matching (Vincent, 2011) sidesteps this: instead of matching the score of the clean data distribution, match the score of a noise-corrupted version, which has a known closed-form solution.

2. The DDPM Framework

Denoising Diffusion Probabilistic Models (Ho et al., 2020) define a forward process that progressively adds noise to data over \(T\) timesteps and a learned reverse process that denoises. The forward process is a Markov chain:

where \(\beta_1, \ldots, \beta_T\) is a variance schedule (a sequence of values controlling how much noise is added at each timestep, determining the rate at which the forward process destroys data structure). A key property of Gaussian noise is that the marginal at any timestep \(t\) can be computed in closed form without iterating through intermediate steps:

where \(\alpha_t = 1 - \beta_t\) and \(\bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s\). This means we can sample \(x_t\) directly from \(x_0\) using the reparameterization:

The reverse process is parameterized as:

where \(\mu_\theta\) is predicted by a neural network. Ho et al. showed that the optimal \(\mu_\theta\) can be expressed in terms of a noise prediction network \(\epsilon_\theta\):

The training objective simplifies to predicting the noise:

This is a mean squared error between the actual noise \(\epsilon\) and the predicted noise \(\epsilon_\theta\), averaged over uniformly sampled timesteps \(t\), data points \(x_0\), and noise realizations \(\epsilon\). In short: teach a network to predict what noise was added, and you have implicitly learned the entire data distribution.

import torch
import torch.nn as nn
import torch.nn.functional as F
import math


class SinusoidalTimeEmbedding(nn.Module):
    """Encode diffusion timestep as a sinusoidal embedding."""
    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=t.device) * -emb)
        emb = t[:, None].float() * emb[None, :]
        return torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)


class NoisePredictor(nn.Module):
    """Simple MLP noise predictor for low-dimensional data.

    For images, replace with a U-Net. For molecules, replace
    with a graph neural network (Section 34.3).
    """
    def __init__(self, data_dim: int, hidden_dim: int = 512, time_dim: int = 128):
        super().__init__()
        self.time_embed = SinusoidalTimeEmbedding(time_dim)
        self.net = nn.Sequential(
            nn.Linear(data_dim + time_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, data_dim)
        )

    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        t_emb = self.time_embed(t)
        return self.net(torch.cat([x, t_emb], dim=-1))


class DDPM:
    """Denoising Diffusion Probabilistic Model.

    Implements the forward noising process, training loss,
    and reverse sampling procedure from Ho et al. (2020).
    """
    def __init__(self, model: nn.Module, T: int = 1000,
                 beta_start: float = 1e-4, beta_end: float = 0.02):
        self.model = model
        self.T = T

        # Linear noise schedule
        self.betas = torch.linspace(beta_start, beta_end, T)
        self.alphas = 1.0 - self.betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0)

    def forward_process(self, x0: torch.Tensor, t: torch.Tensor):
        """Sample x_t from q(x_t | x_0) using reparameterization."""
        alpha_bar_t = self.alpha_bars[t].to(x0.device)
        # Reshape for broadcasting: (batch,) -> (batch, 1, ...)
        while alpha_bar_t.dim() < x0.dim():
            alpha_bar_t = alpha_bar_t.unsqueeze(-1)

        noise = torch.randn_like(x0)
        x_t = torch.sqrt(alpha_bar_t) * x0 + torch.sqrt(1 - alpha_bar_t) * noise
        return x_t, noise

    def training_loss(self, x0: torch.Tensor):
        """Compute the simple denoising loss."""
        batch_size = x0.size(0)
        # Sample random timesteps uniformly
        t = torch.randint(0, self.T, (batch_size,), device=x0.device)
        x_t, noise = self.forward_process(x0, t)
        # Predict the noise that was added
        noise_pred = self.model(x_t, t)
        return F.mse_loss(noise_pred, noise)

    @torch.no_grad()
    def sample(self, shape: tuple, device: torch.device):
        """Generate samples via iterative denoising."""
        x = torch.randn(shape, device=device)  # start from pure noise

        for t in reversed(range(self.T)):
            t_batch = torch.full((shape[0],), t, device=device, dtype=torch.long)
            # Predict and remove noise
            noise_pred = self.model(x, t_batch)

            alpha_t = self.alphas[t].to(device)
            alpha_bar_t = self.alpha_bars[t].to(device)
            beta_t = self.betas[t].to(device)

            # Posterior mean
            x = (1 / torch.sqrt(alpha_t)) * (
                x - (beta_t / torch.sqrt(1 - alpha_bar_t)) * noise_pred
            )
            # Add noise for all steps except the last
            if t > 0:
                x = x + torch.sqrt(beta_t) * torch.randn_like(x)

        return x

3. Noise Schedules and EDM Preconditioning

The step-through above reveals that each reverse step's accuracy depends on the noise level at that timestep, which raises a natural question: does the choice of noise schedule affect how well the model learns across different corruption levels?

The noise schedule \(\{\beta_t\}\) controls how quickly the forward process destroys data. The original DDPM used a linear schedule (\(\beta_1 = 10^{-4}\) to \(\beta_T = 0.02\)), but cosine schedules (Nichol & Dhariwal, 2021) produce more uniform signal-to-noise ratios (the ratio of data signal power to noise power at each timestep, \(\text{SNR}(t) = \bar{\alpha}_t / (1 - \bar{\alpha}_t)\), which determines how much original structure remains visible to the model) across timesteps, improving generation of low-resolution features.

The EDM framework (Karras et al., 2022) unified and simplified diffusion model design by moving to a continuous noise level \(\sigma\) and deriving principled choices for preconditioning, training distribution over \(\sigma\), and sampling schedule. Instead of parameterizing a noise predictor \(\epsilon_\theta\), EDM parameterizes a denoiser \(D_\theta\) that directly predicts the clean data:

where \(F_\theta\) is the raw neural network. The preconditioning functions \(c_{\text{skip}}\), \(c_{\text{out}}\), \(c_{\text{in}}\), and \(c_{\text{noise}}\) normalize \(F_\theta\)'s inputs and outputs to unit variance at every noise level. This makes the network's task uniformly difficult across all \(\sigma\) values, improving training stability and sample quality.

class EDMPreconditioner(nn.Module):
    """EDM preconditioning wrapper (Karras et al., 2022).

    Wraps a raw network F_theta with noise-dependent scaling
    to ensure unit-variance inputs and outputs at all noise levels.
    """
    def __init__(self, network: nn.Module, sigma_data: float = 0.5):
        super().__init__()
        self.network = network
        self.sigma_data = sigma_data

    def forward(self, x_noisy: torch.Tensor, sigma: torch.Tensor):
        """Predict clean data from noisy input at noise level sigma."""
        sigma = sigma.view(-1, *([1] * (x_noisy.dim() - 1)))
        sd2 = self.sigma_data ** 2

        # EDM preconditioning coefficients
        c_skip = sd2 / (sigma ** 2 + sd2)
        c_out = sigma * self.sigma_data / torch.sqrt(sigma ** 2 + sd2)
        c_in = 1 / torch.sqrt(sigma ** 2 + sd2)
        c_noise = 0.25 * torch.log(sigma.squeeze())

        # Apply preconditioning
        F_out = self.network(c_in * x_noisy, c_noise)
        return c_skip * x_noisy + c_out * F_out

    def training_loss(self, x0: torch.Tensor):
        """EDM training loss with log-normal noise sampling."""
        # Sample noise levels from log-normal (EDM recommendation)
        log_sigma = torch.randn(x0.size(0), device=x0.device) * 1.2 - 1.2
        sigma = torch.exp(log_sigma)

        # Add noise
        noise = torch.randn_like(x0)
        x_noisy = x0 + sigma.view(-1, *([1] * (x0.dim() - 1))) * noise

        # Predict clean data
        denoised = self.forward(x_noisy, sigma)

        # Weighted MSE: higher weight for lower noise levels
        weight = (sigma ** 2 + self.sigma_data ** 2) / \
                 (sigma * self.sigma_data) ** 2
        weight = weight.view(-1, *([1] * (x0.dim() - 1)))
        return (weight * (denoised - x0) ** 2).mean()

4. The Continuous-Time SDE Formulation

EDM's shift from discrete timesteps to a continuous noise level \(\sigma\) hints that the underlying process may be better understood in continuous time, and indeed the entire diffusion framework admits an elegant formulation as a stochastic differential equation.

Song et al. (2021) showed that both DDPM and score-based models are discretizations of a continuous-time stochastic differential equation (SDE). The forward process is:

where \(f(x, t)\) is the drift coefficient (the deterministic force pulling the data in a particular direction at each instant), \(g(t)\) is the diffusion coefficient, and \(w\) is a standard Wiener process (a continuous-time random walk, the mathematical formalization of Brownian motion). For the variance-preserving (VP) SDE corresponding to DDPM:

Reversing the Process

The remarkable result of Anderson (1982) is that this SDE has a corresponding reverse-time SDE:

where \(\bar{w}\) is a reverse-time Wiener process. The only unknown quantity is the score \(\nabla_x \log p_t(x)\), which is exactly what the denoising score matching network estimates.

This SDE perspective unifies discrete-time DDPM with continuous-time score matching, and it also reveals that there exists a deterministic counterpart: the probability flow ODE:

This ODE generates the same marginal distributions \(p_t(x)\) as the SDE but without stochasticity. It can be solved with standard ODE solvers, connecting diffusion models to the neural ODE framework from Chapter 33 and enabling exact log-likelihood computation via the instantaneous change-of-variables formula (a continuous-time version of the change-of-variables rule for probability densities, which tracks how the ODE flow stretches or compresses volume and thereby transforms the density at each point).

5. Flow Matching

Flow matching (Lipman et al., 2023; Liu et al., 2023) provides an alternative to diffusion that is both simpler to implement and often faster to sample from. Instead of defining a noising process and learning to reverse it, flow matching directly learns a velocity field \(v_\theta(x, t)\) that transports samples from a noise distribution \(p_0 = \mathcal{N}(0, I)\) at time \(t=0\) to the data distribution \(p_1 = p_{\text{data}}\) at time \(t=1\).

The key idea is conditional flow matching. For each training example \(x_1\), define a simple conditional path from noise to data:

This is an optimal transport path (where optimal transport refers to finding the most efficient way to rearrange one probability distribution into another, minimizing the total distance moved), here a straight line, from a random noise point to the data point. The conditional velocity along this path is:

The flow matching objective trains the network to predict this conditional velocity:

At inference time, we integrate the learned velocity field from \(t=0\) to \(t=1\) using an ODE solver, transforming noise into data. Because the conditional paths are straight lines, the ODE is easy to integrate and typically requires fewer function evaluations than diffusion model sampling.

class FlowMatchingModel:
    """Conditional flow matching for generative modeling.

    Learns a velocity field that transports noise to data
    along optimal transport paths (straight lines).
    """
    def __init__(self, velocity_net: nn.Module, sigma_min: float = 1e-4):
        self.velocity_net = velocity_net
        self.sigma_min = sigma_min  # small noise for numerical stability

    def training_loss(self, x1: torch.Tensor):
        """Compute the conditional flow matching loss."""
        batch_size = x1.size(0)

        # Sample time uniformly in [0, 1]
        t = torch.rand(batch_size, device=x1.device)

        # Sample noise (source distribution)
        x0 = torch.randn_like(x1)

        # Interpolate along the optimal transport path
        t_expanded = t.view(-1, *([1] * (x1.dim() - 1)))
        x_t = (1 - t_expanded) * x0 + t_expanded * x1

        # Target velocity: direction from noise to data
        target_velocity = x1 - x0

        # Predict velocity and compute loss
        pred_velocity = self.velocity_net(x_t, t)
        return F.mse_loss(pred_velocity, target_velocity)

    @torch.no_grad()
    def sample(self, shape: tuple, device: torch.device,
               num_steps: int = 100):
        """Generate samples by integrating the velocity field."""
        x = torch.randn(shape, device=device)  # start from noise
        dt = 1.0 / num_steps

        for i in range(num_steps):
            t = torch.full((shape[0],), i * dt, device=device)
            v = self.velocity_net(x, t)
            x = x + v * dt  # Euler integration step

        return x

    @torch.no_grad()
    def sample_midpoint(self, shape: tuple, device: torch.device,
                        num_steps: int = 50):
        """Generate samples using midpoint method (2nd-order)."""
        x = torch.randn(shape, device=device)
        dt = 1.0 / num_steps

        for i in range(num_steps):
            t = torch.full((shape[0],), i * dt, device=device)
            t_mid = torch.full((shape[0],), (i + 0.5) * dt, device=device)

            # Midpoint method: evaluate velocity at midpoint
            v1 = self.velocity_net(x, t)
            x_mid = x + 0.5 * dt * v1
            v_mid = self.velocity_net(x_mid, t_mid)
            x = x + dt * v_mid

        return x

Figure 34.2.1 compares the sampling trajectories of diffusion models and flow matching. Diffusion follows a curved path through noise space as the reverse SDE iteratively removes noise at varying scales, while flow matching traces nearly straight optimal transport paths from noise to data.

6. Diffusion vs. Flow Matching: When to Use Which

Diffusion models and flow matching solve the same problem with different mathematical machinery. As illustrated in Figure 34.2.1, the geometry of their sampling paths differs fundamentally, and this shapes the practical trade-offs:

Training. Both losses are simple MSE objectives. Flow matching is marginally simpler because it predicts velocity rather than noise, eliminating the need for noise schedule design. The EDM framework reduces this gap by providing principled noise schedule choices for diffusion.

Sampling speed. Flow matching typically requires fewer integration steps (20 to 100) because the optimal transport paths are straight lines, while diffusion models need 50 to 1000 steps along curved reverse-SDE paths. Advanced ODE solvers (DPM-Solver; Lu et al., 2022) narrow this gap to 10 to 25 steps for diffusion.

Conditioning. Both support classifier-free guidance, a technique that steers generation toward a conditioning signal (such as a text prompt) by interpolating between conditional and unconditional score estimates during sampling (Section 34.4). Diffusion models have more mature tooling for guidance (Diffusers library), while flow matching guidance is now mature, with production systems such as Stable Diffusion 3 and Flux (2024) using classifier-free guidance within flow matching pipelines.

Manifold data. Flow matching extends more naturally to non-Euclidean spaces (SE(3), Riemannian manifolds) because it only requires defining interpolation paths, not a noising process. This makes it the preferred choice for protein backbone generation on the SE(3) manifold. Diffusion on manifolds requires wrapping and unwrapping operations that add complexity, though recent work has reduced this overhead.

from diffusers import DDPMScheduler, DDIMScheduler, EulerDiscreteScheduler

# DDPM: 1000 steps, original formulation
ddpm = DDPMScheduler(num_train_timesteps=1000, beta_schedule="linear")

# DDIM: deterministic, 50 steps for fast sampling
ddim = DDIMScheduler(num_train_timesteps=1000, beta_schedule="scaled_linear")

# Euler: for flow matching / EDM-style models
euler = EulerDiscreteScheduler(num_train_timesteps=1000)

# All share the same interface
noisy_sample = ddpm.add_noise(clean_sample, noise, timesteps)
prev_sample = ddpm.step(model_output, timestep, noisy_sample).prev_sample

Exercises

1. (Conceptual) The probability flow ODE and the reverse SDE produce the same marginal distributions \(p_t(x)\) but have different sample trajectories. Explain intuitively why the SDE version (with noise) might produce higher-quality samples than the ODE version (deterministic) even though they share the same marginals. Hint: think about what happens when the score estimate is imperfect.
2. (Coding) Implement both the DDPM and FlowMatchingModel from this section and train both on a 2D Swiss roll dataset (available via sklearn.datasets.make_swiss_roll). Compare: (a) training convergence (loss vs. epochs), (b) sample quality at different numbers of sampling steps (10, 50, 100, 500), and (c) the geometry of the learned transport paths (plot trajectories from noise to data for both models). Which model produces good samples with fewer steps?
3. (Analysis) The EDM framework uses a log-normal distribution \(\log \sigma \sim \mathcal{N}(-1.2, 1.2^2)\) for sampling noise levels during training. Plot this distribution and compare it to the uniform distribution over timesteps used by the original DDPM. At which noise levels does EDM place more training weight? Why does concentrating on these noise levels improve sample quality? Relate your answer to the signal-to-noise ratio \(\text{SNR}(t) = \bar{\alpha}_t / (1 - \bar{\alpha}_t)\).