Diffusion Model Training: Noise Prediction

In our previous lesson, we established the fundamental mechanics of Denoising Diffusion Probabilistic Models (DDPMs). We explored the fixed forward process that systematically corrupts an image with Gaussian noise, and the learned reverse process where a neural network, ${\mathbf{ϵ}}_{\theta }(x_t,t)$, is trained to predict and remove that noise step-by-step. You now understand what the model is trained to do.

Today, we will connect this intuitive goal of "predicting the noise" to its rigorous mathematical foundation. We will derive the training objective for a DDPM, starting from the principled objective for any generative model: maximizing the log-likelihood of the data. You will see how, through a series of elegant mathematical steps, this complex probabilistic objective is transformed into a simple, practical Mean Squared Error loss. This lesson bridges the gap between high-level theory and the code you would actually write to train such a model.

From Maximum Likelihood to a Sum of KL Divergences

The ultimate goal of training a generative model $p_{\theta }(x_0)$ is to make it produce samples that are indistinguishable from the real data distribution $p_{\text{data}}(x_0)$. The standard way to achieve this is by maximizing the log-likelihood of the training data under the model. However, for a DDPM, the likelihood of a given image $x_0$ involves integrating over all possible noise-addition paths, making it intractable:

$$p_{\theta }(x_0)=\int p_{\theta }(x_{0:T})d{x}_{1:T}$$

where $x_{0:T}=(x_0,x_1,...,x_T)$ represents the entire chain.

This is a familiar problem, similar to what we encountered with Variational Autoencoders (VAEs). The solution is also similar: we optimize a lower bound on the log-likelihood, known as the Evidence Lower Bound (ELBO). By applying Jensen's inequality, we can derive this bound, often called the Variational Lower Bound $L_{VLB}$.

The full derivation is quite involved, but it relies on rewriting the bound and leveraging the Markov property of both the forward process ($q$) and our learned reverse process ($p_{\theta }$). This allows us to decompose the total ELBO into a sum of terms, each corresponding to a single timestep in the diffusion process.

The result of this decomposition is:

$$L_{VLB}={\mathbb{E}}_q[\underset{L_T}{\underbrace{D_{KL}(q(x_T|x_0)||p(x_T))}}+\sum _{t=2}^T\underset{L_{t-1}}{\underbrace{D_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))}}-\underset{L_0}{\underbrace{\log p_{\theta }(x_0|x_1)}}]$$

Let's interpret these terms:

• $L_T$: This term measures the KL divergence between the distribution of the fully noised image $q(x_T|x_0)$ and a standard Gaussian prior $p(x_T)=\mathcal{N}(0,\mathbf{I})$. Since the forward process $q$ is designed to converge to this prior, and there are no trainable parameters ($\theta$) in this term, we can ignore it during optimization.
• $L_0$: This is a reconstruction term, representing the log-likelihood of the final denoising step that produces the clean image $x_0$ from $x_1$.
• $L_{t-1}$: This is the heart of the training objective. For each timestep from $T$ down to 2, it's a KL divergence that measures the "distance" between our learned reverse step $p_{\theta }(x_{t-1}|x_t)$ and the true, but conditional, reverse posterior $q(x_{t-1}|x_t,x_0)$. Our goal is to train the network $\theta$ to minimize this divergence at every step.

The following video provides an excellent visual walkthrough of this derivation, from the initial log-likelihood objective to this decomposed sum of KL divergences.

Simplifying the KL Divergence to a Difference of Means

Our objective is now to minimize $L_{t-1}=D_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))$.

As we established in the last lesson, both of these distributions are Gaussian:

1. The true posterior is $q(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1};{\tilde{\mathbf{\mu }}}_t(x_t,x_0),{\tilde{\beta }}_t\mathbf{I})$.
2. Our learned approximation is $p_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mathbf{\mu }}_{\theta }(x_t,t),{\mathbf{\Sigma }}_{\theta }(x_t,t))$.

The DDPM paper makes a key simplifying choice: the variance of the learned reverse step is fixed and not learned. It's set to a constant, ${\sigma }_t^2\mathbf{I}$, typically ${\tilde{\beta }}_t$ or ${\beta }_t$. This means our neural network only needs to learn the mean ${\mathbf{\mu }}_{\theta }(x_t,t)$.

Now for a crucial question from your probability studies.

Socratic Question: The KL divergence between two multivariate Gaussian distributions $P=\mathcal{N}({\mathbf{\mu }}_P,{\mathbf{\Sigma }}_P)$ and $Q=\mathcal{N}({\mathbf{\mu }}_Q,{\mathbf{\Sigma }}_Q)$ has a general, rather complex form. What happens to this formula if we set their covariance matrices to be equal, ${\mathbf{\Sigma }}_P={\mathbf{\Sigma }}_Q={\sigma }^2\mathbf{I}$? How does the expression for $D_{KL}(P||Q)$ simplify?

...

If the covariance matrices are identical, the terms involving the determinant and trace of the matrices cancel out, and the KL divergence reduces to a simple (scaled) L2 norm, or squared Euclidean distance, between their means:

$$D_{KL}(\mathcal{N}({\mathbf{\mu }}_P,{\sigma }^2\mathbf{I})||\mathcal{N}({\mathbf{\mu }}_Q,{\sigma }^2\mathbf{I}))=\frac{1}{2{\sigma }^2}||{\mathbf{\mu }}_P-{\mathbf{\mu }}_Q|{|}^2$$

Applying this to our loss term $L_{t-1}$, we see that minimizing the KL divergence is equivalent to minimizing the squared distance between the predicted mean ${\mathbf{\mu }}_{\theta }$ and the true posterior mean ${\tilde{\mathbf{\mu }}}_t$.

So, our objective simplifies to training the network to predict the correct mean for the denoising step:

$$L_{t-1}\propto {\mathbb{E}}_{x_0,\mathbf{ϵ}}[||{\tilde{\mathbf{\mu }}}_t(x_t,x_0)-{\mathbf{\mu }}_{\theta }(x_t,t)|{|}^2]$$

The Final Step: The Noise Prediction Objective

We've made great progress, boiling the problem down to matching the means of two Gaussian distributions. But as we discussed last lesson, directly predicting the mean ${\mathbf{\mu }}_{\theta }$ is not the most effective strategy. The final insight is to re-parameterize this objective in terms of noise.

Recall the forward process "shortcut" equation: $x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ}$.
We can invert this to express $x_0$ in terms of $x_t$ and the noise $\mathbf{ϵ}$:

$$x_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ})$$

The full expression for the true posterior mean ${\tilde{\mathbf{\mu }}}_t(x_t,x_0)$ is a weighted sum of $x_t$ and $x_0$. If we substitute the equation for $x_0$ above into the formula for ${\tilde{\mathbf{\mu }}}_t$, a remarkable simplification occurs. After some algebra, ${\tilde{\mathbf{\mu }}}_t$ can be rewritten as a function of $x_t$ and the true noise $\mathbf{ϵ}$.

The key idea is to parameterize our network's predicted mean, ${\mathbf{\mu }}_{\theta }(x_t,t)$, using the exact same functional form, but with the network's predicted noise, ${\mathbf{ϵ}}_{\theta }(x_t,t)$:

$${\mathbf{\mu }}_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }(x_t,t)\right)$$

Now, when we compute the difference $||{\tilde{\mathbf{\mu }}}_t-{\mathbf{\mu }}_{\theta }|{|}^2$, the complex terms involving $x_t$ cancel each other out perfectly. The entire expression simplifies to a scaled squared difference between the true noise and the predicted noise.

The authors of the DDPM paper found that ignoring the complex scaling factors in front of this term and simply minimizing the direct Mean Squared Error (MSE) between the noises worked best in practice. This leads to the final, simplified training objective:

$$L_{\text{simple}}(\theta )={\mathbb{E}}_{t\sim [1,T],x_0\sim p_{\text{data}},\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I})}\left[||\mathbf{ϵ}-{\mathbf{ϵ}}_{\theta }(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ},t)|{|}^2\right]$$

This beautifully simple loss function underpins the training of modern diffusion models.

The following resources walk through this final crucial step, showing how the reparameterization leads to the noise prediction objective.

This simplified objective translates into a very direct training algorithm, illustrated by the diagram below.

Conclusion

In this lesson, we have completed the theoretical journey of DDPMs, connecting the high-level goal of generative modeling to a simple and effective training objective.

Key Takeaways:

• From Likelihood to ELBO: Training starts with the goal of maximizing data log-likelihood, which is made tractable by optimizing its Evidence Lower Bound (ELBO).
• ELBO as a Sum of KL Terms: The ELBO for a DDPM can be decomposed into a sum of KL divergences, where each term $L_{t-1}$ measures how well the learned reverse step matches the true posterior at a given timestep.
• KL Simplifies to MSE of Means: By fixing the variance of the reverse process, the KL divergence objective simplifies to minimizing the Mean Squared Error between the predicted mean (${\mathbf{\mu }}_{\theta }$) and the true posterior mean (${\tilde{\mathbf{\mu }}}_t$).
• Reparameterization Yields Noise Prediction: The final and most crucial step is to re-parameterize both means in terms of noise ($\mathbf{ϵ}$ and ${\mathbf{ϵ}}_{\theta }$). This elegantly reduces the objective to a simple MSE loss between the true and predicted noise, $||\mathbf{ϵ}-{\mathbf{ϵ}}_{\theta }|{|}^2$.

This simplification is not just a mathematical convenience; it provides a much more stable and effective learning target for the neural network, which is a major reason for the success of DDPMs.

Preview of the Next Module:

Having covered the foundational principles of modern generative models—VAEs, GANs, and now DDPMs—we are fully equipped to venture into the domain of neural rendering. In our next module, we will explore Neural Radiance Fields (NeRF), a groundbreaking technique that uses a simple MLP not just to generate 2D images, but to represent and render complex 3D scenes from any viewpoint.