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 this https URL.
From gradient flows to Lagrangian mechanics in Wasserstein space
Probability distributions evolving over time appear in domains ranging from cellular biology to fluid dynamics. Traditional modeling often assumes gradient flow dynamics, where the evolution follows the steepest descent of a free energy landscape. However, many important systems exhibit conservative or periodic behavior that gradient flows cannot capture. Lagrangian mechanics reformulates the problem in Wasserstein space, extending the variational framework to handle these richer dynamical classes.
The simulation bottleneck in existing approaches
Existing methods for learning Wasserstein Lagrangian mechanics rely on running a numerical solver at every training step. This simulation-based requirement makes training computationally expensive and limits scalability. The need to integrate equations of motion repeatedly during optimization creates a significant barrier, particularly for high-dimensional or long-horizon dynamics.
Double-Stitch: a simulation-free residual penalty
Double-Stitch avoids solving ODEs during training entirely. Instead, it penalizes the residual of the equation of motion along a learned population path. The key insight is that this residual can be evaluated without ever simulating forward in time, by working directly with the observed snapshots and the learned path geometry.
Derivation from a Clebsch variational principle
The equation of motion is derived from a Clebsch variational principle that does not require gradient velocities. This formulation separates the conservative and dissipative components of the dynamics, allowing the method to capture both gradient-like relaxation and Hamiltonian-style preservation. The residual vanishes exactly when the equation holds, providing a clean optimization target.
Experimental evaluation across domains
Synthetic benchmarks test the method's ability to recover known dynamics from incomplete observations. Single-cell datasets evaluate performance on cellular population evolution, where tracking distribution shifts over time is critical. Ocean vortex datasets assess extrapolation capabilities for geophysical flows. In all cases, Double-Stitch achieves competitive or superior performance compared to gradient-flow baselines and simulation-based Wasserstein Lagrangian methods (WLM).
Training efficiency gains
The most striking improvement is in training speed. Double-Stitch trains 4-14 times faster than WLM while matching or exceeding its performance on most tasks. This dramatic reduction comes from eliminating the per-step simulation requirement, while the Clebsch-based residual penalty provides a stable learning signal without needing numerical integration.
JAX implementation and practical considerations
The authors provide a JAX implementation of Double-Stitch, making the method accessible for research and development. The code operates on discrete snapshot pairs and computes the residual directly from the observed data and the learned path parameters. Users can integrate Double-Stitch into existing optimal transport or Wasserstein framework pipelines with minimal adaptation.
Limitations and future directions
The method's performance depends on the quality of the learned population path and the expressivity of the residual model. For highly turbulent or multi-scale systems, additional inductive biases may be needed. Future work could explore adaptive residual penalties, multi-fidelity approaches combining sparse simulation with the residual framework, and extensions to stochastic settings.