Wasserstein Parallel Transport for Predicting the Dynamics of Statistical Systems

Tristan Luca Saidi; Gonzalo Mena; Larry Wasserman; Florian Gunsilius

stat.MLMar 24, 2026

Wasserstein Parallel Transport for Predicting the Dynamics of Statistical Systems

Authors: Tristan Luca Saidi, Gonzalo Mena, Larry Wasserman, Florian Gunsilius

Organizations: Department of Statistics and Data Science, Carnegie Mellon University · Machine Learning Department, Carnegie Mellon University · Department of Economics, Emory University

Abstract

Many scientific systems, such as cellular populations or economic cohorts, are naturally described by probability distributions that evolve over time. Predicting how such a system would have evolved under different forces or initial conditions is fundamental to causal inference, domain adaptation, and counterfactual prediction. However, the space of distributions often lacks the vector space structure on which classical methods rely. To address this, we introduce a general notion of parallel dynamics at a distributional level. We base this principle on parallel transport of tangent dynamics along optimal transport geodesics and call it ``Wasserstein Parallel Trends''. By replacing the vector subtraction of classic methods with geodesic parallel transport, we can provide counterfactual comparisons of distributional dynamics in applications such as causal inference, domain adaptation, and batch-effect correction in experimental settings. The main mathematical contribution is a novel notion of fanning scheme on the Wasserstein manifold that allows us to efficiently approximate parallel transport along geodesics while also providing the first theoretical guarantees for parallel transport in the Wasserstein space. We also show that Wasserstein Parallel Trends recovers the classic parallel trends assumption for averages as a special case and derive closed-form parallel transport for Gaussian measures. We deploy the method on synthetic data and two single-cell RNA sequencing datasets to impute gene-expression dynamics across biological systems.

Figures & tables

Panel of Figure 1 — Figure 1: Visualizations of Wasserstein Parallel Transport with Gaussian measures in $\mathbb{R}^{2}$ . The left panel illustrates how parallel transport captures deformations to the covariance, while the right panel illustrates how it captures changes in the mean.

Figure 2: Visualization of approximate Wasserstein parallel transport, where the approximate tangent vector is given by the map $p\mapsto{\operatorname{PT}}_{\exp_{p}(\nabla\varphi(p))}(v(p))$ (left), and exact Wasserstein parallel transport, where the true tangent vector is given by the map $p\mapsto({\operatorname{PT}}_{\mathbf{exp}_{\nu}(\nabla\varphi)}(v))(p)$ (right).

Figure 3: Estimated projection of a non-conservative vector field onto the space of gradient fields using the procedure described in Section 3.4

Figure 4: Visualization of counterfactual dynamics prediction procedure described in Algorithm 2 .

Panel of Figure 5 — Figure 5: Difference-in-differences diagram illustrating the parallel trends assumption in the Euclidean case (left) and the Wasserstein case (right). The goal is to reconstruct the counterfactual (denoted $Y_{t}(0)$ on the left and $\mu_{t}^{*}$ on the right).

Panel of Figure 6 — Figure 6: Illustration of the parallel transported dynamics for two systems of time-evolving Gaussian measures using our method and two baselines.

Step	$d=2$					$d=4$
	WPT	WPT -	Brenier	Shift	Gauss PT	WPT	WPT -	Brenier	Shift	Gauss PT
1	0.606	0.142	2.242	0.277	0.118	3.300	0.400	4.669	0.512	0.401
2	1.609	0.121	3.740	0.418	0.101	7.490	0.432	7.765	0.728	0.429
3	2.988	0.124	4.926	0.582	0.097	12.700	0.455	10.283	0.935	0.462
4	4.816	0.145	5.965	0.729	0.139	18.462	0.475	12.477	1.122	0.463
5	7.163	0.114	6.884	0.838	0.090	24.186	0.486	14.463	1.264	0.466

Table 1:

W_{2}

