Score-based constrained generative modeling via Langevin diffusions with boundary conditions
Adam Nordenhög, Akash Sharma
TL;DR
The paper tackles the challenge of enforcing hard domain constraints in score-based diffusion models by introducing a constrained generative framework based on confined Langevin dynamics with specular boundary reflection. It derives the time-reversed dynamics within the constrained domain and develops multiple sampling schemes (projection, symmetrized reflection, penalty, barrier) to preserve constraints, contrasting them with local-time-based reflected SDEs. A corrected score-matching loss that accounts for boundary terms is proposed, enabling tractable training without forcing the score to vanish on the boundary. Empirical results on Gaussian mixtures, maze-like datasets, Flower, and MNIST demonstrate that the proposed constrained diffusion approaches can maintain constraint satisfaction while delivering competitive or superior sample quality compared to unconstrained baselines and other constrained methods.
Abstract
Score-based generative models based on stochastic differential equations (SDEs) achieve impressive performance in sampling from unknown distributions, but often fail to satisfy underlying constraints. We propose a constrained generative model using kinetic (underdamped) Langevin dynamics with specular reflection of velocity on the boundary defining constraints. This results in piecewise continuously differentiable noising and denoising process where the latter is characterized by a time-reversed dynamics restricted to a domain with boundary due to specular boundary condition. In addition, we also contribute to existing reflected SDEs based constrained generative models, where the stochastic dynamics is restricted through an abstract local time term. By presenting efficient numerical samplers which converge with optimal rate in terms of discretizations step, we provide a comprehensive comparison of models based on confined (specularly reflected kinetic) Langevin diffusion with models based on reflected diffusion with local time.
