Forward Process, Transformations, and Gregor Samsa

The Question

I can't help relating diffusion models to a Kafkaesque situation in which Gregor Samsa—the image—awakes from uneasy dreams—the forward process—to find himself transformed into an insect: Gaussian noise. But unlike in The Metamorphosis, diffusion models give us hope of reversing the transformation by studying the nature of the corruption process.

In the “Changing Variables Without Changing Probability” post, we saw how to compute the probability density of a transformation. Namely, given

X=g(Z), \qquad Z=g^{-1}(X),

then

f_X(x) = f_Z(z) \left| \det\left(\frac{\partial z}{\partial x}\right) \right|.

Today, I want to address the open question I left there: how is this transformation related to linear transformations?

Short answer: linear and affine transformations are particular cases of these random-variable transformations. In this post, we'll work with the transformation commonly implemented in a forward diffusion process.

Why I Care

In a standard diffusion forward process, we take a sample from our original data, \(x_0\), and convert it into a noisy sample at time \(t\) using this formula:

x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\varepsilon, \qquad \varepsilon\sim\mathcal{N}(0,I).

We select \(\bar{\alpha}_t\) in a way that, at later timesteps, we have less of our original sample \(x_0\) and more of the generated white noise \(\varepsilon\). As \(\bar{\alpha}_T\to 0\), the conditional distribution of \(x_T\) given \(x_0\) approaches \(\mathcal{N}(0,I)\).

A beneficial property of this formula is that, conditioned on \(x_0\), the resulting \(x_t\) is also normally distributed:

x_t\mid x_0 \sim \mathcal{N}\!\left( \sqrt{\bar{\alpha}_t}\,x_0, (1-\bar{\alpha}_t)I \right).

The conditioning matters. If \(x_0\) is still random and comes from the entire data distribution, then \(x_t\) is generally not Gaussian at a finite timestep.

So it has

\text{mean:}\quad \mu_t=\sqrt{\bar{\alpha}_t}\,x_0,

\text{standard deviation:}\quad \sigma_t=\sqrt{1-\bar{\alpha}_t},

\text{covariance:}\quad \Sigma_t=\sigma_t^2 I=(1-\bar{\alpha}_t)I.

The identity matrix appears because \(\varepsilon\sim\mathcal{N}(0,I)\) has independent, unit-variance coordinates. Multiplying it by the scalar variance gives the simpler isotropic covariance usually chosen in diffusion models. In “When the Score Does Not Point Straight to the Mean.” the covariance was not isotropic, which is why the score could tilt away from the mean.

In general, instead of sampling directly from \(\mathcal{N}(\mu,\sigma^2)\), we can sample from a standard normal \(\mathcal{N}(0,1)\) and use the reparameterization trick:

X=\mu+\sigma\varepsilon, \qquad \varepsilon\sim\mathcal{N}(0,1), \qquad X\sim\mathcal{N}(\mu,\sigma^2).

Let's dive deeper into these ideas.

Setup and Definitions

The following table presents the probability densities of Gaussian distributions in the univariate and multivariate cases.

	Univariate	Multivariate
Standard	\(f_Z(z)=\dfrac{1}{\sqrt{2\pi}} e^{-z^2/2}\)	\(f_Z(z)=\dfrac{1}{(2\pi)^{d/2}} e^{-z^Tz/2}\)
General	\(f_X(x)=\dfrac{1}{\sqrt{2\pi}\,\sigma} e^{-(x-\mu)^2/(2\sigma^2)}\)	\(f_X(x)=\dfrac{1}{(2\pi)^{d/2}\|\Sigma\|^{1/2}} \exp\!\left[-\dfrac{1}{2}(x-\mu)^T \Sigma^{-1}(x-\mu)\right]\)

Here, \(\Sigma\) is the covariance matrix. Covariance matrices are always symmetric positive semidefinite. The presented formula of the multivariate normal also requires \(\Sigma^{-1}\) to exist, which rules out zero-variance directions and makes \(\Sigma\) positive definite. Therefore, we can diagonalize it via its eigendecomposition:

