Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper introduces a fine-tuning framework that integrates constraint guidance through online rollouts into the training of constrained diffusion models, addressing the misalignment between training-time optimization and sampling-time correction. By differentiating through the fixed noise schedule used in the denoising process, the method exposes the model to constraint violations encountered during sampling, thereby enhancing the model's ability to satisfy complex feasibility constraints. Experimental results demonstrate that this approach significantly improves constraint satisfaction while preserving competitive sampling quality compared to existing methods.
Training diffusion models with online rollouts leads to better constraint satisfaction without sacrificing sample quality.
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.