Tree-Based Diffusion Schrödinger Bridge with Applications to Wasserstein Barycenters

Maxence Noble, Valentin De Bortoli, Arnaud Doucet, Alain Durmus

Keywords: Schrödinger bridge, optimal transport, diffusion model, Wasserstein barycenter

Abstract:

Multi-marginal Optimal Transport (mOT), a generalization of OT, aims at minimizing the integral of a cost function with respect to a distribution with some prescribed marginals. In this paper, we consider an entropic version of mOT
with a tree-structured quadratic cost, i.e., a function that can be written as
a sum of pairwise cost functions between the nodes of a tree. To address this
problem, we develop Tree-based Diffusion Schr\"odinger Bridge (TreeDSB), an
extension of the Diffusion Schr\"odinger Bridge (DSB) algorithm. TreeDSB
corresponds to a dynamic and continuous state-space counterpart of the
multimarginal Sinkhorn algorithm. A notable use case of our methodology is to
compute Wasserstein barycenters which can be recast as the solution of a mOT
problem on a star-shaped tree. We demonstrate that our methodology can be applied in high-dimensional settings such as image interpolation and
Bayesian fusion.

Submission Number: 6361