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.

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.

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}
}

• 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)

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