Augmented Bridge Matching
This work solves the coupling preservation issue in stochastic processes for researchers in generative modeling and image translation, representing an incremental improvement over existing flow and bridge matching methods.
The paper addresses the problem that flow and bridge matching processes preserve marginal distributions but not the coupling information between them, which is crucial for tasks like image translation. The authors propose augmenting the velocity field with initial sample information to recover the coupling, demonstrating efficiency on mixture of image translation tasks.
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.