PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Amartya Banerjee, Harlin Lee, Caroline Moosmüller

arXiv:2610.19878v1cs.LGcs.AIstat.ML

TL;DR

Unpaired snapshots leave population evolution between measurements underdetermined, while particle-acceleration penalties can vary across representations of the same distribution curve. NWSP fits a continuous curve with spline-generated samples and a differentiated population-acceleration penalty. It ranks in the top two in eleven of sixteen benchmark settings, with provable representation invariance for a fixed dictionary and ridge parameter.

  • Problem

    Unpaired snapshots do not determine how a population moves between observations, and particle-acceleration penalties can differ for representations of the same marginal curve.

    Source: §1 · p. 1

  • Method

    NWSP maps latent variables to smooth particle-path splines and differentiates a fitted, measure-dependent velocity projection to regularize population acceleration.

    Source: §1 · p. 2, §3.1 · p. 3

  • Results

    NWSP is highly competitive with transport- and flow-based methods, ranking in the top two in eleven of sixteen settings.

    Source: §5 · p. 8

  • Takeaways & Limitations

    For a fixed dictionary and ridge parameter, the population penalty depends only on the marginal curve, and the acceleration prior improves on spline smoothness in thirteen of sixteen settings.

    Source: §3.6 · p. 5, §5 · p. 8

    Finite dictionaries can miss shape changes that preserve mean and variance, and theory does not guarantee learned potentials resolve every such change.

    Source: §3.6 · p. 6

Abstract

from arXiv · show

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.

1 INTRODUCTION

Unpaired snapshots leave the intervening population trajectory underdetermined, and particle-acceleration penalties can depend on unidentifiable choices of particle paths. NWSP uses smooth particle-path splines to fit snapshots while regularizing acceleration of the distribution curve.

1 INTRODUCTION

  • Unpaired observations at a few times leave motion between snapshots undetermined; acceleration regularization favors gradual motion rather than velocity jumps at observation times.In single-cell experiments, destructive assays prevent matching cells measured at successive times.Source: §1 · p. 1
  • NWSP regularizes the intrinsic acceleration of the distribution curve, avoiding the representation dependence of penalties on individual particle paths.The intrinsic Wasserstein acceleration is the rate of change of the minimum-norm transporting velocity and depends only on the curve.Source: §1 · p. 1
  • A neural network maps latent samples to smooth particle-path splines, enabling sampling at any time without numerical integration or cross-time OT couplings.The generator fits snapshots while penalizing differentiated population acceleration; NWSP-N also learns the velocity potentials.Source: §1 · p. 2, Figure 1 · p. 2

1.1 Related work

  • Unlike OT-based spline construction and flow-matching methods (Lipman et al., 2023; Tong et al., 2024a,b), NWSP uses its spline itself to generate samples.Its population-acceleration estimate is indexed by potentials rather than observed particles.Source: §1 · p. 1, §1.1 · p. 2

2 PRELIMINARIES

The preliminaries frame inference as fitting a continuous probability-measure curve to unpaired point clouds. They distinguish intrinsic population motion from particle motions that can change without changing the curve’s marginals.

2 PRELIMINARIES

  • Observations are point clouds at discrete times with no cross-time sample correspondence, so the target is a distributionally fitted measure curve rather than pointwise matching.The curve is assumed to have finite second moments.Source: §2 · p. 3
  • For a regular measure curve, its intrinsic velocity is the minimum-L2(ρt)-norm solution of the continuity equation and is represented by a gradient field in the smooth setting.Intrinsic acceleration is the material derivative of this velocity, and its energy depends only on the full marginal curve.Source: §2 · p. 3
  • Different particle paths can induce identical marginals, so particlewise motion may include changes that leave the population distribution unchanged.A rotating isotropic Gaussian remains stationary as a marginal even though its particles accelerate.Source: §2 · p. 3
  • A population-level prior should assign equal penalties to representations sharing a marginal curve, including paths formed by time-dependent orthogonal latent rotations.This contrasts with particle-acceleration costs, which can vary across such representations.Source: §2 · p. 3, Figure 2 · p. 3

3 NEURAL WASSERSTEIN SPLINES (NWSP)

NWSP generates smooth particle paths, estimates population velocity through potential gradients, and penalizes a differentiated estimate of population acceleration while fitting snapshots. NWSP-N adds learned potentials, and the theoretical guarantees establish representation invariance under stated conditions.

3.1 A spline-based generative model

  • NWSP maps latent draws to spline control points, producing smooth particle paths with analytic positions, velocities, and accelerations at any time without numerical integration.A single forward pass supplies these quantities.Source: §3.1 · pp. 3–4

3.2 From particles to population velocity

  • A finite family of potential gradients estimates population velocity by regressing against generated particle velocities, with minibatch expectations and a ridge fit.The velocity fit uses particle paths supplied by the generator rather than differentiating a learned potential in time.Source: §3.2 · p. 4

