Reflected Schrödinger Bridge Matching

Conference: NeurIPS 2026
arXiv: 2607.03626
OpenReview: https://openreview.net/forum?id=OO3amJIfRg
Code: https://github.com/viktor765/rsbm
Area: Optimization & Theory
Keywords: Schrödinger bridge, reflected diffusion, bridge matching, constrained domain, entropic optimal transport

TL;DR

This paper introduces Reflected Schrödinger Bridge Matching (RSBM), extending the partially simulation-free iterative Markovian fitting framework to reflected Brownian motion on the unit hypercube through coordinate-separable score approximations and a two-step exact bridge midpoint sampler, strictly eliminating out-of-bounds artifacts in bounded image generation with negligible computational overhead.

Background & Motivation

Iterative refinement models, spanning diffusion models and flow matching, have driven rapid advancements in high-dimensional generative modeling, with rectified flows further straightening trajectory paths toward optimal transport. However, deterministic straight flows lack finite-step optimal transport guarantees, whereas Schrödinger Bridges (SB)—the dynamical formulation of entropic optimal transport—provide rigorous measure-preserving couplings between arbitrary empirical distributions. Recent breakthroughs such as Diffusion Schrödinger Bridge Matching (DSBM) and iterative Markovian fitting (IMF / \(\alpha\)-IMF) have significantly accelerated bridge computation by leveraging Doob's \(h\)-transforms and regression losses akin to flow matching, effectively avoiding costly forward-backward SDE simulations in unconstrained Euclidean domains.

Despite these advances, real-world data distributions frequently reside on compact, bounded domains—most notably digital images, whose pixel values are strictly bounded within \([0, 1]^d\). Conventional diffusion and Schrödinger bridge models operating in unconstrained spaces inevitably leak probability mass outside the valid domain during inference, forcing practitioners to resort to heuristic post-hoc clipping that corrupts theoretical consistency and causes severe over-saturation and boundary distortion. Prior attempts to formulate reflected Schrödinger bridges relied on forward-backward SDE (FBSDE) formulations, which require simulating full trajectories for every gradient step and backpropagating through expensive score divergence terms, making them computationally intractable for high-dimensional image-to-image translation.

The fundamental tension lies in enforcing hard boundary constraints via reflected stochastic dynamics while preserving the simulation-free training efficiency of flow matching. This paper resolves this dilemma by restricting the reference process to reflected Brownian motion on the unit hypercube and rigorously establishing the Markovian projection in the reflected path space. Core idea: generalize iterative Markovian fitting to reflected stochastic differential equations on the unit hypercube, exploiting coordinate independence to derive a closed-form truncated Gaussian-sum drift target and introducing an exact two-step bridge midpoint sampler, enabling reflected Schrödinger bridges to train as efficiently as standard flow matching.

Method

Overall Architecture

Reflected Schrödinger Bridge Matching (RSBM) computes the path measure \(\mathbb{P}^*\) that minimizes the Kullback-Leibler divergence with respect to a reflected Brownian motion reference process \(\mathbb{Q}^r\) constrained on \(\mathcal{D} = [0, 1]^d\), subject to boundary marginals \(\pi_0\) and \(\pi_1\). The framework operates via alternating reciprocal and Markovian projections under the continuous online training scheme (\(\alpha\)-RSBM): reciprocal projections draw paired bridge samples from the current Markovian coupling, while Markovian projections train a neural network to approximate the analytical drift adjustment of the reflected reverse and forward SDEs.

The entire pipeline comprises endpoint coupling generation, exact two-step bridge midpoint sampling, bidirectional regression of the analytical score drift, and boundary-preserving Euler-Maruyama numerical integration:

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Marginal Distributions π₀ and π₁<br/>Strictly supported on [0, 1]ᵈ"] --> B["Reflected Reference Process & Markovian Projection<br/>Define RBM reference measure and establish KL minimization"]
    B --> C["Two-Step Reflected Bridge Midpoint Sampling<br/>Sample unreflected endpoint Y₁ + Gaussian bridge + Reflection R"]
    C --> D["Coordinate-Separable Analytical Score Regression<br/>Independent 1D truncated Gaussian sums for drift target v*"]
    D --> E["Bidirectional Online Training & Reflected Numerical SDE<br/>α-RSBM single unified network + Reflected Euler-Maruyama"]
    E --> F["Output: Boundary-preserving image generation without clipping"]

Key Designs

1. Reflected Reference Process & Markovian Projection: Extending Schrödinger Bridges to Compact Domains with Normal Reflection

