Efficient Diffusion Models under Nonconvex Equality and Inequality constraints via Landing

Kijung Jeon, Michael Muehlebach, Molei Tao

Keywords: constrained diffusion models, nonconvex manifold, underdamped Langevin dynamics, landing mechanism

TL;DR: We propose an unified constrained diffusion framework on nonconvex sets that enforces equality and inequality constraints via a projection-free landing step and exploits faster-mixing underdamped dynamics, enabling efficient inference and training.

Abstract: The generative modeling of data in constrained sets is central to scientific and engineering applications with physical, geometric, or safety constraints (e.g., molecular generation, robotics). This article constructs constrained diffusion models on a generic nonconvex feasible sets $\Sigma$, by introducing a unified framework that simultaneously enforces both equality and inequality constraints throughout the diffusion process. Our theory and implementations encompass both overdamped and underdamped dynamics for the forward and backward sampling. The key algorithmic ingredient is a computationally efficient landing mechanism that replaces costly and not-always-well-defined projections onto $\Sigma$, maintaining feasibility without Newton solves and avoiding projection failures. Leveraging underdamped dynamics whose faster mixing reduces the steps needed to reach the prior distribution, the commonly-believed unavoidable heavy forward simulation cost in the constrained diffusion is alleviated. Empirically, this reduces function evaluations, enabling more efficient inference and training while preserving sample quality and substantially lowering memory usage. On equality-only and mixed (equality and inequality) benchmarks, our method shows reasonable sample quality, while substantially reducing computational cost and function evaluations. These results indicate that landing-based enforcement combined with underdamped dynamics provides a practical and scalable recipe for constrained diffusion on nonconvex feasible sets.

Primary Area: generative models

Submission Number: 23119