Reflected Schrödinger Bridge for Constrained Generative Modeling

TL;DR

This work presents Reflected Schrödinger Bridge, a principled framework for constrained generative modeling on bounded domains. By deriving reflected forward-backward SDEs with Neumann and Robin boundary conditions and extending likelihood training to bounded domains, the authors integrate entropy-regularized OT with reflection to enforce boundary effects without ad-hoc thresholding. They connect dynamic and static IPF formulations, provide duality-based convergence guarantees under bounded domains, and demonstrate empirical effectiveness on 2D shapes and image benchmarks, including simplex domains. The approach offers scalable, boundary-aware generative modeling with OT guarantees beyond traditional hypercube constraints, and suggests practical benefits in reducing generation-time costs via OT-informed training.

Abstract

Diffusion models have become the go-to method for large-scale generative models in real-world applications. These applications often involve data distributions confined within bounded domains, typically requiring ad-hoc thresholding techniques for boundary enforcement. Reflected diffusion models (Lou23) aim to enhance generalizability by generating the data distribution through a backward process governed by reflected Brownian motion. However, reflected diffusion models may not easily adapt to diverse domains without the derivation of proper diffeomorphic mappings and do not guarantee optimal transport properties. To overcome these limitations, we introduce the Reflected Schrodinger Bridge algorithm: an entropy-regularized optimal transport approach tailored for generating data within diverse bounded domains. We derive elegant reflected forward-backward stochastic differential equations with Neumann and Robin boundary conditions, extend divergence-based likelihood training to bounded domains, and explore natural connections to entropic optimal transport for the study of approximate linear convergence - a valuable insight for practical training. Our algorithm yields robust generative modeling in diverse domains, and its scalability is demonstrated in real-world constrained generative modeling through standard image benchmarks.

Figures (11)

• Figure 1: Constrained generative modeling via reflected forward-backward SDEs.
• Figure 2: Reflected OU processes (reflected v.s. unconstrained), driven by the same Brownian motion, excluding the reflections. All boundary curves have properly defined unit vectors.
• Figure 3: IPF v.s. aIPF. The approximate (or exact) projections are highlighted through the dotted (or solid) lines.
• Figure 4: Demo of generative samples (top) and vector fields (bottom) based on Reflected SB.
• Figure 5: Samples via reflected SB on MNIST (left), CIFAR10 (middle), and ImageNet 64 (right).

Theorems & Definitions (18)

• Theorem 1
• Proposition 1: Feynman-Kac representation
• Lemma 1: Duality Nutz22_note
• Lemma 2: Convergence of the Dual and Potentials
• Theorem 2: Convergence of Static Couplings
• Proposition 2: Convergence of Dynamic Couplings
• Remark 1
• Proposition 3: Probability Flow ODE
• Lemma 3
• proof