Score-Based Generative Modeling

While Denoising Diffusion Probabilistic Models (DDPMs) offer one perspective on diffusion models, focusing on parameterized forward and reverse Markov chains, Score-Based Generative Modeling provides an alternative, yet deeply connected, viewpoint. This approach centers on directly modeling the gradient of the logarithm of the data probability density function, known as the score function, at various noise levels. Understanding this perspective illuminates the connection between diffusion models and established techniques like score matching and Langevin dynamics.

The Score Function: Gradient of Log-Probability

For a given data distribution $p(\mathbf{x})$, the score function is defined as the gradient of its log-probability with respect to the data $\mathbf{x}$:

$$\nabla_{\mathbf{x}} \log p(\mathbf{x})$$

This vector field is significant because it points in the direction of the steepest ascent of the log-probability density at any given point $\mathbf{x}$. Intuitively, it tells us how to modify $\mathbf{x}$ to make it slightly more probable under the distribution $p(\mathbf{x})$.

In the context of diffusion models, we are interested in the score function of the perturbed data distributions $p_t(\mathbf{x}_t)$ at different noise levels $t$. Recall that $\mathbf{x}_t$ is the data point $\mathbf{x}_0$ after $t$ steps of the forward noising process. The score function $\nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t)$ indicates how to adjust a noisy sample $\mathbf{x}_t$ to increase its likelihood under the marginal distribution at time $t$.

Connecting Scores to the Reverse Diffusion Process

The reverse process in diffusion models aims to denoise a sample $\mathbf{x}_t$ to estimate $\mathbf{x}_{t-1}$. It turns out that the optimal reverse transition, which maximizes the likelihood of the data, is closely related to the score function $\nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t)$.

Specifically, Tweedie's formula provides a link showing that the conditional expectation of the original data point $\mathbf{x}_0$ given the noisy version $\mathbf{x}_t$ involves the score of the noisy distribution $p_t$. The mean of the reverse transition $p_\theta(\mathbf{x}_{t-1} | \mathbf{x}_t)$ depends on estimating this score.

The neural network $\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)$ trained in DDPMs using the simplified objective:

$$L_{simple}(\theta) = \mathbb{E}_{t, \mathbf{x}_0, \boldsymbol{\epsilon}} \left[ \| \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon}, t) \|^2 \right]$$

is essentially learning a re-scaled version of the score function of the noisy data distribution $p_t(\mathbf{x}_t)$. The relationship is approximately:

$$\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t) \approx -\sqrt{1-\bar{\alpha}_t} \nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t)$$

This shows that even when training a DDPM to predict the noise $\boldsymbol{\epsilon}$, the model is implicitly learning the score of the intermediate noisy data distributions.

Training Score Models with Score Matching

Instead of implicitly learning the score via noise prediction, score-based models aim to train a neural network, often denoted as $\mathbf{s}_\theta(\mathbf{x}, t)$, to directly estimate the score function $\nabla_{\mathbf{x}} \log p_t(\mathbf{x})$. A naive approach would be to minimize the squared difference between the model's output and the true score:

$$\mathcal{L}(\theta) = \mathbb{E}_{t} \mathbb{E}_{p_t(\mathbf{x}_t)} \left[ \| \mathbf{s}_\theta(\mathbf{x}_t, t) - \nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t) \|^2_2 \right]$$

However, computing the true score $\nabla_{\mathbf{x}_t} \log p_t(\mathbf{x}_t)$ is generally intractable because $p_t(\mathbf{x}_t)$ involves an unknown normalization constant and requires integrating over all possible $\mathbf{x}_0$.

Score Matching provides a way around this. The original score matching objective, equivalent to the one above, involves the trace of the Hessian of the score model, which can be computationally expensive. More practical variants exist:

