Population dynamics are often observed only as unpaired snapshots, which leave the evolution between measurements underdetermined. Second-order regularization can constrain it, but a penalty on the accelerations of sampled particles is a prior on joint particle paths, and different particle representations of the same marginal curve can have different acceleration costs. We therefore introduce Neural Wasserstein Splines (NWSP), a sample-based method that learns a continuous population curve under a Wasserstein smoothness prior. A neural network maps latent variables to the control points of a cubic spline in time, so samples at any time come from the spline basis without numerical integration. We project the generator's velocity onto gradients of a finite dictionary of scalar potentials and differentiate this measure-dependent projection in time, obtaining a finite-dimensional surrogate for the intrinsic Wasserstein acceleration, trained jointly with distributional fit losses. For a fixed dictionary and ridge parameter, we prove that the population penalty depends only on the marginal curve and bound its sensitivity to moment-estimation errors. Learned potentials (NWSP-N) sharpen acceleration estimates. On various synthetic and single-cell benchmarks, NWSP is highly competitive with transport and flow-based methods, ranking in the top two in 11 of 16 settings. The penalty improves on an unregularized spline in 13 of 16 settings and, while competitive with a particle-acceleration penalty, adds provable representation invariance.
Figures & tables
Figure 1: Learning a smooth population curve from unpaired snapshots. (a) A network maps each latent sample to spline control points (squares), defining a complete path. (b) Ridge regression fits the velocity vM=∇ϕt to particle velocities. Blue/orange indicates lower/higher potential ϕt , scaled separately at each time; shading fades in low-density regions. (c) Differentiating the fitted field gives cloud acceleration aM , whose magnitude colors the same paths. Training balances snapshot fit with mean squared cloud acceleration. NWSP-N also learns the potentials in (b).
Figure 2: Same population, different particle paths. Rotating the latent changes particle paths (solid versus dashed) but preserves every marginal. Dots show the observed snapshots; gray lines show other rotated paths. Particle acceleration thus depends on the representation, whereas the population acceleration aM is unchanged for a fixed dictionary.
Leave-one-out
Leave-alternate-out
Dataset
Spline-only
NWSP
Spline-only
NWSP
Diverging Gauss.
1.084±0.011
1.136±0.017
–
–
Petal
0.153±0.004
0.090±0.014
0.199±0.009
0.086±0.011
Converging Gauss.
0.466±0.025
0.452±0.022
–
–
Embryoid body
0.450±0.003
0.350±0.007
0.457±0.002
0.412±0.007
CITE-seq
30.586±0.052
30.545±0.106
–
–
Table 2: Effect of the acceleration prior. Held-out W1 (mean ± SD over three seeds) for the same generator trained without (Spline-only) or with the population-acceleration penalty. Bold and underline mark the best and second-best mean within each protocol; a dash marks an inapplicable protocol.
d
NWSP (ours)
+16 ridge
+16 untrained
NWSP-N (ours)
2
0.705±0.046
0.683±0.056
0.306±0.025
0.062±0.004
5
0.927±0.025
0.918±0.023
0.899±0.024
0.137±0.016
10
0.962±0.006
0.960±0.008
0.960±0.006
0.199±0.016
30
0.975±0.002
0.975±0.002
0.975±0.002
0.341±0.021
Table 3: Learning the projection space. Relative acceleration error on analytic reference flows (mean ± SD over five seeds; lower is better). Each added block has sixteen potentials; the NWSP-N block is trained for 300 updates. Bold and underline mark the best and second-best means, ties included.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Dataset
d
T
∣H∣loo
∣H∣alt
M0
Circular Gaussian
2
10
8
4
41
Dyngen Cycle
10
15
13
7
101
Dyngen Tree
10
7
5
3
101
Petal
2
5
3
2
41
Converging Gaussian
2
4
2
–
41
Diverging Gaussian
2
4
2
–
41
Appendix
Table 4: Datasets and held-out tasks. State dimension d , snapshots T , held-out targets ∣H∣ under leave-one-out (loo) and leave-alternate-out (alt), and fixed-dictionary size M0 ; NWSP-N adds K=16 potentials. A dash marks an inapplicable protocol.
NWSP (ours)
NWSP-N (ours)
Baseline
wins/ n
median
p
wins/ n
median
p
WLR
9/10
−2.3 %
0.011
8/10
−2.3 %
0.055
4-point
8/10
−4.2 %
0.055
8/10
−4.2 %
0.055
OT-CFM
8/10
−4.2 %
0.055
8/10
−4.9 %
0.055
SB-CFM
8/10
−5.0 %
0.055
8/10
−5.7 %
0.055
CFM
9/10
−12.8 %
0.011
9/10
−13.5 %
0.011
Appendix
Table 7: Dataset-level comparison with each baseline. Relative W1 differences are averaged over protocols within each dataset before counting wins. Columns: win count, median relative change (negative favors ours), and one-sided exact sign-test p (post hoc, uncorrected). Bold and underline compare win counts and medians of our two variants.
NWSP-N (ours): optimization and dictionary size
d
16 learned 300 steps
16 learned 3000 steps
64 learned 300 steps
2
0.062±0.004
0.038±0.005
0.032±0.003
5
0.137±0.016
0.018±0.002
0.112±0.008
10
0.199±0.016
0.019±0.002
0.195±0.012
30
0.341±0.021
0.027±0.003
0.347±0.021
Appendix
Table 8: Optimization and dictionary size for NWSP-N. Relative acceleration error as in Table 3 (mean ± SD over five seeds); the first column repeats the standard configuration, the others use more updates or more learned potentials. Bold and underline mark the best and second-best means.
Dataset
NWSP (ours) Projected acceleration
Raw acceleration
Raw vs projected
Leave-one-out
Diverging Gauss.
1.136±0.017
1.110±0.014
−2.3 %
Petal
0.090±0.014
0.100±0.011
+11.8 %
Converging Gauss.
0.452±0.022
0.498±0.022
+10.1 %
Embryoid body
0.350±0.007
0.361±0.006
+3.2 %
CITE-seq
30.545±0.106
30.539±0.139
−0.0 %
Appendix
Table 9: Population versus particle acceleration. Held-out W1 (mean ± SD over three seeds) with the projected penalty (NWSP) or a raw particle-acceleration penalty, otherwise identical. Last column: change from projected to raw (positive favors NWSP). The two arms ran on different hardware.
Figure 3: Fitted synthetic population trajectories. Each model is fitted to all available snapshots; no time point is withheld. Colors indicate time and snapshot index. These plots illustrate the resulting dynamics rather than held-out generalization, which is measured separately in Table 1 .
Figure 4: Held-out snapshots of (a) Mouse hematopoiesis, (b) Embryoid body (PHATE) and (c) CITE-seq, and each method’s prediction of them; every column is a separate leave-one-out fit that never saw its snapshot, and the number under each panel is that holdout’s W1 (red: the lowest of the column). Clouds use seed zero; scores average the three seeds. Gray clouds are retained neighbors. CITE-seq/Veres PCA and Embryoid-body-standardized PHATE use retained frames only. These views do not establish full-dimensional rankings.
Figure 5: Held-out snapshots of (d) Veres β -cell and each method’s prediction of them, as in Fig. 4 .
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.
Fedor Sergeev, Markus Heinonen, Daniel Waxman +4
Basis Research Institute · ETH Zurich · Helmholtz AI +2
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
In recent years, the machine learning community has increasingly embraced the optimal transport (OT) framework for modeling distributional relationships. In this work, we introduce a sample-based neural solver for computing the Wasserstein geodesic between a source and target distribution, along with the associated velocity field. Building on the dynamical formulation of the optimal transport (OT) problem, we recast the constrained optimization as a minimax problem, using deep neural networks to approximate the relevant functions. This approach not only provides the Wasserstein geodesic but also recovers the OT map, enabling direct sampling from the target distribution. By estimating the OT map, we obtain velocity estimates along particle trajectories, which in turn allow us to learn the full velocity field. The framework is flexible and readily extends to general cost functions, including the commonly used quadratic cost. We demonstrate the effectiveness of our method through experiments on both synthetic and real datasets.
Hailiang Liu, Yan-Han Chen
Department of Mathematics, Iowa State University, Ames, IA 50011, USA. · Department of Statistics, Iowa State University, Ames, IA 50011, USA.