Research map / Image generation / Fast diffusion sampling
End-to-end Learning of Gaussian Mixture Priors for Diffusion Sampler
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Learns Gaussian mixture priors end-to-end inside diffusion-based samplers (DIS, MCD, CMCD, DBS) and adds components iteratively during training, improving ELBO, ESS and mode coverage on synthetic and real-world unnormalized densities.
Problem
Diffusion-based samplers for unnormalized target densities use a fixed Gaussian prior, which poorly matches targets with different support, causes exploration and discretization problems, and can mode-collapse under reverse KL variational training.
Approach
Make the diffusion sampler's prior a learnable Gaussian mixture trained jointly with the control functions and step size by maximizing the extended ELBO, plus an iterative refinement scheme that adds new components with means chosen from MALA-generated candidate samples.
Key results
- On Funnel (d=10), CMCD-GMP reaches ELBO -0.027 vs -0.210 for CMCD and DBS-GMP reaches -0.047 vs -0.180 for DBS.
- On the Fashion target (d=784), DIS-GMP + IMR raises entropic mode coverage to 0.780 from 0.213 for DIS and lowers Sinkhorn distance to 513.8 from 17.7k.
- Learned GMPs improve ELBO over learned Gaussian priors across most real-world tasks, with DBS-GMP best on Cancer, Brownian, Ionosphere and Sonar among diffusion methods.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the fixed Gaussian prior with a learned Gaussian mixture prior | Diffusion-based samplers on Funnel · Funnel | ELBO: -0.211 \scriptstyle \pm 0.002 -> -0.027 \scriptstyle \pm 0.001 | Helped |
| Replace a learned Gaussian prior with a learned Gaussian mixture prior | Diffusion-based samplers on Funnel · Funnel | ESS: 0.567 \scriptstyle \pm 0.023 -> 0.950 \scriptstyle \pm 0.004 | Helped |
| Use a learned Gaussian mixture prior with DBS instead of the fixed prior | DBS on Credit, Seeds, Cancer, Brownian, Ionosphere, and Sonar · Real-world benchmarks | ELBO | Helped |
| Use iterative model refinement to add mixture components (ablation) | DIS-GMP on Fashion · Fashion | EMC: 0.012 \scriptstyle \pm 0.021 -> 0.780 \scriptstyle \pm 0.089 | Helped |
| Use iterative model refinement to add mixture components (ablation) | DIS-GMP on Fashion · Fashion | Sinkhorn distance: 1703.023 \scriptstyle \pm 3.050 -> 513.776 \scriptstyle \pm 13.936 | Helped |
| Use iterative model refinement with the Gaussian mixture prior (ablation) | DIS on Fashion · Fashion | ELBO: -38.873 \scriptstyle \pm 0.175 -> -62.482 \scriptstyle \pm 2.752 | Hurt |
| Increase the number of mixture components (ablation) | DIS · Subset of benchmark tasks | ESS | Helped |
| Increase the number of diffusion steps (ablation) | DIS · Subset of benchmark tasks | ESS | Helped |
| Use a Gaussian mixture prior with DIS on Fashion instead of a learned Gaussian prior | DIS · Fashion | ELBO: -24.712 \scriptstyle \pm 0.253 -> -38.873 \scriptstyle \pm 0.175 | Hurt |
| Use a Gaussian mixture prior with DIS on Fashion instead of a learned Gaussian prior | DIS · Fashion | Sinkhorn distance: 1671.411 \scriptstyle \pm 2.394 -> 1703.023 \scriptstyle \pm 3.050 | Hurt |
Surprising findings
- On the Fashion target, iterative refinement gives worse ELBO and delta log Z than single-Gaussian variants despite better mode coverage and Sinkhorn distance.
- Learned Gaussian priors on the Fashion target (DIS-GP) achieve better ELBO than DIS-GMP, while DIS-GMP is better on mode-related metrics.
Limitations
- Mixture weights are uniformly initialized and fixed during training.
- Higher K makes the control functions more complex, which can worsen ELBO and delta log Z on multimodal targets.
- Iterative refinement relies on MALA-generated candidate samples and prior knowledge of the target support.
Experimental setup
Datasets: Funnel Fashion MNIST (NICE-trained normalizing flow target) Credit Cancer Ionosphere Sonar Seeds Brownian
Benchmarks: Funnel Fashion Credit Cancer Ionosphere Sonar Seeds Brownian
Metrics: ELBO delta log Z ESS Sinkhorn distance W2 entropic mode coverage
Baselines: GVI GMVI MCD CMCD DIS DBS SMC CRAFT FAB
Builds on: Denoising diffusion samplers Monte Carlo Diffusions Controlled Monte Carlo Diffusions Diffusion Bridge Sampler Gaussian mixture variational inference
Closest papers by meaning
- A Mixture-Based Framework for Guiding Diffusion Models (ICML 2025 · similarity 0.91)
- Annealing in variational inference mitigates mode collapse: a theoretical study on Gaussian mixtures (NeurIPS 2026 · similarity 0.90)
- Theoretical insights for diffusion guidance: A case study for Gaussian mixture models (ICML 2024 · similarity 0.89)
- Composing Diffusion Priors with Explicit Physical Context via Generative Gibbs Sampling (NeurIPS 2026 · similarity 0.89)
- Scalable Discrete Diffusion Samplers: Combinatorial Optimization and Statistical Physics (ICLR 2025 · similarity 0.89)
- Efficient Mixture Learning in Black-Box Variational Inference (ICML 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.