3.3 From velocity to population acceleration

  • NWSP differentiates the fitted population velocity along the changing population to approximate intrinsic acceleration, rather than projecting particle acceleration.The changing regression Gram matrix contributes a correction; in the rotating-Gaussian example, particle acceleration is nonzero while intrinsic acceleration vanishes.Source: §3.3 · p. 4

3.4 Training objective and sampling

  • Training combines snapshot distribution-fit losses with normalized acceleration energy, then gradually increases the acceleration weight after fitting the snapshots.The stated distribution losses include debiased Sinkhorn divergence and Gaussian-kernel MMD.Source: §3.4 · pp. 4–5

3.5 NWSP-N: Expanding the basis

  • NWSP-N adds learned scalar potentials, optimizing their gradients to capture population velocity while alternating potential updates with generator updates.The potential objective is normalized to limit domination by times with larger particle speeds; acceleration is not minimized with respect to the potentials.Source: §3.5 · p. 5

3.6 Theoretical guarantees

  • For a fixed time-independent dictionary and positive ridge parameter, lifts inducing the same marginal curve have identical population acceleration energy, subject to regularity condition (R).The result applies to population quantities, not finite batches, and does not extend to the full NWSP-N procedure because its potential-learning objective uses path-dependent particle velocity energy.Source: §3.6 · p. 5
  • The acceleration estimate depends on errors in all four moments A, b, ˙A, and ˙b; omitting the moving-Gram correction adds a field that need not vanish as the dictionary grows.For a regular reference curve with exact moments, the stated Fourier-dictionary rates are O(N^-s) for velocity and O(N^(1-s)) for acceleration; these do not cover the trained generator or learned dictionary.Source: §3.6 · p. 6
  • Linear and quadratic potentials can miss shape changes preserving mean and variance, and the theory does not guarantee that learned potentials resolve every such change.A Gaussian-mixture example reports zero acceleration for a changing distribution under the linear and quadratic potentials.Source: §3.6 · p. 6

4 EXPERIMENTS

Across sixteen held-out prediction settings, NWSP and NWSP-N are competitive with transport- and flow-based baselines, while full-snapshot training and acceleration regularization provide distinct benefits. Controlled-flow experiments show that learned potentials improve acceleration estimates, and tests of the differentiated projection examine representation invariance and its finite-batch limits.

4.1 Experimental setup

  • The ten-dataset benchmark covers sixteen dataset–protocol settings and 66 held-out snapshots in dimensions 2–100, using leave-one-out and leave-alternate-out protocols.Held-out population fit is measured primarily by empirical W1, with W2 and Gaussian-kernel MMD2 also reported; scores average over three seeds.Source: §4.1 · p. 6

4.2 Prediction between observed snapshots

  • NWSP ranks in the top two in eleven of sixteen settings and improves over OT-CFM, SB-CFM, CFM, and [SF]2M in 14, 13, 15, and 15 settings, respectively.NWSP and NWSP-N attain the lowest W1 in five and four settings; inference requires no trajectory integration.Source: §4.2 · p. 6, Table 1 · p. 7
  • Both NWSP variants train on minibatches from all retained observations, whereas exact-transport WLR and four-point implementations subsample five settings; this is a coverage, not runtime, advantage.All methods are scored against complete withheld snapshots, and differing training support prevents attributing every comparison to the acceleration prior.Source: §4.2 · p. 6

4.3 What does the acceleration prior contribute?

  • Adding the population-acceleration penalty lowers mean W1 in thirteen of sixteen settings, with median relative change −2.2%.Gains are heterogeneous: embryoid-body leave-one-out changes from 0.450 to 0.350, while Petal changes from 0.153 to 0.090 and from 0.199 to 0.086 across its two protocols; some settings favor spline-only and Petal remains better served by subdivision.Source: §4.3 · pp. 6–8, Table 2 · p. 7
  • Compared with raw particle acceleration, projected acceleration has lower W1 in eleven settings and higher error in five, with median raw-minus-projected change +0.8%.The ablation supports neither a universal predictive advantage nor equivalence of the penalties; the geometric distinction calls for changing particle paths while holding the marginal curve fixed.Source: §4.3 · p. 8

4.4 Does adaptation improve acceleration estimation?

  • After 300 potential updates, relative acceleration error falls from 0.705 to 0.062 at d = 2 and from 0.975 to 0.341 at d = 30.Learned potentials outperform matched-size additions at every displayed dimension; error nevertheless increases with dimension, and 64 learned potentials do not uniformly improve on 16.Source: §4.4 · p. 8, Table 3 · p. 8
  • A Gaussian latent rotation preserves every marginal while projected energy varies by at most 2.3% through rate four at the largest sample size; on a periodic density curve, the moving-Gram correction reduces absolute L2(ρt) acceleration error from 0.577 to about 7 × 10^-14.Exact invariance is a population statement, not a finite-batch guarantee, and Fourier refinement does not repair the uncorrected differentiation bias.Source: §4.4 · p. 8

