DDPM: Noise Schedules, Parameterizations & the Variational View

"I spent a week deriving a variational bound with a dozen KL terms, and at the end it told me to just predict the noise and square the error. I am told this is what physicists call beauty."

A Loss Function That Simplified Beyond Recognition

In Section 33.1 you built the forward and reverse processes and trained a denoiser with a noise-prediction loss that we asked you to take on faith. This section pays that debt. Here is the payoff worth reading for: we are about to write down a likelihood bound with a dozen KL terms that looks like a graduate exam question, and then watch almost all of it evaporate until what remains is the one-line squared error you already trained against. Stay with the derivation, because the moment it collapses is the moment you understand why diffusion training is both principled and trivial to code. We derive where the loss comes from, the variational lower bound that DDPM shares with the VAE of Chapter 31, and we show how each term of that bound reduces to a tractable quantity. Then we turn to the engineering knobs that separate a model that trains from one that trains well: the schedule, the parameterization, and the loss weighting. By the end you will be able to read the configuration of any modern diffusion model and know what every field means.

1. The Variational Bound, and How It Collapses Advanced

We want to maximize the likelihood the model assigns to real data, $\log p_\theta(x_0)$, but the marginal requires integrating over all the latent noisy states $x_1, \dots, x_T$, which is intractable. The same problem appeared with the VAE, and the same tool solves it: the evidence lower bound (ELBO). Treating the entire forward trajectory as the "encoder" and the reverse process as the "decoder," the bound is

This looks forbidding, but every piece is benign. The last term compares the fully-corrupted $x_T$ to a standard Gaussian; since the forward process is designed to reach exactly that, the term is essentially zero and has no parameters. The first term is a reconstruction at the final denoising step. The interesting work is in the middle sum: each term compares the model's reverse step $p_\theta(x_{t-1} \mid x_t)$ against the true posterior $q(x_{t-1} \mid x_t, x_0)$, which, crucially, is a Gaussian we can write in closed form because we condition on the clean $x_0$. Its mean is

A KL divergence between two Gaussians with the same variance reduces to the squared distance between their means. So each middle term becomes $\| \tilde\mu_t - \mu_\theta \|^2$ up to a constant. Substituting the noise-prediction expression for $\mu_\theta$ from Section 33.1 and simplifying (the algebra is mechanical but lengthy) turns that mean-distance into a distance between the true noise $\epsilon$ and the predicted noise $\epsilon_\theta$, with a per-timestep weight in front. Ho et al. found that simply dropping that weight, setting it to one for every $t$, trained better and gave the now-canonical loss:

2. Noise Schedules: Linear vs Cosine Intermediate

The schedule $\{\beta_t\}$, the variance schedule introduced in Section 33.1 and also called the noise schedule, decides how quickly the forward process destroys the image, and through $\bar\alpha_t$ it sets the signal-to-noise ratio at every step. The original DDPM used a linear schedule, $\beta_t$ rising linearly from $10^{-4}$ to $0.02$. This works but has a flaw at high resolution: it destroys information too quickly near the end, so the last fraction of timesteps carries almost no signal and is wasted. Nichol and Dhariwal proposed a cosine schedule that keeps more signal in the middle of the process by defining $\bar\alpha_t$ directly:

where $s$ is a small offset (about $0.008$) that prevents $\beta_t$ from being too small near $t = 0$. The cosine schedule spends more steps at intermediate noise levels, which is where the model does most of its useful learning, and it noticeably improves sample quality on datasets like ImageNet. Figure 33.2.1 contrasts how $\bar\alpha_t$ decays under the two schedules.

The two schedules are a few lines apart in code. The function below returns the $\bar\alpha_t$ arrays for both, and you can drop either into the q_sample of Section 33.1.

# Two noise schedules expressed through their alpha_bar arrays.
# Linear builds beta_t first then takes the cumulative product; cosine
# defines alpha_bar directly so signal decays gradually at high resolution.
import torch
import math

def linear_alpha_bars(T=1000, b0=1e-4, b1=0.02):
    betas = torch.linspace(b0, b1, T)
    return torch.cumprod(1.0 - betas, dim=0)

def cosine_alpha_bars(T=1000, s=0.008):
    steps = torch.arange(T + 1) / T
    f = torch.cos((steps + s) / (1 + s) * math.pi / 2) ** 2
    ab = f / f[0]
    return ab[1:]                                   # length T, alpha_bar at each step

lin = linear_alpha_bars()
cos = cosine_alpha_bars()
print("linear alpha_bar at t=900:", lin[900].item())   # ~0.0001, nearly all noise
print("cosine alpha_bar at t=900:", cos[900].item())   # larger: signal survives

3. Three Parameterizations: Noise, Image, and Velocity Advanced

