Neural solver for Wasserstein Geodesics and optimal transport dynamics
This work addresses the challenge of modeling distributional relationships in machine learning, offering a flexible framework for optimal transport that can handle general cost functions, though it appears incremental as it builds on existing dynamical formulations.
The paper tackles the problem of computing Wasserstein geodesics and optimal transport dynamics by introducing a sample-based neural solver that recasts the constrained optimization as a minimax problem using deep neural networks, enabling direct sampling from the target distribution and learning of the full velocity field, with effectiveness demonstrated on synthetic and real datasets.
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.