DDPM from Scratch, Part 4: The Loss Function

The final boss. We've assembled the prerequisites (Part 1), the closed-form noising process (Part 2 — the route), and the denoising posterior (Part 3 — the car). What's left is the driver (司机): how do we train the thing? — i.e. where DDPM's loss function comes from. We'll do it two ways: the quick "cheat path," then the principled max-likelihood path. Both land on the same place.

---

Path A — the cheat path (功法逆练)

From Part 3 we already have the reverse posterior:

$$q(x_{t-1} \mid x_t, x_0) = \mathcal{N}\big(x_{t-1};\ \tilde\mu_t(x_t, x_0),\ \tilde\beta_t^2 \mathbf{I}\big).$$

The variance $\tilde\beta_t^2$ is a constant (it depends on $t$ but not on data), so to sample $x_{t-1}$ we only need the mean $\tilde\mu_t$ — and $\tilde\mu_t$ needs $x_t$ (the current, known step) and $x_0$ (unknown at inference). So we let a neural network with parameters $\theta$ predict $x_0$. Since $x_0$ and the noise $\varepsilon$ are interchangeable ($x_t = \sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\varepsilon$), predicting $x_0$ and predicting $\varepsilon$ are the same task — the authors just found predicting $\varepsilon$ trains a bit better.

Following that logic, the answer drops out immediately: minimize the gap between the model's predicted noise $\varepsilon_\theta$ and the true (ground-truth) noise $\varepsilon$ — that is, an MSE loss:

$$\mathcal{L} = \big\|\,\varepsilon - \varepsilon_\theta(x_t, t)\,\big\|_2^2.$$

That's the shortcut — quick, but it skirts the author's actual reasoning. Now let's walk the principled path and confirm it gives the same thing.

---

Path B — the principled path (功法正练)

The objective: maximum likelihood

We want generated images to have high probability of being what we asked for. (Want to generate basketball-playing Kunkun? Then we want the generated image to have a high probability of being that.) So we maximize the log-likelihood:

$$\max_\theta\ \log p_\theta(x_0).$$

From Part 1, the variational bound (the thing we actually optimize) is:

$$L = \mathbb{E}_q\!\left[-\log \frac{p_\theta(x_{0:T})}{q(x_{1:T}\mid x_0)}\right].$$

If the step from $\max \log p_\theta(x_0)$ to this bound is fuzzy, go back to Part 1 — that's where it was cleared.

Set the $\mathbb{E}_q$ aside for now and just play with the log part.

Step 1: expand the ratio with log rules

Substitute the route (Part 2) and reverse (Part 1) product forms:

$$p_\theta(x_{0:T}) = p(x_T)\prod_{t=1}^{T} p_\theta(x_{t-1}\mid x_t), \qquad q(x_{1:T}\mid x_0) = \prod_{t=1}^{T} q(x_t\mid x_{t-1}).$$

So:

$$\log \frac{p_\theta(x_{0:T})}{q(x_{1:T}\mid x_0)} = \log \frac{p(x_T)\prod_{t=1}^{T} p_\theta(x_{t-1}\mid x_t)}{\prod_{t=1}^{T} q(x_t\mid x_{t-1})}.$$

Using $\log(ab) = \log a + \log b$, every product becomes a sum:

$$= \log p(x_T) + \sum_{t=1}^{T}\log p_\theta(x_{t-1}\mid x_t) - \sum_{t=1}^{T}\log q(x_t\mid x_{t-1}).$$

Step 2: Bayes on the $q$ term (telescoping)

Transform $q(x_t\mid x_{t-1})$ with Bayes' rule:

$$q(x_t\mid x_{t-1}) = \frac{q(x_{t-1}\mid x_t)\,q(x_t)}{q(x_{t-1})}.$$

(This is exactly why we sweated through Part 3 — you carry your loot into the next level.) The third sum becomes:

$$\sum_{t=1}^{T}\log q(x_t\mid x_{t-1}) = \sum_{t=1}^{T}\log q(x_{t-1}\mid x_t) + \sum_{t=1}^{T}\log\frac{q(x_t)}{q(x_{t-1})}.$$

That second piece telescopes. Take the 3-term example ($t = 3, 2, 1$):

$$\frac{q(x_3)}{q(x_2)}\cdot\frac{q(x_2)}{q(x_1)}\cdot\frac{q(x_1)}{q(x_0)} = \frac{q(x_3)}{q(x_0)},$$

so in general $\sum_{t=1}^{T}\log\frac{q(x_t)}{q(x_{t-1})} = \log\frac{q(x_T)}{q(x_0)}$. Putting it back:

$$\log \frac{p_\theta(x_{0:T})}{q(x_{1:T}\mid x_0)} = \log p(x_T) + \sum_{t=1}^{T}\log\frac{p_\theta(x_{t-1}\mid x_t)}{q(x_{t-1}\mid x_t)} - \log\frac{q(x_T)}{q(x_0)}.$$

