Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals

2026-10-02

AI Summary

This study addresses the inability of Wasserstein gradient flows to capture conservative dynamics and the prohibitive training costs of existing Lagrangian methods that rely on numerical simulations. To overcome these limitations, this work proposes Double-Stitch, a simulation-free framework. By leveraging the Clebsch variational principle, the method derives velocity-gradient-free equations of motion and achieves efficient learning of Lagrangian mechanics in Wasserstein space by penalizing equation residuals along learned trajectories. Theoretically, vanishing residuals guarantee that the equations of motion are satisfied, thereby entirely eliminating numerical solving steps during training. Experiments demonstrate that Double-Stitch matches or surpasses baseline performance on synthetic, single-cell, and ocean eddy datasets while accelerating training by 4 to 14 times.

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.

Problem

• Population Dynamics
• Wasserstein Space
• Lagrangian Mechanics
• Unpaired Snapshots
• Simulation-Free Learning

Innovation

• Simulation-Free Learning
• Wasserstein Lagrangian Mechanics
• Clebsch Variational Principle
• Population Dynamics
• Residual Penalization