\Sigma=Q\Lambda Q^T,

where \(Q\) has orthonormal columns—the eigenvectors—and \(\Lambda\) is a diagonal matrix of the eigenvalues, which are positive.

I might write a post explaining this better, but take it as a fact if you are unfamiliar with this concept. Luckily, this factorization has a very nice geometric interpretation. For our covariance-matrix decomposition with orthonormal columns in \(Q\), we read the transformation from right to left:

z \xrightarrow{\;Q^T\;} Q^Tz \xrightarrow{\;\Lambda\;} \Lambda Q^Tz \xrightarrow{\;Q\;} Q\Lambda Q^Tz.

First \(Q^T\) changes to the principal-axis basis, \(\Lambda\) scales in that basis, and \(Q\) rotates back to the standard basis. This describes how \(\Sigma\) itself acts on an arbitrary vector. To generate Gaussian samples, we will instead use a covariance factor \(B\) satisfying \(BB^T=\Sigma\).

Minimal Derivation

We'll explore the transformation of the random variable \(Z\sim\mathcal{N}(0,I)\), for invertible \(B\), given by

X=\mu+BZ, \qquad Z\sim\mathcal{N}(0,I).

Notice the connection here with the forward diffusion process presented in the “Why I Care” section:

x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\varepsilon, \qquad \varepsilon\sim\mathcal{N}(0,I).

Conditioned on \(x_0\), this is exactly the affine map \(X=\mu+BZ\), with

Z=\varepsilon, \qquad \mu=\sqrt{\bar{\alpha}_t}\,x_0, \qquad B=\sqrt{1-\bar{\alpha}_t}\,I, \qquad |\det B|=(1-\bar{\alpha}_t)^{d/2}.

We first check the mean:

\mathbb{E}[X] =\mathbb{E}[\mu+BZ] =\mu+B\mathbb{E}[Z] =\mu.

We also find the covariance of \(X\):

\begin{aligned} \operatorname{Cov}(X) &=\operatorname{Cov}(\mu+BZ) \\ &=\operatorname{Cov}(BZ) \\ &=B\operatorname{Cov}(Z)B^T \\ &=BIB^T \\ &=BB^T, \end{aligned}

since \(\mu\) is a constant shift and \(Z\sim\mathcal{N}(0,I)\).

Recall the change-of-variables formula from “The Question”:

X=g(Z), \qquad Z=g^{-1}(X), \qquad f_X(x)=f_Z(g^{-1}(x)) \left| \det\left(\frac{\partial g^{-1}(x)}{\partial x}\right) \right|.

Replacing the values of our transformation,

X=\mu+BZ, \qquad Z=B^{-1}(X-\mu), \qquad \frac{\partial z}{\partial x}=B^{-1}.

So

f_X(x) = f_Z\!\left(B^{-1}(x-\mu)\right) \frac{1}{|\det B|}.

Here we used \(|\det B^{-1}|=1/|\det B|\).

We substitute the formula for the multivariate standard normal from the table:

f_X(x) = \frac{1}{(2\pi)^{d/2}} \exp\!\left( -\frac{1}{2} [B^{-1}(x-\mu)]^T[B^{-1}(x-\mu)] \right) \frac{1}{|\det B|}.

And

\begin{aligned} [B^{-1}(x-\mu)]^T[B^{-1}(x-\mu)] &=(x-\mu)^T(B^{-1})^TB^{-1}(x-\mu) \\ &=(x-\mu)^T(BB^T)^{-1}(x-\mu). \end{aligned}

The term \(BB^T\) matches what we found for the covariance of \(X\), so \(BB^T=\Sigma\). Also,

\det(BB^T)=(\det B)^2=\det\Sigma, \qquad |\det B|=|\Sigma|^{1/2}.

Then

f_X(x) = \frac{1}{(2\pi)^{d/2}|\Sigma|^{1/2}} \exp\!\left[ -\frac{1}{2}(x-\mu)^T\Sigma^{-1}(x-\mu) \right].

