# 4. Deterministic diffusion for Gaussian noise

In diffusion and generative modelling we want to transform one distribution to another. We can achieve this by continuously interpolating between the two distributions, for example by adding varying amounts of Gaussian noise to the data distribution (such as the space of natural images), and having a method for transforming samples from these corresponding distributions.

The time-indexed family of distributions corresponding to the interpolation has a corresponding time-parameterized density function \(\rho(x, t)\), where \(\rho(x, 0)\) is the probability density of the data distribution, and \(\rho(x, t_0)\) is the probability density of the sampling distribution, such as a multivariate Gaussian. By convention, we place the data distribution at \(t=0\), and integrate “backwards in time” from the sampling distribution for some \(t_0 > 0\).

For the purpose of generative modelling, one way of transforming samples is by integrating the probability flow ODE

\[\begin{equation} \label{eq:stochastic_ode} \frac{dX_t}{dt} = v(X_t, t). \end{equation}\]The question then becomes: given the definition of a sequence of distributions (e.g., arising from adding Gaussian noise), with corresponding densities \(\rho(x, t)\), can we find \(v(x, t)\) such that integrating this probability flow ODE yields this target density \(\rho(x, t)\)? (Hint: this is where the *denoiser* that we saw in the introduction will come in.)

## Time-indexed families of distributions

Recall the **continuity equation** from the previous section:

The next corollary says that to find \(v\) such that integrating \(v\) gives some desired \(\rho\), it is sufficient to prove that \(v\) and \(\rho\) satisfy the continuity equation. We state and prove this corollary, and then in this section we look at finding \(v\) for various interpolating distributions (such as adding Gaussian noise).

*calculate*Corollary 4.1: Sufficient to prove continuity

Let \(v : \R^n \times \R \rightarrow \R^n\) and \(\rho : \R^n \times \R \rightarrow \R\) be functions that satisfy the continuity equation \eqref{eq:continuity_equation}. Let \((X_t)_{t \in I}\) for a time interval \(I \subset \R\) be the stochastic process that satisfies

\[\frac{dX_t}{dt} = v(X_t, t)\]with the probability density of \(X_{t_0}\) equal to \(\rho(\cdot, t_0)\) for \(t_0 \in I\).

Then the probability density of \(X_t\) is \(\rho(\cdot, t)\) for \(t \in I\).

*construction*Proof of sufficiency

Let \(\tau(\cdot, t)\) be the probability density of \(X_t\). By construction \(\tau(\cdot, t_0) = \rho(\cdot, t_0)\), and by Theorem 3.3 \(v\) and \(\tau\) satisfy the continuity equation. Our aim is to prove that \(\rho = \tau\), i.e., that there is a unique solution \(\rho\) to the continuity equation involving \(v\).

Write \(\pi = \rho - \tau\). Since \(\partial_t \rho = -\nabla \cdot (\rho v)\) and \(\partial_t \tau = -\nabla \cdot (\tau v)\),

\[\partial_t \pi(x, t) = 0.\]Since \(\pi(x, t_0) = \rho(x, t_0) - \tau(x, t_0) = 0\), we have that \(\pi\) is in fact zero for all \(t\), and thus \(\rho = \tau\) as required.

Thus given functions \(\rho\) and \(v\), call \(v\) a **drift** or **velocity** for \(\rho\) if it obeys the continuity equation \eqref{eq:continuity_equation}.

*edit*Exercise 4.1: Divergence-free flow

Observe that via the substitution \(w(x, t) = \rho(x, t) v(x, t)\), finding \(v\) that satisfies the continuity equation is equivalent to solving

\[\nabla \cdot w(x, t) = - \partial_t \rho(x, t).\]Verify that for \(n \ge 2\) dimensions, there exist non-zero \(w\) such that \(\nabla \cdot w = 0\), and thus multiple solutions to the continuity equation. What do these physically correspond to?

