Neural Wasserstein Splines: Representation-Invariant Second-Order Priors for Trajectory Inference
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.
Related
- A Stieltjes representation of order p for $ \mathrm{tr}[(A+tB)^{-p}]$: Heinävaara's measure, B-splines, and Stahl's theorem
- Wasserstein bounds for invariant measures in the Smoluchowski-Kramers approximation with stable noise
- Wasserstein Kernel Regression with Dependent Distribution-Valued Predictors: Local Recurrence, Joint Inference, and Multivariate Gaussian Geometry
- Deterministic Uniformization for Estimation and Latent Trajectory Inference in Interval-Censored Multistate Triansition Models
- Strong Gaussian Approximation of Constant Learning Rate SGD: Bias and Trajectory Inference
BibTeX
@misc{banerjee2026neural,
title = {{Neural Wasserstein Splines: Representation-Invariant Second-Order Priors for Trajectory Inference}},
author = {Amartya Banerjee and Harlin Lee and Caroline Moosmüller},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.19878},
eprint = {2610.19878},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.19878},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Neural Wasserstein Splines: Representation-Invariant Second-Order Priors for Trajectory Inference
AU - Amartya Banerjee
AU - Harlin Lee
AU - Caroline Moosmüller
PY - 2026
DA - 2026/10/14
PB - arXiv
DO - 10.48550/arxiv.2610.19878
UR - https://arxiv.org/abs/2610.19878
AB - 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.
ER -