Improved Distributional Diffusion Models

_{Official PyTorch Implementation}

Tommaso Martorella^1,2  ·  Alexandre Galashov^3,4  ·  Felix Krause^1,2  ·  Stefan Andreas Baumann^1,2
Valentin De Bortoli³  ·  Arthur Gretton^3,4  ·  Björn Ommer^1,2

¹CompVis @ LMU Munich    ²Munich Center for Machine Learning (MCML)
³Google DeepMind    ⁴Gatsby Unit @ UCL

A single iDDM-XL/2 model, sampled with 4 steps (top) and 50 steps (bottom) on ImageNet 256×256. It reaches FID 4.48 and 2.38, respectively, and is trained from scratch in one stage: no distillation, no teacher, no self-bootstrapping, and no CFG during training.

Note

🚧 Code and pre-trained models are coming soon. We are cleaning up the codebase and will release training and inference code together with checkpoints on Hugging Face. Star ⭐ or watch 👀 the repository to get notified.

🚀 TL;DR

iDDM makes Distributional Diffusion Models (DDMs) practical at scale. DDMs replace the mean-predicting denoiser with a stochastic one, trained with a scoring rule to sample from $p(x_1 \mid x_t)$ instead of regressing to its mean. This strongly helps few-step generation, but has so far been too expensive and too inflexible to scale.

🔥 Contributions

• Deferred particle expansion $\rightarrow$ particles share the transformer trunk and split only in the last layers, so multi-particle training costs ~1.5× flow matching instead of ~4×
• Time-dependent scoring rule schedules $\rightarrow$ instead of fixing the scoring rule hyperparameters for all noise levels, we adapt them along the trajectory: the loss favors diverse samples early, where $p(x_1 \mid x_t)$ is broad, and sharper predictions late, where it has concentrated
• One checkpoint, 4–50 steps $\rightarrow$ FID never degrades with more sampling steps (4.48 → 2.38 for XL/2), so the same model serves both few- and many-step sampling

This repository will contain:

• A simple PyTorch implementation of iDDM
• Pre-trained class-conditional ImageNet 256×256 models (iDDM-B/2 and iDDM-XL/2)
• Sampling and FID evaluation scripts
• A training script for class-conditional ImageNet using PyTorch DDP

For method details, ablations, and comparisons with other few-step methods, see the paper and project page.

🗂️ Pre-trained models

Model	Data	Epochs	FID-50K (4 steps)	FID-50K (50 steps)	Checkpoint
iDDM-B/2	ImageNet 256×256	80	13.13	4.57	coming soon
iDDM-XL/2	ImageNet 256×256	200	4.48	2.38	coming soon

🎓 Citation

If you find our work useful, please cite our paper:

@article{martorella2026iddm,
  title   = {Improved Distributional Diffusion Models},
  author  = {Martorella, Tommaso and Galashov, Alexandre and Krause, Felix and Baumann, Stefan Andreas and De Bortoli, Valentin and Gretton, Arthur and Ommer, Bj{\"o}rn},
  journal = {arXiv preprint arXiv:XXXX.XXXXX},
  year    = {2026}
}

🙏 Acknowledgements

• This work builds on Distributional Diffusion Models with Scoring Rules by De Bortoli et al.
• Some model code is adapted from k-diffusion by Katherine Crowson (MIT)

📄 License

The code in this repository is released under the MIT License.