Yes! This is our exact formula for the general multivariate Gaussian from the table. We just showed how to generate a general multivariate Gaussian from a standard normal. This is why Gaussian distributions are closed under affine transformations: scaling, rotating, and shifting a Gaussian random variable keeps us in the Gaussian family. If there is no shift, the transformation is linear. We can also confirm that the forward expression

x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\varepsilon

means that \(x_t\mid x_0\) is normally distributed with mean \(\sqrt{\bar{\alpha}_t}\,x_0\) and isotropic covariance \((1-\bar{\alpha}_t)I\).

Going back to our diagonalization, we want \(\Sigma=Q\Lambda Q^T\) and \(\Sigma=BB^T\). These are two related but different maps. The diagonal entries of \(\Lambda\) are variances, while those of \(\Lambda^{1/2}\) are standard deviations. Applying \(\Sigma\) to an arbitrary vector uses \(Q^T\), then \(\Lambda\), then \(Q\). To generate samples, we only need a factor \(B\) such that \(BB^T=\Sigma\). Because standard Gaussian noise has no preferred orientation, \(Q^TZ\overset{d}{=}Z\), so the first rotation does not change its distribution. This is why we can begin by scaling. One possible factor is

B=Q\Lambda^{1/2},

such that

\begin{aligned} BB^T &=Q\Lambda^{1/2}(Q\Lambda^{1/2})^T \\ &=Q\Lambda^{1/2}\Lambda^{1/2}Q^T \\ &=Q\Lambda Q^T. \end{aligned}

So our transformation becomes

X=\mu+Q\Lambda^{1/2}Z.

Geometrically, we generate noise centered at zero, stretch it by \(\Lambda^{1/2}\), rotate it by \(Q\), and shift everything by \(\mu\).

Translation does not change volume, and the orthogonal matrix satisfies \(|\det Q|=1\), so only the scaling changes it:

|\det B| =|\det Q|\det(\Lambda^{1/2}) =\prod_i\sqrt{\lambda_i}.

From a linear-algebra or geometric point of view, \(B\) tells us what happens to points and space, as shown in the beautiful videos of the 3Blue1Brown linear algebra series. From probability, we know that the same event must have the same probability after the transformation, so we must compensate for the change in space produced by \(B\). In effect,

\text{volume or area scales by }|\det B| \quad\Longleftrightarrow\quad \text{density scales by }\frac{1}{|\det B|}.

Visualizing the Transformation

First, transform one fixed sample of standard Gaussian noise by scaling, rotating, and shifting it. Then use the same linear map on a grid cell to see why area and density must change in opposite directions.

What I Think I Learned

The key idea of this post is the connection between random-variable transformations and their geometric interpretation using linear algebra. We saw the specific affine transformation implemented in the forward process and why we use the Jacobian in the transformed density: we need to account for the change in area or volume introduced by the transformation so the probability mass remains the same.

In that quest, we talked about the reparameterization trick, how to generate Gaussian noise from standard Gaussian noise, and mentioned some properties of the covariance matrix. I know it was a lot, but I hope that writing and reading these ideas is helpful for solidifying the building blocks of state-of-the-art generative models.

Open Questions

If \(B\) is not unique—for example, a Cholesky factor can produce the same \(\Sigma\)—does the choice matter?
Conditioned on \(x_0\), \(x_{t-1}\) and \(x_t\) are jointly Gaussian. How does conditioning give us the reverse step's mean and covariance?
Can we get Gregor back?

References

Changing Variables Without Changing Probability—the earlier derivation of the change-of-variables formula used in this post.
Denoising Is Not Inversion—the earlier distinction between reversing corruption and recovering the exact original sample.
3Blue1Brown, Linear transformations and matrices—a visual introduction to reading matrices as transformations of space.
3Blue1Brown, The determinant—the area-scaling interpretation of \(|\det B|\).
3Blue1Brown, Eigenvectors and eigenvalues—a visual resource for the eigendecomposition used here.
Ho, Jain & Abbeel, Denoising Diffusion Probabilistic Models (2020)—the forward-process marginal used in this post.