5 DISCUSSION AND CONCLUSION

NWSP separates population sampling from motion regularization, using spline paths and a differentiated projection to regularize marginal acceleration rather than unidentifiable particle paths. The paper reports competitive benchmark performance and identifies limits to its guarantees and scaling.

5 DISCUSSION AND CONCLUSION

  • NWSP ranks in the top two in eleven of sixteen settings, improves on spline smoothness in thirteen, and learned potentials reduce controlled-flow relative acceleration error from 0.705 to 0.062 at d = 2.NWSP-N attains the lowest W1 in four settings.Source: §5 · p. 8
  • The invariance guarantee is for frozen dictionaries, while Fourier rates assume exact moments and regular positive densities; NWSP identifies marginal curves, not individual cell histories.The dictionary reaches 5186 fixed features at d = 100, shifting the scaling bottleneck toward dictionary estimation; sparser dictionaries and matched runtime and memory comparisons remain future directions.Source: §5 · p. 8

AI use statement

  • Generative AI tools assisted with drafting and text revision, and the authors take responsibility for the final manuscript and its claims.Source: AI use statement · p. 9

A IMPLEMENTATION AND EVALUATION

The appendix specifies the datasets, fitting choices, and evaluation protocol used to compare NWSP and its baselines. Predictions are assessed with held-out distributional metrics, with dataset-specific processing and transport subsampling noted.

A IMPLEMENTATION AND EVALUATION

  • The benchmark uses nine datasets from prior benchmarks plus a tenth 100-dimensional CITE-seq time course, with dataset dimensions and held-out tasks tabulated.The table records state dimension, snapshot count, held-out targets under leave-one-out and leave-alternate-out, and fixed dictionary size; NWSP-N adds 16 potentials.Source: Table 4 · p. 11, A IMPLEMENTATION AND EVALUATION · p. 11

A.1 Hyperparameters and evaluation protocol

  • NWSP selects its acceleration weight from six values including zero, while NWSP-N uses the same grid except for a two-value CITE-seq grid.On embryoid body, Mouse, and Veres, exact-transport baselines use three independent 1,000-point retained-snapshot subsamples; CITE-seq uses full-support transport.Source: §A.1 · p. 13

A.2 Prediction metrics

  • Predictions are evaluated against withheld snapshots using exact empirical W1, W2, and Gaussian-kernel MMD2, with holdouts and three seed averages summarized by one sample standard deviation.Embryoid body scores use standardized PHATE coordinates and are not comparable with scores for other datasets.Source: §A.2 · p. 13

B COMPUTATIONAL COST

The computational discussion distinguishes data-use and scaling properties from measured speed, then presents theoretical guarantees for the finite-dimensional acceleration surrogate. Those guarantees cover representation invariance, moment-error stability, and Fourier approximation under stated conditions.

B COMPUTATIONAL COST

  • NWSP trains from minibatches of full snapshots without cross-time transport couplings, whereas exact-transport spline methods required 1,000-point subsamples on three benchmark datasets.This is a difference in usable data, not a measured speedup; the paper reports no matched runtime or memory comparisons against baselines.Source: B COMPUTATIONAL COST · p. 13

C.1 Representation invariance

  • For a fixed dictionary and ridge parameter, lifts satisfying condition (R) that induce the same marginal curve yield identical projected fields and population acceleration energies.The result also applies to a frozen learned dictionary.Source: §C.1 · p. 14

C.2 Stability of the two solves

  • The stability theorem bounds acceleration-estimation error using perturbations in all four moments, but its Gram-matrix condition was not measured in the reported experiments.The condition holds when Gram-matrix error is below the ridge parameter.Source: §C.2 · p. 15

C.3 Fourier approximation

  • On the torus, under the stated smoothness, positive-density, and exact-moment assumptions, Fourier velocity error is O(N^-s) and integrated acceleration energy error is O(N^(1-s)).Here N is the frequency cutoff, not the number of dictionary columns.Source: §C.3 · p. 15

C.4 Why the moving-Gram correction cannot be dropped

The moving-Gram correction contributes a term that need not vanish, and the periodic check shows a large error when it is omitted. A separate example shows that representation invariance does not guarantee that a finite dictionary captures every distributional change.

