Flow models II: Score Matching Techniques

March 2025

Let $p$ be any smooth probability density function. Its score is the gradient of its log-density: $\nabla \log p(x)$ . This object is of paramount importance in many fields, like physics, statistics, and machine learning. In particular, it is the key to sample from $p$ using the Langevin dynamics, and we've seen it to be the key when reversing diffusion processes. In this note, we survey the classical technique used for learning the score of a density from its samples: score matching.

Vanilla Score Matching
Tweedie's formula for denoising
Back to diffusions
Conclusion
References

Vanilla Score Matching

The Fisher Divergence

The L2-distance between the scores of two probability densities is often called the Fisher divergence:

\mathrm{fisher}(\rho_1 \mid \rho_2) = \int \rho_1(x)|\nabla\log\rho_1(x) - \nabla\log\rho_2(x)|^2dx.

Since our goal is to learn $\nabla\log p(x)$ , it is natural to choose a parametrized family of functions $s_\theta$ and to optimize $\theta$ so that the divergence

\int p(x)|\nabla\log p(x) - s_\theta(x)|^2dx

is as small as possible. However, this optimization problem is intractable, due to the presence of the explicit form of $p$ inside the integral. This is where Score Matching techniques come into play.

The score matching trick

Let $p$ be a smooth probability density function supported over $\mathbb{R}^d$ and let $X$ be a random variable with density $p$ . The following elementary identity is due to Hyvärinen, 2005; it is the basis for score matching estimation in statistics.

Let

s : \mathbb{R}^d \to \mathbb{R}^d

be any smooth function with sufficiently fast decay at

\infty

, and

X \sim p

. Then,

\mathbb{E}[\vert \nabla \log p(X) - s(X)\vert^2] = c + \mathbb{E}\left[|s(X)|^2 + 2 \nabla \cdot s(X)\right]

where

c

is a constant not depending on

s

Proof. We start by expanding the square norm: $\begin{aligned}\int p(x)|\nabla \log p(x) - s(x)|^2 dx &= \int p(x)|\nabla \log p(x)|^2 dx + \int p(x)|s(x)|^2 dx - 2\int \nabla \log p(x)\cdot p(x)s(x) dx. \end{aligned}$ The first term does not depend on $s$ , it is our constant $c$ . For the last term, we use $\nabla \log p = \nabla p / p$ then we use the integration-by-parts formula:

2\int \nabla \log p(x)\cdot p(x)s(x) dx = 2\int \nabla p(x) \cdot s(x) dx = -2 \int p(x)( \nabla \cdot s(x))dx

and the identity is proved.

The loss we want to minimize is thus

\theta \in \argmin_\theta \mathbb{E}[\vert \nabla \log p(X) - s_{\theta}(X)\vert^2] = \argmin_\theta \mathbb{E}[|s_{\theta}(X)|^2 + 2 \nabla \cdot (s_{\theta}(X))].

Practically, learning the score of a density $p$ from samples $x^1, \dots, x^n$ is done by minimizing the empirical version of the right-hand side of (3):

\ell(\theta) = \frac{1}{n}\sum_{i=1}^n |s_{\theta}(x^i)|^2 + 2 \nabla \cdot (s_{\theta}(x^i)).

How do we empirically optimize (6) in the context of diffusion models?

In the context of diffusion models, we have a whole family of densities $p_t$ and we want to learn the score for every $t \in [0,T]$ .

First, we need not solve this optimization problem for every $t$ . We could obviously discretize $[0,T]$ with $t_1, \dots, t_N$ and only solve for $\theta_{t_i}$ independently, but it is actually smarter and cheaper to approximate the whole function $(t,x) \to \nabla \log p_t(x)$ by a single neural network (a U-net, in general). That is, we use a parametrized family $s_\theta(t,x)$ . This enforces a form of time-continuity which seems natural. Now, since we want to aggregate the losses at each time, we solve the following problem:

\argmin_\theta \int_0^T w(t)\mathbb{E}[|s_{\theta}(t, X_t)|^2 + 2 \nabla \cdot (s_{\theta}(t, X_t))]dt

