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
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
Figure 1: Visualizations of Wasserstein Parallel Transport with Gaussian measures in R2 . 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↦PTexpp(∇φ(p))(v(p)) (left), and exact Wasserstein parallel transport, where the true tangent vector is given by the map p↦(PTexpν(∇φ)(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 .
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 Yt(0) on the left and μt∗ on the right).
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: W2 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.
npcs
mid ( ∼ day 65)
late ( ∼ 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: W2 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 npcs=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 npcs=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: W2 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.
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.
Junda Ying, Zhiwei Zeng, Peijie Zhou +1
Beijing International Center for Mathematical Research, Peking University · Center for Quantitative Biology, Peking University · Center for Machine Learning Research, Peking University +4
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
Learning a transport model that maps a source distribution to a target distribution is a canonical problem in machine learning, but scientific applications increasingly require models that can generalize to source and target distributions unseen during training. We introduce distribution-conditioned transport (DCT), a framework that conditions transport maps on learned embeddings of source and target distributions, enabling generalization to unseen distribution pairs. DCT also allows semi-supervised learning for distributional forecasting problems: because it learns from arbitrary distribution pairs, it can leverage distributions observed at only one condition to improve transport prediction. DCT is agnostic to the underlying transport mechanism, supporting models ranging from flow matching to distributional divergence-based models (e.g. Wasserstein, MMD). We demonstrate the practical performance benefits of DCT on synthetic benchmarks and four applications in biology: batch effect transfer in single-cell genomics, perturbation prediction from mass cytometry data, learning clonal transcriptional dynamics in hematopoiesis, and modeling T-cell receptor sequence evolution.
Nic Fishman, Gokul Gowri, Paolo L. B. Fischer +3
Harvard University · MIT · Beth Israel Deaconess Medical Center +2