让扩散模型在训练时就学会满足约束,提升生成效果。
Integration Matters: Rollout-Based Training for Constrained Diffusion Models

- 训练时引入在线滚动优化,让模型提前感知生成过程中的约束违反
- 在多个任务上约束满足率显著提升,且采样质量不下降
- 适合需要严格满足规则的生成场景,如科学模拟或安全应用
约束生成模型旨在生成符合复杂可行性约束的同时保持数据分布忠实性。现有方法通常通过训练时优化或采样时修正来施加约束。训练时优化针对训练分布诱导的状态,与采样时实际遇到的状态差异较大;采样时修正则引入分布偏移,需昂贵调参,尤其在少步采样下更明显。本文提出一种微调框架,将通过在线滚动获得的约束引导融入训练过程,通过对固定噪声调度的反向传播,使训练与采样对齐,暴露模型于去噪轨迹中出现的违反情况,从而实现扩散学习与采样过程的一致性。多任务实验表明,该方法在保持竞争性采样质量的同时,显著提升了约束满足率。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。