Efficient Diffusion Models under Nonconvex Equality and Inequality constraints via Landing
This work provides a practical and scalable solution for generative modeling under physical, geometric, or safety constraints, which is important for scientific and engineering applications like molecular generation and robotics.
The paper proposes a unified framework for constrained diffusion models that handles nonconvex equality and inequality constraints via a landing mechanism, replacing costly projections. The method reduces function evaluations and memory usage while achieving comparable sample quality to state-of-the-art baselines on benchmarks with equality and mixed constraints.
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.