between predicted and true Gaussian CF (under Wasserstein Parallel Trends) across dimensions using

5000

samples. WPT denotes Wasserstein parallel transport with projection ( Algorithm 2 ), WPT - omits the projection step, Brenier applies the control OT map without transport, Shift applies the mean-shift baseline, and Gauss PT fits Gaussian parameters from samples and applies closed-form parallel transport. Lower is better; bold marks the best method at each step and underlining marks the second-best.

$n_{\mathrm{pcs}}$	mid ( $\sim$ day 65)		late ( $\sim$ day 110)
	WPT -	Mean-shift	WPT -	Mean-shift
2	1.262	1.230	0.817	1.367
3	1.262	1.298	0.833	1.433
5	1.294	1.376	0.896	1.621
7	1.389	1.474	1.121	1.707
10	1.537	1.578	1.428	1.851

Table 2:

W_{2}

between predicted and observed human organoid gene expressions per non-initial time point for WPT - and mean-shift baseline as a function of the number of principal components used.

Figure 7: UMAP ( McInnes et al., 2018 ) visualizations of chimp (top row), predicted human (middle row) and observed human (bottom row) organoid dynamics over time using $n_{\text{pcs}}=15$ .

Figure 8: UMAP ( McInnes et al., 2018 ) visualizations of female (top row), predicted male (middle row) and observed male (bottom row) microglial dynamics over time using $n_{\text{pcs}}=15$ . The panel on the left uses only cells from the untreated (sham) group, while the panel on the right uses only cells from the treated (SNI) group.

Table 3:

W_{2}

between predicted and observed male microglial gene expressions per non-initial timepoint for WPT - and mean-shift baseline as a function of the number of principal components used.

Batch Effect Correction

See all 9 papers

Dynamic Generalized Gromov-Wasserstein Optimal Transport
Sep 17, 2026Junda Ying, Zhiwei Zeng, Peijie Zhou +1Flow Matching Trajectory Inference
Beijing International Center for Mathematical Research, Peking University · Center for Quantitative Biology, Peking University · Center for Machine Learning Research, Peking University +4
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
Distribution-Conditioned Transport
Mar 5, 2026Nic Fishman, Gokul Gowri, Paolo L. B. Fischer +3OOD Generalization
Harvard University · MIT · Beth Israel Deaconess Medical Center +2

Sep 17, 2026J/K move · Enter open · S save

Dynamic Generalized Gromov-Wasserstein Optimal Transport

Junda Ying, Zhiwei Zeng, Peijie Zhou, Lei Zhang

Beijing International Center for Mathematical Research, Peking University · Center for Quantitative Biology, Peking University · Center for Machine Learning Research, Peking University+4

Flow Matching Trajectory Inference Optimal Transport

Open paper AskPDF arXiv

Gromov--Wasserstein optimal transport (GW-OT) extends classical optimal transport by introducing structure-aware transport cost. This is particularly relevant for spatial transcriptomics, where dynamical reconstruction should preserve tissue structure in addition to matching expression patterns. While static formulations have been widely used for such structure-aware alignment, a general dynamic formulation for reconstructing continuous trajectories is still missing. We introduce Travelling Pair Dynamical Alignment and Trajectory Estimation (TP-DATE), a theoretical and computational framework to generalize GW-OT dynamically in a simulation-free manner. We formulate a broad class of static and dynamic Quadratic-form OT (QOT) through path actions and prove the static dynamic equivalence. We further develop travelling-pair flow matching, which allows interacting conditional paths and marginalizes their interactions into a single vector field. On synthetic and real spatial transcriptomics data, TP-DATE better preserves spatial structure and improves continuous 3D dynamics reconstruction.

Abstract

Batch Effect Correction

Explore similar work

Dynamic Generalized Gromov-Wasserstein Optimal Transport

Wasserstein Residuals: Learning Gradient Flows from Population Dynamics

Distribution-Conditioned Transport