🤖 AI Summary
This work addresses the problem of efficiently learning Wasserstein geodesics and their associated optimal transport velocity fields directly from samples to model the dynamic transport process between a source and a target distribution. Building upon the dynamical formulation of optimal transport, the constrained optimization problem is reformulated as a minimax game, enabling joint approximation of the geodesic, optimal map, and full velocity field using deep neural networks. The proposed method constitutes the first purely sample-driven neural solver capable of handling general cost functions—including the quadratic cost—without requiring explicit density estimation. Experiments on both synthetic and real-world datasets demonstrate that the approach accurately reconstructs geodesics and velocity fields and enables direct sampling from the target distribution, thereby validating its effectiveness and broad applicability.
📝 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.