Standard diffusion and bridge models assume unbounded Brownian motion, leading to boundary leakage when dealing with bounded supports. RSBM formulates the reference measure \(\mathbb{Q}^r\) as a reflected stochastic differential equation governed by boundary local time: $\(d X_t = \mu(t, X_t) dt + \sigma(t) d W_t + n(X_t) d L_t\)$ where \(n(x)\) denotes the inward-pointing unit normal at the boundary \(\partial \mathcal{D}\), and the continuous scalar process \(L_t\) acts only when \(X_t \in \partial \mathcal{D}\) to prevent escaping the domain. The authors prove a fundamental theoretical result (Proposition 4.1): for any bridge mixture \(\Pi = \Pi_{0,1} \mathbb{Q}^r_{\mid 0,1}\), its Markovian projection \(\mathbb{M}^* = \operatorname{proj}_{\mathcal{M}^r}(\Pi)\) preserves all marginal distributions \(\mathbb{M}^*_t = \Pi_t\) for all \(t \in [0, 1]\), and uniquely minimizes the path-space KL divergence: $\(\operatorname{D_{KL}}(\Pi \mid\mid \mathbb{M}^*) = \frac{1}{2} \int_0^1 \frac{1}{\sigma^2(t)} \mathbb{E}_{\Pi_{0,t}} \left[ \left\| \sigma^2(t) \mathbb{E}_{\Pi_{1 \mid 0, t}} [\nabla \log \mathbb{Q}^r_{1 \mid t}(X_1 \mid X_t) \mid X_0, X_t] - v^*(t, X_t) \right\|^2 \right] dt\)$ By applying Girsanov's theorem under finite entropy conditions, this result circumvents the notoriously difficult task of characterizing the domain of the infinitesimal generator on bounded manifolds, providing a rigorous mathematical foundation for training reflected drift fields via least-squares regression.

2. Two-Step Reflected Bridge Midpoint Sampling: Exact and Computationally Free Bridge Interpolation via Inverse Image Folding

In the reciprocal projection step, intermediate bridge states must be sampled conditionally on endpoint pairs: \(X_t \sim \mathbb{Q}^r_{t \mid 0, 1}(\cdot \mid X_0=x_0, X_1=x_1)\). Simulating an SDE with boundary reflection between endpoints is computationally prohibitive and introduces numerical time-discretization errors. The authors exploit the trajectory equivalence between reflected Brownian motion \(X_t\) and the folded unconstrained Brownian motion \(R(W_t)\), where \(R\) is the hypercube reflection operator. This yields an exact two-step bridge sampler (Algorithm 2): First, Bayes' rule is applied over the discrete preimage set \(\mathcal{I}(x_1) = \{y \in \mathbb{R}^d \mid R(y) = x_1\}\) to sample an unfolded virtual endpoint \(Y_1 \sim \sum_{y_1 \in \mathcal{I}(x_1)} \mathbb{Q}_{1 \mid 0}(y_1 \mid x_0) \delta_{y_1}\); second, a standard unconstrained Gaussian bridge is sampled between \(x_0\) and \(Y_1\) via \(Y_t \sim \mathcal{N}(t Y_1 + (1-t) x_0, t(1-t)\sigma^2 I_d)\); finally, the true bridge point is recovered by applying the reflection operator \(X_t = R(Y_t)\). This closed-form scheme eliminates all intermediate numerical SDE stepping, guaranteeing zero sampling bias while matching the execution speed of standard Gaussian bridges.

3. Coordinate-Separable Analytical Score Regression: Mitigating the Curse of Dimensionality with Truncated Gaussian Sums

The regression target requires evaluating the score of the reflected transition density \(\nabla \log \mathbb{Q}^r_{1 \mid t}(x_1 \mid x_t)\). While in free space this degenerates to a simple linear displacement \((x_1 - x_t) / (1-t)\), reflection across hypercube boundaries yields an infinite sum of mirrored Gaussian modes. In arbitrary domains, expanding this sum requires \(K^d\) terms, which is computationally intractable for high-dimensional images. RSBM observes that for orthogonal hypercubes under independent Brownian increments, the transition kernel factorizes completely across dimensions: $\(\mathbb{Q}^r_{t \mid s}(x_t \mid x_s) = \prod_{i=1}^d \mathbb{Q}^{r, i}_{t \mid s}(x_t^i \mid x_s^i) = \prod_{i=1}^d \sum_{y_t^i \in \mathcal{I}(x_t^i)} \mathbb{Q}^i_{t \mid s}(y_t^i \mid x_s^i)\)$ Consequently, the gradient of the log-density decouples into independent 1D scalar calculations: $\((\nabla_{x_s} \log \mathbb{Q}^r_{t \mid s}(x_t \mid x_s))_i = \frac{\sum_{y_t^i \in \mathcal{I}(x_t^i)} \frac{d}{d x_s^i} \mathbb{Q}^i_{t \mid s}(y_t^i \mid x_s^i)}{\sum_{y_t^i \in \mathcal{I}(x_t^i)} \mathbb{Q}^i_{t \mid s}(y_t^i \mid x_s^i)}\)$ Truncating the 1D sum to 11 reflection images per dimension achieves near-perfect numerical precision, collapsing the computational complexity from exponential down to strictly linear \(\mathcal{O}(K \cdot d)\) without requiring spatial discretization or score divergence calculations.

