Wasserstein Parallel Trends

• Wasserstein Parallel Trends is a distribution-level generalization of the classical parallel trends assumption, comparing full evolving distributions rather than just means.
• It leverages geodesic parallel transport in the 2-Wasserstein space to align tangent dynamics between treated and control units for counterfactual prediction.
• This framework has practical implications in causal inference, domain adaptation, and biology by capturing changes in shape, spread, and multimodality of outcomes.

Searching arXiv for the specified paper and closely related work on Wasserstein parallel transport and trend assumptions. Wasserstein Parallel Trends is a distribution-level analogue of the classical parallel trends assumption from Difference-in-Differences. Instead of assuming that means evolve in parallel across treated and control groups, it assumes that the entire distributional dynamics evolve in parallel once one transports the control group’s tangent dynamics to the treated group’s baseline along a Wasserstein geodesic. In the formulation introduced in "Wasserstein Parallel Transport for Predicting the Dynamics of Statistical Systems" (Saidi et al., 24 Mar 2026), the method replaces vector subtraction by geodesic parallel transport on the $2$-Wasserstein space, thereby enabling counterfactual comparisons of evolving probability measures in causal inference, domain adaptation, and batch-effect correction. The same work states that Wasserstein Parallel Trends recovers the classic parallel trends assumption for averages as a special case and derives closed-form parallel transport for Gaussian measures (Saidi et al., 24 Mar 2026).

1. Classical antecedents and the move from mean trends to distributional dynamics

In standard DiD, the identifying condition is a statement about first moments: untreated outcome trends are equal in expectation. Wasserstein Parallel Trends generalizes this by replacing equality of mean changes with equality of tangent dynamics after parallel transport. The key assumption is written in terms of probability measures $\mu_t$ and $\nu_t$ and their tangent velocities: $$\nabla\varphi_{t_i}^* = {PT}_{\nu_{t_i} \rightarrow \mu_{t_i}^*}(\nabla \varphi_{t_i}), \qquad i=1,\dots,T-1.$$ Here $\nabla\varphi_{t_i}$ is the tangent velocity of the control trajectory at time $t_i$, and $\nabla\varphi_{t_i}^*$ is the tangent velocity of the counterfactual treated trajectory (Saidi et al., 24 Mar 2026).

This shifts the target of comparison from averages to full evolving distributions. The paper explicitly presents the result as a genuine distributional generalization: it captures changes not only in location but also in spread, shape, multimodality, and other structural features of the distribution (Saidi et al., 24 Mar 2026).

A useful antecedent in the DiD literature is entropy balancing of pre-intervention outcome trends. "Reducing bias in difference-in-differences models using entropy balancing" (Cefalu et al., 2020) argues that researchers may still recover credible causal effects by reweighting the comparison group so that its pre-treatment outcome trends match those of the treated group. That paper does not present the method as Wasserstein optimal transport, but it explicitly frames the method as close in spirit to distributional balancing and describes the weighting problem as a “transport” idea in which mass is moved in the comparison group so its pre-trend distribution matches the treated group’s pre-trend distribution as closely as possible, subject to moment constraints (Cefalu et al., 2020). This suggests that Wasserstein Parallel Trends belongs to a broader shift from balancing static levels to balancing dynamic structure.

2. Geometric formulation on Wasserstein space

The framework is built on the Riemannian geometry of the $2$-Wasserstein space $(\mathcal P_2(M), W_2)$. The basic distance is

$$W_2(\mu,\nu) = \inf_{\gamma\in \Gamma_{\mu,\nu}} \left(\int_{M\times M} d_M(x,y)^2\,\gamma(dx,dy)\right)^{1/2}.$$

A measure-valued curve $\mu_t$0 evolves under a velocity field $\mu_t$1 through the continuity equation

$\mu_t$2

Among all such velocity fields, the optimal one is the tangent vector to the Wasserstein geodesic, and the Benamou–Brenier theorem gives

$\mu_t$3