1. Denoising Score Matching (DSM): This method perturbs the data slightly with known noise and trains the model to predict the score of the perturbed distribution. Interestingly, the objective function for DSM under Gaussian noise turns out to be closely related to the noise prediction objective used in DDPMs. It trains the score network $\mathbf{s}_\theta(\mathbf{x}, t)$ by minimizing an objective related to recovering clean data from noisy versions, effectively connecting score estimation to denoising.
2. Sliced Score Matching (SSM): This technique avoids computing the full Hessian trace by projecting the score matching objective onto random directions, making it more scalable for high-dimensional data.

Training a score-based model $\mathbf{s}_\theta(\mathbf{x}_t, t)$ involves sampling timesteps $t$, sampling data points $\mathbf{x}_0$, generating corresponding noisy samples $\mathbf{x}_t$ using the forward process, and then optimizing $\theta$ using a suitable score matching objective like DSM.

Generating Samples via Langevin Dynamics

Once a time-dependent score model $\mathbf{s}_\theta(\mathbf{x}, t)$ has been trained, we can generate new samples by simulating the reverse diffusion process. A common algorithm for this is Annealed Langevin Dynamics.

Langevin dynamics is an iterative method originally used in physics to sample from a probability distribution $p(\mathbf{x})$ using its score function $\nabla_{\mathbf{x}} \log p(\mathbf{x})$. The update rule combines a gradient ascent step on the log probability (following the score) with Gaussian noise injection:

$$\mathbf{x}_{i+1} = \mathbf{x}_i + \frac{\eta}{2} \nabla_{\mathbf{x}} \log p(\mathbf{x}_i) + \sqrt{\eta} \mathbf{z}_i$$

Here, $\eta$ is the step size and $\mathbf{z}_i \sim \mathcal{N}(0, \mathbf{I})$ is Gaussian noise.

For score-based generative models, we apply this idea across decreasing noise levels (annealing). We start with a sample from the prior distribution, typically pure Gaussian noise $\mathbf{x}_T \sim \mathcal{N}(0, \mathbf{I})$, and iteratively refine it using the learned score function $\mathbf{s}_\theta(\mathbf{x}_t, t)$ for $t = T, T-1, \dots, 1$. A typical update step looks like:

$$\mathbf{x}_{t-1} = \mathbf{x}_t + \alpha_t \mathbf{s}_\theta(\mathbf{x}_t, t) + \sqrt{2 \alpha_t} \mathbf{z}_t$$

where $\alpha_t > 0$ is a step size that depends on the noise level $t$, and $\mathbf{z}_t \sim \mathcal{N}(0, \mathbf{I})$. This process gradually moves the sample away from noise and towards regions of high probability under the data distribution, guided by the learned score field. More sophisticated sampling procedures, often involving predictor-corrector steps inspired by numerical methods for solving SDEs, are commonly used to improve sample quality and efficiency. These samplers often resemble the reverse process sampler used in DDPMs.

Unifying Perspectives: DDPMs and Score-Based Models

The score-based perspective highlights the fundamental role of the score function in generative modeling via diffusion. It reveals that:

• DDPMs, while trained to predict noise, implicitly learn the score function. Their sampling process is a discretized version of simulating the reverse-time SDE governed by the score.
• Score-based models explicitly train a network to estimate the score using score matching techniques. Sampling is performed using algorithms like annealed Langevin dynamics, which also approximate the reverse-time SDE.

The Stochastic Differential Equation (SDE) formulation, discussed previously, provides a continuous-time framework that naturally encompasses both DDPMs and score-based models. Both approaches can be seen as different ways to discretize and implement the same underlying continuous diffusion process.

Summary

Score-Based Generative Modeling provides a powerful theoretical lens for understanding diffusion models. By focusing on the score function $\nabla_{\mathbf{x}} \log p_t(\mathbf{x}_t)$, it connects diffusion processes to score matching for training and Langevin dynamics for sampling. This perspective not only clarifies the mechanisms behind DDPMs but also opens avenues for designing new models and samplers based on directly estimating and utilizing the score function of the evolving data distribution under noise. The close relationship between noise prediction in DDPMs and score estimation underscores the deep connections between these seemingly different formulations.