where $w(t)$ is a weighting function (for example, $w(t)$ can be higher for $t\approx 0$ , since we don't really care about approximating $p_T$ as precisely as $p_0$ ).

In the preceding formulation we cannot exactly compute the expectation with respect to $p_t$ , but we can approximate it with our samples $x_t^i$ . Additionnaly, we need to approximate the integral, for instance we can discretize the time steps with $t_0=0 < t_1 < \dots < t_N = T$ . Our objective function becomes

\ell(\theta) =\frac{1}{n}\sum_{t \in \{t_0, \dots, t_N\}} w(t)\sum_{i=1}^n |s_{\theta}(t, x_t^i)|^2 + 2 \nabla\cdot(s_{\theta}(t, x_t^i))

which looks computable… except it's not ideal. Suppose we perform a gradient descent on $\theta$ to find the optimal $\theta$ for time $t$ . Then at each gradient descent step, we need to evaluate $s_{\theta}$ as well as its divergence; and then compute the gradient in $\theta$ of the divergence in $x$ , in other words to compute $\nabla_\theta \nabla_x \cdot s_\theta$ . In high dimension, this can be too costly.

Tweedie's formula for denoising

Denoising Score Matching

Fortunately, there is another way to perform score matching when $p_t$ is the distribution of a random variable with gaussian noise added, as in our setting. We'll present this result in a fairly abstract setting; we suppose that $p$ is a density function, and $q = p*g$ where $g$ is an other density. The following result is due to Vincent, 2010.

Denoising Score Matching Objective

Let $s:\mathbb{R}^d \to \mathbb{R}^d$ be a smooth function. Let $X$ be a random variable with density $p$ , and let $\varepsilon$ be an independent random variable with density $g$ . We call $p_{\mathrm{noisy}}$ the distribution of $X_{\mathrm{noisy}}=X + \varepsilon$ . Then, $\mathbb{E}[\vert \nabla \log p_{\mathrm{noisy}}(X+\varepsilon) - s(X+\varepsilon)\vert^2] = c + \mathbb{E}[|\nabla \log g(\varepsilon) - s(X+\varepsilon)|^2]$ where $c$ is a constant not depending on $s$ .

Proof. Note that $p_{\mathrm{noisy}} = p * g$ . By expanding the square, we see that $\mathbb{E}[\vert \nabla \log p_{\mathrm{noisy}}(X+\varepsilon) - s(X+\varepsilon)\vert^2]$ is equal to

c + \int p_{\mathrm{noisy}}(x)|s(x)|^2dx -2\int \nabla p_{\mathrm{noisy}}(x)\cdot s(x)dx.

Now by definition, $p_{\mathrm{noisy}}(x) = \int p(y)g(x-y)dy$ , hence $\nabla p_{\mathrm{noisy}}(x) = \int p(y)\nabla g(x-y)dy$ , and the last term above is equal to $\begin{aligned} -2\int \int p(y)\nabla g(x-y)\cdot s(x)dxdy &= -2\int \int p(y)g(x-y)\nabla \log g(x-y)\cdot s(x)dydx\\ &= -2\mathbb{E}[\nabla \log g(\varepsilon)\cdot s(X + \varepsilon)]. \end{aligned}$ But then, upon adding and subtracting the term $\mathbb{E}[|\nabla \log g(\varepsilon)|^2]$ which does not depend on $s$ , we get another constant $c'$ such that

\mathbb{E}[\vert \nabla \log p_{\mathrm{noisy}}(X) - s(X)\vert^2] = c' + \mathbb{E}[|\nabla \log g(\varepsilon) - s(X + \varepsilon)|^2].

Tweedie's formula

It turns out that the Denoising Score Matching objective is just an avatar of a deep, not so-well-known result, called Tweedie's formula. Herbert Robbins is often credited with the first discovery of this formula in the context of exponential (Poisson) distributions; Maurice Tweedie extended it, and Bradley Efron popularized it in his excellent paper on selection bias.

Tweedie's formula

Let $X$ be a random variable with density $p$ , and let $\varepsilon$ be an independent random variable with distribution $N(0, \sigma^2)$ . We still note $p_{\mathrm{noisy}}$ the distribution of $X_{\mathrm{noisy}} = X + \varepsilon$ . Then, $\mathbb{E}[X \mid X_{\mathrm{noisy}}] = X_{\mathrm{noisy}} + \sigma^2 \nabla \ln p_{\mathrm{noisy}}(X_{\mathrm{noisy}}).$ Equivalently, $\mathbb{E}\left[\left. \varepsilon \right| X_{\mathrm{noisy}}\right] = -\sigma^2 \nabla \ln p_{\mathrm{noisy}}(X_{\mathrm{noisy}}).$

Proof. The joint distribution of $(X, X_{\mathrm{noisy}})$ is $p(x)g(z-x)$ , hence the conditional expectation of $X$ given $X_{\mathrm{noisy}} = z$ is

\frac{\int x p(x)g(z-x)dx}{q(z)} = z - \frac{\int (z-x)g(z-x)p(x)dx}{p_{\mathrm{noisy}}(z)}.

Note that $p_{\mathrm{noisy}}(z) = \int p(x)g(z-x)dx$ . Now, since $g$ is the density of $N(0,\sigma^2 I_d)$ , we have $\nabla g(x) = -x/\sigma^2 g(x)$ , and the term above is equal to

z +\sigma^2 \frac{\int \nabla_z g(z-x)p(x)dx}{p_{\mathrm{noisy}}(z)} = z + \sigma^2 \nabla_z \ln \int g(z-x)p(x)dx

which is the desired result. For the second formula, just note that $X_{\mathrm{noisy}} = \mathbb{E}[X_{\mathrm{noisy}}|X_{\mathrm{noisy}}] = \mathbb{E}[X \mid X_{\mathrm{noisy}}] + \mathbb{E}[\varepsilon \mid X_{\mathrm{noisy}}]$ then simplify.

Since $X_{\mathrm{noisy}}$ is centered around $X$ , the classical ("frequentist") estimate for $X$ given $X_{\mathrm{noisy}}$ is $X_{\mathrm{noisy}}$ . Tweedie's formula corrects this estimate by adding a term accounting for the fact that $X$ is itself random: for instance, if $X_{\mathrm{noisy}}$ lands in a region where $X$ is extremely unlikely to live in, then it is probable that the noise is responsible for this, and the estimate for $X$ .

Now, what's the link between this and Denoising Score Matching ? Well, the conditional expectation of any $X$ given $Z$ minimizes the $L^2$ -distance, in the sense that $\mathbb{E}[X \mid Z] = f(Z)$ where $f$ minimizes $\mathbb{E}[|f(Z) - X|^2]$ . Formula (15) says that $\nabla \ln p_{\mathrm{noisy}}$ , the quantity we need to estimate, is nothing but the « best denoiser » up to a scaling $-\sigma^2$ :

\nabla \ln p_{\mathrm{noisy}} \in \arg \min \mathbb{E}[|f(X + \varepsilon) - (-\varepsilon/\sigma^2)|^2].

Replacing $f$ with a neural network, we exactly find the Denoising Score Matching objective in (10).

I find this result to give some intuition on score matching and how to parametrize the network.

If $s(x)$ is « noise predictor » trained with the objective $\mathbb{E}[|s(X+\varepsilon) - \varepsilon|^2]$ , then $-s(x)/\sigma^2$ is a proxy for $\nabla \ln p_{\mathrm{noisy}}$ . This is called « noise prediction » training.
If $s(x)$ had rather been trained on a pure denoising objective $\mathbb{E}[|s(X+\varepsilon) - X|^2]$ , then $x - s(x)$ is a noise predictor for $X$ , hence $\sigma^{-2} (s(x) - x)$ is a proxy for $\nabla \ln p_{\mathrm{noisy}}$ . This is called « data prediction » training, or simply denoising training.

Both formulations are found in the litterature, as well as affine mixes of both.

Back to diffusions

Let us apply this to our setting. Remember that $p_t$ is the density of $\alpha_t X_0 + \varepsilon_t$ where $\varepsilon_t \sim \mathscr{N}(0,\bar{\sigma}_t^2)$ , hence in this case $g(x) = (2\pi\bar{\sigma}_t^2)^{-d/2}e^{-|x|^2 / 2\bar{\sigma}_t^2}$ and $\nabla \log g(x) = - x / \bar{\sigma}^2_t$ .

Data prediction model

The « data prediction » (or denoising) parametrization of the neural network would minimize the objective

\int_0^T w(t)\mathbb{E}[|s_{\theta}(t, X_t) - \alpha_t X_0|^2]dt.

Noise prediction model

Alternatively, the « noise prediction » parametrization of the neural network would minimize the objective

\int_0^T w(t)\mathbb{E}[|s_{\theta}(t, X_t) - \varepsilon|^2]dt

where I noted $\varepsilon$ instead of $X_1$ to emphasize we're predicting noise.

Since we have access to samples $(x^i, \varepsilon^i, \tau^i)$ where $\tau \sim w$ , we get the empirical version:

\hat{\ell}(\theta) = \frac{1}{n}\sum_{i=1}^n \left[|\varepsilon^i - s_\theta(\alpha_\tau x^i + \bar{\sigma}_\tau \varepsilon^i)|^2\right].

Up to the constants and the choice of the drift $\mu_t$ and variance $w_t$ , this is exactly the loss function (14) from the DDPM paper. Once this is done, the proxy for $\nabla \ln p_t$ would be

\nabla \ln p_t(x) \approx \frac{x - s_{\theta}(t,x)}{\bar{\sigma}_t^2}.

Plugging this back into the SDE (DDPM) sampling formula, we get

dY_t = \mu_{T-t}Y_t + 2\frac{w^2_{T-t}}{\bar{\sigma}_{T-t}^2}(Y_t - s_{\theta}(T-t,Y_t))dt + \sqrt{2w_{T-t}}dB_t

dY_t = \left(\mu_{T-t} + 2\frac{w^2_{T-t}}{\bar{\sigma}_{T-t}^2}\right)Y_t - 2\frac{w^2_{T-t}}{\bar{\sigma}_{T-t}^2}s_{\theta}(T-t,Y_t)dt + \sqrt{2w_{T-t}}dB_t.

Conclusion

We are now equipped with learning $\nabla \ln p_t$ when $p_t$ is the distribution of something like $\alpha_t X_0+\sigma_t \varepsilon$ . The diffusion models we presented earlier provide such distributions given the drift and diffusion coefficients, $\mu_t$ and $\sigma_t$ . But in general, we should be able to directly choose $\alpha$ and $\sigma$ without relying on the (intricate) maths behind the Fokker-Planck equation. This is where the flow matching formulation enters the game.

References

Hyvärinen's paper on Score Matching.

Efron's paper on Tweedie's formula - a gem in statistics.