Schr\"odinger Bridge Flow for Unpaired Data Translation

TL;DR

This paper introduces a new algorithm for computing the Schr"odinger Bridge, a dynamic entropy-regularised optimal transport, that avoids training multiple diffusion models, and demonstrates its effectiveness on unpaired data translation tasks.

Contribution

The paper presents a novel algorithm for Schr"odinger Bridge computation that simplifies high-dimensional optimal transport without multiple model training.

Findings

Effective unpaired data translation demonstrated

Algorithm avoids training multiple diffusion models

Performs well on various data translation tasks

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\"odinger 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\"odinger 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\"odinger Bridge Flow, whose only stationary point is the Schr\"odinger Bridge. We demonstrate the performance of our algorithm on a variety of unpaired data translation tasks.