Step 3: regroup and restore the $x_0$ conditions

Collect the $x_T$ pieces, $\log p(x_T) - \log q(x_T) = \log\frac{p(x_T)}{q(x_T)}$, and recall that every $q$ here secretly carries the $x_0$ condition (Part 3's whole point). Adding $x_0$ back:

$$= \log\frac{p(x_T)}{q(x_T\mid x_0)} + \sum_{t=1}^{T}\log\frac{p_\theta(x_{t-1}\mid x_t)}{q(x_{t-1}\mid x_t, x_0)} + \underbrace{\log q(x_0 \mid x_0)}_{=\,\log 1\,=\,0}.$$

That last term is the running joke: "Given that Kunkun practiced for two and a half years, what's the probability that Kunkun practiced for two and a half years?" It's 1, and $\log 1 = 0$, so it vanishes.

Step 4: why the paper's $L_0$ split is unnecessary

The paper (Eq. 5) writes the bound as three labeled pieces:

$$L = \mathbb{E}_q\Big[\underbrace{D_{\mathrm{KL}}(q(x_T\mid x_0)\,\|\,p(x_T))}_{L_T} + \sum_{t>1}\underbrace{D_{\mathrm{KL}}(q(x_{t-1}\mid x_t, x_0)\,\|\,p_\theta(x_{t-1}\mid x_t))}_{L_{t-1}} \underbrace{-\,\log p_\theta(x_0\mid x_1)}_{L_0}\Big].$$

The author splits the $t=1$ term out as a special $L_0$. But that's showmanship — it's unnecessary. Pull $t=1$ out of our unified sum:

$$\log\frac{p_\theta(x_0\mid x_1)}{q(x_0\mid x_1, x_0)} = \log p_\theta(x_0\mid x_1),$$

because $q(x_0\mid x_1, x_0) = 1$ (the Kunkun joke again — you already know $x_0$). So the special $L_0$ term is just the $t=1$ case of the general sum. Keep it folded into $\sum_{t=1}^{T}$ — it's cleaner, and it's what most real implementations do anyway.

Step 5: recognize the KL divergences

Restore the negative sign (we minimize $-\log$) and the expectation $\mathbb{E}_q$. A term of the form $\mathbb{E}_{q}\big[\log\frac{q}{p}\big] = \int q\log\frac{q}{p}$ is a KL divergence (Part 1). So:

$$L = \mathbb{E}_q\Big[\underbrace{D_{\mathrm{KL}}\big(q(x_T\mid x_0)\,\|\,p(x_T)\big)}_{\text{(i)}} + \sum_{t\ge1}\underbrace{D_{\mathrm{KL}}\big(q(x_{t-1}\mid x_t, x_0)\,\|\,p_\theta(x_{t-1}\mid x_t)\big)}_{\text{(ii)}}\Big].$$

• Term (i) has both ends known: $q(x_T\mid x_0)$ from the forward process, and $p(x_T)$ defined as fixed Gaussian noise. No neural network involved → it's a constant → drop it.
• Term (ii) is a KL between two Gaussians — and that's the heart of the loss.

(A small honesty note: the outer $\mathbb{E}_q$ doesn't fully cancel into the KL, because the $q(x_{1:T}\mid x_0)$ expectation has only "used up" the $x_{t-1}, x_t$ marginals — the other latents $x_1,\dots,x_{T-1}$ are still being integrated. So an expectation lingers. It doesn't affect the final result.)

---

Computing the KL between two Gaussians

Let $p = \mathcal{N}(\mu_1, \sigma_1^2)$ (this is $q(x_{t-1}\mid x_t,x_0)$) and $q = \mathcal{N}(\mu_2, \sigma_2^2)$ (this is $p_\theta(x_{t-1}\mid x_t)$).

$$D_{\mathrm{KL}}(P\|Q) = \int p(x)\log\frac{p(x)}{q(x)}\,dx.$$

Trick: keep the leading factor as $p(x)$ — don't expand it. Only expand the log of the ratio:

$$\log\frac{p(x)}{q(x)} = \log\frac{\frac{1}{\sigma_1\sqrt{2\pi}}e^{-\frac{(x-\mu_1)^2}{2\sigma_1^2}}}{\frac{1}{\sigma_2\sqrt{2\pi}}e^{-\frac{(x-\mu_2)^2}{2\sigma_2^2}}} = \log\frac{\sigma_2}{\sigma_1} - \frac{(x-\mu_1)^2}{2\sigma_1^2} + \frac{(x-\mu_2)^2}{2\sigma_2^2}.$$

So:

$$D_{\mathrm{KL}} = \int p(x)\left[\log\frac{\sigma_2}{\sigma_1} - \frac{(x-\mu_1)^2}{2\sigma_1^2} + \frac{(x-\mu_2)^2}{2\sigma_2^2}\right]dx.$$

Evaluate term by term as expectations under $p$:

Term 1 — expectation of a constant is itself: $$\int p(x)\log\tfrac{\sigma_2}{\sigma_1}\,dx = \log\frac{\sigma_2}{\sigma_1}.$$

Term 2 — $\int p(x)(x-\mu_1)^2\,dx = \operatorname{Var}_p(X) = \sigma_1^2$, so: $$-\frac{1}{2\sigma_1^2}\cdot\sigma_1^2 = -\frac12.$$

Term 3 — expand $(x-\mu_2)^2 = x^2 - 2\mu_2 x + \mu_2^2$ and use: $$\int p(x)\,x\,dx = \mu_1, \qquad \int p(x)\,x^2\,dx = \mathbb{E}[X^2] = \mu_1^2 + \sigma_1^2.$$ The second identity comes from the variance definition (a quick college-level reminder): $$ \sigma_1^2 = \mathbb{E}[(X-\mathbb{E}X)^2] = \mathbb{E}[X^2] - 2\mu_1^2 + \mu_1^2 = \mathbb{E}[X^2] - \mu_1^2 \ \Rightarrow\ \mathbb{E}[X^2] = \sigma_1^2 + \mu_1^2. $$ Therefore: $$\int p(x)(x-\mu_2)^2\,dx = (\mu_1^2+\sigma_1^2) - 2\mu_2\mu_1 + \mu_2^2 = \sigma_1^2 + (\mu_1-\mu_2)^2,$$ and Term 3 $= \dfrac{\sigma_1^2 + (\mu_1-\mu_2)^2}{2\sigma_2^2}$.

Putting it together — the Gaussian KL formula:

$$\boxed{\,D_{\mathrm{KL}}\big(\mathcal{N}(\mu_1,\sigma_1^2)\,\|\,\mathcal{N}(\mu_2,\sigma_2^2)\big) = \log\frac{\sigma_2}{\sigma_1} + \frac{\sigma_1^2 + (\mu_1-\mu_2)^2}{2\sigma_2^2} - \frac12.\,}$$

---

Collapsing to MSE

Here's the payoff: almost everything is dead weight. In DDPM we set the model's variance equal to the posterior variance — i.e. $\sigma_1 = \sigma_2$ (both equal the constant $\tilde\beta_t$). Then:

• $\log\dfrac{\sigma_2}{\sigma_1} = \log 1 = 0$,
• $\dfrac{\sigma_1^2}{2\sigma_2^2} = \dfrac12$, which cancels the $-\dfrac12$,
• leaving only $\dfrac{(\mu_1-\mu_2)^2}{2\sigma_2^2}$.

Since $\sigma_2^2$ is a constant, the loss is proportional to:

$$L \ \propto\ (\mu_1 - \mu_2)^2 = \big(\tilde\mu_t(x_t, x_0) - \mu_\theta(x_t, t)\big)^2.$$

Both means are functions of $(x_t, x_0)$, and $p_\theta$ has the same functional form as $q$ — the only difference is that $q$ uses the true $x_0$ while $p_\theta$ uses the model's predicted $\hat{x}_0$. Plugging the Part 3 expression for $\tilde\mu_t$ into both and converting $x_0 \to \varepsilon$ via $\hat{x}_0 = (x_t - \sqrt{1-\bar\alpha_t}\,\hat\varepsilon)/\sqrt{\bar\alpha_t}$, the whole thing simplifies to a noise-prediction MSE:

$$\boxed{\,L \ \propto\ \big\|\,\varepsilon - \varepsilon_\theta(x_t, t)\,\big\|_2^2.\,}$$

The principled path lands exactly where the cheat path did. Forward and reverse "cultivation" agree — which is why DDPM training is, in the end, just MSE on predicted noise.

---

Takeaways

• Two routes, one answer. The shortcut (predict $\varepsilon$, match with MSE) and the principled max-likelihood derivation give the same loss.
• Drop the constants. Of the variational bound's terms, $L_T$ is network-free (constant), the variances cancel, and only the mean-difference term survives — which becomes MSE on noise.
• Don't over-engineer the reading. The paper's separate $L_0$ term is unnecessary — it folds into the general sum. Most implementations just use the plain noise-MSE. When learning diffusion, think simple, not complicated — the more you try to over-formalize every term, the scarier (and less clear) it gets.

That's the driver — and with it, the four-level run through the DDPM main paper is complete. 🎉 Congratulations on clearing all four gates: prerequisites → noising → denoising → loss. Anything left unclear can go in a bonus "DLC" episode.

---

Notes taken from a video tutorial on DDPM (Part 4 of the series).

---

References

• Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. NeurIPS 2020. arXiv:2006.11239