Generative modeling within constrained sets is essential for scientific and engineering applications involving physical, geometric, or safety requirements (e.g., molecular generation, robotics). We present a unified framework for constrained diffusion models on generic nonconvex feasible sets $Σ$ that simultaneously enforces equality and inequality constraints throughout the diffusion process. Our framework incorporates both overdamped and underdamped dynamics for forward and backward sampling. A key algorithmic innovation is a computationally efficient landing mechanism that replaces costly and often ill-defined projections onto $Σ$, ensuring feasibility without iterative Newton solves or projection failures. By leveraging underdamped dynamics, we accelerate mixing toward the prior distribution, effectively alleviating the high simulation costs typically associated with constrained diffusion. Empirically, this approach reduces function evaluations and memory usage during both training and inference while preserving sample quality. On benchmarks featuring equality and mixed constraints, our method achieves comparable sample quality to state-of-the-art baselines while significantly reducing computational cost, providing a practical and scalable solution for diffusion on nonconvex feasible sets.
翻译:在受约束集合内进行生成建模对于涉及物理、几何或安全要求(例如分子生成、机器人技术)的科学与工程应用至关重要。我们提出了一个在一般非凸可行集 $Σ$ 上,同时在整个扩散过程中强制执行等式与不等式约束的约束扩散模型统一框架。该框架综合了过阻尼与欠阻尼动力学以实现前向与反向采样。一项关键的算法创新是计算高效的降落机制,该机制替代了对 $Σ$ 进行昂贵且常定义模糊的投影操作,从而避免了迭代牛顿求解或投影失败,确保了可行性。通过利用欠阻尼动力学,我们加速了向先验分布的混合过程,有效缓解了约束扩散通常伴随的高昂模拟成本。实验表明,该方法在保持样本质量的同时,减少了训练和推理期间的函数评估次数与内存使用。在涉及等式与混合约束的基准测试中,我们的方法在显著降低计算成本的同时,实现了与最先进基线相当的样本质量,为非凸可行集上的扩散提供了一种实用且可扩展的解决方案。