Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals

Fedor Sergeev; Markus Heinonen; Daniel Waxman; Tim Cooijmans; Ricardo Baptista; Dmitry Batenkov; Eli Bingham

cs.LGOct 2, 2026

Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals

Authors: Fedor Sergeev, Markus Heinonen, Daniel Waxman, Tim Cooijmans, Ricardo Baptista, Dmitry Batenkov, Eli Bingham

Organizations: Basis Research Institute · ETH Zurich · Helmholtz AI · MIT ·

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.

Figures & tables

Figure 1: Simulation-free fitting across three dynamical regimes. Left: conservative, damped and gradient-flow examples, with analytic densities, energies and Double-Stitch fits. Right: schematic training procedures for Double-Stitch and WLM. Grey curves show kinetic and total energy for the oscillators; the gradient-flow panel shows free energy. Grey contours show analytic densities.

Family	Method	Inputs	Extrapolation	Encoded flow		Simulation-free
Family	Method	Inputs	Extrapolation	Gradient	Lagrangian	Simulation-free
Flow-based	Flow matching, bridges	$X_{i}$	✗	✗	✗	✓
Flow-based	Curly-FM	$X_{i}$ , reference process	✗ *	✗	✗ *	✓
Gradient flow	JKOnet family	$X_{i}$	✓	✓	✗	✓/✗ ‡
Gradient flow	Stitching	$X_{i}$	✓	✓	✗	✓
Lagrangian flow	WLM	$X_{i},\ (v_{0})$	✓	✓ †	✓	✗

Table 1: Comparison of methods for learning population dynamics. Double-Stitch is a simulation-free method that extrapolates in time and learns both gradient and second-order Lagrangian flows.

	Method	LPO, $W_{1}$	L1O, $W_{1}$	L2O, $W_{2}$
Transport, bridges & neural flows	Neural SDE	$0.820\phantom{{}\pm 0.000}$	—	—
	TrajectoryNet	$0.930\phantom{{}\pm 0.000}$	$0.848\phantom{{}\pm 0.000}$	$1.980\phantom{{}\pm 0.000}$
	SB-FBSDE	$0.840\phantom{{}\pm 0.000}$	—	—
	NLSB	$0.780\phantom{{}\pm 0.000}$	$0.970\phantom{{}\pm 0.000}$	—
	OT-CFM	$0.770\phantom{{}\pm 0.000}$	$0.790\pm 0.068$	—
	WLF-OT	$0.740\phantom{{}\pm 0.000}$	$0.742\pm 0.012$	—

Table 2: Double-Stitch matches performance of the best gradient and Lagrangian methods on the EB dataset (

W_{1}

and

W_{2}

, mean ± std, 5 seeds). Lower is better. Dashes indicate unreported cells.

V(x)

and

V(x,t)

denote static and time-dependent potentials. Stitching and WLM rows are our reruns ( Section D.2 ); Section E.2 also lists their published values.

	Method	Interp. pushfwd.	Forecast pushfwd.	Forecast rollout
Action & continuity	AM	$0.376\phantom{{}\pm 0.000}$	—	—
	UAM	$0.383\phantom{{}\pm 0.000}$	—	—
	sAM	$0.378\phantom{{}\pm 0.000}$	—	—
	DICE	$0.244\phantom{{}\pm 0.000}$	—	—
Transport & bridges	Vanilla-SB	$0.214\phantom{{}\pm 0.000}$	—	—
	OT-CFM	$0.204\phantom{{}\pm 0.000}$	—	—

Table 3: Double-Stitch outperforms gradient methods and matches the best Lagrangian method on ocean vortex interpolation (

W_{1}

, mean ± std, 3 seeds). Lower is better. Dashes indicate unreported cells.

V(x)

and

V(x,t)

denote static and time-dependent potentials. Stitching and WLM rows are our reruns ( Section D.2 ); Section E.3 also lists WLM’s published values.

Figure 5

Appendix figures & tables22 assets

Supplementary material from the paper’s appendix.

Appendix

Indices
Symbol	Meaning
$i$	Snapshot, $i=0,\dots,K$
$n$	Sample within a snapshot, $n=1,\dots,N_{i}$
$k,l$	Learned centers, $k,l=1,\dots,P$
$j$	Time step, $j=0,\dots,M$
Problem setting

Appendix

Table 5: Notation

Figure 3: EB snapshots over time. Shared UMAP.

Figure 4: Ocean vortex snapshots. Color indicates time. Circular points are used for training, diamonds for testing.