What happens in \(n=1\) dimension?

Given a sequence of distributions with probability density \(\rho(x, t)\) where \(t\) ranges over some interval \(I \subset \R\), there are thus many stochastic processes \((X_t)_{t \in I}\) with marginal density given by \(\rho\), since there are many distinct \(v\) which solve \(\eqref{eq:continuity_equation}\).

There are even more \(v\) if we only care about the initial and final distributions at time \(t=0\) and \(t=t_0\).

Often one will want to find \(v\) such that the resulting map has some optimal property, for example, having low curvature which allows for low-step numerical integration (this is the idea behind consistency diffusion). Transportation theory looks at the more general problem of finding mappings between distributions that optimize some objective, and this is a deep mathematical field that we do not cover here.

## Convolved distributions

Consider the canonical diffusion example: adding Gaussian noise to an initial distribution. Let \(\pi : \R^n \rightarrow \R\) be the density function of the data distribution, corresponding to \(t=0\). Let

\[\begin{equation} \label{eq:gaussian_noise_t} \tau(x, t) = \frac{1}{(2 \pi)^{n/2} t} \exp \left\{ - \vert \vert x \vert \vert^2 / 2 t^2 \right\} \end{equation}\]be the density of a multivariate Gaussian with variance \(t^2 1_n\). Then the convolution

\[\begin{equation} \label{eq:spatial_convolution_with_time} \rho(x, t) = (\pi * \tau(\cdot, t))(x) := \int_{\R^n} \pi(y) \tau(x - y, t) dy \end{equation}\]corresponds to the time-indexed family of distributions given by adding Gaussian noise with variance \(t^2\) to the initial distribution given by \(\pi\). (The convolution formula above of course holds for general \(\tau\), not just Gaussian noise.)

*edit*Exercise 4.2: Gaussian noise?

Verify that \(w(x, t) = x / t\) satisfies the continuity equation for \(\tau\) defined in \eqref{eq:gaussian_noise_t}.

The main result is that if we have \(w\) that satisfies the continuity equation for \(\tau\), we can derive \(v\) that satisfies the continuity equation for the convolved distribution \(\rho\).

*calculate*Theorem 4.2: Drift for convolution

Let \(w : \R^n \times \R \rightarrow \R^n\) and \(\tau : \R^n \times \R \rightarrow \R\) satisfy the continuity equation

\[\partial_t \tau(x, t) = - \nabla \cdot (\tau(x, t) w(x, t)).\]Let \(\pi : \R^n \rightarrow \R\) be an “initial distribution”. Let \(v : \R^n \times \R \rightarrow \R^n\) be the minimizer of

\[\begin{equation} \label{eq:convolved_drift_loss} \underset{\substack{ z \sim \pi(z) \\ y \sim \tau(y, t) }}{\E} \vert\vert v(y+z, t) - w(y, t) \vert\vert^2 . \end{equation}\]Then \(v\) and \(\rho\) satisfy the continuity equation, where \(\rho\) is the spatial convolution of \(\pi\) with \(\tau\) defined by \(\eqref{eq:spatial_convolution_with_time}\).

This is kinda a cool result. Let us pause and take stock of what it is saying.

The first thing to note is due to a similar result to Exercise 1.5 in the introduction, the definition of \(v\) in equation \(\eqref{eq:convolved_drift_loss}\) is equivalent to the following definition:

\[\begin{equation} \label{eq:explicit_drift_for_convolved} v(x, t) = \underset{\substack{ z \vert y + z = x }}{\E} w(x - z, t) := \frac{ \int_{\R^n} \pi(z) \tau(x-z, t) w(x-z, t) dz }{\int_{\R^n} \pi(z) \tau (x-z, t) dz} , \end{equation}\]provided we restrict \(v\) to \(\text{supp}(\rho) := \{ (x, t) : \rho(x, t) > 0 \}\).

