Augmented Bridge Matching
Valentin De Bortoli Guan-Horng Liu Tianrong Chen Evangelos A. Theodorou Weilie Nie
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.