Panel of Figure 5 — Figure 5: Canonical and general velocities can produce the same snapshots. Left: translation with the canonical velocity (solid arrows). Right: translation and rotation (dashed arcs) with a general velocity. Color indicates time.

Figure 6: Snapshots are not identifiable without the initial velocity. Left: simulated trajectories with their individual initial velocities pass through the same points and are therefore not identifiable from them. Right: shared circular potential.

Figure 7: Potentials are not identifiable from only snapshots. Left: simulated trajectories with their shared initial velocities pass through the same points and are therefore not identifiable from them. Right: the potentials that produced the corresponding trajectories.

Method	Source / settings	$t_{1}$	$t_{2}$	$t_{3}$	$t_{4}$	Mean
Neural SDE	Li et al. 2020	$0.690\phantom{{}\pm 0.000}$	$0.910\phantom{{}\pm 0.000}$	$0.850\phantom{{}\pm 0.000}$	$0.810\phantom{{}\pm 0.000}$	$0.820\phantom{{}\pm 0.000}$
TrajectoryNet	Tong et al. 2020	$0.730\phantom{{}\pm 0.000}$	$1.060\phantom{{}\pm 0.000}$	$0.900\phantom{{}\pm 0.000}$	$1.010\phantom{{}\pm 0.000}$	$0.930\phantom{{}\pm 0.000}$
SB-FBSDE	Chen et al. 2022	$0.560\phantom{{}\pm 0.000}$	$0.800\phantom{{}\pm 0.000}$	$1.000\phantom{{}\pm 0.000}$	$1.000\phantom{{}\pm 0.000}$	$0.840\phantom{{}\pm 0.000}$
NLSB	Koshizuka & Sato 2023	$0.680\phantom{{}\pm 0.000}$	$0.840\phantom{{}\pm 0.000}$	$0.810\phantom{{}\pm 0.000}$	$0.790\phantom{{}\pm 0.000}$	$0.780\phantom{{}\pm 0.000}$
OT-CFM	Tong et al. 2024	$0.780\phantom{{}\pm 0.000}$	$0.760\phantom{{}\pm 0.000}$	$0.770\phantom{{}\pm 0.000}$	$0.750\phantom{{}\pm 0.000}$	$0.770\phantom{{}\pm 0.000}$
WLF-OT	Neklyudov et al. 2024	$0.650\phantom{{}\pm 0.000}$	$0.780\phantom{{}\pm 0.000}$	$0.760\phantom{{}\pm 0.000}$	$0.750\phantom{{}\pm 0.000}$	$0.740\phantom{{}\pm 0.000}$

Appendix

Table 6: Per-snapshot performance on EB LPO (

W_{1}

, mean ± std, 5 seeds, lower is better).

Method	Source / settings	$t_{1}$	$t_{2}$	$t_{3}$	Mean
TrajectoryNet	Kapuśniak et al. 2024	—	—	—	$0.848\phantom{{}\pm 0.000}$
NLSB	Kapuśniak et al. 2024	—	—	—	$0.970\phantom{{}\pm 0.000}$
OT-CFM	Kapuśniak et al. 2024	—	—	—	$0.790\pm 0.068$
WLF-OT	Kapuśniak et al. 2024	—	—	—	$0.742\pm 0.012$
WLF-SB	Kapuśniak et al. 2024	—	—	—	$0.746\pm 0.016$
Reg. CNF	Kapuśniak et al. 2024	—	—	—	$0.825\pm 0.429$

Appendix

Table 7: Per-snapshot performance on EB L1O (

W_{1}

, mean ± std, 5 seeds, lower is better).

Method	Source / settings	$t_{1}$	$t_{3}$	Mean
TrajectoryNet	Tong et al. 2020	$2.030\pm 0.040$	$1.930\pm 0.080$	$1.980\phantom{{}\pm 0.000}$
Vanilla-SB	Vargas et al. 2021	$1.490\pm 0.060$	$1.550\pm 0.030$	$1.520\phantom{{}\pm 0.000}$
MMSB	Shen et al. 2025	$1.270\pm 0.030$	$1.570\pm 0.050$	$1.420\phantom{{}\pm 0.000}$
DMSB	Chen et al. 2023	$1.130\pm 0.080$	$1.450\pm 0.160$	$1.290\phantom{{}\pm 0.000}$
JKOnet ∗ , $V(x)$	Terpin et al. 2024	$1.150\pm 0.030$	$2.530\pm 0.010$	$1.840\phantom{{}\pm 0.000}$
JKOnet ∗ , $V(x,t)$	Terpin et al. 2024	$4.410\pm 1.500$	$2.770\pm 0.200$	$3.590\phantom{{}\pm 0.000}$

