Diagnosing and Improving Diffusion Models by Estimating the Optimal Loss Value

TL;DR

Diffusion training loss is not zero at the optimum — it equals an averaged conditional variance that depends on the dataset and the diffusion schedule but not on the model. The paper derives this optimal loss in closed form under a unified clean-data-prediction / VE-process formulation, builds a scalable estimator (cDOL) that converges on ImageNet-scale data, and uses it as a diagnostic to (a) compare existing diffusion variants on a common ground, (b) design a noise schedule and loss weight from the loss-gap rather than the raw loss, and (c) propose that scaling laws should fit the gap, not the loss. The recipe gives FID improvements of 0.2–4.6 across EDM and Flow Matching on CIFAR-10 / ImageNet-64, and recovers cleaner power laws on 120M–1.5B EDM2 runs.

Key claims

The optimal diffusion loss is non-zero and equals an averaged conditional variance of clean data given the noised observation; the optimum depends only on the dataset and the diffusion settings, not on the model architecture or parameterization [§3.1, Thm. 1, Eq. 12].
A direct empirical estimator has $O(N^2)$ cost; the proposed cDOL estimator with bias-correcting coefficient $c \in (0, M-1)$ is consistent and unbiased in expectation, equivalent in expectation to the $O(N^2)$ estimator with subset size $M+c$ [§3.3, Thm. 2, Eq. 17].
The optimal loss curve has a dataset-dependent critical point $\sigma_{\rm crit}$ below which the loss is near zero; $\sigma_{\rm crit}$ shrinks with dataset resolution and size (CIFAR-10 lowest, ImageNet-64 between FFHQ-64 and CIFAR-10) [§3.4, Fig. 1d].
On CIFAR-10 the gap between actual and optimal loss reveals that most diffusion variants already near-optimally fit large noise scales — the remaining room is at intermediate $\sigma$ — contradicting the conventional reading of the raw loss curve [§4.1, Fig. 2a].
The loss gap’s correlation with FID flips sign around the critical point: gap and FID correlate negatively near $\sigma_{\rm crit}$ but positively immediately to the left of it, so it’s worth trading off near-critical-point fit for cleaner fit at smaller noise [§4.1, Fig. 2a insets].
Among prediction targets evaluated under the FM process, $\mathbf{F}$ -prediction (EDM-style) achieves the closest fit across all noise scales; $\boldsymbol{\epsilon}$ excels at small $\sigma$ , $\mathbf{x}$ at large $\sigma$ , but neither alone matches across the range [§4.1, Fig. 2b].
The proposed training schedule sets loss weight $w \propto 1/L^*$ above the critical point with a cutoff plus an additional weight bump below it, and the noise schedule samples proportional to the weight-calibrated loss gap [§4.2].
The proposed schedule improves EDM and Flow Matching FID on CIFAR-10 (1.79 → 1.75 EDM, 2.07 → 1.77 FM cond) and ImageNet-64 (2.44 → 2.25 EDM cond, 3.06 → 2.29 FM cond) [Table 2].
Subtracting the estimated optimal loss from the actual training loss yields a better-fit power law across EDM2 model sizes 120M–1.5B on ImageNet-64/-512, suggesting the gap (not raw loss) is the right scaling-law metric [§1, §5].

Method

The paper unifies clean-data / noise / vector-field / score prediction across VP, VE, and FM diffusion processes under a single VE clean-data prediction formulation with precondition coefficients $(c_{\rm in}, c_{\rm out}, c_{\rm skip}, c_{\rm noise})$ [§4.1, Table 1]. Under this unified view, the optimal loss for clean-data prediction is $L^*(\sigma) = \mathbb{E}_{p(\mathbf{x}_\sigma)}\left[\mathrm{Var}_{p(\mathbf{x}_0 | \mathbf{x}_\sigma)}[\mathbf{x}_0]\right]$ — an averaged conditional variance [Thm. 1, §3.1].