Thus the drift of the convolved distribution at a point \((x, t)\) is equal to the average of the drift \(w(x-z, t)\) where \(z \sim \pi(z)\). This has an intuitive interpretation. Let’s treat the density \(\rho\) as corresponding to a fluid, that is generated by sampling particles \(z \sim \pi(z)\), and then for each particle following the flow given by \(\tau\). We can colour each particle by its original position \(z\). If we look at a given position \(x\) at time \(t\), then we will see a mix of particles of different colours from their original starting positions. What is the flow of the fluid at this point? It will be the weighted average of the flows of the individual particles, where the weighting is given by the amount of each colour (corresponding to the particle origin \(z\)), and the direction of the flow for this colour and at this point given by \(w(x-z, t)\).

Let’s prove Theorem 4.2.

*construction*Proof of drift for convolution formula

Our aim is to show that that \(v\) and \(\rho\) satisfy the continuity equation. Indeed

\[\begin{align*} \partial_t \rho (x, t) &= \partial_t (\pi * \tau)(x, t) \\ &= \partial_t \int \pi(z) \tau(x - z, t) dz \\ & = \int \pi(z) \partial_t \tau(x - z, t) dz \\ & = - \int \pi(z) \nabla_x \cdot (\tau(x - z, t) w(x-z, t)) dz \qquad\text{by \eqref{eq:continuity_equation}} \\ &= - \nabla_x \cdot \int \pi(z) \tau(x-z, t) w(x-z, t) dz \\ &= -\nabla_x \cdot \left[ v(x, t) \int \pi(z) \tau(x-z, t) dz \right] \qquad\text{by \eqref{eq:explicit_drift_for_convolved}} \\ &= -\nabla \cdot (v(x, t) \cdot (\pi * \tau) (x, t)) \\ &= -\nabla \cdot (v(x, t) \rho (x, t)) \end{align*}\]which is the continuity equation, as required.

To obtain a generative diffusion model, we want to convert this into an algorithm for learning the drift \(v\) by a neural network. One option is to directly parameterize \(v\) as a neural network with parameters \(\theta\), and then use equation \(\eqref{eq:convolved_drift_loss}\) as a loss function and perform stochastic gradient over samples \(z, y\) to learn \(\theta\). An important unspecified aspect is the distribution over time \(t\). Although the true theoretical minimizer of \(\eqref{eq:convolved_drift_loss}\) is independent of the distribution of \(t\), when we have only a neural function approximator, and limited compute budget, the distribution over \(t\) makes a large difference. We look at this and other practical details later on.

It is actually more common to learn a *denoiser* \(D\), which is conceptually simpler and numerically better conditioned than directly learning the drift \(v\). The denoiser predicts the “unnoised ground truth” value (i.e., at \(t=0\)) given a noisy sample (i.e., as convolved with \(\tau\)).

### Gaussian noise

In the case of adding Gaussian noise, the denoiser can be defined as follows.

*calculate*Definition 4.3: Denoiser, for additive noise

Given an initial distribution with density function \(\pi(z)\), and \(\sigma \ge 0\), the **denoiser** is the function \(D(x, \sigma)\) that minimizes

By Exercise 1.5 we can equivalently define

\[\begin{equation*} D(x, \sigma) = \underset{z \vert y+z=x}{\E} z := \frac{ \int \pi(z) n(x-z, \sigma^2) z dz }{ \int \pi(z) n(x-z, \sigma^2) dz } \end{equation*}\]where \(n(\cdot, \sigma^2)\) is the density of a Gaussian with variance \(\sigma^2\).

The following exercise shows how to convert a denoised prediction to a drift.

*edit*Exercise 4.3: Deterministic diffusion

Recall from Exercise 4.2 that \(w(y, t) = y / t\) is a drift for a Gaussian with variance \(t^2\). By substituting the expressions for \(v(y, t)\) and \(w(x, t)\) into \(\eqref{eq:convolved_drift_loss}\), show that

