Neural solver for Wasserstein Geodesics and optimal transport dynamics
Authors: Hailiang Liu, Yan-Han Chen
Organizations: Department of Mathematics, Iowa State University, Ames, IA 50011, USA. · Department of Statistics, Iowa State University, Ames, IA 50011, USA.
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.
Algorithm 2 Training Procedure for Solving Problem ( 3.1 ).
Figure 1. Phase 1 of Synthetic-1: Samples of the source distribution ρa (open circle at the bottom left corner) are transported to the upper right corner. Each gray line illustrates the learned trajectory Gθ(t;⋅)#ρa . The contour shows the true log density of ρb .
Figure 2. Phase 2 of Synthetic-1: learned velocity fields at selected time points. The arrows illustrate the direction and magnitude of the learned velocity.
Figure 3. Phase 1 of Synthetic-2: Samples of the source distribution ρa (open circle at the center) are transported to the four corners. Each gray line illustrates the learned trajectory Gθ(t;⋅)#ρa . The contour shows the true log density of ρb .
Figure 4. Phase 2 of Synthetic-2: learned velocity fields at selected time points. The arrows illustrate the direction and magnitude of the learned velocity.
Figure 5. Phase 1 of Synthetic-3: 2-dimensional projection of the 10-dimensional distributions. Samples of the source distribution ρa (open circle at the bottom left corner) are transported to the upper right corner. Each gray line illustrates the learned trajectory Gθ(t;⋅)#ρa . The contour shows the true log density of ρb .
Figure 6. Phase 2 of Synthetic-3: learned velocity fields at selected time points. The arrows illustrate the direction and magnitude of the learned velocity.
Figure 7. Phase 1 of Synthetic-4: Left panel presents the analytical trajectories ( 4.1 ) of 20 samples from ρa (open circle at the bottom left corner). In the right panel, those 20 samples are transported by the learned map, where the gray lines illustrate the learned trajectories Gθ(t;⋅)#ρa . The contour in the right panel shows the true log density of ρb .
Figure 8. Phase 2 of Synthetic-4: learned velocity fields at selected time points. The arrows illustrate the direction and magnitude of the learned velocity.
Figure 9. Phase 1 of Real data (MNIST): Left panel contains 100 samples from the source distribution ρa (digit 9). Right panel consists of transported samples Gθ(1;⋅)#ρa , which is meant to match ρb (digit 6).
Figure 10. Phase 1 of Real data case (MNIST): Selected time slices of the transformation process from ρa (digit 9) to Gθ(1;⋅)#ρa , which is meant to match ρb (digit 6).
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 11. Synthetic-1, Phase 1: The left panel shows the estimated W2(ρa,ρb) over training iterations; the dash line refers to the analytical value. The right panel shows the estimated W1(Gθ(1;⋅)#ρa,ρb) over iterations.
Figure 12. Synthetic-1, Phase 2: The left panel shows the loss of objective function value. The right panel shows the mean squared error between the learned velocity field vη and the analytical one v∗ .
Figure 13. Synthetic-2, Phase 1: The left panel shows the estimated W2(ρa,ρb) over training iterations; the dash line refers to the analytical value. The right panel shows the estimated W1(Gθ(1;⋅)#ρa,ρb) over iterations.
Figure 14. Synthetic-2, Phase 2: Loss of the objective function value.
Figure 15. Synthetic-3, Phase 1: The left panel shows the estimated W2(ρa,ρb) over training iterations; the dash line refers to the analytical value. The right panel shows the estimated W1(Gθ(1;⋅)#ρa,ρb) over iterations.
Figure 16. Synthetic-3, Phase 2: The left panel shows the loss of objective function value. The right panel shows the mean squared error between the learned velocity field vη and the analytical one v∗
Figure 17. Synthetic-4: The left panel shows the estimated cost over training iterations. The right panel shows the estimated W1(Gθ(1;⋅)#ρa,ρb) over iterations.
Figure 18. Synthetic-4: Loss of the objective function value.
Figure 19. Realistic case: The left panel shows the estimated W2(ρa,ρb) over training iterations. The right panel shows the estimated W1(Gθ(1;⋅)#ρa,ρb) over iterations.
Figure 20. Realistic case, Phase 2: Loss of the objective function value.
We study optimal transport between distributions when only a subset of the target population is relevant to the source. The goal is to map the source distribution to a reweighted target distribution whose density ratio with respect to the full target lies in [0,c]. This constraint enables the transported distribution to select an appropriate target subpopulation while limiting how strongly any region of the target can be emphasized. We propose a dynamic neural optimal transport framework for this setting using two neural networks. The first parameterizes a time-dependent potential whose spatial gradient defines the velocity field and whose terminal value enforces the density-ratio constraint. The second parameterizes an interpolation between the source distribution and the learned terminal distribution. Unlike prior neural dynamic optimal transport methods, which typically assume a fixed terminal distribution, our formulation learns the terminal distribution within the support of the target. We demonstrate the method on positive-unlabeled learning, unpaired domain translation, and trajectory inference, where meaningful alignments exist between source and target distributions, but direct distribution matching is inappropriate because the target contains additional or irrelevant components
Austin J. Brockmeier, Bilal Riaz
Department of Electrical and Computer Engineering, University of Delaware, Newark, DE, 19711 · Department of Computer and Information Sciences, University of Delaware, Newark, DE, 19711
We propose an implicit neural formulation of optimal transport that eliminates adversarial min--max optimization and multi-network architectures commonly used in existing approaches. Our key idea is to parameterize a single potential in the Kantorovich dual and reformulate the associated c-transform as a proximal fixed-point problem. This yields a stable single-network framework in which dual feasibility is enforced exactly through proximal optimality conditions rather than adversarial training. Despite the inner fixed-point computation, gradients can be computed without differentiating through the fixed-point iterations, enabling efficient training without requiring implicit differentiation. We further establish convergence of stochastic gradient descent. The resulting framework is efficient, scalable, and broadly applicable: it simultaneously recovers forward and backward transport maps and naturally extends to class-conditional settings. Experiments on high-dimensional Gaussian benchmarks, physical datasets, and image translation tasks demonstrate strong transport accuracy together with improved training stability and favorable computational and memory efficiency.
Yesom Park, Eric Gelphman, Stanley Osher +1
Department of Mathematics, University of California, Los Angeles · Department of Applied Mathematics and Statistics, Colorado School of Mines
We introduce Potential Matching Optimal Transport (PMOT), a potential-flow framework for general p-cost optimal transport with cp(x,y)=∥x−y∥p. PMOT parameterizes the CNF velocity field with a scalar potential in the generalized Benamou--Brenier form for the chosen exponent p. It trains the potential gradient with a self-induced matching loss along straight bridges determined by the model's own endpoints, while allowing flexible terminal distribution matching. Our main result establishes zero-loss exactness: under the stated regularity, exact terminal matching, and uniqueness assumptions, any zero-loss solution satisfies the generalized Benamou--Brenier optimality system and recovers the corresponding p-optimal transport map and dynamics. On synthetic benchmarks, PMOT learns p-specific maps that agree with the corresponding p-matched OT references. It also remains competitive as a likelihood-based density model on high-dimensional tabular data, and an MMD-based color transformation experiment demonstrates flexible sample-based terminal matching.
Lishuo Zhang, Ruizhi Huang, Yang Yu +1
School of Mathematical Sciences, Shanghai Jiao Tong University · Institute of Natural Sciences, MOE-LSC, Shanghai Jiao Tong University