Integration Matters: Rollout-Based Training for Constrained Diffusion Models

Constrained generative models aim to produce samples that satisfy complex feasibility constraints while remaining faithful to the underlying data distribution. Existing methods typically enforce constraints through either training-time optimization or sampling-time correction. Training-time optimization approaches optimize over 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 costly tuning, particularly in few-step sampling regimes. We propose IMCD (Integration Matters for Constrained Diffusion), a fine-tuning framework that integrates constraint guidance into training through online rollouts. IMCD aligns training with sampling by differentiating through the fixed noise schedule used to numerically integrate the denoising process. This exposes the model to constraint 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.