Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals
Abstract
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.
Related
- WTF?! Simulation-Free Reinforcement Learning with Wasserstein-Tilted Flow Maps
- Hyperbolic phase-field coupling with explicit Lagrangian hydrodynamics for the simulation of damage in brittle materials
- Low Residuals, Wrong Value: Learning the Viscosity Solution of Hamilton-Jacobi-Isaacs Equations with PINNs
- Special Lagrangian cones in Deep Learning
- Minimax Rates for Conditional Generative Learning under Wasserstein Loss
BibTeX
@misc{sergeev2026simulationfree,
title = {{Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals}},
author = {Fedor Sergeev and Markus Heinonen and Daniel Waxman and Tim Cooijmans and Ricardo Baptista and Dmitry Batenkov and Eli Bingham},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.03679},
eprint = {2610.03679},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.03679},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals
AU - Fedor Sergeev
AU - Markus Heinonen
AU - Daniel Waxman
AU - Tim Cooijmans
AU - Ricardo Baptista
AU - Dmitry Batenkov
AU - Eli Bingham
PY - 2026
DA - 2026/10/02
PB - arXiv
DO - 10.48550/arxiv.2610.03679
UR - https://arxiv.org/abs/2610.03679
AB - 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.
ER -