Improved Distributional Diffusion Models
Tommaso Martorella, Alexandre Galashov, Felix Krause, Stefan Andreas Baumann, Valentin De Bortoli, Arthur Gretton, Björn Ommer
LMU Munich · Munich Center for Machine Learning (MCML) · Google DeepMind · UCL
Published September 29, 2026 · on the feed September 30, 2026 · ▲ 5 · arXiv 2609.37147
The abstract
Distributional Diffusion Models (DDMs) replace the standard mean-prediction denoiser with a distributional denoiser trained via a scoring rule objective, learning a stochastic approximation to p(x_1 mid x_t) rather than its conditional mean. However, scaling DDMs to modern image-generation settings faces two obstacles: (i) multi-particle training incurs overhead that scales with the number of particles, (ii) DDMs use globally fixed scoring rule hyperparameters, forcing a single trade-off across sampling budgets. We mitigate these limitations by deferring particle expansion to late transformer layers, and the hyperparameter trade-off by introducing time-dependent scoring rule schedules informed by the dynamical regimes of~Biroli2024. Combined with a DiT-based latent setup, these changes make DDM training practical on class-conditional ImageNet-256^2, achieving 4.48 FID at 4 steps and 2.38 at 50 steps with DiT-XL/2, from a single model trained from scratch in one stage, without a teacher, self-distillation or JVPs. The result is a stochastic few-step generator whose FID does not degrade as the sampling budget grows from 4 to 50 NFE, and the same recipe transfers to text-to-image generation. Code and pre-trained models available at https://github.com/CompVis/iDDM.
The Brief
2:13A short podcast summary of the paper. Written and voiced by AI.
Transcript
Improved Distributional Diffusion Models. By Tommaso Martorella and six others, from LMU Munich, the Munich Center for Machine Learning, and two other institutions. Published on the twenty-ninth of September. Its code repository is a placeholder: a readme, and nothing to run yet.
Distributional diffusion models replace the standard denoiser, which predicts a mean, with a distributional denoiser trained through a scoring rule objective. Rather than the conditional mean of the clean sample given a noisy one, they learn a stochastic approximation to the whole conditional distribution. Scaling them to modern image generation, however, meets two obstacles. Training with multiple particles carries overhead that grows with the number of particles. And the scoring rule hyperparameters are fixed globally, which forces one trade-off across every sampling budget. A setting that suits a handful of steps is stuck with the model when it is given many, and the reverse.
Each obstacle gets its own remedy. For the cost, particle expansion is deferred to the late transformer layers. For the trade-off, the fixed hyperparameters give way to scoring rule schedules that depend on time, informed by the dynamical regimes described in work by Biroli. Rather than one compromise held for the whole of sampling, the scoring rule is allowed to change as noise levels change.
These changes are combined with a latent setup based on DiT. Together they make training practical on class-conditional ImageNet at two hundred and fifty-six pixels square. The model is trained from scratch, in a single stage, as one model. There is no teacher, no self-distillation, and no JVPs. Nothing is borrowed from another model, before or after.
With DiT-XL/2, the model reaches an FID of four point four eight at four steps, and two point three eight at fifty steps. The result is a stochastic few-step generator whose FID does not degrade as the sampling budget grows from four to fifty NFE. The same recipe also transfers to text-to-image generation. Code and pre-trained models are available.
One model, one stage of training, and quality that holds from four steps to fifty. The two fixes divide the work cleanly. Deferring particles is what makes training affordable at this scale, and time-dependent schedules are what release the model from choosing a single trade-off for every budget.
Paper Quest · 7 chapters · about 39 min
Play this paper
Learn it by doing: make the calls the authors made, guess the result before you see it, and defend it in a viva graded on your own words. Chapter 1 is right here.
Chapter 1 🧩 · The problem
A generator that samples instead of averaging, held back by the cost of its own samples.
You train an image generator whose denoiser returns samples instead of a single average guess, which helps when few sampling steps leave the clean image very uncertain.
The road ahead
- 1 · The problem ~6 min A denoiser that returns samples instead of one average guess needs several samples per training example, and each sample costs a full network pass.
- 2 · What came before ~3 min
- 3 · The key idea ~8 min Share one pass for most of the network, split into particles only near the end, and let the scoring rule's settings change with the noise level.
- 4 · Choose your own research ~3 min
- 5 · Build the method ~8 min The noisy input runs once through early layers, splits into particles in late layers, is scored by a time-dependent rule, and then generates images.
- 6 · The experiment ~6 min You’ll guess: As the sampling budget grows from 4 to 50 NFE, what happens to the image quality, measured by FID, of the improved distributional generator?
- 7 · The viva ~5 min Three questions from the examiners, answered in your own words. Win it for the paper’s trophy.
Free, with Google or an email address. Your stars and progress are kept.