Iterated Schrödinger bridge approximation to Wasserstein Gradient Flows

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Agarwal, Medha, Harchaoui, Zaid, Mulcahy, Garrett, Pal, Soumik
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866910489046417408
author Agarwal, Medha
Harchaoui, Zaid
Mulcahy, Garrett
Pal, Soumik
author_facet Agarwal, Medha
Harchaoui, Zaid
Mulcahy, Garrett
Pal, Soumik
contents We introduce a novel discretization scheme for Wasserstein gradient flows that involves successively computing Schrö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ödinger bridge with the same marginals at temperature $ε$ and the joint distribution of a stationary Langevin diffusion at times zero and $ε$ is of the order $o(ε^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ö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.
format Preprint
id arxiv_https___arxiv_org_abs_2406_10823
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Iterated Schrödinger bridge approximation to Wasserstein Gradient Flows
Agarwal, Medha
Harchaoui, Zaid
Mulcahy, Garrett
Pal, Soumik
Probability
Machine Learning
49N99, 49Q22, 60J60
We introduce a novel discretization scheme for Wasserstein gradient flows that involves successively computing Schrö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ödinger bridge with the same marginals at temperature $ε$ and the joint distribution of a stationary Langevin diffusion at times zero and $ε$ is of the order $o(ε^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ö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.
title Iterated Schrödinger bridge approximation to Wasserstein Gradient Flows
topic Probability
Machine Learning
49N99, 49Q22, 60J60
url https://arxiv.org/abs/2406.10823