Efficient Diffusion Models under Nonconvex Equality and Inequality constraints via Landing

TL;DR

This paper introduces a unified, efficient framework for constrained diffusion models that enforces complex constraints throughout the process, reducing computational costs while maintaining high sample quality.

Contribution

The paper proposes a landing mechanism and underdamped dynamics to efficiently handle nonconvex constraints in diffusion models, avoiding costly projections and accelerating mixing.

Findings

Reduces function evaluations and memory usage during training and inference.

Achieves comparable sample quality to state-of-the-art methods on benchmark tasks.

Significantly lowers computational costs for constrained diffusion modeling.

Abstract

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 $\Sigma$ 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 $\Sigma$, 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.