# Neural Optimal Transport with Lagrangian Costs

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

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 63 references that are in the map.
- First posted on arXiv 1 Jun 2024.

## Builds on (most cited first, 30 of 39)

- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Multisample Flow Matching: Straightening Flows with Minibatch Couplings](https://mlsanity.com/paper/2304.14772/) (2023)
- [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)
- [TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics](https://mlsanity.com/paper/2002.04461/) (2020)
- [Geometric Dataset Distances via Optimal Transport](https://mlsanity.com/paper/2002.02923/) (2020)
- [OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal Transport](https://mlsanity.com/paper/2006.00104/) (2020)
- [Large-Scale Optimal Transport and Mapping Estimation](https://mlsanity.com/paper/1711.02283/) (2017)
- [Gromov-Wasserstein Alignment of Word Embedding Spaces](https://mlsanity.com/paper/1809.00013/) (2018)
- [Wasserstein-2 Generative Networks](https://mlsanity.com/paper/1909.13082/) (2019)
- [A Machine Learning Framework for Solving High-Dimensional Mean Field Game and Mean Field Control Problems](https://mlsanity.com/paper/1912.01825/) (2019)
- [Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark](https://mlsanity.com/paper/2106.01954/) (2021)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization](https://mlsanity.com/paper/2012.05942/) (2020)
- [Entropic estimation of optimal transport maps](https://mlsanity.com/paper/2109.12004/) (2021)
- [Robust Optimal Transport with Applications in Generative Modeling and Domain Adaptation](https://mlsanity.com/paper/2010.05862/) (2020)
- [Tutorial on amortized optimization](https://mlsanity.com/paper/2202.00665/) (2022)
- [Generative Modeling with Optimal Transport Maps](https://mlsanity.com/paper/2110.02999/) (2021)
- [Plugin Estimation of Smooth Optimal Transport Maps](https://mlsanity.com/paper/2107.12364/) (2021)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [Rates of Estimation of Optimal Transport Maps using Plug-in Estimators via Barycentric Projections](https://mlsanity.com/paper/2107.01718/) (2021)
- [Deep Generalized Schr"odinger Bridge](https://mlsanity.com/paper/2209.09893/) (2022)
- [Multi-marginal optimal transport: theory and applications](https://mlsanity.com/paper/1406.0026/) (2014)
- [From optimal transport to generative modeling: the VEGAN cookbook](https://mlsanity.com/paper/1705.07642/) (2017)
- [Neural Monge Map estimation and its applications](https://mlsanity.com/paper/2106.03812/) (2021)
- [Generalized Schr"odinger Bridge Matching](https://mlsanity.com/paper/2310.02233/) (2023)
- [Optimal Transport for Diffeomorphic Registration](https://mlsanity.com/paper/1706.05218/) (2017)
- [The Monge Gap: A Regularizer to Learn All Transport Maps](https://mlsanity.com/paper/2302.04953/) (2023)
- [On amortizing convex conjugates for optimal transport](https://mlsanity.com/paper/2210.12153/) (2022)
- [Learning Riemannian Manifolds for Geodesic Motion Skills](https://mlsanity.com/paper/2106.04315/) (2021)

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

- [Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control](https://mlsanity.com/paper/2409.08861/) (2024)
- [Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport](https://mlsanity.com/paper/2410.00844/) (2024)
- [Conditional simulation via entropic optimal transport: Toward non-parametric estimation of conditional Brenier maps](https://mlsanity.com/paper/2411.07154/) (2024)
- [Integrating Dynamical Systems Modeling with Spatiotemporal scRNA-seq Data Analysis](https://mlsanity.com/paper/2503.11347/) (2025)
- [Generalized Flow Matching for Transition Dynamics Modeling](https://mlsanity.com/paper/2410.15128/) (2024)
- [Zero-Shot Offline Imitation Learning via Optimal Transport](https://mlsanity.com/paper/2410.08751/) (2024)
- [Differentiable Cost-Parameterized Monge Map Estimators](https://mlsanity.com/paper/2406.08399/) (2024)
- [Optimal Stochastic Trace Estimation in Generative Modeling](https://mlsanity.com/paper/2502.18808/) (2025)
- [Lagrangian Flow Matching: A Least-Action Framework for Principled Path Design](https://mlsanity.com/paper/2605.15419/) (2026)
- [Flows and Diffusions on the Neural Manifold](https://mlsanity.com/paper/2507.10623/) (2025)
- [HOTA: Hamiltonian framework for Optimal Transport Advection](https://mlsanity.com/paper/2507.17513/) (2025)
- [Experts-Guided Unbalanced Optimal Transport for ISP Learning from Unpaired and/or Paired Data](https://mlsanity.com/paper/2512.05635/) (2025)
- [Twisted Schr"odinger Bridge Matching](https://mlsanity.com/paper/2607.16987/) (2026)
- [Sum-of-Squares Programming for Ma-Trudinger-Wang Regularity of Optimal Transport Maps](https://mlsanity.com/paper/2412.13372/) (2024)
- [Hyperparameter Trajectory Inference with Conditional Lagrangian Optimal Transport](https://mlsanity.com/paper/2603.01771/) (2026)
- [Reliable Modeling of Distribution Shifts via Displacement-Reshaped Optimal Transport](https://mlsanity.com/paper/2605.04965/) (2026)

## Abstract (arXiv)

We investigate the optimal transport problem between probability measures when the underlying cost function is understood to satisfy a least action principle, also known as a Lagrangian cost. These generalizations are useful when connecting observations from a physical system where the transport dynamics are influenced by the geometry of the system, such as obstacles (e.g., incorporating barrier functions in the Lagrangian), and allows practitioners to incorporate a priori knowledge of the underlying system such as non-Euclidean geometries (e.g., paths must be circular). Our contributions are of computational interest, where we demonstrate the ability to efficiently compute geodesics and amortize spline-based paths, which has not been done before, even in low dimensional problems. Unlike prior work, we also output the resulting Lagrangian optimal transport map without requiring an ODE solver. We demonstrate the effectiveness of our formulation on low-dimensional examples taken from prior work. The source code to reproduce our experiments is available at https://github.com/facebookresearch/lagrangian-ot.