Title:	Reflected Schrödinger Bridge Matching
DNr:	Berzelius-2025-371
Project Type:	LiU Berzelius
Principal Investigator:	Pierre Nyquist <pierren@kth.se>
Affiliation:	Kungliga Tekniska högskolan
Duration:	2025-10-20 – 2026-05-01
Classification:	10106
Keywords:

Abstract

This project seeks to explore the application of Schrödinger Bridge Matching to the setting of constrained domains with reflecting boundary conditions. Diffusion Schrödinger bridges on such domains were first explored in [1] using forward-backward SDE theory with score(-like) networks, using methods adapted from [2]. These methods yield high flexibility with respect to the domain and reference dynamics; the authors of [2] show that a range of 2-dimensional domains, including non-convex ones, are workable for the model. Additionally, the method extends to high-dimensional settings, where the learned models show decent performance on tasks such as generating MNIST, CIFAR10, and ImageNet64 from pure noise. However, the method suffers from slow training due to its simulation-based behavior. Namely, just like its non-reflected predecessor from [2], it requires generating full trajectories from the current model during training. The main remedy for this is using cached trajectories, which instead biases the learning objective. This is unfortunate, as one would always prefer to have a simulation-free training algorithm, like in the seminal works on denoising diffusion models, as well as flow-matching. In the unreflected setting, this has given rise to a new type of method called Diffusion Schrödinger bridge matching [3], inspired by flow-matching. This is a simulation-free method that iteratively regresses a learned score network against the Doob h-transform of the reference process, where samples are drawn from bridge processes between the coupled marginal samples. In this work, we seek to implement Schrödinger bridge matching for reflected domains with a Brownian motion reference process. This will be done using extensions of the training methods of [4]. [4] showed how to accurately compute the score of a reflected Brownian motion on the unit cube. To leverage this for the bridge matching problem, we must additionally compute the Doob h-transform term and also develop an efficient method of sampling bridges on the reflected domain. Our preliminary work has solved both of these questions. What remains is implementing this as a deep learning study, over a multitude of datasets and tasks. Especially, we would like to extend the demonstrations in [1] to include coupling of high-dimensional data, as done in [3] for the unreflected case. In addition, we would like to explore whether this framework could enable conditional generation (e.g. prompt-based) using Schrödinger bridges. To the best of our knowledge, this has not been done yet, but the constrained setting may allow methodologies such as classifier-free guidance to be adopted without adding ad-hoc projections in the sampling stage, as shown in [4] for diffusion models. [1] Deng, Wei, et al. "Reflected Schrödinger Bridge for Constrained Generative Modeling." arXiv:2401.03228 (2024). [2] Chen, Tianrong, et al. "Likelihood training of Schrödinger bridge using forward-backward SDEs theory." ICLR (2021). [3] Shi, Yuyang, et al. "Diffusion Schrödinger bridge matching." NeurIPS (2023). [4] Lou, Aaron, and Stefano Ermon. "Reflected diffusion models." ICML (2023).