C.4 Why the moving-Gram correction cannot be dropped

  • The moving-Gram correction cannot generally be dropped: its additional term need not vanish, even when the velocity fit is exact.This matters because omitting the term leaves a nonvanishing bias in the stated Fourier limit.Source: §C.4 · p. 16
  • At cutoff N = 16, the periodic check reports absolute acceleration error 7.3 × 10^-14 with the correction and 0.577 without it.The passage distinguishes this absolute norm from the relative-error metric used in learned-potential experiments.Source: §C.4 · p. 17

C.5 A finite-dictionary resolution example

  • A time-varying Gaussian mixture can change its distribution while preserving its first two moments, so a linear-and-quadratic dictionary can return zero estimated acceleration.The example motivates learning potentials but does not show that the implemented learning rule resolves every such change.Source: §C.5 · p. 17

D SUPPLEMENTARY PREDICTION RESULTS

The supplementary tables report held-out W2 and MMD2 for the same fitted models, plus dataset-level relative W1 comparisons. Their summaries describe observed score differences rather than independent replications across protocols.

D SUPPLEMENTARY PREDICTION RESULTS

  • Table 5 reports held-out W2 for the same fitted models, splits, and conventions as the main W1 table.The caption identifies W2 as the metric but provides no score values in the supplied passage.Source: Table 5 · p. 18
  • Table 6 reports held-out MMD2 averaged over five Gaussian bandwidths for the same fitted models, splits, and conventions as the main W1 table.The supplied caption does not include the table’s score values.Source: Table 6 · p. 19
  • Table 7 aggregates relative W1 differences within each dataset before counting wins, but these protocol-level summaries are not independent replications.Its columns are win count, median relative change, and one-sided exact sign-test p-values, which are post hoc and uncorrected.Source: D SUPPLEMENTARY PREDICTION RESULTS · p. 18, Table 7 · p. 19

E ACCELERATION ESTIMATION AND REGULARIZATION

The section evaluates acceleration estimation against analytic references and under alternative particle representations, then compares projected and raw acceleration penalties on held-out W1. It also reports limits of potential adaptation and an optimization and dictionary-size comparison.

E ACCELERATION ESTIMATION AND REGULARIZATION

  • The latent-rotation check preserves the complete marginal curve and tests finite-sample estimation of the same population quantity, not held-out prediction or distributional stationarity.Rapid rotations make cancellations in the differentiated moments difficult to estimate.Source: E ACCELERATION ESTIMATION AND… · p. 20
  • Analytic-reference tests cover rotating three-well gradient flows, blind Gaussian mixtures, and rotating lifts of constant clouds, with relative errors averaged over 21 time points.Reference accelerations are evaluated analytically and checked numerically; reported tables use five-seed means and standard deviations.Source: E ACCELERATION ESTIMATION AND… · p. 20
  • Table 8 compares the standard NWSP-N configuration with runs using more updates or more learned potentials.It reports relative acceleration error as five-seed mean ± SD, but the supplied caption gives no error values.Source: Table 8 · p. 20
  • Adaptation controls show that an untrained block beats the trained block on the one-dimensional blind mixture, while long constant-rate runs can degrade or diverge.An offline Action-Matching network overtakes NWSP-N between 300 and 3000 steps, reaching errors 3.4×10^-4–1.0×10^-3 at 10^4 steps; this does not test either estimator inside a moving generator.Source: E ACCELERATION ESTIMATION AND… · p. 20

E.1 Raw versus projected acceleration

  • Raw acceleration yields lower W1 in 5 of 16 settings and higher W1 in 11, with a median change of +0.8% and mean change of +2.4% relative to NWSP.Six settings differ by less than 1%; the raw arm ran on different hardware, which can shift individual scores by about a seed spread.Source: §E.1 · p. 21, Table 9 · p. 21

F QUALITATIVE RESULTS

The qualitative figures distinguish fitted synthetic dynamics from held-out single-cell predictions. Their plots provide visual comparisons, while quantitative comparisons use full evaluation coordinates rather than the two-dimensional projections.

F QUALITATIVE RESULTS

  • Quantitative comparisons use full evaluation coordinates, not the two-dimensional plotting projections.Figure 4 uses retained-neighbor clouds; CITE-seq/Veres PCA and Embryoid-body standardized PHATE use retained frames only.Source: F QUALITATIVE RESULTS · p. 22, Figure 4 · p. 23
  • Figure 3 shows fitted synthetic population trajectories using all available snapshots, so it illustrates fitted dynamics rather than held-out generalization.The held-out evaluation is reported separately in Table 1.Source: Figure 3 · p. 22, F QUALITATIVE RESULTS · p. 22
  • Figures 4 and 5 show held-out single-cell snapshots and method predictions, with each Figure 4 column representing a separate leave-one-out fit.The number below each Figure 4 panel is that holdout’s W1; red marks the column’s lowest value, but these views do not establish full-dimensional rankings.Source: Figure 4 · p. 23, Figure 5 · p. 24
Loading 2610.19878v1…