Mathematical Underpinnings: Score Matching and ODEs

Mathematical foundations provide a thorough understanding of diffusion models and link them to other areas of generative modeling. These foundations cover discrete-time formulations like DDPM and DDIM. This connection, particularly through score matching and differential equations, explains why these models work and enables advanced techniques, including custom noise schedules and sophisticated samplers.

Score Matching: Guiding the Reverse Process

At the core of diffusion models lies the idea of reversing a process that gradually adds noise to data. The main element needed for this reversal is knowing, at any noise level $t$, which direction to step in to make the data $\mathbf{x}_t$ slightly less noisy, moving it closer to the original data distribution. This direction is mathematically captured by the gradient of the log-probability density of the noisy data, known as the score function:

$$\nabla_{\mathbf{x}_t} \log q(\mathbf{x}_t)$$

The score function points towards regions of higher data density for the noisy data distribution $q(\mathbf{x}_t)$ at time $t$. If we had access to this score function for all $t$, we could use it to guide a reverse process, starting from pure noise and gradually removing it to generate realistic data.

Score matching is a principle for training a model, let's call it $s_\theta(\mathbf{x}_t, t)$, to approximate this true score function. The goal is to minimize the difference between the model's output and the true score:

$$\mathbb{E}_{t}\mathbb{E}_{\mathbf{x}_t \sim q(\mathbf{x}_t)} \left[ || s_\theta(\mathbf{x}_t, t) - \nabla_{\mathbf{x}_t} \log q(\mathbf{x}_t) ||^2_2 \right]$$

While directly computing the true score $\nabla_{\mathbf{x}_t} \log q(\mathbf{x}_t)$ is often intractable, techniques like denoising score matching show that this objective is equivalent to training a model to denoise samples.

Specifically, for Gaussian noise perturbations commonly used in diffusion models, it can be shown that the denoising objective used in DDPMs, where a model $\epsilon_\theta(\mathbf{x}_t, t)$ is trained to predict the noise $\epsilon$ added to obtain $\mathbf{x}_t$ from $\mathbf{x}_0$, is implicitly learning the score function. The relationship is remarkably simple:

$$s_\theta(\mathbf{x}_t, t) \approx \nabla_{\mathbf{x}_t} \log q(\mathbf{x}_t) \approx -\frac{\epsilon_\theta(\mathbf{x}_t, t)}{\sigma_t}$$

Here, $\sigma_t$ represents the standard deviation of the noise at time $t$. This fundamental connection reveals that the neural network in a DDPM, while framed as predicting noise, is effectively learning the score of the noise-perturbed data distribution. This perspective provides a solid theoretical grounding for the model's architecture and objective function.

Continuous-Time View: Stochastic Differential Equations (SDEs)

The discrete steps of the DDPM forward process $q(\mathbf{x}_t | \mathbf{x}_{t-1})$ can be seen as approximations of a continuous-time process described by a Stochastic Differential Equation (SDE). An SDE describes how a variable evolves over continuous time under the influence of both deterministic forces (drift) and random fluctuations (diffusion).

A common SDE formulation corresponding to many DDPM variants is the Variance Preserving (VP) SDE:

$$d\mathbf{x} = f(\mathbf{x}, t) dt + g(t) d\mathbf{w}$$

where:

• $d\mathbf{x}$ is the infinitesimal change in the data $\mathbf{x}$ over an infinitesimal time step $dt$.
• $f(\mathbf{x}, t)$ is the drift coefficient, determining the deterministic part of the change. For the VP-SDE often related to DDPMs, this is typically $f(\mathbf{x}, t) = -\frac{1}{2}\beta(t)\mathbf{x}$, pushing the data towards the origin.
• $g(t)$ is the diffusion coefficient, scaling the magnitude of the random noise added. For the VP-SDE, $g(t) = \sqrt{\beta(t)}$, where $\beta(t)$ is a continuous-time noise schedule.
• $d\mathbf{w}$ represents an infinitesimal step of a standard Wiener process (Brownian motion), providing the stochasticity.

This SDE describes a continuous noising process where data gradually drifts towards zero mean while accumulating Gaussian noise, with the rate controlled by $\beta(t)$. The discrete forward process $q(\mathbf{x}_t | \mathbf{x}_0)$ of a DDPM can be recovered by solving this SDE from $t=0$ to $t=T$.

