Table of Contents
Fetching ...

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.

Score-based constrained generative modeling via Langevin diffusions with boundary conditions

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.

Paper Structure

This paper contains 43 sections, 3 theorems, 73 equations, 22 figures, 8 tables.

Key Result

Proposition 4.1

Let $s_{\theta} \in C^{1}([0,T] \times \mathbb{R}^{d} ; \mathbb{R}^{d})$. Then, the score matching loss function for CLD (cld_pos_for)-(cld_vel_for) is given by: where $C >0$ is a constant and $\mathcal{L}_{t}^{(X,V)}$ is the time-marginal of joint law of $(X,V)$ from (cld_pos_for)-(cld_vel_for).

Figures (22)

  • Figure 1: A simple comparison : Original data; Data generated by unconstrained diffusion model; Data generated by projected diffusion model.
  • Figure 2: Trajectory of $X_t$ of confined Langevin dynamics \ref{['cld_pos_for']}-\ref{['cld_vel_for']}.
  • Figure 3: Trajectory of $X_t$ of reflected (overdamped) Langevin dynamics.
  • Figure 4: Collision step if $\mathcal{Y}_1 \notin \bar{G}$ i.e. if auxiliary step is outside domain $G$.
  • Figure 5: Description of [A$_c$OA$_c$] (for $b \equiv 0$) and [CBBK] (for $b = -x$).
  • ...and 17 more figures

Theorems & Definitions (6)

  • Proposition 4.1
  • Proposition 4.2
  • Proposition 4.3
  • proof : Proof of Proposition \ref{['prop_cld']}
  • proof : Proof of Proposition \ref{['ref_lossProp_2']}
  • proof : Proof of Proposition \ref{['ref_lossProp_3']}