Score Matching & Score-Based Generative Models

Why It Matters for Modern Models

Score functions underlie diffusion models: the network trained at each noise level is, up to scaling, an estimate of the score of the noised data distribution (Song & Ermon, 2019).
This is why diffusion training is a regression problem: the best possible noise prediction is $-\sigma\, \nabla_{\tilde x} \log q_\sigma(\tilde x)$ , a rescaled score.
Energy-based models and diffusion models meet here: both can be trained through the score, which does not depend on the intractable normalizing constant.

Questions to Think About

Why does a denoiser that predicts $\mathbb E[x \mid \tilde x]$ give the score through Tweedie's formula, $\nabla \log q_\sigma(\tilde x) = (\mathbb E[x \mid \tilde x] - \tilde x)/\sigma^2$ ?
Why can score matching avoid the partition function entirely, and what does a model trained only on gradients of log p fail to tell you about p itself?

Interactive Demo

The interactive demo for this concept is in the newer notebook.

Open the notebook demo

Core Math

If you prefer intuition first, try the demo in the newer notebook, then come back for the math.

Key Equation

s(x) = \nabla_x \log p(x)

The score of a density is the gradient of its logarithm with respect to the input:

s(x) = \nabla_x \log p(x)

If $p(x) = e^{-E(x)}/Z$ , the score is $-\nabla_x E(x)$ : the normalizing constant $Z$ drops out.

Score matching fits a model $s_\theta(x) \approx \nabla_x \log p_{\text{data}}(x)$ without knowing $Z$ or the true score. After integration by parts, the objective needs only samples from the data (Hyvärinen, 2005):

\mathcal{L}_{SM} = \mathbb{E}_{p_{\text{data}}}\left[ \frac{1}{2}\|s_\theta(x)\|^2 + \text{tr}(\nabla_x s_\theta(x)) \right]

The trace term is expensive in high dimensions, which motivates the denoising form.

Denoising score matching adds Gaussian noise, $\tilde{x} = x + \sigma\epsilon$ with $\epsilon \sim \mathcal N(0, I)$ , and regresses onto the score of the noise kernel (Vincent, 2011):

\mathcal{L}_{DSM} = \mathbb{E}_{x, \tilde{x}}\left[ \|s_\theta(\tilde{x}) - \nabla_{\tilde{x}} \log q(\tilde{x}|x)\|^2 \right], \qquad \nabla_{\tilde{x}} \log q(\tilde{x}|x) = -\frac{\tilde x - x}{\sigma^2} = -\frac{\epsilon}{\sigma}

The regression target $-\epsilon/\sigma$ is noisy, but the minimizer is its conditional mean given $\tilde x$ , which equals the score of the noised data distribution, $\nabla_{\tilde x} \log q_\sigma(\tilde x)$ .