4. Bidirectional Online Training & Reflected Numerical SDE: Unified Parameterization and Boundary Integration

To prevent the compounding discretization errors of alternating forward and backward model refitting in standard IMF, RSBM adapts the \(\alpha\)-IMF online update protocol. A single shared U-Net backbone \(v_\theta(t, X_t; \operatorname{dir})\) is trained simultaneously across both forward and backward time flows. A directional embedding is concatenated with the sinusoidal time embedding and injected via FiLM modulation into the GroupNorm layers across residual blocks. During inference, sample trajectories are evolved via Euler-Maruyama discretization, where coordinates exceeding boundaries are folded back into \([0, 1]\) via the modulo reflection operator \(R(y)\) at each step, with a hard projection applied exclusively at the final step \(t=1\) to eliminate floating-point truncation drift.

Loss & Training

The model is trained end-to-end using the bidirectional symmetrized regression objective: $\(\mathcal{L}(\theta; \Pi) = \frac{1}{2} \left( \mathcal{L}_\to(\theta; \Pi) + \mathcal{L}_\leftarrow(\theta; \Pi) \right)\)$ where the forward drift loss is given by: $\(\mathcal{L}_\to(\theta; \Pi) = \int_0^1 \frac{1}{\sigma^2(t)} \mathbb{E}_{\Pi_{t, 1}} \left[ \left\| \sigma^2(t) \nabla \log \mathbb{Q}^r_{1 \mid t}(X_1 \mid X_t) - v_\theta(t, X_t; \to) \right\|^2 \right] dt\)$ and \(\mathcal{L}_\leftarrow\) is defined symmetrically by reversing the temporal direction. Training proceeds across two phases: 1. Pretraining: Endpoints are drawn independently from the marginal product \((X_0, X_1) \sim \pi_0 \otimes \pi_1\) to establish an initial transport vector field (50k steps for MNIST \(\leftrightarrow\) EMNIST; 75k steps for AFHQ). 2. Finetuning: At each online optimization step, approximate endpoints \(\hat{X}_1\) and \(\hat{X}_0\) are sampled using the current exponential moving average (EMA, decay rate 0.999) model, followed by single-step gradient updates on the empirical bridge loss (100k steps for MNIST \(\leftrightarrow\) EMNIST; 20k steps for AFHQ). The Adam optimizer is employed with learning rates of \(1 \times 10^{-4}\) and \(2 \times 10^{-4}\) and unit gradient clipping.

Key Experimental Results

Main Results

The method was evaluated on two high-dimensional unpaired image-to-image translation benchmarks: MNIST ↔ EMNIST (\(28 \times 28\) grayscale digits and letters) and AFHQv2 cat ↔ wild (\(64 \times 64\) RGB animal faces). The primary baseline is the unconstrained counterpart \(\alpha\)-DSBM, whose generated samples require explicit clipping to compute valid evaluation metrics. Quantitative metrics include Fréchet Inception Distance (FID) for generative fidelity, per-dimension Mean Square Distance (MSD) for source-target correspondence, and Learned Perceptual Image Patch Similarity (LPIPS) in AlexNet feature space.

Translation Task	Model	MSD ↓	LPIPS ↓	FID ↓	Wall-clock Time (GPU)
MNIST → EMNIST	\(\alpha\)-DSBM (clipped)	0.416	—	6.41	31.5 h (1× A100)
	\(\alpha\)-RSBM (Ours)	0.356	—	5.22	32.0 h (1× A100)
EMNIST → MNIST	\(\alpha\)-DSBM (clipped)	0.415	—	8.58	31.5 h (1× A100)
	\(\alpha\)-RSBM (Ours)	0.362	—	8.23	32.0 h (1× A100)
AFHQ: cat → wild	\(\alpha\)-DSBM (clipped)	0.131	0.239	27.6	108.0 h (4× A100)
	\(\alpha\)-RSBM (Ours)	0.084	0.248	26.1	109.0 h (4× A100)
AFHQ: wild → cat	\(\alpha\)-DSBM (clipped)	0.200	0.302	31.5	108.0 h (4× A100)
	\(\alpha\)-RSBM (Ours)	0.082	0.244	33.8	109.0 h (4× A100)

Out-of-Bounds Pixel Statistics

To quantify the severity of domain leakage in unconstrained models, the authors evaluated the violation statistics of \(\alpha\)-DSBM prior to clipping:

