# A Computational Framework for Solving Wasserstein Lagrangian Flows

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

## Key facts

- Cited by 46 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 86 references that are in the map.
- First posted on arXiv 16 Oct 2023.

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

- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [A survey of the Schr"odinger problem and some of its connections with optimal transport](https://mlsanity.com/paper/1308.0215/) (2013)
- [Rectified Flow: A Marginal Preserving Approach to Optimal Transport](https://mlsanity.com/paper/2209.14577/) (2022)
- [Multisample Flow Matching: Straightening Flows with Minibatch Couplings](https://mlsanity.com/paper/2304.14772/) (2023)
- [Scalable Gradients for Stochastic Differential Equations](https://mlsanity.com/paper/2001.01328/) (2020)
- [Likelihood Training of Schr"odinger Bridge using Forward-Backward SDEs Theory](https://mlsanity.com/paper/2110.11291/) (2021)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [Ab-Initio Solution of the Many-Electron Schr"odinger Equation with Deep Neural Networks](https://mlsanity.com/paper/1909.02487/) (2019)
- [TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics](https://mlsanity.com/paper/2002.04461/) (2020)
- [Solving Schr"odinger Bridges via Maximum Likelihood](https://mlsanity.com/paper/2106.02081/) (2021)
- [Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measures](https://mlsanity.com/paper/1508.07941/) (2015)
- [On the relation between optimal transport and Schr"odinger bridges: A stochastic control viewpoint](https://mlsanity.com/paper/1412.4430/) (2014)
- [Action Matching: Learning Stochastic Dynamics from Samples](https://mlsanity.com/paper/2210.06662/) (2022)
- [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)
- [Simulation-free Schr"odinger bridges via score and flow matching](https://mlsanity.com/paper/2307.03672/) (2023)
- [Diffusion Bridge Mixture Transports, Schr"odinger Bridge Problems and Generative Modeling](https://mlsanity.com/paper/2304.00917/) (2023)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks](https://mlsanity.com/paper/2106.00774/) (2021)
- [Variational Wasserstein gradient flow](https://mlsanity.com/paper/2112.02424/) (2021)
- [Accelerating Langevin Sampling with Birth-death](https://mlsanity.com/paper/1905.09863/) (2019)

## Built on by (most cited first, 30 of 46)

- [Metric Flow Matching for Smooth Interpolations on the Data Manifold](https://mlsanity.com/paper/2405.14780/) (2024)
- [Schr"odinger Bridge Flow for Unpaired Data Translation](https://mlsanity.com/paper/2409.09347/) (2024)
- [Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold](https://mlsanity.com/paper/2408.14608/) (2024)
- [Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport](https://mlsanity.com/paper/2410.00844/) (2024)
- [Doob's Lagrangian: A Sample-Efficient Variational Approach to Transition Path Sampling](https://mlsanity.com/paper/2410.07974/) (2024)
- [Dynamic Conditional Optimal Transport through Simulation-Free Flows](https://mlsanity.com/paper/2404.04240/) (2024)
- [Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling](https://mlsanity.com/paper/2505.13413/) (2025)
- [Curly Flow Matching for Learning Non-gradient Field Dynamics](https://mlsanity.com/paper/2510.26645/) (2025)
- [Geometry-Aware Generative Autoencoders for Warped Riemannian Metric Learning and Generative Modeling on Data Manifolds](https://mlsanity.com/paper/2410.12779/) (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)
- [Efficient Trajectory Inference in Wasserstein Space Using Consecutive Averaging](https://mlsanity.com/paper/2405.19679/) (2024)
- [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)
- [Score-based Neural Ordinary Differential Equations for Computing Mean Field Control Problems](https://mlsanity.com/paper/2409.16471/) (2024)
- [Improving Neural Optimal Transport via Displacement Interpolation](https://mlsanity.com/paper/2410.03783/) (2024)
- [WFR-FM: Simulation-Free Dynamic Unbalanced Optimal Transport](https://mlsanity.com/paper/2601.06810/) (2026)
- [Scalable Simulation-free Entropic Unbalanced Optimal Transport](https://mlsanity.com/paper/2410.02656/) (2024)
- [Branched Schr"odinger Bridge Matching](https://mlsanity.com/paper/2506.09007/) (2025)
- [A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots](https://mlsanity.com/paper/2605.08550/) (2026)
- [Conditional Lagrangian Wasserstein Flow for Time Series Imputation](https://mlsanity.com/paper/2410.07550/) (2024)
- [Generalized Flow Matching for Transition Dynamics Modeling](https://mlsanity.com/paper/2410.15128/) (2024)
- [Symplectic Neural Flows for Modeling and Discovery](https://mlsanity.com/paper/2412.16787/) (2024)
- [Multi-Marginal Flow Matching with Adversarially Learnt Interpolants](https://mlsanity.com/paper/2510.01159/) (2025)
- [Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots](https://mlsanity.com/paper/2605.00545/) (2026)
- [Lagrangian Flow Matching: A Least-Action Framework for Principled Path Design](https://mlsanity.com/paper/2605.15419/) (2026)
- [Minimum-Excess-Work Guidance](https://mlsanity.com/paper/2505.13375/) (2025)
- [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)
- [Learning Explicit Single-Cell Dynamics Using ODE Representations](https://mlsanity.com/paper/2510.02903/) (2025)

## Abstract (arXiv)

The dynamical formulation of the optimal transport can be extended through various choices of the underlying geometry (kinetic energy), and the regularization of density paths (potential energy). These combinations yield different variational problems (Lagrangians), encompassing many variations of the optimal transport problem such as the Schr"odinger bridge, unbalanced optimal transport, and optimal transport with physical constraints, among others. In general, the optimal density path is unknown, and solving these variational problems can be computationally challenging. We propose a novel deep learning based framework approaching all of these problems from a unified perspective. Leveraging the dual formulation of the Lagrangians, our method does not require simulating or backpropagating through the trajectories of the learned dynamics, and does not need access to optimal couplings. We showcase the versatility of the proposed framework by outperforming previous approaches for the single-cell trajectory inference, where incorporating prior knowledge into the dynamics is crucial for correct predictions.