# Diffusion Models - modeling noise?

In this post about diffusion models, IIUC, we want to use a neural network to approximate the mean of the reverse diffusion: $$p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t))$$. That is, we train a neural network to predict the quantity $$\tilde{\boldsymbol{\mu}}_t = \frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_t \Big)$$.

Instead, according to the article, since $$\mathbf{x}_t$$ is available at training time, we predict $$\boldsymbol{\epsilon}_t$$ instead.

But according to the forward diffusion process, we have:

\begin{aligned} \mathbf{x}_t &= \sqrt{\alpha_t}\mathbf{x}_{t-1} + \sqrt{1 - \alpha_t}\boldsymbol{\epsilon}_{t-1} & \text{ ;where } \boldsymbol{\epsilon}_{t-1}, \boldsymbol{\epsilon}_{t-2}, \dots \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \\ &= \sqrt{\alpha_t \alpha_{t-1}} \mathbf{x}_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}} \bar{\boldsymbol{\epsilon}}_{t-2} & \text{ ;where } \bar{\boldsymbol{\epsilon}}_{t-2} \text{ merges two Gaussians (*).} \\ &= \dots \\ &= \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon} \\ q(\mathbf{x}_t \vert \mathbf{x}_0) &= \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t)\mathbf{I}) \end{aligned}

So $$\boldsymbol{\epsilon}_t$$ is a random sample from the standard normal. I don't fully understand this; why are we trying to train a neural network to predict on the noise? If this is a random draw from a normal that is independent of any covariates, isn't this also available at training time, and so we can just directly compute a sample?

TLDR; the noise $$\boldsymbol{\epsilon}_\theta$$ is the only thing we need learn so we can model reverse diffusion process. We can use $$\boldsymbol{\epsilon}_\theta$$ to get mean and variance of the Gaussian distributions we are looking for!

I also had the exact same question when I encountered the DDPM paper! I think we need to clarify a few things to understand the purpose of predicting the noise!

1- Why are we learning the noise?

• So in diffusion models (including DDPM), given the forward diffusion (i.e., $$q(x_t | x_{t-1})$$), our goal is to learn how to reverse that process (i.e., $$p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)$$) with a neural network.
• It turns out after doing lots of math, $$p_\theta$$ is Guassian in the form of $$p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t))$$. Hence, all we need to do is to find $$\mathbf{\mu_\theta}$$ and $$\mathbf{\Sigma_\theta}$$
• Now, we can do the math and see: \begin{aligned} \boldsymbol{\mu}_\theta(\mathbf{x}_t, t) &= \color{black}{\frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t) \Big)} \end{aligned} where $$\mathbf{\alpha}_t$$ depends on noise values and non-trainable. So the only part we actually need to learn, is $$\boldsymbol{\epsilon}_\theta$$.
• Once we learn $$\boldsymbol{\epsilon}_\theta$$, we can put it in the equation and get $$\mathbf{p}_\theta$$

2- If the noise is independent of data, why/how can we learn it?

• Well, I think given the algorithm 1 below, all our NN has to do is to predict $$\mathcal{N}(0, I)$$. But that's not possible. We can't predict random noise.

• However, there's a catch! I think the confusion comes from the fact that we don't want to learn a random noise. We want to learn the noise component, given a noisy input! So the input to our model is image + noise, and we want to learn the noise part.

So let's see a pseudo-code of DDPM for clarification:

def train_loss(denoise_model, x_0, t):
noise = torch.randn_like(x_0)

x_noisy = q_sample(x_0=x_0, t=t, noise=noise)
predicted_noise = denoise_model(x_noisy, t)

loss = F.l2_loss(noise, predicted_noise)

return loss



We can see that the model is trained to denoise a noisy input in the code.

• But why is noise added at each timestamp during the reverse process? Commented Jan 24, 2023 at 9:14
• So during the forward process, we add noise step by step. In the reverse process, we are trying to do the exact opposite. So at each step, we try to denoise. The reason that we add noise at each step is that it will be extremely difficult to directly find a mapping between noise and real data. GANs try to do the same thing and their training is very unstable/difficult. Commented Jan 24, 2023 at 23:32
• What will happen if we don't add this noise at every backward step? Commented Jan 26, 2023 at 11:33
• I don't think I fully understand what you mean. If we're trying to get rid of the random noise component in the reverse process, why do we add a random noise component in every step of the backward process? Commented Apr 23, 2023 at 10:33
• @NathanB I also didn't understand why the noise was added at each step. I wrote an answer in stats.stackexchange.com/a/624237/394779 Commented Aug 17, 2023 at 14:03

