Integration Matters: Rollout-Based Training for Constrained Diffusion Models

2026-07-15

AI Summary

Existing diffusion models often suffer from low constraint satisfaction rates or degraded sample quality when generating samples under complex feasibility constraints, primarily due to distributional mismatches between training and sampling phases. This work proposes a trajectory-aware fine-tuning framework based on online rollouts, which integrates constraint guidance during training and, for the first time, incorporates a numerical integration perspective into the diffusion process. By end-to-end differentiating denoising trajectories under a fixed noise schedule, the method explicitly exposes constraint violations, thereby aligning the training and sampling distributions. Combining constraint-aware guidance, differentiable noise scheduling, and an online rollout mechanism, the approach significantly improves constraint satisfaction across multiple tasks while maintaining generation quality on par with current state-of-the-art methods.

Abstract

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.

Problem

• constrained generation
• diffusion models
• distribution shift
• sampling-time correction
• training-time optimization

Innovation

• rollout-based training
• constrained diffusion models
• differentiable integration
• constraint satisfaction
• denoising trajectory alignment