DDIM-GMM: Improved DDIM Sampling with Moment Matching Gaussian Mixtures
Prasad Gabbur
Denoising Diffusion Implicit Models (DDIM) perform fast sampling from diffusion models using a unimodal Gaussian reverse kernel. We propose DDIM-GMM, which replaces this kernel with a multimodal Gaussian mixture, constrained so that the forward marginals retain the same first and second order moments as the DDPM forward marginals. This allows DDIM-GMM to be used as a drop-in replacement for DDIM with no retraining. We provide two implementations: one for Stable Diffusion / latent diffusion models and one for HuggingFace Diffusers.
pip install torch numpy tqdm pillow
For the Diffusers implementation, additionally install:
pip install diffusers
Stable Diffusion / Latent Diffusion
Located in ddim_gmm_stable_diffusion/. See ddim_gmm_stable_diffusion/README.md for full details.
from ddim_gmm import DDIMSampler, GMM
gmm_params = GMM(gpu=0)
gmm_params.initialize(
dim=3 * 64 * 64,
n_components=16,
n_steps=50,
scale=1.0,
uniform_priors=True,
orthonormal=True,
upper_bound_vars=True
)
sampler = DDIMSampler(model=diffusion_model, gmm=True, gmm_params=gmm_params)
sampler.make_schedule(ddim_num_steps=50, ddim_eta=0.0)
samples, _ = sampler.sample(steps=50, batch_size=4, shape=(3, 64, 64), eta=0.0)
Located in ddim_gmm_diffusers/. See ddim_gmm_diffusers/README.md for full details, including how to integrate the scheduler into your Diffusers installation.
from diffusers.schedulers.scheduling_ddim import DDIMScheduler, GMM
gmm_params = GMM(device='cuda')
gmm_params.initialize(
dim=3 * 256 * 256,
n_components=16,
n_steps=50,
scale=1.0,
uniform_priors=True,
orthonormal=True,
upper_bound_vars=True
)
scheduler = DDIMScheduler.from_pretrained("google/ddpm-cat-256")
scheduler.set_gmm_params(gmm_params=gmm_params)
scheduler.set_timesteps(50)
@article{gabbur2026ddimgmm,
title={Improved {DDIM} Sampling with Moment Matching {G}aussian Mixtures},
author={Prasad Gabbur},
journal={Transactions on Machine Learning Research},
year={2026},
url={https://openreview.net/forum?id=CdSPjfmrQN}
}
MIT — see LICENSE for details.