Direct evaluation requires nested expectations: an outer expectation over $p(\mathbf{x}_\sigma)$ and an inner posterior $p(\mathbf{x}_0 | \mathbf{x}_\sigma)$ that can’t be sampled. The authors apply Bayes’ rule to rewrite the posterior expectation as a self-normalized importance-sampled ratio over the dataset, then introduce two scalability tricks: (i) self-normalized importance sampling with a shared random subset $\mathcal{S}$ of size $M$ (SNIS) [§3.3, Eq. 15]; (ii) DOL — always include the data sample that generated $\mathbf{x}_\sigma$ in $\mathcal{S}$ , because at small $\sigma$ that sample dominates the importance weight; (iii) cDOL — down-weight the always-included sample’s contribution by a coefficient $c$ to correct the bias DOL introduces. Theorem 2 shows cDOL with subset size $M$ matches the SNIS expectation at subset size $M + c$ [Eq. 17].

Applied to CIFAR-10 / FFHQ-64 / ImageNet-64, cDOL with $c \approx M/2$ is most accurate; SNIS ( $c=0$ ) has high variance from missed dominating samples; DOL ( $c=M-1$ ) under-estimates [Fig. 1c]. The estimator output is the per- $\sigma$ optimal loss curve.

For training-schedule design, the loss weight is $w(\sigma) = 1/L^*(\sigma)$ with a cutoff below $\sigma_{\rm crit}$ (to avoid divergence) and an additive bump targeting the small- $\sigma$ interval where loss-gap correlates positively with FID; the noise schedule is then chosen so the per-step training intensity is proportional to the weight-calibrated loss gap $w(\sigma)(L(\sigma) - L^*(\sigma))$ [§4.2].

Results

CIFAR-10, EDM: 1.79 → 1.75 (cond), 1.98 → 1.94 (uncond) FID [Table 2].
CIFAR-10, FM (with EDM sampler): 2.07 → 1.77 (cond), 2.24 → 2.03 (uncond) [Table 2].
ImageNet-64, EDM: 2.44 → 2.25 (cond) [Table 2].
ImageNet-64, FM (with EDM sampler): 3.06 → 2.29 (cond) [Table 2].
ImageNet-256 LightningDiT: 0.06 FID improvement (from §1 summary) — small in absolute terms but at a stronger baseline.
Scaling-law fit: across EDM2 sizes 120M / 200M / 350M / 700M / 1.5B on ImageNet-64 and ImageNet-512, the loss-gap power law fit is reported as cleaner than the raw-loss power law [§1, §5].
Diagnostic finding: on CIFAR-10 the actual loss of every mainstream variant is near-optimal at large $\sigma$ — the data-fitting room is at intermediate noise, the opposite of what the raw-loss curve suggests [Fig. 2a].

Why it’s interesting

This is a measurement instrument paper rather than a new model: it gives a way to compute a quantity (the per-step optimal loss) that the diffusion-training literature has been working without, and uses it both as a diagnostic and as a recipe for noise-schedule design. It directly contests an assumption baked into a lot of the wiki’s diffusion-training-efficiency work — that large- $\sigma$ regions of the noise schedule are where models still underfit. Under cDOL the room-for-improvement is in fact concentrated at intermediate $\sigma$ for every mainstream variant, which is a constraint on what new schedules (e.g. the schedule modifications in Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis Dual-Timestep Scheduling or The Geometry of Noise: Why Diffusion Models Don't Need Noise Conditioning‘s noise-conditioning-free regime) can reasonably claim to fix. It also complements The Design Space of Tri-Modal Masked Diffusion Models‘s scaling-law work: that paper fits Kaplan-form scaling laws on raw training loss for masked diffusion; this paper argues the right quantity to fit a power law against is the loss gap, not the loss — which, if it holds in the continuous-token DiT setting, changes the shape of every diffusion scaling-law curve currently in Hyperparameter scaling laws.