We begin by highlighting key facts and details about diffusion models:

• The forward diffusion process is a step-wise deformation of the data distribution caused by the injection of additive noise at each time step.
• The forward process assumes the latent dimension is exactly equal to the data dimension $$x_{t+1} \equiv x_{t}$$
• The structure of the latent space of each time step is fixed as a linear Gaussian model (not learned) centered on the output of the previous time-step.
• The transition probability kernel for a single time-step is $$q(x_{t} \vert x_{t-1}) \: = \mathcal{N} \big( \: \sqrt{\smash[b]{\alpha_t}} \cdot x_{t-1} , \: (1 \: - \: \alpha_t) \cdot \mathcal{I} \: \big)$$ where moments depend on the parameter $$\alpha_{t}$$ with mean $$\mu=\sqrt{\smash[b]{\alpha_t}}$$ centered on the previous timepoint and covariance $$\Sigma=(1 \: - \: \alpha_t)$$
• To sample from the diffusion kernel distribution $$x_{t} \sim q(x_{t} \vert x_{t-1})$$, we could use the re-parameterization trick, mathematically expressed as $$x_{t} \: = \: \mu_{t} \cdot x_{t-1} \: + \: \sigma_{t} \cdot \epsilon_{t}$$

These modeling constraints on the diffusion process labor the point that all that the encoding process $$q(x_{t} \vert x_{t-1})$$ requires is learning the $$\mu_{t}$$ and $$\sigma_{t}^2$$ of each latent Gaussian distribution.

Because each of these independent and sequential time-steps are analytically tractable (Gaussian), in consequence, so is the full chain as well. It can then be shown that the KL-divergence requiring optimization of the posterior is

$$\underbrace{ \mathbb{E}_{ q(x_{1} \vert x_{0}) } \: \Big\lbrack \: \log p_{\theta}(x_{0} \vert x_{1}) \: \Big\rbrack }_{\text{reconstruction term}} \: + \: \underbrace{ D_{KL} \: \big( \: q(x_{T} \vert x_{0}) \: \Vert \: p(x_{T}) \: \big) }_{\text{prior matching term}} \: - \: \sum_{t=2}^T \underbrace{ \mathbb{E}_{ q(x_{t} \vert x_{0}) } \: \Big\lbrack \: D_{KL} \: \big( q(x_{t-1} \vert x_{t}, x_{0}) \: \Vert \: p_{\theta}(x_{t-1} \vert x_{t}) \big) \: \Big\rbrack }_{\text{denoising matching term}}$$

With some math, the main computational challenge comes from the third KL-divergence term which is shown to be equal to

$$D_{KL} \big( \: q( x_{t-1} \vert x_{t}, x_{0}) \: \big\Vert \: p_{\theta}(x_{t-1} \vert x_{t} ) \: \big) \: = \frac {1} {2 \sigma_{q}^2 (t)} \Big\lbrack \: \big\Vert \: \mu_{\theta} \: - \: \mu_{q} \: \big\Vert_{2}^2 \: \Big\rbrack$$

where $$\mu_{q}$$ and $$\mu_{\theta}$$ are shorthand for $$\mu_{q}(x_{t}, x_{0})$$ and $$\mu_{\theta}(x_{t}, t)$$ respectively. Examination of this result explains why diffusion models train a parameterized neural network model to learn to predict the original training image $$x_{0}$$ from noisy image $$x_{t}$$ and sampled time index $$t$$. By extension, a simple application of the re-parameterization trick permits modeling the noise instead $$\epsilon_{\theta}(x_{t}, t)$$.

In summary, with a few modeling constraints, the forward diffusion (encoding) process serves as a way to calculate the parameters of the true posterior for the reverse (decoding) process for each time-step $$p_{\theta}(x_{t-1} \vert x_{t})$$.

The following excerpts are taken from my book on variational inference. Learn more on the topic by visiting https://www.thevariationalbook.com/