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.
Department of Electrical and Computer Engineering, University of Delaware, Newark, DE, 19711 · Department of Computer and Information Sciences, University of Delaware, Newark, DE, 19711