Metropolis Sampling for Constrained Diffusion Models
This work addresses the need for efficient and flexible constrained generative modelling in fields such as natural sciences and robotics, representing an incremental improvement over prior methods.
The paper tackles the problem of applying denoising diffusion models to domains with arbitrary constraints, which existing methods either handle inefficiently or only for convex subsets. It introduces a Metropolis sampling-based noising scheme that achieves substantial gains in computational efficiency and empirical performance, demonstrated across applications like geospatial modelling, robotics, and protein design.
Denoising diffusion models have recently emerged as the predominant paradigm for generative modelling on image domains. In addition, their extension to Riemannian manifolds has facilitated a range of applications across the natural sciences. While many of these problems stand to benefit from the ability to specify arbitrary, domain-informed constraints, this setting is not covered by the existing (Riemannian) diffusion model methodology. Recent work has attempted to address this issue by constructing novel noising processes based on the reflected Brownian motion and logarithmic barrier methods. However, the associated samplers are either computationally burdensome or only apply to convex subsets of Euclidean space. In this paper, we introduce an alternative, simple noising scheme based on Metropolis sampling that affords substantial gains in computational efficiency and empirical performance compared to the earlier samplers. Of independent interest, we prove that this new process corresponds to a valid discretisation of the reflected Brownian motion. We demonstrate the scalability and flexibility of our approach on a range of problem settings with convex and non-convex constraints, including applications from geospatial modelling, robotics and protein design.