Appendix

Table 8: Per-snapshot performance on EB L2O (

W_{2}

, mean ± std, 5 seeds, lower is better).

Figure 8: Double-Stitch predictions on EB L1O (

V(x)

, seed

0

). Shared UMAP. Pale points are ground truth, dark points are predictions.

Figure 9: Double-Stitch predictions on EB L2O (

V(x)

, seed

0

). Shared UMAP. Pale points are ground truth, dark points are predictions.

Method	Source / settings	$t_{1}$	$t_{3}$	$t_{5}$	$t_{7}$	Mean
AM	Neklyudov et al. 2023	$0.358\pm 0.009$	$0.353\pm 0.020$	$0.447\pm 0.013$	$0.346\pm 0.016$	$0.376\phantom{{}\pm 0.000}$
UAM	Neklyudov et al. 2023	$0.291\pm 0.008$	$0.378\pm 0.021$	$0.514\pm 0.010$	$0.349\pm 0.009$	$0.383\phantom{{}\pm 0.000}$
sAM	Neklyudov et al. 2023	$0.358\pm 0.009$	$0.354\pm 0.010$	$0.451\pm 0.016$	$0.350\pm 0.009$	$0.378\phantom{{}\pm 0.000}$
DICE	Blickhan et al. 2025	$0.369\pm 0.019$	$0.252\pm 0.020$	$0.240\pm 0.011$	$0.115\pm 0.009$	$0.244\phantom{{}\pm 0.000}$
Vanilla-SB	Vargas et al. 2021	$0.112\pm 0.061$	$0.178\pm 0.073$	$0.262\pm 0.043$	$0.305\pm 0.032$	$0.214\phantom{{}\pm 0.000}$
OT-CFM	Tong et al. 2024	$0.148\pm 0.004$	$0.227\pm 0.008$	$0.191\pm 0.012$	$0.250\pm 0.018$	$0.204\phantom{{}\pm 0.000}$

Appendix

Table 9: Per-snapshot performance on the small ocean vortex (

W_{1}

, mean ± std, 3 seeds, lower is better).

Method	Source / settings	Rollout $0\!\to\!10$	Pushfwd. $9\!\to\!10$	Direct
SBIRR	Shen et al. 2025	$\phantom{0}1.062\pm 0.111$	$0.567\pm 0.014$	—
Snap-MMD	Berlinghieri et al. 2025	$\phantom{0}0.896\pm 0.141$	$0.473\pm 0.013$	—
WLM, published	Guan et al. 2026a	$\phantom{0}0.180\pm 0.038$	$0.065\pm 0.005$	—
Stitching, $V(x)$	Main settings	$14.370\pm 6.123$	$1.388\pm 0.119$	$0.295\pm 0.032$
Stitching, $V(x,t)$	Main settings	$\phantom{0}0.224\pm 0.017$	$0.195\pm 0.022$	$0.212\pm 0.018$
Stitching, $V(x)$	Original settings	$\phantom{0}0.149\pm 0.007$	$0.205\pm 0.006$	$0.206\pm 0.001$

Appendix

Table 10: Per-snapshot performance on the large ocean vortex (

W_{1}

, mean ± std, 3 seeds, lower is better). Dashes indicate unreported or unavailable cells.

Figure 10: Double-Stitch predictions on the small ocean vortex ( $V(x)$ , seed $0$ ). Left: learned trajectories (used for direct evaluation). Center: pushforward predictions (used for pushforward evaluation). Right: learned potential.

Figure 11: Double-Stitch learns all three Gaussian-mixture flows. Left: ground-truth conservative, damped and gradient densities. Center: Stitching, which fails on the conservative flow and underperforms on the damped flow. Right: Double-Stitch, which succeeds on all three. Dotted lines mark the held-out times.

Interpolation
Method	Settings	Direct	Pushforward	Rollout
Stitching, $V(x)$	Main settings	$0.366\pm 0.015$	$1.515\pm 0.273$	$\phantom{0}3.858\pm 0.912$
Stitching, $V(x,t)$	Main settings	$0.125\pm 0.022$	$0.166\pm 0.035$	$1.089\pm 0.475$
Stitching, $V(x)$	Original settings	$0.301\pm 0.016$	$4.750\pm 0.535$	$32.764\pm 2.684$
Stitching, $V(x,t)$	Original settings	$0.104\pm 0.021$	$0.104\pm 0.020$	$\phantom{0}0.104\pm 0.017$
WLM	Main settings	—	$0.064\pm 0.015$	$\phantom{0}0.176\pm 0.098$

Appendix