The tangent space at $\mu_t$4 is

$\mu_t$5

with inner product

$\mu_t$6

The Wasserstein tangent vector is the minimal-norm solution to the continuity equation, and it is the limit of gradient fields (Saidi et al., 24 Mar 2026).

Parallel transport is defined through the Wasserstein Levi-Civita connection: $\mu_t$7 where $\mu_t$8 is orthogonal projection onto $\mu_t$9. A field is parallel along $\nu_t$0 if

$\nu_t$1

equivalently,

$\nu_t$2

In this geometry, the “difference” between two evolving distributions is not a Euclidean subtraction but a transported tangent field (Saidi et al., 24 Mar 2026).

The same paper also defines Wasserstein exponential and logarithmic maps. On $\nu_t$3,

$\nu_t$4

These maps supply the local coordinates used in counterfactual construction (Saidi et al., 24 Mar 2026).

3. Fanning schemes and computational approximation

Exact Wasserstein parallel transport is hard to compute. The central algorithmic contribution of the 2026 paper is a fanning scheme: approximate Wasserstein parallel transport by transporting tangent fields along the base manifold $\nu_t$5 using Jacobi fields, then projecting back onto the Wasserstein tangent space (Saidi et al., 24 Mar 2026).

The construction begins with the Lagrangian particle paths induced by the OT map. The paper uses a result of Gigli to link Wasserstein parallel transport to parallel transport of tangent vectors on $\nu_t$6 along those Lagrangian paths. On $\nu_t$7, parallel transport is then locally approximated by a Jacobi field. For a geodesic $\nu_t$8, the relevant Jacobi field $\nu_t$9 is the unique solution of

$$\nabla\varphi_{t_i}^* = {PT}_{\nu_{t_i} \rightarrow \mu_{t_i}^*}(\nabla \varphi_{t_i}), \qquad i=1,\dots,T-1.$$0

with prescribed initial conditions, and the approximation result is

$$\nabla\varphi_{t_i}^* = {PT}_{\nu_{t_i} \rightarrow \mu_{t_i}^*}(\nabla \varphi_{t_i}), \qquad i=1,\dots,T-1.$$1

for sufficiently small $$\nabla\varphi_{t_i}^* = {PT}_{\nu_{t_i} \rightarrow \mu_{t_i}^*}(\nabla \varphi_{t_i}), \qquad i=1,\dots,T-1.$$2, under a lower bound on the injectivity radius and compactness assumptions (Saidi et al., 24 Mar 2026).

The paper then lifts this one-step approximation to Wasserstein space. If $$\nabla\varphi_{t_i}^* = {PT}_{\nu_{t_i} \rightarrow \mu_{t_i}^*}(\nabla \varphi_{t_i}), \qquad i=1,\dots,T-1.$$3 is the measure after moving a small step along the geodesic from $$\nabla\varphi_{t_i}^* = {PT}_{\nu_{t_i} \rightarrow \mu_{t_i}^*}(\nabla \varphi_{t_i}), \qquad i=1,\dots,T-1.$$4 to $$\nabla\varphi_{t_i}^* = {PT}_{\nu_{t_i} \rightarrow \mu_{t_i}^*}(\nabla \varphi_{t_i}), \qquad i=1,\dots,T-1.$$5, and

$$\nabla\varphi_{t_i}^* = {PT}_{\nu_{t_i} \rightarrow \mu_{t_i}^*}(\nabla \varphi_{t_i}), \qquad i=1,\dots,T-1.$$6

then

$$\nabla\varphi_{t_i}^* = {PT}_{\nu_{t_i} \rightarrow \mu_{t_i}^*}(\nabla \varphi_{t_i}), \qquad i=1,\dots,T-1.$$7

Under strong regularity, this simplifies to

$$\nabla\varphi_{t_i}^* = {PT}_{\nu_{t_i} \rightarrow \mu_{t_i}^*}(\nabla \varphi_{t_i}), \qquad i=1,\dots,T-1.$$8

