Improved Distributional Diffusion Models
Tommaso Martorella, Alexandre Galashov, Felix Krause, Stefan Andreas Baumann, Valentin De Bortoli, Arthur Gretton, Björn Ommer
cs.CV, cs.LG
2026-09-29
Deferred expansion and scheduled scoring rules scale DDMs to DiT-XL: 4.48 FID at 4 steps, 2.38 at 50 on ImageNet-256, from scratch, no teacher.
In the few-step regime, p(x1|xt) is wide. Standard diffusion and flow matching fit only the conditional mean E[x1|xt]. On a coarse grid that mean is a thin summary of a fat posterior, and samples smear. The usual fixes are better ODE solvers, distilling a many-step teacher, or training a few-step generator with consistency or flow-map objectives.
Distributional Diffusion Models (DDMs) from De Bortoli et al. (2025) change the target. The denoiser also takes auxiliary noise ξ and is trained with a generalized energy score so that it samples from an approximation to p(x1|xt). Setting λ=0 and β=2 recovers ordinary MSE. Two costs kept DDMs off modern image-scale training. Each example uses m particles, so the forward pass grows linearly with m. A single global (λ, β) also freezes the fidelity-diversity trade-off at every noise level, so one checkpoint cannot serve both 4-step and 50-step budgets.
iDDM changes two pieces and keeps a latent DiT backbone.
Deferred population expansion. Early transformer blocks run once on xt. Only the last few blocks are replicated into m particles, with ξ injected there. Particles share those late weights and differ only in ξ. On DiT-XL (28 layers, expansion from layer 24, m=4) a training step costs about 36% of naive DDM, or about 1.5× flow matching. On 4×A100 the measured XL step cost is 1.41× FM. Among ξ injections (channel concat, residual add, AdaNorm, register tokens), fixed-inner-width concat wins on 4-step FID; AdaNorm and tokens collapse in the few-step regime. With CFG, fewer distributional layers work better: 2-4 on DiT-B, and the B-scale default is ℓstart=10 of 12.
Time-dependent scoring rules. The posterior is broad at the noise end and tight at the data end, so a fixed (λ, β) is a bad compromise. Biroli et al. (2024) split the reverse process into speciation and collapse, plus a geometric separation threshold nearer the data. The schedule is more distributional early and more regression-like late. The default is linear, s(t)=1-t with βmin=0.1, close to a piecewise dyn-reg schedule without estimating the thresholds. Training times are drawn from a jit distribution that oversamples the noisy end, where the distributional loss matters more.
Inference draws one ξ per Euler step. Training is single-stage from scratch, with no teacher and no JVPs.
Controlled DiT-B, ImageNet-256, 400k steps, all retrained. FID uses 50k samples. CFG is swept at evaluation and is not used in training. Shortcut's 50-step column is 64 steps, because that model only supports powers of two.
| Method | 50-step FID | 4-step FID |
| Shortcut | 9.85 | 15.02 |
| IMM | 7.57 | 24.85 |
| MFM | 7.72 | 9.17 |
| Naive DDM | 10.04 | 43.68 |
| iDDM | 4.57 | 13.13 |
Versus naive DDM, 4-step FID falls from 43.68 to 13.13 and 50-step from 10.04 to 4.57. At matched training compute (iDDM 260k vs FM 400k), 4-step FID is 14.33 vs 26.70; FM is still slightly better at 50 steps (4.97 vs 5.62). At 4 steps, FID, FD-DINOv2, KID, Inception Score, precision, and recall all beat both FM and naive DDM.
DiT-XL/2, 200 epochs:
| Method | Training | 50-step FID | 4-step FID |
| MeanFlow-XL/2 | 240 ep | 3.29 | 2.93 |
| iMF-XL/2 (48 layers) | 800 ep | 1.43 | 1.51 |
| IMM-XL/2 | 3840 ep | 1.82 | 2.51 |
| MFM (multi-stage) | 800+80+28 | 5.74 | 1.97 |
| Naive DDM-XL/2 | 200 ep | 4.71 | 23.25 |
| iDDM-XL/2 | 200 ep | 2.38 | 4.48 |
MeanFlow's FID rises from 2.93 at 4 steps to 3.29 at 50; MFM rises from 1.97 to 5.74. iDDM goes from 4.48 to 2.38. One checkpoint covers 4 to 50 steps. Per-step cost is 1.41× FM, against 1.96× for MeanFlow and 6.20× for iMF. Total training compute is about 1.67× for MeanFlow and 17.6× for iMF.
The same recipe transfers to a 1.6B text-to-image model. MS-COCO FID at 4 steps is 41.05 vs 78.20 for a matched FM baseline; at 8 steps 19.40 vs 26.33; at 16 steps the two are tied (11.27 vs 11.20). CLIP, GenEval, and HPS v2.1 also degrade less at 4 steps. Spatial on T2I-CompBench at 16 steps stays weaker than FM (0.12 vs 0.29).
For a from-scratch, single-stage, stochastic generator whose one checkpoint serves 4 through 50 steps, this is the only recipe in Table 4 that hits all three. Distillation (sCD, ADD) is stronger at 1-2 steps and needs a teacher plus a second stage. iMF and IMM post better FID at much higher training cost. MeanFlow and MFM win at 4-8 steps and get worse if you spend more sampling compute.
The extra 1.4× per-step training cost is real. Code and checkpoints are public.
The authors flag that the regime anchors come from a mean-field analysis and are approximate at finite scale; smooth schedules beat hard steps; distilled and fastforward methods still win at 1-2 steps. The optimal (λ(t), β(t)) is open. In the B-scale ablation, scheduling only β with λ=1 and a cut at t=0.1 reaches 13.06 4-step FID, better than the default linear schedule, and was rejected because variance is unidentified for t>0.1. The schedule remains a heuristic.
The T2I 512² model is a 20k-step finetune of a 256² checkpoint. Spatial and some other CompBench axes stay behind FM at 16 steps. The iMF comparison is not matched in depth or training budget. Deferred expansion helps with more distributional layers when CFG is off, and the opposite when CFG is on; that interaction has no clean theory yet.