\[\begin{equation} \label{eq:denoiser_drift_for_adding_gaussian} v(x, t) := \frac{x - D(x, t)}{t} \end{equation}\]is a drift for the family of distributions that arises from adding Gaussian noise with variance \(t^2\) to samples from \(\pi(z)\).

The drift of equation \eqref{eq:denoiser_drift_for_adding_gaussian} is particularly suggestive. For large \(t\), \(D(x, t) \approx \E z\), the unconditional expectation of \(z \sim \pi(z)\), and so \(v(x, t) \approx (x - \E z) / t\) points to the mean of the distribution \(\pi\). As shown in the figure below, at large \(t\) the trajectories are tangent to the line that goes approximately through the origin (the mean of the signed Bernoulli random variable at time \(t=0\)). As we follow this trajectory from large values of \(t\) to smaller values, finer details about the original distribution are resolved, and the trajectories curve towards these; the initial Gaussian noise is gradually and continuously transformed into the distribution \(\pi\) at time \(t=0\). This might be one clue as to why diffusion models perform so well: they are able to incrementally add information about the original distribution, rather than resolving it all at once.

To obtain samples, we numerically integrate the drift to obtain trajectories. An important aspect affecting numerical integration error is the curvature of the trajectories.

*edit*Exercise 4.4: Curvature of diffusion trajectories

By using Exercise 3.2, show that the curvature for a trajectory \(\phi(t)\) in the setup of Exercise 4.3 is

\[\begin{equation} \partial_{tt} \phi(t) = - \frac{1}{t} \big( \partial_t + v(\phi(t), t) \cdot \nabla \big) D(\phi(t), t). \end{equation}\]Again, the material derivative (which gives the rate of change along a trajectory) makes an appearance, this time operating on the denoiser \(D\). In particular this implies that if \(D\) is constant along a trajectory, then the curvature is zero, and this allows one-step integration. This idea occurs in consistency diffusion [15]**Consistency models**

Song, Yang and Dhariwal, Prafulla and Chen, Mark and Sutskever, Ilya

arXiv preprint arXiv:2303.01469, 2023.

*edit*Exercise 4.5: Denoising at low noise levels

Show that when the data distribution \(\pi\) is a discrete distribution, i.e., the support of \(\pi\) consists of a discrete finite set of points \(\mathcal{A} \subset \R^n\), that for Gaussian noise,

\[D(x, \sigma) \rightarrow (1 + o_\sigma(1)) \, \underset{z \in \mathcal{A}}{\text{argmin}} ||z - x||^2\]where \(o_\sigma(1) = o_{\mathcal{A};\, \sigma \rightarrow 0}(1) \rightarrow 0\) as \(\sigma \rightarrow 0\) (this error term can depend on \(\mathcal{A}\), but is independent of \(x\)).

In other words, when \(\sigma\) is small, \(D(x, \sigma) \approx z(x)\) where \(z(x) \in \mathcal{A}\) is the closest point to \(x\) in the support of \(\pi\).

We can visualize the implication of Exercise 4.5.

## Final remarks

In this section we only looked at *deterministic sampling* achieved by following a deterministic ODE. If our denoiser is “perfect”, i.e., gives the true conditional expectation of the data distribution, and we follow these trajectories precisely (rather than having numerical integration errors), then we will correctly sample from the data distribution. However this is not always the case, and there are alternative ways of sampling which do not involve following a deterministic ODE, but rather involve following a *stochastic differential equation*. We will look at this in a later section.

We also only derived the flow for Gaussian noise with variance \(t^2\). More general formulations are possible, and we will also look at these in later sections.

**Table of contents**

- Home
- Motivating example
- Stochastic processes
- Probability flow
- Deterministic diffusion for Gaussian noise
- Numerical integration
- Reparameterizations of time and space
- Stochastic calculus
- Diffusion via SDEs, and score functions