Double-Stitch: Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Mechanics

Forum topic · 小凯 · 2026-10-06

Summary

Dynamics of cells, organisms, and fluids are often modeled as probability distributions evolving over time. Reconstructing such evolution from unpaired snapshots typically assumes Wasserstein gradient flows, which cannot capture conservative or periodic dynamics. Lagrangian mechanics in Wasserstein space can model both, but existing learning methods are simulation-based, running a numerical solver at every training step and making training expensive. This paper introduces Double-Stitch, a simulation-free method that learns Wasserstein Lagrangian mechanics by penalizing the residual of the equation of motion along a learned population path. The equation is derived from a Clebsch variational principle that does not require gradient velocities, and the authors prove the residual is exactly zero only when the equation holds. On synthetic, single-cell, and ocean eddy datasets, Double-Stitch matches or outperforms gradient-flow methods and simulation-based WLM on most tasks while training 4-14x faster. A JAX implementation is available at https://github.com/BasisResearch/stitching. arXiv: 2610.03679.

Paper Overview

Field: Machine Learning Authors: Fedor Sergeev, Markus Heinonen, Daniel Waxman, Tim Cooijmans, Ricardo Baptista, Dmitry Batenkov, Eli Bingham Published: 2026-10-02 arXiv: 2610.03679

Summary

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 regimes, but existing methods for learning it are simulation-based: they run a numerical solver at every training step, which makes training expensive.

Key Contributions

Double-Stitch, a simulation-free method that learns Wasserstein Lagrangian mechanics by penalizing the residual of the equation of motion along a learned population path.
The equation is derived from a Clebsch variational principle that does not require gradient velocities.
A proof that the residual is exactly zero only when the equation of motion holds.

Results

Evaluated on synthetic, single-cell, and ocean eddy datasets, Double-Stitch matches or outperforms gradient-flow methods and simulation-based WLM on most tasks, while training 4-14x faster than WLM.

Code

A JAX implementation is available at https://github.com/BasisResearch/stitching.

Double-Stitch: Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Mechanics

Summary

Paper Overview

Summary

Key Contributions

Results

Code

Tags