Table 11: Double-Stitch and WLM sensitivity to the prediction mode on ocean vortices (

W_{1}

, mean ± std, 3 seeds, lower is better). Dashes indicate unavailable cells.

	EB			Oceans
Setting	LPO	L1O	L2O	Interp.	Forecast
$P=100$	$28.8\pm 0.2$	$84.7\pm 3.3$	$29.0\pm 2.3$	$104.7\pm 1.0$	$104.4\pm 3.0$
$P=400$	$63.3\pm 0.3$	$186.8\pm 2.1$	$62.8\pm 0.6$	$252.7\pm 6.0$	$277.1\pm 2.9$
$P=1000$	$127.3\pm 2.9$	$385.5\pm 4.9$	$125.6\pm 0.4$	$512.4\pm 16.9$	$532.5\pm 7.5$

Appendix

Table 12: Double-Stitch training time versus number of particles across datasets (seconds, mean ± std, 3 seeds).

	EB			Oceans
Setting	LPO $W_{1}$	L1O $W_{1}$	L2O $W_{2}$	Interp. $W_{1}$	Forecast $W_{1}$
OT-path, coarsest (ref.)	$0.618\pm 0.001$	$0.733\pm 0.003$	$0.795\pm 0.006$	$0.065\pm 0.004$	$0.134\pm 0.003$
OT-path, Silverman	$0.614\pm 0.001$	$0.729\pm 0.004$	$0.795\pm 0.010$	$0.077\pm 0.021$	$0.131\pm 0.013$
Endpoint OT, Silverman	$0.777\pm 0.004$	$0.814\pm 0.011$	$0.814\pm 0.011$	$0.504\pm 0.274$	$0.578\pm 0.056$
McCann-path, Silverman	$0.640\pm 0.004$	$0.736\pm 0.001$	$0.791\pm 0.002$	$0.121\pm 0.009$	$0.127\pm 0.009$

Appendix

Table 13: Double-Stitch sensitivity to initialization across datasets (mean ± std, 3 seeds, lower is better).

	EB			Oceans
Setting	LPO $W_{1}$	L1O $W_{1}$	L2O $W_{2}$	Interp. $W_{1}$	Forecast $W_{1}$
$V(x)$	$0.618\pm 0.001$	$0.733\pm 0.003$	$0.794\pm 0.005$	$0.065\pm 0.004$	$0.134\pm 0.003$
$V(x,t)$	$0.558\pm 0.004$	$0.716\pm 0.007$	$0.800\pm 0.004$	$0.073\pm 0.009$	$0.119\pm 0.019$

Appendix

Table 14: Double-Stitch sensitivity to a time-dependent potential across datasets (mean ± std, 3 seeds, lower is better).

	EB			Oceans
Setting	LPO $W_{1}$	L1O $W_{1}$	L2O $W_{2}$	Interp. $W_{1}$	Forecast $W_{1}$
$\lambda=0.3$	$0.591\pm 0.005$	$0.724\pm 0.005$	$0.805\pm 0.005$	$0.061\pm 0.002$	$0.134\pm 0.008$
$\lambda=1$ (ref.)	$0.618\pm 0.001$	$0.733\pm 0.003$	$0.795\pm 0.006$	$0.065\pm 0.004$	$0.134\pm 0.003$
$\lambda=3$	$0.633\pm 0.003$	$0.739\pm 0.006$	$0.786\pm 0.004$	$0.071\pm 0.011$	$0.129\pm 0.004$

Appendix

Table 15: Double-Stitch sensitivity to the residual-loss weight across datasets (mean ± std, 3 seeds, lower is better).

	EB			Oceans
Setting	LPO $W_{1}$	L1O $W_{1}$	L2O $W_{2}$	Interp. $W_{1}$	Forecast $W_{1}$
$\gamma_{0}=0.1$	$0.610\pm 0.003$	$0.728\pm 0.004$	$0.793\pm 0.003$	$0.062\pm 0.001$	$0.113\pm 0.018$
$\gamma_{0}=1$ (ref.)	$0.618\pm 0.001$	$0.733\pm 0.003$	$0.795\pm 0.006$	$0.065\pm 0.004$	$0.134\pm 0.003$
$\gamma_{0}=10$	$0.740\pm 0.003$	$0.812\pm 0.003$	$0.827\pm 0.003$	$0.148\pm 0.005$	$0.127\pm 0.007$

Appendix

Table 16: Double-Stitch sensitivity to the initial friction across datasets (mean ± std, 3 seeds, lower is better).

	EB			Oceans
