Track trials, findings, and clinical evidence as it lands.
Search papers, labs, and topics across Lattice.
This paper introduces Double-Stitch, a simulation-free method for learning population dynamics governed by Lagrangian mechanics in Wasserstein space, which can capture conservative and periodic behaviors missed by Wasserstein gradient flows. The method penalizes the residual of the learned equation of motion along a population path derived from a Clebsch variational principle, avoiding expensive numerical solvers during training. Experiments on synthetic, single-cell, and ocean vortex data demonstrate competitive or superior performance to existing methods, with a significant speedup in training time.
Learning conservative and periodic population dynamics is now 4-14x faster without needing numerical solvers.
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.