By composing many small steps of size $$\nabla\varphi_{t_i}^* = {PT}_{\nu_{t_i} \rightarrow \mu_{t_i}^*}(\nabla \varphi_{t_i}), \qquad i=1,\dots,T-1.$$9, the paper defines an approximate transport operator $\nabla\varphi_{t_i}$0 and proves

$\nabla\varphi_{t_i}$1

This is the main approximation guarantee for the multi-step scheme (Saidi et al., 24 Mar 2026).

The algorithm itself is presented in explicit steps: compute the Brenier map $\nabla\varphi_{t_i}$2; set $\nabla\varphi_{t_i}$3, which generates the Wasserstein geodesic; discretize with $\nabla\varphi_{t_i}$4; and then, at each step, form $\nabla\varphi_{t_i}$5, define $\nabla\varphi_{t_i}$6, set $\nabla\varphi_{t_i}$7, compute the Jacobi field $\nabla\varphi_{t_i}$8, and project the resulting field back to the tangent space,

$\nabla\varphi_{t_i}$9

For the projection $t_i$0, the paper also gives an RKHS-based Helmholtz-Hodge projection method and proves consistency: $t_i$1 under boundedness and density assumptions and if $t_i$2 for $t_i$3 (Saidi et al., 24 Mar 2026).

4. Theoretical consequences and special cases

A central structural result is that Wasserstein Parallel Trends implies equality of mean trends in $t_i$4: $t_i$5 Thus the distributional statement strictly generalizes the classical equibias condition. The paper also uses the subresult that Wasserstein parallel transport preserves constant vector fields: $t_i$6 These statements formalize the relationship between classical mean-parallel trends and the Wasserstein formulation (Saidi et al., 24 Mar 2026).

For Gaussian measures, the framework admits a closed form. If $t_i$7 and $t_i$8, with affine tangent vector

$t_i$9

then the parallel transported field along the Gaussian Wasserstein geodesic has the form

$\nabla\varphi_{t_i}^*$0

with $\nabla\varphi_{t_i}^*$1 governed by the continuous Lyapunov equation

$\nabla\varphi_{t_i}^*$2

where

$\nabla\varphi_{t_i}^*$3

and

$\nabla\varphi_{t_i}^*$4

This gives an explicit affine transport rule and shows how covariance deformation is transported in tandem with mean motion (Saidi et al., 24 Mar 2026).

The paper also proves a stability theorem: $\nabla\varphi_{t_i}^*$5 For iterative counterfactual prediction, the one-step error satisfies

$\nabla\varphi_{t_i}^*$6

and summing over time yields the cumulative bound reported in the paper (Saidi et al., 24 Mar 2026).

5. Applications in causal inference, domain adaptation, and single-cell biology

The primary causal-inference use is a distributional DiD. Instead of subtracting control changes from treated changes in means, the method transports the control trajectory to the treated baseline and uses that as the counterfactual trajectory. This allows treatment effects that alter the shape of the outcome distribution, not just its mean, to be detected (Saidi et al., 24 Mar 2026).

The same geometry is presented as a dynamic counterpart to OT-based domain adaptation. Traditional OT adaptation aligns source and target distributions at a single time point; here the goal is to transfer temporal dynamics from one domain to another by parallel-transporting the tangent flow. In controlled biological experiments, baseline differences may come from batch effects or technical shifts, and the paper argues that if these discrepancies are approximately stable across time, Wasserstein Parallel Trends gives a principled way to align the dynamics across baselines while respecting the distributional geometry (Saidi et al., 24 Mar 2026).

