
\begin{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 \citep{reflected_diffusion_model} 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 Schrödinger 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. 
\end{abstract}




\section{Introduction}

Iterative refinement is key to the unprecedented success of diffusion models. They exhibit statistical efficiency \citep{Koehler_Heckett_Risteski} and reduced dimensionality dependence \citep{dim_free_doucet}, driving innovation in image, audio, video, and molecule synthesis \citep{SGMS_beat_GAN, 
%DiffWave, 
imagen_video, 3d_generation, Gaussian_SB}. However, diffusion models do not inherently guarantee optimal transport properties \citep{Lavenant_Santambrogio_22} and often result in slow inference \citep{DDPM, Progressive_distillation, DPMsolver}. Furthermore, the consistent reliance on Gaussian priors imposes limitations on the application potential and sacrifices the efficiency when the data distribution significantly deviates from the Gaussian prior.

The predominant method for fast inference originates from the field of optimal transport (OT). Notably, the (static) iterative proportional fitting (IPF) algorithm \citep{Kullback_68, IPF_95} addresses this challenge by employing alternating projections onto each marginal distribution. This algorithm has showcased impressive performance in low-dimensional contexts \citep{Chen16, Pavon_CPAM_21, Caluya21}. In contrast, the Schr\"odinger bridge (SB) problem \citep{leonard_14} introduces a principled framework for the dynamic treatment of entropy-regularized optimal transport (EOT) \citep{Villani03, 
% OT_applied_math, 
Compute_OT}. Recent advances \citep{DSB, forward_backward_SDE} have pushed the frontier of IPFs to  {(ultra-)}high-dimensional generative models {using deep neural networks (DNNs)} and have generated straighter trajectories; {Additionally, SBs based on Gaussian process \citep{SBP_max_llk} demonstrates great promise in robustness and scalability; Bridge matching methods \citep{SB_matching, Peluchetti23} also offers promising alternatives for solving complex dynamic SB problems.}


Real-world data, such as pixel values in images, often exhibits bounded support. To address this challenge, a common practice involves the use of thresholding techniques \citep{DDPM} to guide the sampling process towards the intended domain of simple structures. \citet{reflected_diffusion_model} introduced reflected diffusion models that employ reflected Brownian motion on constrained domains such as hypercubes and simplex. However, constrained domains on general Euclidean space with optimal transport guarantee are still not well developed. Moreover, \citet{reflected_diffusion_model} relies on a uniform prior based on variance-exploding (VE) SDE to derive closed-form scores, and the popular variance-preserving (VP) SDE is not fully exploited.



% \begin{figure*}[H]
% \centering
% \label{checkboard}
% % \vspace{-0.1in}
%   {\includegraphics[scale=0.5]{figures/rFB-SDE_demo_v6.png}}
% % \vspace{-0.15in}
% \caption{Constrained generative modeling via reflected forward-backward SDEs.}
% \end{figure*}

\begin{figure*}
\centering
% \vspace{-0.1in}
  {\includegraphics[scale=0.5]{figures/rFB-SDE_demo_v6.png}}
% \vspace{-0.15in}
\caption{Constrained generative modeling via reflected forward-backward SDEs.}
\label{checkboard}
\end{figure*}


To bridge this gap, we propose the \emph{Reflected Schr\"odinger Bridge} (SB) to model the transport between any smooth distributions with bounded support. We derive novel reflected forward-backward stochastic differential equations (reflected FB-SDEs) with Neumann and Robin boundary conditions and extend the divergence-based likelihood training to ensure its confinement within any smooth boundaries. We further establish connections between reflected FB-SDEs and EOT on bounded domains, where the latter facilitates the theoretical understanding by analyzing the convergence of the dual, potentials, and couplings on bounded domains. Notably, our analysis provides the first non-geometric approach to study the uniform-in-time stability w.r.t. the marginals and is noteworthy in its own right. We empirically validate our algorithm on 2D examples and standard image benchmarks, showcasing its promising performance in generative modeling over constrained domains. The flexible choices on the priors allow us to choose freely between VP-SDE and VE-SDE.


\section{Related Works}
\label{related_works}
% \paragraph{Constrained Sampling} \citet{sebastien_bubeck} studied the convergence of Langevin Monte Carlo on bounded domains and showed a polynomial sample time from a log-concave distribution, which is later extended to non-convex settings \citep{Andrew_Lamperski_21_COLT}. \citet{Sampling_Hamiltonian} demonstrated the constrained sampling with ill-conditioned and non-smooth distributions using Hamiltonian Monte Carlo. 

\paragraph{Constrained Sampling} \citet{sebastien_bubeck} studied the convergence of Langevin Monte Carlo within bounded domains. His work revealed a polynomial sample time for log-concave distributions, which is later extended to non-convex settings by \citet{Andrew_Lamperski_21_COLT}. %reflected_non_convex 
Furthermore, the exploration of constrained sampling in challenging scenarios with ill-conditioned and non-smooth distributions was explored by \citet{Sampling_Hamiltonian}, who leveraged Hamiltonian Monte Carlo techniques. Other constrained sampling works include proximal Langevin dynamics \citep{Proximal_LMC} and mirrored Langevin dynamics \citep{mirror_LMC}. 

\paragraph{Constrained Generation} \citet{Riemannian_score_generative, Riemannian_Diffusion} studied the extension of diffusion models on Riemannian manifolds, and the convergence is further analyzed by \citet{diffusion_manifold}. This groundwork subsequently motivated follow-up research, including implicit score-matching loss via log-barrier methods and reflected Brownian motion \citep{Diffusion_constrained} and Schr\"odinger bridge \citep{Riemannian_diffusion_SB} on the Riemannian manifold. Alternatively, drawing inspiration from the popular thresholding technique in real-world diffusion applications, \citet{reflected_diffusion_model} proposed to train explicit score-matching loss based on reflected Brownian motion, which demonstrated compelling empirical performance. Mirror diffusion models \citep{Mirror_Diffusion} studied constrained generation on convex sets and found interesting applications in watermarked generations. \citet{diffusion_bridge_qiang} employed Doob's h-transform to learn diffusion bridges on various constrained domains. The study of reflected Schr\"odinger bridge was initiated by \citet{Caluya_reflected_SB} in the control community and has shown remarkable performance in low-dimensional problems.
