Reflected Schrödinger Bridge Matching

Abstract: Recent advances in generative modeling have enabled the efficient computation of Schrödinger bridges (SB) in high-dimensional settings by leveraging partially simulation-free training methods inspired by flow matching. However, these have not covered SBs with reflecting dynamics, a useful model choice with built-in guarantees that generated samples stay in the data domain. Existing alternatives for reflected SBs instead rely on more complex training based on forward--backward SDE theory, requiring expensive higher-order derivatives and sampling entire paths during training. In this article, we introduce a partially simulation-free framework that allows reflected SBs to be trained similarly to flow matching, using a new sampling method and regression target. We demonstrate our results by coupling pairs of well-known high-dimensional image datasets. Using reflected dynamics incurs negligible additional wall-clock time during both training and inference while maintaining or slightly improving generative performance.

• The paper introduces RSBM, a novel method that extends iterative Markovian fitting to reflected stochastic dynamics for bounded domains.
• It employs reciprocal and Markovian projections to efficiently enforce pixel value and physical support constraints during generative modeling.
• RSBM achieves superior performance metrics (MSD, FID, LPIPS) while incurring negligible computational overhead compared to non-reflected baselines.

Reflected Schrödinger Bridge Matching: Efficient Constrained Generative Modeling

Background and Motivation

The Schrödinger Bridge (SB) framework offers an entropy-regularized approach to probabilistic optimal transport (OT) between arbitrary marginal distributions. Recent advances in high-dimensional generative modeling, notably diffusion models and flow matching, have enabled practical computation of SBs in unconstrained domains. However, scenarios requiring rigorous domain constraints (e.g., pixel value bounds or physical support constraints) motivate extending SB methods to reflected stochastic dynamics. In particular, reflected SBs provide sample guarantees by enforcing domain boundary conditions throughout the generative process.

Current implementations of reflected SBs (e.g., forward-backward SDE frameworks) are computationally demanding, requiring full path sampling during training and expensive higher-order derivative computations. This paper introduces Reflected Schrödinger Bridge Matching (RSBM), a framework leveraging the simulation-free characteristics of flow matching and iterative Markovian fitting (IMF), thereby enabling efficient training and inference of reflected SBs on domains such as the unit hypercube.

Figure 1: Visualization of a reflected Schrödinger bridge between two 2D distributions, illustrating constrained sample evolution between pixel-based marginals.

Algorithmic Framework

RSBM generalizes non-reflected IMF and $\alpha$-IMF to reflected reference processes, specifically reflected Brownian motion (RBM) on $[0,1]^d$. The algorithm iteratively alternates between two core steps:

• Reciprocal projection: Sampling bridge paths between endpoints coupled from the marginal distributions, ensuring path-wise domain constraint.
• Markovian projection: Learning the drift correction via a neural network using regression against the Doob $h$-transform score of the reflected process transition density. Backward and forward objective losses are symmetrized with conditional directional embeddings, enabling bidirectional parametrization.

The reference SDE is discretized with Euler–Maruyama and incorporates reflections at each step (per dimension for $[0,1]^d$). Efficient bridge sampling exploits the invariance properties of RBM, leveraging a mixture of reflected Gaussian distributions.

The loss function minimizes the KL divergence between bridge mixtures and the reflected Markovian projected process, with explicit expectations over the score functions: $$L(\theta ; \Pi) = \frac{1}{2} \left( L_{\to}(\theta ; \Pi) + L_{\gets}(\theta ; \Pi) \right)$$ where each term involves regression against the gradient of the log reflected transition density (Doob $h$-transform).

Efficient Sampling and Score Computation

Bridge sampling with reflection is performed using a multi-stage approach:

• Sample possible endpoint reflections per dimension by weighting with transition Gaussian likelihoods.
• Sample intermediate bridge points conditionally, then reflect as needed.
• For high-dimensional domains and independent reference processes, the transition density and score computation are factorized to remain tractable.

Score computation for reflected transitions is crucial. For RBM on $[0,1]^d$, the transition density at $x_t$ given $x_s$ is constructed as a sum of Gaussian densities over all candidate reflections, leveraging the symmetry and independence in each dimension.

Experimental Evaluation

RSBM is empirically evaluated on high-dimensional image translation tasks: MNIST $\leftrightarrow$ EMNIST and AFHQ Cat $[0,1]^d$0 Wild. The method guarantees that all generated samples lie strictly within the pixel value domain (unlike the non-reflected baseline $[0,1]^d$1-DSBM, which produces significant out-of-bounds values).

Strong numerical results are observed:

• MSD (mean square distance): Reflected samples exhibit significantly lower MSD from source images compared to non-reflected baselines, reflecting superior source-to-target likeness preservation.
• FID (Fréchet Inception Distance): $[0,1]^d$2-RSBM matches or marginally improves generative quality relative to $[0,1]^d$3-DSBM, with lower FID in most cases except the wild-to-cat transfer direction.
• LPIPS: $[0,1]^d$4-RSBM generally outperforms or closely matches the non-reflected baseline.

Figure 2: Qualitative samples for cat $[0,1]^d$5 wild translation. RSBM samples always honor domain constraints, while non-reflected samples must be clipped post-generation due to out-of-bounds values.

Figure 3: MNIST $[0,1]^d$6 EMNIST transfers, demonstrating sharper source-target similarity and rigorous pixel support preservation in RSBM outputs.

RSBM incurs only negligible computational overhead relative to non-reflected bridge matching, with nearly identical wall-clock time for both training and inference.

Theoretical Implications

The paper establishes convergence guarantees for reflected IMF, extending theoretical results from prior works to RBM-constrained domains. Fundamental properties, including preservation of marginal distributions and minimal KL divergence within the Markov class of reflected stochastic processes, are rigorously proved. The Doob $[0,1]^d$7-transform score regression, implemented efficiently and tractably for RBM, enables high-fidelity generative modeling on bounded domains.

The limitations of current RSBM implementation—restriction to RBM on $[0,1]^d$8 and zero drift—arise from analytical tractability. Extensions to more general domains and reflected SDEs with nontrivial drift present substantial challenges, largely in bridge sampling and $[0,1]^d$9-transform approximation.

Practical Implications and Future Directions

RSBM enables generative modeling in domains where strict support constraints are non-negotiable. Applications span image translation with pixel bounds, physical simulation (e.g., molecular trajectories constrained within physical boundaries), and any generative task requiring unassailable support enforcement.

Future work may address:

• Neural bridge samplers and $h$0-transform approximators for non-analytic domains or reference dynamics.
• PDE-based score estimation for complex reflected SDEs.
• Generalization to domains admitting diffeomorphic mapping to $h$1, or even nonconvex supports.
• Integration with classifier-free guidance and other diffusion model enhancements.
• Exploration of high-dimensional settings with interacting particle dynamics and empirical measure constraints.

Conclusion

Reflected Schrödinger Bridge Matching provides an efficient, simulation-free framework for constrained generative modeling, extending flow matching and iterative bridge fitting to support-reflected stochastic dynamics. The method achieves rigorous sample support, high generative quality, and computational efficiency. While currently restricted to RBM on hypercubes, the theoretical and empirical advances offer a principled foundation for further extension to more complex domains and reference processes, with broad implications for constrained generative modeling in AI.