# A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 56 references that are in the map.
- First posted on arXiv 8 May 2026.

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

- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [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)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)
- [TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics](https://mlsanity.com/paper/2002.04461/) (2020)
- [Action Matching: Learning Stochastic Dynamics from Samples](https://mlsanity.com/paper/2210.06662/) (2022)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Metric Flow Matching for Smooth Interpolations on the Data Manifold](https://mlsanity.com/paper/2405.14780/) (2024)
- [The Schr"odinger Bridge between Gaussian Measures has a Closed Form](https://mlsanity.com/paper/2202.05722/) (2022)
- [A Computational Framework for Solving Wasserstein Lagrangian Flows](https://mlsanity.com/paper/2310.10649/) (2023)
- [Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold](https://mlsanity.com/paper/2408.14608/) (2024)
- [Towards a mathematical theory of trajectory inference](https://mlsanity.com/paper/2102.09204/) (2021)
- [Learning diffusion at lightspeed](https://mlsanity.com/paper/2406.12616/) (2024)
- [Curly Flow Matching for Learning Non-gradient Field Dynamics](https://mlsanity.com/paper/2510.26645/) (2025)
- [Reconciling emergences: An information-theoretic approach to identify causal emergence in multivariate data](https://mlsanity.com/paper/2004.08220/) (2020)
- [Identifying Drift, Diffusion, and Causal Structure from Temporal Snapshots](https://mlsanity.com/paper/2410.22729/) (2024)
- [Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action](https://mlsanity.com/paper/2505.11823/) (2025)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [DICE: Discrete inverse continuity equation for learning population dynamics](https://mlsanity.com/paper/2507.05107/) (2025)
- [Joint trajectory and network inference via reference fitting](https://mlsanity.com/paper/2409.06879/) (2024)
- [Oh SnapMMD! Forecasting Stochastic Dynamics Beyond the Schr"odinger Bridge's End](https://mlsanity.com/paper/2505.16082/) (2025)
- [Conditional Lagrangian Wasserstein Flow for Time Series Imputation](https://mlsanity.com/paper/2410.07550/) (2024)
- [Gradient-flow SDEs have unique transient population dynamics](https://mlsanity.com/paper/2505.21770/) (2025)

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

- [Measure-to-measure Regression with Transformers](https://mlsanity.com/paper/2605.28075/) (2026)
- [DisKO: Deep Koopman Learning in Distribution Space from Unpaired Snapshots](https://mlsanity.com/paper/2609.34629/) (2026)
- [Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals](https://mlsanity.com/paper/2610.03679/) (2026)
- [Discrete Action Matching: Learning Stochastic Dynamics from Samples via State Graphs](https://mlsanity.com/paper/2610.05071/) (2026)
- [Neural Wasserstein Splines: Representation-Invariant Second-Order Priors for Trajectory Inference](https://mlsanity.com/paper/2610.19878/) (2026)

## Abstract (arXiv)

The population dynamics of molecules, cells, and organisms are governed by a number of unknown forces. In the last decade, population dynamics have predominantly been modeled with Wasserstein gradient flows. However, since gradient flows minimize free energy, they fail to capture important dynamical properties, such as periodicity. In this work, we propose a change in perspective by considering dynamics that minimize a population-level action under a damped Wasserstein Lagrangian. By deriving the corresponding Hamiltonian equations of motion, we formalize Wasserstein Lagrangian Mechanics, a structured class of second-order dynamics that encompasses classical mechanics, quantum mechanics, and gradient flows. We then propose WLM as the first algorithm that learns these second-order dynamics from observed marginals, without specifying the Lagrangian. By directly learning the population mechanics, WLM can both forecast and interpolate unseen marginals, and outperforms existing gradient flow and flow matching methods across a wide range of dynamics, including vortex dynamics, embryonic development, and flocking.