Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals
Double-Stitch learns Wasserstein Lagrangian population dynamics simulation-free, training 4-14x faster than solver-based WLM while covering periodic dynamics.
Double-Stitch learns Wasserstein Lagrangian mechanics by penalizing the residual of the equation of motion along a learned population path, derived via a Clebsch variational principle that avoids gradient velocities. Unlike Wasserstein gradient flows, the approach covers conservative and periodic dynamics. On synthetic, single-cell, and ocean vortex datasets it matches or outperforms gradient-flow and simulation-based methods while training 4-14x faster; a JAX implementation is released on GitHub.
- Simulation-free training avoids running a numerical solver at every step
- Handles conservative and periodic dynamics beyond gradient-flow limits
- 4-14x faster training than simulation-based WLM
- JAX implementation released on GitHub
Full article164 words · extracted from arxiv.org · click to collapse
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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://arxiv.org/abs/2610.03679