Neural Wasserstein Splines: Representation-Invariant Second-Order Priors for Trajectory Inference

Amartya Banerjee; Harlin Lee; Caroline Moosmüller

cs.LGOct 14, 2026

Neural Wasserstein Splines: Representation-Invariant Second-Order Priors for Trajectory Inference

Authors: Amartya Banerjee, Harlin Lee, Caroline Moosmüller

Organizations: Dept. of Computer Science, UNC Chapel Hill · SDSS, UNC Chapel Hill · Dept. of Mathematics, UNC Chapel Hill

Abstract

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 $v_{M}=\nabla\phi_{t}$ to particle velocities. Blue/orange indicates lower/higher potential $\phi_{t}$ , scaled separately at each time; shading fades in low-density regions. (c) Differentiating the fitted field gives cloud acceleration $a_{M}$ , 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

a_{M}

is unchanged for a fixed dictionary.

	Leave-one-out		Leave-alternate-out
Dataset	Spline-only	NWSP	Spline-only	NWSP
Diverging Gauss.	$\mathbf{1.084}\pm 0.011$	$\underline{1.136}\pm 0.017$	–	–
Petal	$\underline{0.153}\pm 0.004$	$\mathbf{0.090}\pm 0.014$	$\underline{0.199}\pm 0.009$	$\mathbf{0.086}\pm 0.011$
Converging Gauss.	$\underline{0.466}\pm 0.025$	$\mathbf{0.452}\pm 0.022$	–	–
Embryoid body	$\underline{0.450}\pm 0.003$	$\mathbf{0.350}\pm 0.007$	$\underline{0.457}\pm 0.002$	$\mathbf{0.412}\pm 0.007$
CITE-seq	$\underline{30.586}\pm 0.052$	$\mathbf{30.545}\pm 0.106$	–	–

Table 2: Effect of the acceleration prior. Held-out

W_{1}

(mean

\pm

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$ $\pm\,0.046$	$0.683$ $\pm\,0.056$	$\underline{0.306}$ $\pm\,0.025$	$\mathbf{0.062}$ $\pm\,0.004$
5	$0.927$ $\pm\,0.025$	$0.918$ $\pm\,0.023$	$\underline{0.899}$ $\pm\,0.024$	$\mathbf{0.137}$ $\pm\,0.016$
10	$0.962$ $\pm\,0.006$	$\underline{0.960}$ $\pm\,0.008$	$\underline{0.960}$ $\pm\,0.006$	$\mathbf{0.199}$ $\pm\,0.016$
30	$\underline{0.975}$ $\pm\,0.002$	$\underline{0.975}$ $\pm\,0.002$	$\underline{0.975}$ $\pm\,0.002$	$\mathbf{0.341}$ $\pm\,0.021$

Table 3: Learning the projection space. Relative acceleration error on analytic reference flows (mean

\pm

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$	$\|\mathcal{H}\|_{\mathrm{loo}}$	$\|\mathcal{H}\|_{\mathrm{alt}}$	$M_{0}$
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

|\mathcal{H}|

under leave-one-out (loo) and leave-alternate-out (alt), and fixed-dictionary size

M_{0}

; 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	$\mathbf{9/10}$	$\mathbf{-2.3}$ %	0.011	$\underline{8/10}$	$\mathbf{-2.3}$ %	0.055
4-point	$\mathbf{8/10}$	$\mathbf{-4.2}$ %	0.055	$\mathbf{8/10}$	$\mathbf{-4.2}$ %	0.055
OT-CFM	$\mathbf{8/10}$	$\underline{-4.2}$ %	0.055	$\mathbf{8/10}$	$\mathbf{-4.9}$ %	0.055
SB-CFM	$\mathbf{8/10}$	$\underline{-5.0}$ %	0.055	$\mathbf{8/10}$	$\mathbf{-5.7}$ %	0.055
CFM	$\mathbf{9/10}$	$\underline{-12.8}$ %	0.011	$\mathbf{9/10}$	$\mathbf{-13.5}$ %	0.011

Appendix

Table 7: Dataset-level comparison with each baseline. Relative

W_{1}

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\pm 0.004$	$\underline{0.038}\pm 0.005$	$\mathbf{0.032}\pm 0.003$
5	$0.137\pm 0.016$	$\mathbf{0.018}\pm 0.002$	$\underline{0.112}\pm 0.008$
10	$0.199\pm 0.016$	$\mathbf{0.019}\pm 0.002$	$\underline{0.195}\pm 0.012$
30	$\underline{0.341}\pm 0.021$	$\mathbf{0.027}\pm 0.003$	$0.347\pm 0.021$

Appendix

Table 8: Optimization and dictionary size for NWSP-N. Relative acceleration error as in Table 3 (mean

\pm

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.

Leave-one-out
Dataset	NWSP (ours) Projected acceleration	Raw acceleration	Raw vs projected
Diverging Gauss.	$\underline{1.136}$ $\,\pm\,0.017$	$\mathbf{1.110}$ $\,\pm\,0.014$	$-2.3$ %
Petal	$\mathbf{0.090}$ $\,\pm\,0.014$	$\underline{0.100}$ $\,\pm\,0.011$	$+11.8$ %
Converging Gauss.	$\mathbf{0.452}$ $\,\pm\,0.022$	$\underline{0.498}$ $\,\pm\,0.022$	$+10.1$ %
Embryoid body	$\mathbf{0.350}$ $\,\pm\,0.007$	$\underline{0.361}$ $\,\pm\,0.006$	$+3.2$ %
CITE-seq	$\underline{30.545}$ $\,\pm\,0.106$	$\mathbf{30.539}$ $\,\pm\,0.139$	$-0.0$ %

Appendix

Table 9: Population versus particle acceleration. Held-out

W_{1}

(mean

\pm

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 $W_{1}$ (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 $\beta$ -cell and each method’s prediction of them, as in Fig. 4 .

Trajectory Inference

Trajectory inference reconstructs continuous paths and underlying dynamics from sparse, noisy, or incomplete observations.

See all 37 papers

Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals
Oct 2, 2026Fedor Sergeev, Markus Heinonen, Daniel Waxman +4Dynamical Systems Residual Dynamics Learning
Basis Research Institute · ETH Zurich · Helmholtz AI +2
Wasserstein Residuals: Learning Gradient Flows from Population Dynamics
Jul 6, 2026Markus Heinonen, Yair Shenfeld, Ricardo Baptista +4Dynamical Systems Wasserstein Gradient Flows
Basis Research Institute · Aalto University · Brown University +2
Neural solver for Wasserstein Geodesics and optimal transport dynamics
Feb 25, 2026Hailiang Liu, Yan-Han ChenOptimal Transport
Department of Mathematics, Iowa State University, Ames, IA 50011, USA. · Department of Statistics, Iowa State University, Ames, IA 50011, USA.

Oct 2, 2026J/K move · Enter open · S save

Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals

Fedor Sergeev, Markus Heinonen, Daniel Waxman, Tim Cooijmans+3

Basis Research Institute · ETH Zurich · Helmholtz AI+2

Dynamical Systems Residual Dynamics Learning

Open paper AskPDF arXiv

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.

Abstract

Trajectory Inference

Explore similar work

Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals

Wasserstein Residuals: Learning Gradient Flows from Population Dynamics

Neural solver for Wasserstein Geodesics and optimal transport dynamics