Mathematical Foundations: Stochastic Differential Equations

Diffusion models operate by progressively adding noise to data and then learning to reverse this process. While models like DDPM are often presented using discrete time steps, a more general and powerful mathematical framework arises when we consider the continuous-time limit of this noising process. This leads us to the language of Stochastic Differential Equations (SDEs).

Understanding the SDE formulation provides deeper insights into why diffusion models work and unifies various discrete-time diffusion approaches under a single mathematical structure. It also opens doors to more flexible noise scheduling and sampling techniques.

From Discrete Steps to Continuous Processes

Recall the discrete forward process in DDPM, where noise is added incrementally at each step $t$. If we consider infinitesimally small time steps, this sequence of transformations converges to a continuous stochastic process. An SDE describes the evolution of a variable over continuous time, incorporating both deterministic change (drift) and random fluctuations (diffusion).

A general Itô SDE takes the form:

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

Here:

• $\mathbf{x}_t$ represents the state (e.g., an image) at time $t$.
• $d\mathbf{x}_t$ is the infinitesimal change in $\mathbf{x}_t$.
• $\mathbf{f}(\mathbf{x}_t, t)$ is the drift function, describing the deterministic component of the change.
• $g(t)$ is the diffusion coefficient, scaling the magnitude of the random noise.
• $d\mathbf{w}_t$ represents an infinitesimal step of a standard Wiener process (Brownian motion), modeling the random fluctuations. It's essentially Gaussian noise scaled by $\sqrt{dt}$.

The Forward SDE: Data to Noise

In the context of diffusion models, the forward process transforms complex data $\mathbf{x}_0$ into a simple noise distribution (typically Gaussian) as time $t$ progresses from $0$ to $T$. This "information destruction" process can be modeled by a specific SDE. A common choice, corresponding to the Variance Preserving (VP) SDE often linked to DDPM, is:

$$d\mathbf{x}_t = -\frac{1}{2} \beta(t) \mathbf{x}_t dt + \sqrt{\beta(t)} d\mathbf{w}_t$$

Here, $\beta(t)$ is a positive, time-dependent function often called the noise schedule.

• The drift term $-\frac{1}{2} \beta(t) \mathbf{x}_t dt$ pushes the state towards the origin (mean zero).
• The diffusion term $\sqrt{\beta(t)} d\mathbf{w}_t$ continuously adds Gaussian noise, with the variance controlled by $\beta(t)$.

As $t$ increases from $0$ to $T$, the influence of the initial data $\mathbf{x}_0$ diminishes, and $\mathbf{x}_T$ approaches a standard Gaussian distribution, irrespective of $\mathbf{x}_0$.

The Reverse SDE: Noise to Data

The generative power of diffusion models comes from reversing this process. We start with a sample $\mathbf{x}_T$ from the simple noise distribution and evolve it backward in time from $T$ to $0$ to generate a data sample $\mathbf{x}_0$. A remarkable result from stochastic calculus (Anderson, 1982) states that the reverse trajectory of a diffusion process defined by a forward SDE also follows an SDE, provided we know the score function of the marginal distributions $p_t(\mathbf{x}_t)$.

The reverse SDE corresponding to the forward process above is given by:

$$d\mathbf{x}_t = \left[ -\frac{1}{2} \beta(t) \mathbf{x}_t - \beta(t) \nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t) \right] dt + \sqrt{\beta(t)} d\bar{\mathbf{w}}_t$$

Here:

• $d\bar{\mathbf{w}}_t$ represents a standard Wiener process running backward in time.
• $\nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t)$ is the score function of the data distribution at time $t$. This term represents the gradient of the log-probability density with respect to the data $\mathbf{x}_t$. Intuitively, it points in the direction where the data density increases most rapidly.

This reverse SDE tells us how to infinitesimally adjust the current state $\mathbf{x}_t$ to make it slightly more likely under the data distribution $p_t$. The drift term now includes the score function, effectively guiding the process away from noise and towards plausible data structures.

The Role of Score Matching

The central challenge in using the reverse SDE for generation is that we don't know the true score function $\nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t)$ for the intermediate distributions $p_t$. This is where neural networks come in. We train a time-dependent neural network, often denoted as $\mathbf{s}_\theta(\mathbf{x}_t, t)$, to approximate the true score function:

$$\mathbf{s}_\theta(\mathbf{x}_t, t) \approx \nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t)$$

This network is typically trained using objectives derived from score matching or objectives equivalent to those used in DDPMs (like the $L_{simple}$ objective mentioned in the chapter introduction, which implicitly learns the score). Once trained, $\mathbf{s}_\theta(\mathbf{x}_t, t)$ can be plugged into the reverse SDE:

$$d\mathbf{x}_t = \left[ -\frac{1}{2} \beta(t) \mathbf{x}_t - \beta(t) \mathbf{s}_\theta(\mathbf{x}_t, t) \right] dt + \sqrt{\beta(t)} d\bar{\mathbf{w}}_t$$

Simulating this SDE backward in time, starting from $\mathbf{x}_T \sim \mathcal{N}(0, \mathbf{I})$, allows us to generate new data samples $\mathbf{x}_0$.

Diagram illustrating the forward SDE destroying data structure over time and the learned reverse SDE reconstructing data from noise by following the estimated score function.

Significance of the SDE Framework

Viewing diffusion models through the lens of SDEs offers several advantages:

1. Unification: It provides a common mathematical framework encompassing various discrete-time models like DDPM and related approaches derived from different discretizations of the SDEs.
2. Flexibility: It naturally handles continuous time, allowing for more flexible noise schedules $\beta(t)$ and sampling strategies with fixed discrete steps.
3. Theoretical Grounding: It connects diffusion models to established theories of stochastic processes and score matching, facilitating deeper analysis.
4. Advanced Sampling: SDE solvers (numerical methods for simulating SDEs) can be employed for potentially faster or more accurate generation compared to simple reverse-stepping of discrete models.

This continuous-time perspective sets the stage for understanding score-based generative modeling and advanced techniques like DDIM, which uses properties of the underlying SDEs for efficient sampling. We will build on these foundations as we examine the implementation details and improvements of diffusion models in the subsequent sections.