Augmented Bridge Matching

TL;DR

This paper introduces a modification to flow and bridge matching methods to preserve the original coupling information between data points, enhancing their ability to learn complex distribution transfers, especially in image translation tasks.

Contribution

The authors propose an augmentation of the velocity field in flow and bridge matching to recover coupling information, which is lost under standard methods.

Findings

Enhanced coupling preservation in distribution transfer models

Improved learning of image translation tasks

Demonstrated effectiveness of augmentation in experiments

Abstract

Flow and bridge matching are a novel class of processes which encompass diffusion models. One of the main aspect of their increased flexibility is that these models can interpolate between arbitrary data distributions i.e. they generalize beyond generative modeling and can be applied to learning stochastic (and deterministic) processes of arbitrary transfer tasks between two given distributions. In this paper, we highlight that while flow and bridge matching processes preserve the information of the marginal distributions, they do \emph{not} necessarily preserve the coupling information unless additional, stronger optimality conditions are met. This can be problematic if one aims at preserving the original empirical pairing. We show that a simple modification of the matching process recovers this coupling by augmenting the velocity field (or drift) with the information of the initial sample point. Doing so, we lose the Markovian property of the process but preserve the coupling information between distributions. We illustrate the efficiency of our augmentation in learning mixture of image translation tasks.