Neural solver for Wasserstein Geodesics and optimal transport dynamics

TL;DR

This paper presents a neural network-based method for computing Wasserstein geodesics and optimal transport maps, enabling efficient sampling and velocity field estimation between distributions.

Contribution

We introduce a sample-based neural solver that recasts the OT problem as a minimax optimization, allowing direct computation of geodesics, OT maps, and velocity fields.

Findings

Effective computation of Wasserstein geodesics demonstrated on synthetic data.

Accurate estimation of OT maps and velocity fields shown on real datasets.

Flexible extension to various cost functions including quadratic cost.

Abstract

In recent years, the machine learning community has increasingly embraced the optimal transport (OT) framework for modeling distributional relationships. In this work, we introduce a sample-based neural solver for computing the Wasserstein geodesic between a source and target distribution, along with the associated velocity field. Building on the dynamical formulation of the optimal transport (OT) problem, we recast the constrained optimization as a minimax problem, using deep neural networks to approximate the relevant functions. This approach not only provides the Wasserstein geodesic but also recovers the OT map, enabling direct sampling from the target distribution. By estimating the OT map, we obtain velocity estimates along particle trajectories, which in turn allow us to learn the full velocity field. The framework is flexible and readily extends to general cost functions, including the commonly used quadratic cost. We demonstrate the effectiveness of our method through experiments on both synthetic and real datasets.