Research Log

Averaging Paths to the Peaks of Our Mountains

The Question

How is the marginal score related to the posterior mean of the clean sample \(\mathbb{E}[x_0 \mid x_t]\)?

Short answer: the marginal score points from the noisy observation \(x_t\) toward the scaled posterior mean \(\alpha_t\mathbb{E}[x_0 \mid x_t]\), divided by the noise variance:

\[ \nabla_{x_t}\log p_t(x_t) = \frac{\alpha_t\mathbb{E}[x_0 \mid x_t]-x_t}{\sigma_t^2}. \]

Why I Care

In previous posts, we've been discussing the score function and the posterior conditional mean. We've seen that the score function helps find the direction that increases the log-density. On the other hand, we proved that, under MSE, the posterior conditional mean \(\mathbb{E}[x_0 \mid x_t]\) is the best estimator we could have to recover \(x_0\) after seeing a specific noisy sample \(x_t\).

The conditional Gaussian score is easy to write down when \(x_0\) is known. But during generation, \(x_0\) is exactly what we do not have. We only have \(x_t\), so what we really need is the marginal score.

Today, we'll reward our brains with a delightful connection between these two concepts while we continue our quest to master SotA generative models from first principles.

Setup and Definitions

Let's revisit our forward noising formula at time \(t\):

\[ x_t=\alpha_t x_0+\sigma_t\epsilon, \qquad \epsilon\sim\mathcal{N}(0,I), \qquad \epsilon\perp x_0. \]

Here:

  • \(x_0\) is the clean data variable.
  • \(x_t\) is the noisy observation at time \(t\).
  • \(\alpha_t\) controls how much clean signal remains.
  • \(\sigma_t\) controls how much Gaussian noise is added.

In the “Score of a Gaussian” post, we derived the score function as:

\[ \nabla_{x_t}\log p(x_t \mid x_0) = -\frac{x_t-x_0}{\sigma^2}. \]

Under this Gaussian forward process, the conditional distribution is:

\[ p(x_t \mid x_0) = \mathcal{N}(\alpha_t x_0,\sigma_t^2I). \]

Its conditional score is therefore:

\[ \nabla_{x_t}\log p(x_t \mid x_0) = -\frac{x_t-\alpha_t x_0}{\sigma_t^2}. \]
\((1)\)

Finally, let's recall two identities:

\[ \nabla_x\log f(x) = \frac{\nabla_x f(x)}{f(x)}, \]
\((2)\)

and, similarly,

\[ \nabla_x f(x) = f(x)\nabla_x\log f(x). \]
\((3)\)

Minimal Derivation

By the law of total probability, the marginal density of \(x_t\) is obtained by taking a weighted average over the conditional densities associated with all possible clean values \(x_0\) that could have produced the observed \(x_t\). Namely:

\[ p_t(x_t) = \int p(x_t \mid x_0)p_{\text{data}}(x_0)\,dx_0. \]

We also know, by the identity in (2), that:

\[ \nabla_{x_t}\log p_t(x_t) = \frac{1}{p_t(x_t)}\nabla_{x_t}p_t(x_t). \]

Let's substitute the definition of \(p_t(x_t)\):

\[ \nabla_{x_t}\log p_t(x_t) = \frac{1}{p_t(x_t)} \nabla_{x_t} \int p(x_t \mid x_0)p_{\text{data}}(x_0)\,dx_0. \]

Assuming the usual regularity conditions, we can move the derivative inside the integral:

\[ \nabla_{x_t}\log p_t(x_t) = \frac{1}{p_t(x_t)} \int \nabla_{x_t} \left[ p(x_t \mid x_0)p_{\text{data}}(x_0) \right]\,dx_0. \]

With respect to \(x_t\), \(p_{\text{data}}(x_0)\) is constant:

\[ \nabla_{x_t}\log p_t(x_t) = \frac{1}{p_t(x_t)} \int \nabla_{x_t}p(x_t \mid x_0) p_{\text{data}}(x_0)\,dx_0. \]

Using (3) from Setup and Definitions:

\[ \nabla_{x_t}p(x_t \mid x_0) = p(x_t \mid x_0) \nabla_{x_t}\log p(x_t \mid x_0). \]

So:

\[ \nabla_{x_t}\log p_t(x_t) = \frac{1}{p_t(x_t)} \int p(x_t \mid x_0) \nabla_{x_t}\log p(x_t \mid x_0) p_{\text{data}}(x_0)\,dx_0. \]

Let's move the normalizing factor inside the integral:

\[ \nabla_{x_t}\log p_t(x_t) = \int \nabla_{x_t}\log p(x_t \mid x_0) \frac{ p(x_t \mid x_0)p_{\text{data}}(x_0) }{ p_t(x_t) }\,dx_0. \]

And we could use Bayes' rule to simplify things:

\[ p(x_0 \mid x_t) = \frac{ p(x_t \mid x_0)p_{\text{data}}(x_0) }{ p_t(x_t) }. \]

So:

\[ \nabla_{x_t}\log p_t(x_t) = \int \nabla_{x_t}\log p(x_t \mid x_0) p(x_0 \mid x_t)\,dx_0. \]

Nice! This is, by definition, the conditional expectation of the score function:

\[ \nabla_{x_t}\log p_t(x_t) = \mathbb{E}\left[ \nabla_{x_t}\log p(x_t \mid x_0) \mid x_t \right]. \]
\((4)\)

Equation (4) is the general identity. The clean posterior-mean formula is the Gaussian bonus.

We'll now use (1) from Setup and Definitions, and the math becomes straightforward:

\[ \begin{aligned} \nabla_{x_t}\log p_t(x_t) &= \mathbb{E}\left[ -\frac{x_t-\alpha_t x_0}{\sigma_t^2} \,\middle|\,x_t \right] \\ &= -\frac{1}{\sigma_t^2} \mathbb{E}\left[ x_t-\alpha_t x_0 \mid x_t \right] \qquad \text{(pull out the constant)} \\ &= -\frac{1}{\sigma_t^2} \left( \mathbb{E}[x_t \mid x_t] - \alpha_t\mathbb{E}[x_0 \mid x_t] \right) \qquad \text{(by linearity)} \\ &= -\frac{1}{\sigma_t^2} \left( x_t-\alpha_t\mathbb{E}[x_0 \mid x_t] \right) \qquad \text{(given }x_t\text{, }x_t\text{ is fixed)} \\ &= \frac{ \alpha_t\mathbb{E}[x_0 \mid x_t]-x_t }{ \sigma_t^2 }. \end{aligned} \]

Reading the same identity in the other direction gives us:

\[ \mathbb{E}[x_0 \mid x_t] = \frac{ x_t+\sigma_t^2\nabla_{x_t}\log p_t(x_t) }{ \alpha_t }, \qquad \alpha_t\ne 0. \]

So, for a fixed noise schedule, knowing the posterior mean gives us the marginal score, and knowing the marginal score gives us the posterior mean.

This inverse reading is a classical result: with \(\alpha_t = 1\), it is Tweedie's formula. Robbins (1956) published it, crediting Maurice Tweedie, and Efron (2011) gives the modern treatment. Identity (4) is also the backbone of denoising score matching (Vincent, 2011): a model can learn the marginal score from conditional targets that we can compute in closed form.

Visualizing the Average Arrow

The derivation says that the marginal score is a posterior-weighted average of local conditional scores. This playground lets us move \(x_t\) and watch those arrows and their weights change. To make that average visible, we'll replace the integral with only three equally likely clean origins.

What I Think I Learned

I find myself in a valley surrounded by mountains. I know that I got to my current position walking from one of the peaks of the many possible mountains, but which one?

The key insight from formula (4) is that, at my current position, every possible peak gives me a local uphill arrow. The score averages those arrows, weighted by how probable each clean origin looks after observing \(x_t\). A full path appears only if I move, recompute the score, and keep going.

Open Questions

  • MSE targets the posterior mean. What estimator would MAE target instead?
  • How do neural networks help us when \(\mathbb{E}[x_0 \mid x_t]\) is not tractable?
  • If the posterior mean and score determine each other, why might we prefer to train for one rather than the other?

References