Setting	LPO $W_{1}$	L1O $W_{1}$	L2O $W_{2}$	Interp. $W_{1}$	Forecast $W_{1}$
$P=100$	$0.667\pm 0.013$	$0.766\pm 0.021$	$0.833\pm 0.003$	$0.065\pm 0.004$	$0.122\pm 0.018$
$P=400$	$0.629\pm 0.009$	$0.735\pm 0.002$	$0.794\pm 0.003$	$0.079\pm 0.004$	$0.134\pm 0.003$
$P=1000$	$0.618\pm 0.001$	$0.733\pm 0.003$	$0.795\pm 0.006$	$0.099\pm 0.013$	$0.098\pm 0.008$

Appendix

Table 17: Double-Stitch sensitivity to the number of particles across datasets (mean ± std, 3 seeds, lower is better).

Dynamical Systems

Dynamical systems studies how states evolve over time, including stability, feedback, chaos, and transitions in physical, biological, and engineered systems.

See all 478 papers

May 8, 2026

A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots

The population dynamics of molecules, cells, and organisms are governed by a number of unknown forces. In the last decade, population dynamics have predominantly been modeled with Wasserstein gradient flows. However, since gradient flows minimize free energy, they fail to capture important dynamical properties, such as periodicity. In this work, we propose a change in perspective by considering dynamics that minimize a population-level action under a damped Wasserstein Lagrangian. By deriving the corresponding Hamiltonian equations of motion, we formalize Wasserstein Lagrangian Mechanics, a structured class of second-order dynamics that encompasses classical mechanics, quantum mechanics, and gradient flows. We then propose WLM as the first algorithm that learns these second-order dynamics from observed marginals, without specifying the Lagrangian. By directly learning the population mechanics, WLM can both forecast and interpolate unseen marginals, and outperforms existing gradient flow and flow matching methods across a wide range of dynamics, including vortex dynamics, embryonic development, and flocking.

Vincent Guan, Lazar Atanackovic, Kirill Neklyudov

University of British Columbia · Mila - Quebec AI Institute · University of Alberta +4

Dynamical Systems

Ask PDF

Jul 6, 2026

Wasserstein Residuals: Learning Gradient Flows from Population Dynamics

Reconstructing population dynamics is a central problem in the physical and data sciences. Often, the dynamics are modeled as a Wasserstein gradient flow (WGF): a curve of distributions driven by an energy functional. Though there are multiple mathematical characterizations of a WGF, the dominant algorithmic approach relies on the Jordan--Kinderlehrer--Otto (JKO) scheme. JKO-based methods are inflexible to time discretisation and require solving costly optimal transport problems. We take a residual approach, enforcing the continuity equations via a non-negative loss function whose minimum is the WGF. Combined with a data-fitting divergence, this gives a single global objective. This perspective unifies several existing methods and leads to a new particle-based method, stitching, that is simulation-free and robust to large gaps between observations. We demonstrate that the stitching method achieves state-of-the-art performance across trajectory inference benchmarks. For code see github.com/BasisResearch/wasserstein-residuals.

Markus Heinonen, Yair Shenfeld, Ricardo Baptista +4

Basis Research Institute · Aalto University · Brown University +2

Dynamical Systems Wasserstein Gradient Flows Residual Learning Trajectory Inference Optimal Transport

Ask PDF

Oct 24, 2026

Learning Population Dynamics with Variational Doob-Wasserstein Flow

This paper introduces a novel simulation-free framework for learning unbalanced population dynamics on the Riemannian Wasserstein-Fisher-Rao (WFR) manifold, termed Variational Doob-Wasserstein Flow (VDWF), built upon the theory of the unbalanced mean-field Schrödinger bridge (MFSB). First, by independently learning a geometry model equipped with a data-dependent metric, our model captures the intrinsic geometric structure of population dynamics. Furthermore, leveraging Doob's h-transform, the proposed method incorporates boundary information to improve flow-based trajectory inference. For unbalanced transport, a growth model is employed to learn the dynamics while constructing intermediate mass samples in the WFR space without simulation. Additionally, we analyze the convergence rate of the proposed framework using a uniform Log-Sobolev inequality (LSI). Finally, we validate the proposed approach on multiple real-world datasets, demonstrating competitive performance in learning both balanced and unbalanced population dynamics compared with state-of-the-art methods.

Weizhu Qian

School of Computer Science and Technology Soochow University, China

Optimal Transport Dynamical Systems Schrödinger Bridges

Ask PDF

Abstract

Dynamical Systems

Explore similar work

A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots

Wasserstein Residuals: Learning Gradient Flows from Population Dynamics

Learning Population Dynamics with Variational Doob-Wasserstein Flow