Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Diffusion Models for Constrained Domains

Nic Fishman, Leo Klarner, Valentin De Bortoli, Emile Mathieu, Michael Hutchinson

TL;DR

This work presents two distinct noising processes based on (i) the logarithmic barrier metric and (ii) the reflected Brownian motion induced by the constraints, and derives new tools to define such models in this framework.

Abstract

Denoising diffusion models are a novel class of generative algorithms that achieve state-of-the-art performance across a range of domains, including image generation and text-to-image tasks. Building on this success, diffusion models have recently been extended to the Riemannian manifold setting, broadening their applicability to a range of problems from the natural and engineering sciences. However, these Riemannian diffusion models are built on the assumption that their forward and backward processes are well-defined for all times, preventing them from being applied to an important set of tasks that consider manifolds defined via a set of inequality constraints. In this work, we introduce a principled framework to bridge this gap. We present two distinct noising processes based on (i) the logarithmic barrier metric and (ii) the reflected Brownian motion induced by the constraints. As existing diffusion model techniques cannot be applied in this setting, we derive new tools to define such models in our framework. We then demonstrate the practical utility of our methods on a number of synthetic and real-world tasks, including applications from robotics and protein design.

Diffusion Models for Constrained Domains

TL;DR

This work presents two distinct noising processes based on (i) the logarithmic barrier metric and (ii) the reflected Brownian motion induced by the constraints, and derives new tools to define such models in this framework.

Abstract

Denoising diffusion models are a novel class of generative algorithms that achieve state-of-the-art performance across a range of domains, including image generation and text-to-image tasks. Building on this success, diffusion models have recently been extended to the Riemannian manifold setting, broadening their applicability to a range of problems from the natural and engineering sciences. However, these Riemannian diffusion models are built on the assumption that their forward and backward processes are well-defined for all times, preventing them from being applied to an important set of tasks that consider manifolds defined via a set of inequality constraints. In this work, we introduce a principled framework to bridge this gap. We present two distinct noising processes based on (i) the logarithmic barrier metric and (ii) the reflected Brownian motion induced by the constraints. As existing diffusion model techniques cannot be applied in this setting, we derive new tools to define such models in our framework. We then demonstrate the practical utility of our methods on a number of synthetic and real-world tasks, including applications from robotics and protein design.
Paper Structure (15 sections, 2 theorems, 15 equations, 6 figures, 3 algorithms)

This paper contains 15 sections, 2 theorems, 15 equations, 6 figures, 3 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Riemannian manifolds.
  4. Constrained manifolds.
  5. Continuous diffusion models.
  6. Inequality-Constrained Diffusion Models
  7. Log-barrier diffusion models
  8. Barrier Langevin dynamics.
  9. Time-reversal.
  10. Sampling.
  11. Reflected diffusion models
  12. Skorokhod problem.
  13. Time-reversal.
  14. Sampling.
  15. Likelihood evaluation.

Key Result

Proposition 3.1

For any $t > 0$, the distribution of $\bar{\mathbf{B}}_t$ admits a density w.r.t. the Lebesgue measure denoted $p_t$. In addition, we have for any $x \in \mathrm{int}(\mathcal{M})$ and $x_0 \in \partial \mathcal{M}$ where we recall that $\mathbf{n}$ is the outward normal to $\mathcal{M}$.

Figures (6)

  • Figure 1: The behaviour of different types of noising processes considered in this work defined on the unit interval. $\mathbf{B}_t$: Unconstrained (Euclidean) Brownian motion. $\hat{\mathbf{B}}_t$: Log-barrier forward noising process. $\bar{\mathbf{B}}_t$: Reflected Brownian motion. All sampled with the same initial point and driving noise. Black line indicates the boundary.
  • Figure 2: A convex polytope defined by six constraints $\{f_i\}_{i \in \c{I}}$, along with the log barrier potential, and 'straight trajectories' under the log-barrier metric and under the Euclidean metric with and without reflection at the boundary.
  • Figure 3: Convergence of the Barrier Langevin dynamics on the unit interval to the uniform distribution.
  • Figure 4: Illustrative diagram of the barrier method and the change of metric.
  • Figure 5: Left: Convergence of the reflected Brownian motion on the unit interval to the uniform distribution. Right: Value of $\abs{\mathbf{k}}_t$ for the trajectory samples on the left through time.
  • ...and 1 more figures

Theorems & Definitions (2)

  • Proposition 3.1: burdzy2004heat
  • Theorem 3.2