Dataset	Violating Samples (%)	Mean Violating Pixels (%)	Conditional Violation Magnitude	Max Image Violation
MNIST / EMNIST	100.0%	34% – 46%	0.002 – 0.003	0.008 – 0.013
AFHQ (64×64)	~30.0% (>1% violation)	~2.0% (all channels)	0.020	0.030 (mean max)

Key Findings

• Superior Trajectory Correspondence: Across all translation settings, \(\alpha\)-RSBM substantially reduces the Mean Square Distance (MSD) compared to \(\alpha\)-DSBM (e.g., dropping from 0.200 to 0.082 on AFHQ wild \(\to\) cat, a 59% reduction), demonstrating that reflected dynamics yield significantly tighter transport paths with superior background and color consistency.
• Virtually Negligible Computational Overhead: Unlike prior FBSDE-based reflected bridges that require full-path trajectory caches and score divergence backpropagation, \(\alpha\)-RSBM introduces virtually zero runtime penalty over standard \(\alpha\)-DSBM (32.0h vs 31.5h on MNIST; 109h vs 108h on AFHQ, a <1.5% wall-clock increase).
• Complete Elimination of Boundary Leakage: While 100% of samples from the unconstrained baseline suffer from out-of-bounds pixel artifacts on MNIST/EMNIST and 15% exceed a 5% channel violation threshold on AFHQ, \(\alpha\)-RSBM inherently guarantees that all generated samples reside strictly within the physical support domain \([0, 1]^d\).

Highlights & Insights

• Bridging Reflected Dynamics and Flow Matching: Elegantly bridges bounded stochastic differential equations with modern simulation-free regression by expressing boundary interactions as expectations over discrete mirror preimages.
• Dimensional Factorization Beats Curse of Dimensionality: Overcomes exponential complexity on the unit hypercube by exploiting the orthogonal separability of the reflected Gaussian kernel, turning high-dimensional score computation into 1D scalar evaluations.
• Optimal Balance between Theoretical Rigor and Scalability: Avoids both the heuristic boundary hacking of unconstrained flows and the computational bottlenecks of FBSDE formulations, delivering an industrial-grade framework for constrained optimal transport.

Limitations & Future Work

• Zero-Drift Reference Process: The current theoretical proof and closed-form score calculation require a zero-drift reflected Brownian motion, precluding variance-preserving (VP) SDEs commonly employed in standard diffusion models.
• Hypercube Domain Restriction: The analytical tractability relies heavily on the orthogonal coordinate separability of \([0, 1]^d\). Extending to non-orthogonal domains (such as probability simplices or spherical manifolds) breaks coordinate independence, requiring neural bridge approximations or PDE numerical solvers.
• Future Directions: Exploring diffeomorphic coordinate transformations to map convex geometries to the hypercube, and developing adaptive step-size ODE/SDE solvers for faster generation.

Related Work & Insights

• vs. Unconstrained Schrödinger Bridge Matching (DSBM / \(\alpha\)-DSBM): DSBM (Shi et al., 2023) and \(\alpha\)-DSBM (De Bortoli et al., 2024) operate only in unconstrained Euclidean space, producing severe out-of-bounds artifacts that require heuristic post-hoc clipping; RSBM natively preserves boundary constraints with full mathematical rigor.
• vs. Forward-Backward Reflected Bridges (Deng et al., 2024): While Deng et al. handle more general geometries via FBSDE theory, their method demands full-path simulation and divergence computations per iteration; RSBM offers a vastly more scalable alternative with negligible overhead for hypercube domains.
• vs. Reflected Diffusion Models (Lou & Ermon, 2023): Lou & Ermon restricted unconditional score-based diffusion to hypercubes for one-way denoising; RSBM extends reflected dynamics to bidirectional two-marginal Schrödinger bridges for unpaired domain translation.

Rating

• Novelty: ⭐⭐⭐⭐ [First scalable adaptation of flow-matching-style bridge fitting to reflected Brownian motion with closed-form sampling]
• Experimental Thoroughness: ⭐⭐⭐⭐ [Evaluated on both grayscale and RGB unpaired translation with rigorous violation and timing audits]
• Writing Quality: ⭐⭐⭐⭐⭐ [Clear mathematical exposition with rigorous stochastic calculus and measure-theoretic proofs]
• Value: ⭐⭐⭐⭐ [Provides an efficient, principled blueprint for constrained generative modeling and entropic optimal transport]

Related Papers

• [ICLR 2026] A Schrödinger Eigenfunction Method for Long-Horizon Stochastic Optimal Control
• [NeurIPS 2026] SplineFlow: Flow Matching for Dynamical Systems with B-Spline Interpolants
• [NeurIPS 2026] Compatible Likelihoods for Flow Matching on Manifolds
• [NeurIPS 2026] Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion
• [NeurIPS 2026] Path-Guided Flow Matching for Dataset Distillation