Research map / Image generation / Fast diffusion sampling
Constrained Diffusion Models via Dual Training
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.
Casts diffusion model training under distribution constraints as a KL-constrained problem, solves it with a Lagrangian dual training algorithm, and shows the optimal model samples from a lambda-weighted mixture of the target and constraint distributions.
Problem
Diffusion models reproduce dataset biases (e.g., underrepresented classes) and fine-tuning on new data overfits or degrades pretrained generation; there is no principled way to impose distribution requirements on training.
Approach
Formulate diffusion training as constrained distribution optimization with KL-divergence constraints to desired distributions, show strong duality implies the optimal model is a mixture distribution weighted by dual variables, and train via a dual ascent algorithm over score networks with theoretical TV-distance optimality bounds.
Key results
- On MNIST with minority digits 4, 5, 7, constrained model FID 13.4 vs unconstrained 15.9.
- On CelebA with male minority, constrained model FID 11.6 vs unconstrained 19.6.
- On MNIST fine-tuning with digit 9 held out, constrained model FID 25.2 vs unconstrained 45.9, and it samples the new class while retaining old classes.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add minority-class distribution constraints for digits 4, 5, and 7 | Diffusion model trained on MNIST · MNIST | FID: 15.9 -> 13.4 | Helped |
| Add a minority-class distribution constraint for male samples | Diffusion model trained on Celeb-A · Celeb-A | FID: 19.6 -> 11.6 | Helped |
| Add minority-class distribution constraints for Cassette player, French horn, and Golf ball | Latent diffusion model trained on ImageNet · ImageNet (10-class subset) | FID: 36.0 -> 27.3 | Helped |
| Add a constraint to preserve the pretrained model's generation distribution | Fine-tuned diffusion model · MNIST fine-tuning (new digit 9) | FID: 45.9 -> 25.2 | Helped |
Limitations
- Theory relies on disjoint supports for the closed-form dual solution and on bounded data and score-error assumptions.
- Parametrized dual problem is non-convex, so optimality guarantees are only approximate.
- Experiments are limited to small or subset-based datasets and unconditional models.
Experimental setup
Datasets: MNIST CelebA ImageNet (10-class subset)
Benchmarks: MNIST CelebA ImageNet (10-class subset)
Metrics: FID class sampling frequency
Baselines: unconstrained diffusion model pretrained diffusion model fine-tuning
Builds on: denoising diffusion probabilistic models latent diffusion constrained learning
Closest papers by meaning
- Composition and Alignment of Diffusion Models using Constrained Learning (NeurIPS 2025 · similarity 0.91)
- Fast constrained sampling in pre-trained diffusion models (NeurIPS 2025 · similarity 0.90)
- The Emergence of Reproducibility and Consistency in Diffusion Models (ICML 2024 · similarity 0.90)
- Constrained Diffusion with Trust Sampling (NeurIPS 2024 · similarity 0.89)
- Training Data Attribution in Diffusion Models via Mirrored Unlearning and Noise-Consistent Skew (NeurIPS 2026 · similarity 0.89)
- Scalable Discrete Diffusion Samplers: Combinatorial Optimization and Statistical Physics (ICLR 2025 · 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.