Research Log

Denoising Is Not Inversion

The Question

After seeing \(x_t\), could we recover exactly \(x_0\)?

Short answer: in general, no. Once nonzero noise has been added, the best we can do under mean-squared error (MSE) is estimate the posterior conditional mean:

\[ \mathbb{E}[x_0 \mid x_t] \]

Why I Care

The goal of a diffusion model is to generate clean data from an easy-to-sample distribution, like Gaussian noise. In the previous post, we saw how we can use the score function to move back toward cleaner regions. In this post, we'll highlight that moving toward those cleaner regions is not exact reversal, but our best guess of \(x_0\) from what we have, \(x_t\).

Setup and Definitions

One common misconception in diffusion models is to think about the reverse process — the denoising process — as exactly undoing the forward process where noise was added to a clean image. Many tutorials present an image that is first corrupted by adding Gaussian noise, and then we see arrows from the final noise back to the clean image.

This makes sense intuitively and can trick you into thinking that you have mastered diffusion models. But then you see the math and read the papers and you wonder: "Where is this coming from?" "Why did we add noise again?" among many other questions.

So let's try to step back and see how we can actually denoise an observed noisy sample.

Let's work with a simplified scalar model:

\[ x_0 \sim \mathcal{N}(\mu, \sigma_0^2), \qquad \varepsilon \sim \mathcal{N}(0, 1), \qquad \varepsilon \perp x_0 \]

In words, the noise \(\varepsilon\) is independent of \(x_0\).

Then the forward process is defined as:

\[ x_t = \alpha_t x_0 + \sigma_t \varepsilon \]

Once we observe \(x_t\), we do not know how much of it came from \(x_0\) and how much came from the noise \(\varepsilon\). Many possible clean values could have produced the same noisy observation, each with a different noise realization. So denoising is not algebraic inversion; it is conditional guessing.

Let's also define \(g(x_t)\) as a function that estimates \(x_0\) from \(x_t\), our denoiser.

We will now minimize a squared-error objective:

\[ L(g) = \mathbb{E}\left[(x_0 - g(x_t))^2\right] \]

This is the \(x_0\)-prediction MSE view of denoising. Many practical diffusion models train a network to predict the noise or the score instead, but these targets are closely connected through the Gaussian forward process.

Also, the Gaussian setup gives us diffusion context, but the MSE posterior-mean argument below is more general. It does not really need Gaussianity; it needs the relevant conditional expectations to exist.

Minimal Derivation

\[ L(g) = \mathbb{E}\left[(x_0 - g(x_t))^2\right] \]

Let's condition on \(x_t\):

\[ L(g) = \mathbb{E}\left[ \mathbb{E}\left[(x_0 - g(x_t))^2 \mid x_t\right] \right] \]

Let's define:

\[ m(x_t) := \mathbb{E}[x_0 \mid x_t] \]

the posterior mean.

Algebra trick: subtract and add the posterior mean inside the square.

\[ x_0 - g(x_t) = (x_0 - m(x_t)) + (m(x_t) - g(x_t)) \]

Then:

\[ \begin{aligned} \mathbb{E}\left[(x_0 - g(x_t))^2 \mid x_t\right] &= \mathbb{E}\left[ \left((x_0 - m(x_t)) + (m(x_t) - g(x_t))\right)^2 \mid x_t \right] \\ &= \mathbb{E}\left[(x_0 - m(x_t))^2 \mid x_t\right] \\ &\quad + 2\mathbb{E}\left[ (x_0 - m(x_t))(m(x_t) - g(x_t)) \mid x_t \right] \\ &\quad + \mathbb{E}\left[(m(x_t) - g(x_t))^2 \mid x_t\right] \end{aligned} \]

First term:

\[ \mathbb{E}\left[(x_0 - m(x_t))^2 \mid x_t\right] = \operatorname{Var}(x_0 \mid x_t) \]

by definition.

Cross-term:

\[ \mathbb{E}\left[ (x_0 - m(x_t))(m(x_t) - g(x_t)) \mid x_t \right] = (m(x_t) - g(x_t)) \mathbb{E}\left[x_0 - m(x_t) \mid x_t\right] \]

Once we condition on \(x_t\), both \(m(x_t)\) and \(g(x_t)\) are fixed, so the remaining expectation is:

\[ \begin{aligned} \mathbb{E}\left[x_0 - \mathbb{E}[x_0 \mid x_t] \mid x_t\right] &= \mathbb{E}[x_0 \mid x_t] - \mathbb{E}[x_0 \mid x_t] \\ &= 0 \end{aligned} \]

So:

\[ \mathbb{E}\left[ (x_0 - m(x_t))(m(x_t) - g(x_t)) \mid x_t \right] = 0 \]

Last term:

\[ \mathbb{E}\left[(m(x_t) - g(x_t))^2 \mid x_t\right] = (m(x_t) - g(x_t))^2 \]

Since \(m\) and \(g\) depend only on \(x_t\), there's no other randomness.

So:

\[ \mathbb{E}\left[(x_0 - g(x_t))^2 \mid x_t\right] = \operatorname{Var}(x_0 \mid x_t) + (m(x_t) - g(x_t))^2 \]
\[ L(g) = \mathbb{E}\left[\operatorname{Var}(x_0 \mid x_t)\right] + \mathbb{E}\left[(m(x_t) - g(x_t))^2\right] \]

Where:

\[ \mathbb{E}\left[\operatorname{Var}(x_0 \mid x_t)\right] \]

is irreducible uncertainty, and

\[ \mathbb{E}\left[(m(x_t) - g(x_t))^2\right] \]

is the estimator mismatch.

What I Think I Learned

We found that:

\[ L(g) = \mathbb{E}\left[\operatorname{Var}(x_0 \mid x_t)\right] + \mathbb{E}\left[(m(x_t) - g(x_t))^2\right] \]

We cannot do anything about the first term since it doesn't depend on \(g\). For the second term, the optimal outcome is when:

\[ m(x_t) - g(x_t) = 0 \quad \Rightarrow \quad g^*(x_t) = m(x_t) \]

But remember that \(m\) is the posterior conditional mean:

\[ m(x_t) = \mathbb{E}[x_0 \mid x_t] \]

Under MSE, we just proved that the best function we could use to predict a clean observation from a noisy point is the posterior conditional mean \(\mathbb{E}[x_0 \mid x_t]\). We also observed that the first term, \(\mathbb{E}[\operatorname{Var}(x_0 \mid x_t)]\), is randomness we cannot remove.

The name for what we just proved is the MMSE property of conditional expectation. Blitzstein and Hwang prove the general version in chapter 9 of Introduction to Probability.

The key insight of this post is that we should interpret this MSE denoising objective as learning the best conditional estimator, and not as exactly reverting the noise that was added. This makes sense intuitively: the best way to reduce the expected squared error is to make \(g(x_t)\) the average clean sample we expect after observing \(x_t\).

Open Questions

  • How does this connect back to the score function?
  • What changes in multiple dimensions or multiple steps?
  • How does this "many possible clean \(x_0\) values" view connect to conditional expectation more formally?
  • Can we compute \(\mathbb{E}[x_0 \mid x_t]\) analytically for real image data? What if not?

References