Reversing Time with ODEs: The Probability Flow

The real power of the SDE perspective comes when considering the reverse process. It's a known result from SDE theory (attributed to Anderson, 1982) that the forward SDE has a corresponding reverse-time SDE that maps noise back to data. This reverse SDE also has drift and diffusion terms, and crucially, its drift term depends directly on the score function $\nabla_{\mathbf{x}_t} \log q_t(\mathbf{x}_t)$ of the marginal distributions $q_t(\mathbf{x})$ generated by the forward SDE.

The reverse SDE for the VP-SDE example is:

$$d\mathbf{x} = \left[ f(\mathbf{x}, t) - g(t)^2 \nabla_{\mathbf{x}_t} \log q_t(\mathbf{x}_t) \right] dt + g(t) d\overline{\mathbf{w}}$$

Here, $d\overline{\mathbf{w}}$ is a reverse-time Wiener process, and we integrate time backwards from $T$ to $0$. This SDE tells us how to stochastically remove noise using the score.

Furthermore, Song et al. (NeurIPS 2020) showed that associated with any diffusion process defined by an SDE, there exists a deterministic process described by an Ordinary Differential Equation (ODE) whose trajectories share the same marginal densities $q_t(\mathbf{x})$ over time. This is often called the Probability Flow ODE. For the VP-SDE, this ODE is:

$$d\mathbf{x} = \left[ f(\mathbf{x}, t) - \frac{1}{2} g(t)^2 \nabla_{\mathbf{x}_t} \log q_t(\mathbf{x}_t) \right] dt$$

Notice the absence of the stochastic term $d\mathbf{w}$. This ODE provides a deterministic path from noise to data.

Connecting ODEs, Score Models, and Sampling

This ODE is immensely significant:

1. Deterministic Generation: If we start with a sample $\mathbf{x}_T$ from the noise distribution $q_T(\mathbf{x})$ (usually a standard Gaussian) and solve this ODE backwards in time from $t=T$ to $t=0$, the resulting $\mathbf{x}_0$ is a sample from the data distribution.

2. Requires the Score: The ODE requires the score function $\nabla_{\mathbf{x}_t} \log q_t(\mathbf{x}_t)$.

3. Practical Implementation: We don't know the true score, but we have trained a model $s_\theta(\mathbf{x}_t, t)$ (or equivalently, $\epsilon_\theta(\mathbf{x}_t, t)$) to approximate it! By plugging our learned score model $s_\theta$ into the ODE, we get a practical way to generate samples:

   $$d\mathbf{x} \approx \left[ f(\mathbf{x}, t) - \frac{1}{2} g(t)^2 s_\theta(\mathbf{x}_t, t) \right] dt$$

   Substituting $f(\mathbf{x}, t) = -\frac{1}{2}\beta(t)\mathbf{x}$, $g(t)^2 = \beta(t)$, and $s_\theta(\mathbf{x}_t, t) \approx -\epsilon_\theta(\mathbf{x}_t, t)/\sigma_t$ (where $\sigma_t^2$ is the variance at time $t$) connects this directly back to the operations performed in diffusion model sampling.

4. Foundation for Advanced Samplers: DDIM sampling can be interpreted as a particular numerical method (specifically, a first-order discretization) for solving this Probability Flow ODE. This explains its deterministic nature and ability to use larger step sizes than DDPM. Moreover, this ODE formulation allows us to apply more sophisticated numerical ODE solvers (like Runge-Kutta methods, DPM-Solver, UniPC, etc.) to potentially achieve higher accuracy with fewer function evaluations (i.e., fewer model inference steps), leading to faster sampling, which we will explore in Chapter 6.

Diagram illustrating the relationship between the forward SDE, the score function, the learned score model, the reverse SDE, and the probability flow ODE used for sampling.

This perspective, linking discrete diffusion steps to underlying continuous-time SDEs and ODEs via score matching, is a foundation for understanding many advanced techniques. It clarifies the role of the noise schedule (related to $\beta(t)$ or $g(t)$) and motivates the search for better ODE solvers for faster and more accurate sampling. As we proceed to discuss custom noise schedules and learned variances, keep this SDE/ODE framework in mind as the continuous-time ideal that discrete implementations aim to approximate effectively.