The empirical evaluation includes synthetic Gaussian experiments and two single-cell RNA sequencing datasets. On synthetic time-evolving Gaussian measures, where ground-truth parallel transport is available in closed form, the paper compares WPT, WPT$\nabla\varphi_{t_i}^*$7, Brenier baseline, Mean-shift baseline, and Gaussian PT. Evaluation uses $\nabla\varphi_{t_i}^*$8 between predicted and true counterfactual distributions at each time step. The reported findings are conditional: when the covariance is stationary, mean-shift performs reasonably; when covariance changes over time, WPT outperforms naive baselines in capturing shape deformation; in low dimension, WPT$\nabla\varphi_{t_i}^*$9 is strong; the Gaussian baseline is best when the data are truly Gaussian; and in higher dimensions performance degrades, especially for the projection step (Saidi et al., 24 Mar 2026).

The biological datasets are: human versus chimp cerebral organoids, with $2$0, and mouse microglial cells over time in female versus male, both sham and SNI conditions. Each cell is treated as a point in $2$1 and each time point as an empirical distribution over cells. The tasks are to impute chimp developmental dynamics onto human initial conditions, and to impute female dynamics onto male and compare against observed male trajectories, separately for sham and SNI. Performance is measured by Wasserstein distance between predicted and observed point clouds. The paper reports that WPT$2$2 usually beats or matches the mean-shift baseline, with stronger gains as the prediction horizon increases, and that UMAP plots show predicted trajectories qualitatively following the observed developmental progression (Saidi et al., 24 Mar 2026).

6. Relation to Wasserstein parallel transport theory, limitations, and common confusions

Wasserstein Parallel Trends is not identical to the general mathematical problem of parallel transport in Wasserstein space. Rather, it is a counterfactual assumption and computational framework built on that geometry. Earlier work had already developed several notions of parallel transport: along Wasserstein geodesics and tangent cones (Shen, 2016), via intrinsic metric constructions on $2$3 (Lott, 2017), through stochastic regular curves on $2$4 (Ding et al., 2023), by a relaxed measured-valued tangent space on $2$5 (Bertucci, 10 Dec 2025), and by lifting the problem to the diffeomorphism group and constructing stochastic parallel transport on the smooth Wasserstein space $2$6 (Martin, 2 Dec 2025).

Those works differ substantially in geometric setup. The 2016 and 2017 papers focus on geodesic transport between tangent cones or linear tangent spaces and on agreement with earlier formal computations (Shen, 2016, Lott, 2017). The 2023 paper develops intrinsic Itô formulae, stochastic regular curves, and stochastic parallel translations, with a special existence result on $2$7 (Ding et al., 2023). The 2025 papers broaden the landscape further: one introduces tangent vectors as probability-measure-valued fields and defines Lagrangian parallel transport by freezing the tangent label along a path (Bertucci, 10 Dec 2025), while the other uses the principal-bundle structure $2$8 to establish existence and uniqueness of stochastic parallel transport along diffusions (Martin, 2 Dec 2025). Within this broader literature, Wasserstein Parallel Trends uses parallel transport as an operational device for counterfactual prediction rather than as an end in itself.

Several limitations are explicit in the 2026 formulation. Statistical theory is incomplete: the convergence and error guarantees are population-level approximation results, not full finite-sample statistical guarantees. Regularity requirements are strong: parallel transport is only well-defined along regular geodesics, and the stability theory assumes substantial smoothness and boundedness of densities and velocities. High-dimensional estimation is hard: the Helmholtz projection via RKHS regression suffers from slow rates and deteriorates with dimension. Much of the detailed stability analysis is done on the flat torus $2$9, avoiding boundary complications. In genomics applications, the authors present the results as evidence that Wasserstein Parallel Trends is a plausible assumption, not as definitive causal claims (Saidi et al., 24 Mar 2026).

A common confusion is to treat the method as a mere distributional rebranding of mean-parallel trends. The theory does not do that: its core object is the transported tangent velocity, not the mean difference. Another confusion is to equate it with reweighting-based DiD corrections. That connection is conceptually suggestive, especially because entropy balancing on pre-intervention trends also targets dynamic comparability (Cefalu et al., 2020), but the formal machinery is different: Wasserstein Parallel Trends replaces vector subtraction by geodesic parallel transport on the space of measures (Saidi et al., 24 Mar 2026).