# A scalable deep learning approach for solving high-dimensional dynamic optimal transport

arXiv: https://arxiv.org/abs/2205.07521 · HTML: https://mlsanity.com/paper/2205.07521/

## Key facts

- Cited by 5 papers in the map (data through 7 Oct 2026).
- Builds on 29 of its 56 references that are in the map.
- First posted on arXiv 16 May 2022.

## Builds on (most cited first, 29 of 29)

- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Deep Learning in Neural Networks: An Overview](https://mlsanity.com/paper/1404.7828/) (2014)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [DGM: A deep learning algorithm for solving partial differential equations](https://mlsanity.com/paper/1708.07469/) (2017)
- [Wasserstein GAN](https://mlsanity.com/paper/1701.07875/) (2017)
- [The Deep Ritz method: A deep learning-based numerical algorithm for solving variational problems](https://mlsanity.com/paper/1710.00211/) (2017)
- [Optimal Transport for Domain Adaptation](https://mlsanity.com/paper/1507.00504/) (2015)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [Stochastic Optimization for Large-scale Optimal Transport](https://mlsanity.com/paper/1605.08527/) (2016)
- [Deep learning-based numerical methods for high-dimensional parabolic partial differential equations and backward stochastic differential equations](https://mlsanity.com/paper/1706.04702/) (2017)
- [Deep Network Approximation for Smooth Functions](https://mlsanity.com/paper/2001.03040/) (2020)
- [TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics](https://mlsanity.com/paper/2002.04461/) (2020)
- [Improving GANs Using Optimal Transport](https://mlsanity.com/paper/1803.05573/) (2018)
- [Large-Scale Optimal Transport and Mapping Estimation](https://mlsanity.com/paper/1711.02283/) (2017)
- [A Machine Learning Framework for Solving High-Dimensional Mean Field Game and Mean Field Control Problems](https://mlsanity.com/paper/1912.01825/) (2019)
- [ANODE: Unconditionally Accurate Memory-Efficient Gradients for Neural ODEs](https://mlsanity.com/paper/1902.10298/) (2019)
- [A Primer on Neural Network Models for Natural Language Processing](https://mlsanity.com/paper/1510.00726/) (2015)
- [Optimal Transport for Multi-source Domain Adaptation under Target Shift](https://mlsanity.com/paper/1803.04899/) (2018)
- [Neural Network Approximation: Three Hidden Layers Are Enough](https://mlsanity.com/paper/2010.14075/) (2020)
- [Discretize-Optimize vs. Optimize-Discretize for Time-Series Regression and Continuous Normalizing Flows](https://mlsanity.com/paper/2005.13420/) (2020)
- [A Survey on Optimal Transport for Machine Learning: Theory and Applications](https://mlsanity.com/paper/2106.01963/) (2021)
- [A Geometric View of Optimal Transportation and Generative Model](https://mlsanity.com/paper/1710.05488/) (2017)
- [Advancements in Image Classification using Convolutional Neural Network](https://mlsanity.com/paper/1905.03288/) (2019)
- [Unnormalized Optimal Transport](https://mlsanity.com/paper/1902.03367/) (2019)
- [AE-OT-GAN: Training GANs from data specific latent distribution](https://mlsanity.com/paper/2001.03698/) (2020)
- [Multilevel Optimal Transport: a Fast Approximation of Wasserstein-1 distances](https://mlsanity.com/paper/1810.00118/) (2018)
- [A Fast Proximal Gradient Method and Convergence Analysis for Dynamic Mean Field Planning](https://mlsanity.com/paper/2102.13260/) (2021)

## Built on by (most cited first, 5 of 5)

- [Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport](https://mlsanity.com/paper/2410.00844/) (2024)
- [Modeling Cell Dynamics and Interactions with Unbalanced Mean Field Schr"odinger Bridge](https://mlsanity.com/paper/2505.11197/) (2025)
- [Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action](https://mlsanity.com/paper/2505.11823/) (2025)
- [Integrating Dynamical Systems Modeling with Spatiotemporal scRNA-seq Data Analysis](https://mlsanity.com/paper/2503.11347/) (2025)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)

## Abstract (arXiv)

The dynamic formulation of optimal transport has attracted growing interests in scientific computing and machine learning, and its computation requires to solve a PDE-constrained optimization problem. The classical Eulerian discretization based approaches suffer from the curse of dimensionality, which arises from the approximation of high-dimensional velocity field. In this work, we propose a deep learning based method to solve the dynamic optimal transport in high dimensional space. Our method contains three main ingredients: a carefully designed representation of the velocity field, the discretization of the PDE constraint along the characteristics, and the computation of high dimensional integral by Monte Carlo method in each time step. Specifically, in the representation of the velocity field, we apply the classical nodal basis function in time and the deep neural networks in space domain with the H1-norm regularization. This technique promotes the regularity of the velocity field in both time and space such that the discretization along the characteristic remains to be stable during the training process. Extensive numerical examples have been conducted to test the proposed method. Compared to other solvers of optimal transport, our method could give more accurate results in high dimensional cases and has very good scalability with respect to dimension. Finally, we extend our method to more complicated cases such as crowd motion problem.