Sampling Our Way Down, One Gaussian Step at a Time

The Question

In some diffusion-model approaches, why can each reverse step be sampled from a Gaussian distribution?

Short answer: because the forward process adds Gaussian noise, and when the amount of noise added per step is very small, we can approximate the conditional reverse distribution with a Gaussian. In the simple setup of this post,

p(x_{t-\Delta t}\mid x_t) \approx \mathcal{N}\!\left( x_{t-\Delta t}; \mu(x_t,t), h \right).

Why I Care

Sampling in diffusion models is a huge topic with many subtleties. In this post, I want to address one of those subtleties that, from what I have seen, causes a lot of confusion. It is clear that the forward process ultimately transforms the input image into pure noise. What is harder to follow is why we use a Gaussian distribution in the reverse process, i.e., during sampling. I think this confusion comes from not differentiating between marginal, joint, and conditional distributions.

Today, I want to explore a setup that lets us approximate \(p(x_{t-\Delta t}\mid x_t)\) with a Gaussian distribution. If we can show this under our setup, then, for an observed \(x_t\), the reverse step reduces to sampling a possible \(x_{t-\Delta t}\) from a Gaussian with the correct conditional mean and variance. Starting from the corresponding terminal distribution and repeating accurate reverse steps can then produce \(x_0\), a sample from our original distribution.

I want to highlight that the marginal distribution, e.g., \(p(x_{t-\Delta t})\), is not necessarily Gaussian. It is the conditional distribution that is approximately Gaussian. More specifically, once we know the image at the next step, the possible previous images that could have generated it are distributed approximately as a Gaussian.

This also explains why we inject noise in some sampling approaches, like DDPM, which might seem odd when cleaning up the image.

Setup and Definitions

Let

Z:=x_t, \qquad Y:=x_{t-\Delta t},

such that our forward process is

Z=Y+\sqrt{h}\,\varepsilon, \qquad \varepsilon\sim\mathcal{N}(0,1),

where

\(h\) is the variance injected over one step, so \(\sqrt{h}\) scales the noise; in a continuous-time interpretation, \(h\) is proportional to \(\Delta t\),
\(\varepsilon\) is independent of \(Y\), and
\(\Delta t\) is one discrete step among equally spaced steps from \(0\) to \(T\).

With this setup, after conditioning on \(Y=y\), the only randomness left in \(Z\) comes from \(\varepsilon\). Therefore,

Z\mid Y=y\sim\mathcal{N}(y,h).

Its conditional density is

\[ p_{Z\mid Y}(z\mid y) = \frac{1}{\sqrt{2\pi h}} \exp\left(-\frac{(z-y)^2}{2h}\right). \] \((1)\)

Minimal Derivation

Let \(p_Y(y)\) be the density of the previous state. It does not need to be Gaussian. Under mild regularity conditions around the observed \(z\), Bayes' rule gives:

p_{Y\mid Z}(y\mid z) = \frac{p_{Z\mid Y}(z\mid y)p_Y(y)}{p_Z(z)}.

We can substitute Equation \((1)\) from our setup:

p_{Y\mid Z}(y\mid z) = \frac{ \dfrac{1}{\sqrt{2\pi h}} \exp\left(-\dfrac{(z-y)^2}{2h}\right)p_Y(y) }{p_Z(z)}.

Since \(z\) is fixed, we can write this as a proportionality:

p_{Y\mid Z}(y\mid z) \propto \frac{1}{\sqrt{2\pi h}} \exp\left(-\frac{(z-y)^2}{2h}\right)p_Y(y).

Now let \(h\) become small. This is the key step. As \(h\) shrinks,

\exp\left(-\frac{(z-y)^2}{2h}\right)

matters only when \(z\) and \(y\) are close. If \(z-y\) is large, dividing its square by a small \(h\) makes the exponent strongly negative, pushing the expression toward zero. Therefore, we have a density that is almost zero everywhere except inside the tiny window where \(z-y\approx 0\).

