# Integration Matters: Rollout-Based Training for Constrained Diffusion Models

arXiv: https://arxiv.org/abs/2607.14398 · HTML: https://mlsanity.com/paper/2607.14398/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 44 of its 54 references that are in the map.
- First posted on arXiv 15 Jul 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 44)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [FiLM: Visual Reasoning with a General Conditioning Layer](https://mlsanity.com/paper/1709.07871/) (2017)
- [Video Diffusion Models](https://mlsanity.com/paper/2204.03458/) (2022)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [MotionDiffuse: Text-Driven Human Motion Generation with Diffusion Model](https://mlsanity.com/paper/2208.15001/) (2022)
- [DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models](https://mlsanity.com/paper/2305.16381/) (2023)
- [Flexible Diffusion Modeling of Long Videos](https://mlsanity.com/paper/2205.11495/) (2022)
- [FreeDoM: Training-Free Energy-Guided Conditional Diffusion Model](https://mlsanity.com/paper/2303.09833/) (2023)
- [Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control](https://mlsanity.com/paper/2409.08861/) (2024)
- [INTERACTION Dataset: An INTERnational, Adversarial and Cooperative moTION Dataset in Interactive Driving Scenarios with Semantic Maps](https://mlsanity.com/paper/1910.03088/) (2019)
- [Motion Planning Diffusion: Learning and Planning of Robot Motions with Diffusion Models](https://mlsanity.com/paper/2308.01557/) (2023)
- [Manifold Preserving Guided Diffusion](https://mlsanity.com/paper/2311.16424/) (2023)
- [CFG++: Manifold-constrained Classifier Free Guidance for Diffusion Models](https://mlsanity.com/paper/2406.08070/) (2024)
- [Dynamical Regimes of Diffusion Models](https://mlsanity.com/paper/2402.18491/) (2024)
- [Diffusion-Based Planning for Autonomous Driving with Flexible Guidance](https://mlsanity.com/paper/2501.15564/) (2025)
- [TFG: Unified Training-Free Guidance for Diffusion Models](https://mlsanity.com/paper/2409.15761/) (2024)
- [Input Perturbation Reduces Exposure Bias in Diffusion Models](https://mlsanity.com/paper/2301.11706/) (2023)
- [Physics-Informed Diffusion Models](https://mlsanity.com/paper/2403.14404/) (2024)
- [Reflected Diffusion Models](https://mlsanity.com/paper/2304.04740/) (2023)
- [Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching](https://mlsanity.com/paper/2504.11713/) (2025)

## Abstract (arXiv)

Constrained generative models aim to produce samples that satisfy complex feasibility constraints while remaining faithful to the data distribution. Existing constrained generation methods typically enforce constraints either through training-time optimization or sampling-time correction. Training-time optimization approaches optimize on states induced by the training distribution, which can differ substantially from those encountered during sampling. Sampling-time correction methods instead modify the sampling process at inference, introducing distribution shift and requiring expensive tuning, particularly for few-step sampling. We propose a fine-tuning framework that incorporates constraint guidance obtained through online rollout into the training process, which aligns training with sampling by differentiating through the fixed noise schedule used to numerically integrate the denoising process. This exposes the model to violations that arise along the denoising trajectory and aligns diffusion learning with the sampling process. Experiments across multiple tasks show that our method improves constraint satisfaction while maintaining competitive sampling quality compared to prior methods.