Iterated Schr\"odinger bridge approximation to Wasserstein Gradient Flows

TL;DR

This paper introduces a new discretization method for Wasserstein gradient flows using iterated Schr"odinger bridges, which avoids score functions and is suitable for particle-based algorithms, with proven convergence properties.

Contribution

The paper proposes a novel scheme for Wasserstein gradient flows based on iterated Schr"odinger bridges, providing a rigorous convergence analysis and connections to machine learning models.

Findings

The scheme avoids the use of score functions.

Convergence to Wasserstein gradient flows is proven for certain classes.

Provides a probabilistic framework for understanding self-attention in transformers.

Abstract

We introduce a novel discretization scheme for Wasserstein gradient flows that involves successively computing Schr\"{o}dinger bridges with the same marginals. This is different from both the forward/geodesic approximation and the backward/Jordan-Kinderlehrer-Otto (JKO) approximations. The proposed scheme has two advantages: one, it avoids the use of the score function, and, two, it is amenable to particle-based approximations using the Sinkhorn algorithm. Our proof hinges upon showing that relative entropy between the Schr\"{o}dinger bridge with the same marginals at temperature $\epsilon$ and the joint distribution of a stationary Langevin diffusion at times zero and $\epsilon$ is of the order $o(\epsilon^2)$ with an explicit dependence given by Fisher information. Owing to this inequality, we can show, using a triangular approximation argument, that the interpolated iterated application of the Schr\"{o}dinger bridge approximation converge to the Wasserstein gradient flow, for a class of gradient flows, including the heat flow. The results also provide a probabilistic and rigorous framework for the convergence of the self-attention mechanisms in transformer networks to the solutions of heat flows, first observed in the inspiring work SABP22 in machine learning research.