# Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 27 of its 50 references that are in the map.
- First posted on arXiv 2 Oct 2026.

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

- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Characterizing possible failure modes in physics-informed neural networks](https://mlsanity.com/paper/2109.01050/) (2021)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [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)
- [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)
- [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)
- [Neural Conservation Laws: A Divergence-Free Perspective](https://mlsanity.com/paper/2210.01741/) (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)
- [Deep Momentum Multi-Marginal Schr"odinger Bridge](https://mlsanity.com/paper/2303.01751/) (2023)
- [Multi-marginal Schr"odinger Bridges with Iterative Reference Refinement](https://mlsanity.com/paper/2408.06277/) (2024)
- [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)
- [Identifying Drift, Diffusion, and Causal Structure from Temporal Snapshots](https://mlsanity.com/paper/2410.22729/) (2024)
- [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)
- [Oh SnapMMD! Forecasting Stochastic Dynamics Beyond the Schr"odinger Bridge's End](https://mlsanity.com/paper/2505.16082/) (2025)
- [A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots](https://mlsanity.com/paper/2605.08550/) (2026)
- [Gradient-flow SDEs have unique transient population dynamics](https://mlsanity.com/paper/2505.21770/) (2025)
- [Generative Path-Finding Method for Wasserstein Gradient Flow](https://mlsanity.com/paper/2604.11519/) (2026)
- [Wasserstein Residuals: Learning Gradient Flows from Population Dynamics](https://mlsanity.com/paper/2607.04738/) (2026)

## Abstract (arXiv)

The dynamics of cells, organisms, and fluids are often modeled as probability distributions evolving over time. Reconstructing and extrapolating this evolution from unpaired snapshots requires assumptions about the underlying process. Wasserstein gradient flows are a common choice, but they cannot describe conservative or periodic dynamics. Lagrangian mechanics in Wasserstein space covers both, but existing methods for learning it are simulation-based: they run a numerical solver at every training step, which makes training expensive. We propose Double-Stitch, a simulation-free method that learns these mechanics by penalizing the residual of the equation of motion along a learned population path. We derive this equation from a Clebsch variational principle that does not require gradient velocities, and show that the residual vanishes exactly when the equation holds. We test Double-Stitch on synthetic, single-cell and ocean vortex datasets and find that it matches or outperforms gradient-flow methods and simulation-based WLM on most tasks, while training 4-14 times faster than WLM. We provide a JAX implementation of Double-Stitch at https://github.com/BasisResearch/stitching.