DDPM Training & the Noise Prediction Objective

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. The Variational Lower Bound
5. Simplification to L-simple
6. The Noise Prediction Objective
7. Training Algorithm
8. Loss Weighting Strategies
9. Variance Prediction (Improved DDPM)
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Derive the variational lower bound (VLB) for diffusion models.
2. Show how the VLB decomposes into per-timestep KL divergence terms.
3. Prove the simplified noise prediction objective is equivalent.
4. Explain different loss weighting strategies and their effects.
5. Derive the learned variance parameterization (Improved DDPM).

Notation

• $\beta_t$ — noise schedule at step $t$
• $\bar{\alpha}_t = \prod_{s=1}^t (1-\beta_s)$ — cumulative signal retention
• $\boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)$ — noise prediction network
• $\mathcal{L}_{\text{simple}}$ — simplified DDPM loss

Core Intuition

Training a diffusion model means learning to reverse the noise addition process. The model sees a noisy image $\mathbf{x}_t$ at a random timestep $t$ and must predict what noise was added. This simple objective — predicting $\boldsymbol{\epsilon}$ — emerges naturally from the evidence lower bound (ELBO) on the log-likelihood.

The Variational Lower Bound

Start from the marginal log-likelihood:

Decompose using the Markov structure:

Key insight: Both $q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0)$ and $p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)$ are Gaussian, so each $L_t$ is a KL between Gaussians (closed form).

Simplification to L-simple

The posterior $q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0)$ is Gaussian with:

Parameterize the model mean as:

The KL divergence $L_t$ reduces to:

The Noise Prediction Objective

DDPM simplified loss (Ho et al., 2020): Drop the weighting prefactor:

Steps:

1. Sample $\mathbf{x}_0 \sim p_{\text{data}}$, $t \sim \text{Uniform}(1, T)$, $\boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I})$.
2. Construct noisy input: $\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon}$.
3. Predict noise: $\hat{\boldsymbol{\epsilon}} = \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t)$.
4. Loss: $\|\boldsymbol{\epsilon} - \hat{\boldsymbol{\epsilon}}\|^2$.

Training Algorithm

Algorithm (DDPM Training):

Repeat until convergence:

1. $\mathbf{x}_0 \sim p_{\text{data}}(\mathbf{x})$
2. $t \sim \text{Uniform}\{1, \ldots, T\}$
3. $\boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$
4. Compute gradient of $\|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon}, t)\|^2$
5. Update $\theta$ with optimizer step

Key properties:

• No adversarial training (unlike GANs).
• Stable optimization (simple MSE loss).
• Each training step only evaluates one random timestep $t$.

Loss Weighting Strategies

Uniform weighting ($\mathcal{L}_{\text{simple}}$): All timesteps weighted equally. Emphasizes large $t$ (high noise) where MSE is large.

SNR weighting: Weight by signal-to-noise ratio:

Min-SNR weighting (Hang et al., 2023): Clip the SNR weight:

Prevents early timesteps (low noise, high SNR) from dominating the loss.

P2 weighting: Emphasize perceptually important timesteps (mid-range noise).

Variance Prediction (Improved DDPM)

Standard DDPM: Fixed variance $\sigma_t^2 = \beta_t$ or $\sigma_t^2 = \tilde{\beta}_t = \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\beta_t$.

Improved DDPM (Nichol & Dhariwal, 2021): Learn the variance as interpolation in log-space:

where $v = \text{sigmoid}(\mathbf{v}_\theta(\mathbf{x}_t, t))$ is predicted by the network.

Training: Use a hybrid loss: $\mathcal{L}_{\text{simple}}$ for the mean + $\mathcal{L}_{\text{VLB}}$ for the variance (with stop-gradient on the mean prediction).

Result: Better log-likelihoods and improved sample quality, especially with fewer sampling steps.

Common Pitfalls

Pitfall 1. Using the VLB loss directly for training. The prefactors in $L_t$ create highly non-uniform gradient magnitudes across timesteps. $\mathcal{L}_{\text{simple}}$ (uniform weighting) trains better empirically.

Pitfall 2. Forgetting that $t$ is sampled uniformly. Each training step only sees one timestep — the network must generalize across all $T$ timesteps from random sampling.

Pitfall 3. Confusing $\mathbf{x}_0$-prediction with $\boldsymbol{\epsilon}$-prediction. They're equivalent but have different numerical properties: $\boldsymbol{\epsilon}$-prediction is better at high noise, $\mathbf{x}_0$-prediction at low noise.

Summary

• DDPM training maximizes a variational lower bound on $\log p(\mathbf{x}_0)$.
• The VLB decomposes into per-timestep KL terms between Gaussians.
• Simplification: predict the added noise $\boldsymbol{\epsilon}$ with an MSE loss.
• Uniform loss weighting works well; min-SNR weighting improves convergence.
• Learned variance (Improved DDPM) improves log-likelihood and few-step sampling.

Exercises

Exercise 1. Derive $\tilde{\boldsymbol{\mu}}_t$ (equation 3) from $q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0)$ using Bayes' rule on two Gaussians.

Exercise 2. Show that $\mathcal{L}_{\text{simple}}$ (equation 6) differs from $\mathcal{L}_{\text{VLB}}$ only by a time-dependent weighting factor.

Exercise 3. Prove that $\boldsymbol{\epsilon}$-prediction and $\mathbf{x}_0$-prediction are equivalent parameterizations.

Exercise 4. Compute the SNR at $t=0, T/4, T/2, 3T/4, T$ for a linear noise schedule with $\beta_1=10^{-4}, \beta_T=0.02$.

Exercise 5. Explain why learning the variance requires a hybrid loss with stop-gradient.