Noise Prediction

This note covers the actual, simplified training objective diffusion models use in practice — not directly predicting the reverse process's distribution parameters, but the much simpler, more tractable task of predicting the noise that was added at a given step.

The Simplified Training Objective

\(\boldsymbol\epsilon\) is the actual random noise that was added (known exactly, since it was generated by the forward process). \(\boldsymbol\epsilon_\theta(\mathbf{x}_t,t)\) is the network's predicted noise, given the noisy image \(\mathbf{x}_t\) and the timestep \(t\). This is simply mean squared error (see Mean Squared Error) between the true and predicted noise — a remarkably simple, stable regression loss for what's ultimately a sophisticated generative model.

Why Predicting Noise (Not the Clean Image) Works So Well

Predicting the noise \(\boldsymbol\epsilon\) turns out to be mathematically equivalent to predicting the reverse process's mean \(\boldsymbol\mu_\theta\) from the previous note (via a direct algebraic substitution using the forward process's closed-form formula) — but it's an empirically easier, better-behaved quantity for the network to learn, and it produces higher-quality results in practice. This specific reformulation — from the original DDPM paper — is a large part of why modern diffusion training is so remarkably simple and stable.

The Full Training Step, Assembled

1. Sample a real image \(\mathbf{x}_0\) from the training data.
2. Sample a random timestep \(t\) uniformly from \(1\) to \(T\).
3. Sample random noise \(\boldsymbol\epsilon\), and compute \(\mathbf{x}_t\) via the forward process's closed-form formula from Diffusion Forward Process.
4. Feed \(\mathbf{x}_t\) and \(t\) into the network; compute its predicted noise \(\boldsymbol\epsilon_\theta(\mathbf{x}_t,t)\).
5. Compute MSE loss between the true \(\boldsymbol\epsilon\) and predicted \(\boldsymbol\epsilon_\theta\); backpropagate and update.

Code

import torch
import torch.nn.functional as F

def training_step(model, x0, T, alpha_bar):
    batch_size = x0.shape[0]
    t = torch.randint(0, T, (batch_size,))          # random timestep per example
    epsilon = torch.randn_like(x0)                     # random noise to add

    sqrt_alpha_bar_t = alpha_bar[t].sqrt().view(-1, 1, 1, 1)
    sqrt_one_minus_alpha_bar_t = (1 - alpha_bar[t]).sqrt().view(-1, 1, 1, 1)
    x_t = sqrt_alpha_bar_t * x0 + sqrt_one_minus_alpha_bar_t * epsilon

    predicted_epsilon = model(x_t, t)   # the network's job: predict the noise that was added
    loss = F.mse_loss(predicted_epsilon, epsilon)   # a simple, stable regression loss
    return loss

Common Mistakes

• Assuming the network predicts the clean image directly — it predicts the noise, which is then used (via the formulas in Diffusion Reverse Process) to compute an estimate of the slightly-less-noisy previous step.
• Forgetting the network must also receive the timestep \(t\) as input — the same network is used across every noise level, so it needs to know which specific noise level it's currently dealing with to predict accurately.

Interview Relevance

Q: "What does a diffusion model's neural network actually predict during training, and why is this a convenient choice?" It predicts the random noise that was added to produce the current noisy image at a given timestep — a simple mean-squared-error regression target, since the true noise is exactly known (it was generated during the forward process). This is mathematically equivalent to predicting the reverse process's distribution parameters directly, but empirically trains more easily and stably as a straightforward regression problem.

Practice Question

Why does the training loop sample a random timestep \(t\) for each training example, rather than training sequentially through all timesteps in order?

Related DL Notes

• Diffusion Reverse Process
• Mean Squared Error
• U-Net (Diffusion)