Raghuram Shankar
← Diffusion Language Models

Part 3 · The Maths I: Continuous Diffusion

The reverse process and the ELBO

Generation needs q(xt−1 | xt), which is intractable — it implicitly marginalizes over the entire data distribution. Two facts rescue us. First, for small βt the true reversal is itself nearly Gaussian, so a Gaussian model pθ(xt−1 | xt) is well-specified. Second, conditioned additionally on the clean image, the reversal is exactly Gaussian with closed form:

q(xt−1 | xt, x0) = 𝒩( μ̃t, β̃t·I ),    μ̃t = √ᾱt−1·βt1−ᾱt·x0  +  √αt·(1−ᾱt−1)1−ᾱt·xt,    β̃t = 1−ᾱt−11−ᾱt·βt

μ̃t, β̃t the mean and variance of the true reverse step · it is conditioned on x0, which the model does not have — that gap is what the network is trained to close

The standard variational argument (the same one behind VAEs) then bounds the negative log-likelihood — the NELBO — and the bound splits cleanly per timestep:

−log pθ(x0) ≤ 𝔼q[ −log pθ(x0 | x1) ]  +  ∑t=2T 𝔼q KL( q(xt−1 | xt, x0) ‖ pθ(xt−1 | xt) )  +  KL( q(xT | x0) ‖ 𝒩(0, I) )

pθ the model's reverse process · q the fixed forward process · KL the divergence between them at each step · the sum is a bound on the negative log likelihood, so minimizing it maximizes likelihood

Read the three terms left to right: a reconstruction term for the final tiny step; the denoising-matching sum — the workhorse, demanding that at every noise level the model's one-step denoiser match the tractable posterior above; and the prior term, which contains no parameters at all and is ≈0 whenever the schedule truly destroys the signal. Every KL is between Gaussians, hence closed-form.

One reparameterization turns this into the loss everyone actually ships. Predict the noise ε rather than the mean:

μθ(xt, t) = 1√αt·( xt − βt√(1−ᾱt)·εθ(xt, t) )

εθ the network's prediction of the noise that was added · αt the per-step signal retained, ᾱt the cumulative version · predicting noise and predicting the mean are equivalent, related by this rearrangement

after which the KL terms collapse into weighted squared errors on ε. DDPM famously drops the weights (it trains better) and arrives at:

Lsimple = 𝔼t, x0, ε ‖ ε − εθ( √ᾱt·x0 + √(1−ᾱt)·ε,  t ) ‖²

ε the noise actually drawn · εθ the network's guess at it · the expectation runs over timesteps t, clean data x0, and noise draws — a plain regression loss

“Guess the noise I added” — that single regression, averaged over all noise levels, trains the entire generative hierarchy.

The score view. A direct computation from the closed form gives ∇xt log q(xt | x0) = −ε / √(1−ᾱt), so the trained network estimates the score (gradient of the log-density) of each noisy marginal: sθ(xt) = −εθ(xt, t) / √(1−ᾱt). Denoising is gradient ascent on log-probability; this identification (score-based generative modeling, developed extensively in Ermon’s group) unifies diffusion with stochastic differential equations and is where samplers, guidance, and much else are derived from.

Faded derivation: forward kernel → Lsimple
1Compose the per-step Gaussians into a one-jump marginal q(xt | x0). What are its mean and variance?
q(xt | x0) = 𝒩( √ᾱt·x0, (1−ᾱt)·I ),  ᾱt = ∏s≤t(1−βs)

Gaussians compose: means multiply, variances accumulate; variance preservation keeps the total at I.

2Use Bayes to get the reverse step conditioned on the clean image, q(xt−1 | xt, x0). What form does it take?
𝒩( μ̃t, β̃t·I ),  μ̃t = √ᾱt−1·βt1−ᾱt·x0 + √αt·(1−ᾱt−1)1−ᾱt·xt,  β̃t = 1−ᾱt−11−ᾱt·βt

Product of two Gaussians in xt−1 is Gaussian; complete the square to read off mean and variance.

3Write the NELBO and split it per timestep. Which three kinds of term appear?
𝔼q[−log pθ(x0|x1)] + ∑t≥2 𝔼q KL( q(xt−1|xt,x0) ‖ pθ(xt−1|xt) ) + KL( q(xT|x0) ‖ 𝒩(0,I) )

Reconstruction + denoising-matching (the workhorse) + parameter-free prior (≈0 by schedule design).

4Reparameterize the model mean μθ in terms of a noise prediction εθ.
μθ(xt, t) = 1√αt·( xt − βt√(1−ᾱt)·εθ(xt, t) )

Solve xt = √ᾱtx0 + √(1−ᾱt)ε for x0 and substitute into μ̃t.

5Collapse the Gaussian KLs and drop the weights. What loss remains?
Lsimple = 𝔼t, x0, ε ‖ ε − εθ( √ᾱt·x0 + √(1−ᾱt)·ε, t ) ‖²

KL between Gaussians with shared variance is a scaled squared distance between means; DDPM found the unweighted version trains better.

Try to produce each step on paper before revealing it. When you can go 1→5 unaided, you own this derivation. (Steps start hidden; “worked” shows everything.)