Then, within that tiny window,

p_Y(y)\approx p_Y(z),

such that

p_{Y\mid Z}(y\mid z) \propto \frac{1}{\sqrt{2\pi h}} \exp\left(-\frac{(z-y)^2}{2h}\right)p_Y(z).

Because \(z\) is fixed, \(p_Y(z)\) is constant with respect to \(y\). It multiplies every possible \(y\) by the same amount, so it disappears when we normalize the density. We therefore obtain

p_{Y\mid Z}(y\mid z) \approx \frac{1}{\sqrt{2\pi h}} \exp\left(-\frac{(y-z)^2}{2h}\right),

which is the density of

Y\mid Z=z\approx\mathcal{N}(z,h).

Approximating \(p_Y(y)\) as constant inside the tiny window gives us this zeroth-order Gaussian shape. It does not yet determine the reverse mean accurately enough for actual sampling. Let

\mu(z):=\mathbb{E}[Y\mid Z=z].

Then the more useful small-step approximation is

Y\mid Z=z\approx\mathcal{N}(\mu(z),h),

where \(\mu(z)\) is close to, but generally not exactly, \(z\). In this smooth setting, \(\mu(z)=z+O(h)\). That small displacement gives the reverse step its denoising direction. I will connect it to the score in a follow-up.

Visualizing the Tiny Window

The distinction between the marginal and the conditional is easier to see than to say. In the playground below, move the observed \(z=x_t\) and shrink \(h\). The marginal \(p_Y(y)\) stays multimodal, while the Gaussian likelihood isolates a tiny window around \(z\). After normalization, the exact reverse conditional approaches the zeroth-order local Gaussian \(\mathcal{N}(z,h)\). The dashed curve isolates this shape approximation; the exact conditional can still be slightly shifted because its mean need not equal \(z\).

What I Think I Learned

With the right setup, we can make the conditional reverse distribution approximately Gaussian. This is huge! We know how to sample from a Gaussian, so we have identified the local form of the reverse step. To actually reverse the process, we still need its conditional mean and variance—exactly the questions left open below.

We also reviewed the importance of telling apart a marginal from a conditional distribution. In our case, the marginal \(p(x_{t-\Delta t})\) might not be Gaussian. But with our setup, the conditional \(p(x_{t-\Delta t}\mid x_t)\) is indeed approximately Gaussian.

Moving away from the math, the result should make intuitive sense. If we add only a little noise to generate \(x_t\), then only a small neighborhood of possible \(x_{t-\Delta t}\) images could plausibly have generated it. These nearby images lie around the peak of the conditional distribution. Most other possible images have very low probability of having generated \(x_t\), so they fall in the tails of the umbrella.

Open Questions

I have been talking about the score in previous posts. A natural follow-up is to connect it to this sampling approach.
To sample from a Gaussian, we still need the mean and variance. Where could we get these?

References

Ho, Jain & Abbeel, Denoising Diffusion Probabilistic Models (2020)—the sampler this post is trying to justify: every reverse step draws from a Gaussian, and noise is injected on the way down.
Sohl-Dickstein et al., Deep Unsupervised Learning using Nonequilibrium Thermodynamics (2015), section 2.1—makes the same small-step argument, citing Feller: for a small enough step, the reverse conditional keeps the functional form of the forward one.
Anderson, Reverse-Time Diffusion Equation Models (1982)—the continuous-time result my discrete approximation is a shadow of, and where the reverse drift picks up a score term.
Song et al., Score-Based Generative Modeling through Stochastic Differential Equations (2021)—uses Anderson's reversal to put the score inside the sampler; the follow-up direction for the first open question.
Nakkiran et al., Step-by-Step Diffusion: An Elementary Tutorial—covers this same small-step reasoning, and treats the per-step drift and noise scales I skipped here.