Schrödinger Bridge Flow for Unpaired Data Translation

TL;DR

The paper addresses learning entropic transport maps between unpaired distributions by recasting OT as a dynamic Schrödinger Bridge and introducing a discretised flow of path measures called Schrödinger Bridge Flow. It develops the $\alpha$-IMF family, showing convergence to the SB, and proposes an online, parametric variant, $\alpha$-DSBM, that avoids heavy inner minimisations yet retains SB properties. Theoretical results link non-parametric and parametric updates via $\alpha$-IMF, and experiments on Gaussian and image translation benchmarks demonstrate faster convergence and competitive performance with reduced hyperparameter loads. The approach offers a practical, diffusion-based framework for unpaired data translation with strong transport-theoretic guarantees, while noting remaining limitations around sampling-free representations and scalability.

Abstract

Mass transport problems arise in many areas of machine learning whereby one wants to compute a map transporting one distribution to another. Generative modeling techniques like Generative Adversarial Networks (GANs) and Denoising Diffusion Models (DDMs) have been successfully adapted to solve such transport problems, resulting in CycleGAN and Bridge Matching respectively. However, these methods do not approximate Optimal Transport (OT) maps, which are known to have desirable properties. Existing techniques approximating OT maps for high-dimensional data-rich problems, such as DDM-based Rectified Flow and Schrödinger Bridge procedures, require fully training a DDM-type model at each iteration, or use mini-batch techniques which can introduce significant errors. We propose a novel algorithm to compute the Schrödinger Bridge, a dynamic entropy-regularised version of OT, that eliminates the need to train multiple DDM-like models. This algorithm corresponds to a discretisation of a flow of path measures, which we call the Schrödinger Bridge Flow, whose only stationary point is the Schrödinger Bridge. We demonstrate the performance of our algorithm on a variety of unpaired data translation tasks.

Figures (26)

• Figure 1: Illustration of the SB Flow and comparison with IMF. $\mathbb{Pbb}^\star$ is the SB, $(\hat{\mathbb{P}}^n)_{n \in \mathbb{N}}$ the IMF sequence and $(\hat{\mathbb{P}}^s)_{s \geq 0}$ the flow we consider. See \ref{['sec:euclidean_flow']} for the analysis of this example.
• Figure 2: Evolution of the covariance during online and iterative DSBM finetuning for forward and backward networks. The finetuning starts after 10K steps of training a bridge matching model. For the iterative case, we alternate between forward and backward updates with varying frequencies, i.e. changing after 1K, 2.5K and 5K steps.
• Figure 3: Left: FID and Mean Squared Distance (MSD) on EMNIST to MNIST translation before and after finetuning with different values of $\varepsilon$. Right: AFHQ-64 samples after the finetuning. For both, we use a bidirectional model with online finetuning. More results are in \ref{['sec:experimental_details_mnist']} and \ref{['sec:experimental_details_afhq']}.
• Figure 5: Samples from the original distributions $\pi_0$ (left) and $\pi_1$ (right) are shown in red, while sample paths from $\mathbb{P} = (\pi_0 \otimes \pi_1) \mathbb{Q}_{|0,1}$ are shown in blue.
• Figure 6: Samples from the original distributions are shown in red, while sample paths from $\mathbb{M} = \mathrm{proj}_{\mathcal{M}}(\mathbb{P})$ are shown in blue.

Theorems & Definitions (24)

• theorem 1
• proposition 1
• proposition 2
• proposition 3
• proof
• proposition 4
• proof
• proposition 5
• proof
• proposition 6