The network has to output something from which the reverse step can be computed, but that something is not unique. Because $x_t$, $x_0$, and $\epsilon$ are linked by the single closed-form equation, predicting any one of them determines the others. There are three standard choices, and each is the right one in a different regime.

• Noise prediction ($\epsilon$-prediction): the network outputs the noise $\epsilon_\theta(x_t, t)$. This is the DDPM default and works well across most of the schedule.
• Image prediction ($x_0$-prediction): the network outputs the clean image directly. This is more stable at very high noise levels: when $x_t$ is almost pure noise, predicting $\epsilon$ amounts to echoing back the input (the target $\epsilon$ is most of what the network already sees), so a tiny error in the prediction barely changes the loss yet badly corrupts the recovered image, whereas predicting the faint underlying $x_0$ gives the network a target that actually carries signal.
• Velocity prediction ($v$-prediction): the network predicts $v = \sqrt{\bar\alpha_t}\,\epsilon - \sqrt{1-\bar\alpha_t}\,x_0$, a blend that interpolates between the two. It behaves like image prediction at high noise and noise prediction at low noise, giving the best of both, and is the choice used in progressive distillation and many modern systems.

The three are exact algebraic re-encodings of each other; given one and the noisy $x_t$, you can recover the others. The reason the choice matters is the implied loss weighting: predicting $\epsilon$, $x_0$, or $v$ puts different emphasis on different noise levels, which changes what the network spends its capacity on. The conversion utilities are short.

# Convert among the noise, image, and velocity parameterizations.
# All three are exact re-encodings linked by the closed-form equation,
# so a velocity prediction must round-trip back to the original noise.
import torch

def to_x0(xt, eps, ab):                              # ab is alpha_bar_t, shape (B,1,1,1)
    """Recover the clean image from x_t and predicted noise."""
    return (xt - (1 - ab).sqrt() * eps) / ab.sqrt()

def v_target(x0, eps, ab):
    """The velocity target used by v-prediction."""
    return ab.sqrt() * eps - (1 - ab).sqrt() * x0

def eps_from_v(xt, v, ab):
    """Recover noise from a velocity prediction (for the reverse step)."""
    return ab.sqrt() * v + (1 - ab).sqrt() * xt

# Sanity check: round-trip consistency on random tensors
x0 = torch.randn(2, 1, 4, 4); eps = torch.randn_like(x0)
ab = torch.tensor([0.3, 0.7]).view(-1, 1, 1, 1)
xt = ab.sqrt() * x0 + (1 - ab).sqrt() * eps
v = v_target(x0, eps, ab)
assert torch.allclose(eps_from_v(xt, v, ab), eps, atol=1e-5)
print("parameterizations are mutually consistent")

4. The Network: A Time-Conditioned U-Net Intermediate

For images, the denoiser $\epsilon_\theta$ is almost always a U-Net: an encoder that downsamples the image through convolutional blocks, a bottleneck, and a decoder that upsamples back to full resolution, with skip connections joining matching levels. This is the same architecture you met for segmentation in Chapter 24, repurposed: instead of predicting a per-pixel class, it predicts per-pixel noise. Two ingredients make it diffusion-specific. First, the timestep $t$ is embedded (with the sinusoidal encoding of Section 33.1) and injected into every block, usually by adding it to the feature maps, so the same weights behave differently at different noise levels. Second, self-attention layers are inserted at the lower-resolution levels so the network can relate distant regions, which a pure convolution cannot do in one layer. The skip connections matter especially here: they carry the high-frequency detail from the encoder straight to the decoder, exactly the detail that the bottleneck would otherwise lose, and that the denoiser must restore at low noise levels.

5. Putting It Together: The DDPM Training Step Intermediate

With the schedule chosen, the parameterization fixed, and the U-Net in hand, the training step is the same three-line idea from Section 33.1, now on real images. The function below is the complete DDPM loss for a batch, and it is what every image-diffusion codebase runs at its core.

# The complete DDPM training objective L_simple for a batch of images.
# Draw a random timestep, corrupt with the closed form, predict the noise,
# and return the mean squared error; this is the core of every codebase.
import torch

def ddpm_loss(unet, x0, alpha_bars, T=1000):
    """One DDPM training step on a batch of images x0 in [-1, 1]."""
    B = x0.size(0)
    t = torch.randint(0, T, (B,), device=x0.device)        # random step per image
    ab = alpha_bars[t].view(B, 1, 1, 1)                     # signal level
    noise = torch.randn_like(x0)                            # the target
    xt = ab.sqrt() * x0 + (1 - ab).sqrt() * noise           # closed-form corruption
    pred = unet(xt, t).sample                               # predicted noise
    return ((pred - noise) ** 2).mean()                     # L_simple

# Inside a normal training loop:
# loss = ddpm_loss(unet, batch, alpha_bars)
# loss.backward(); optimizer.step()