Transport G-Computation: A Distributional Approach to Longitudinal Causal Inference via Optimal Transport

Technical Summary: Transport G-Computation (TGC)

Problem Statement
Longitudinal causal inference using observational health data faces significant challenges, primarily time-dependent confounding, feedback loops between treatment assignment and evolving patient states, and heterogeneous post-intervention outcome distributions. Conventional longitudinal g-computation methods typically focus on estimating mean-level effects (e.g., average treatment effects). However, in many clinical contexts, interventions may alter the shape, dispersion, or multimodality of outcome distributions rather than just the mean. Furthermore, parametric g-computation relies heavily on the correct specification of conditional transition models; when data exhibit nonlinearity, multimodality, heteroscedasticity, or complex feedback structures, these models often fail, leading to systematic bias. There is a need for a framework that can estimate distributional causal effects while robustly handling complex dynamic dependencies without relying on rigid global parametric assumptions.

Methodology: Transport G-Computation (TGC)
The authors propose Transport G-Computation (TGC), a distributional framework that integrates longitudinal g-computation with optimal transport (OT) theory. The core objective is to estimate the Wasserstein Causal Effect (WCE), defined as the second-order Wasserstein distance ($W_2$) between the terminal potential outcome distributions under two static intervention regimes ($A_t \equiv 1$ and $A_t \equiv 0$).

The TGC algorithm proceeds as follows:

1. Identification: Under standard assumptions (Consistency, Sequential Ignorability, and Positivity), the terminal distribution is identified via the longitudinal g-formula.
2. Local Transition Estimation: Instead of fitting a global parametric model for the conditional transition kernel $\mathbb{P}(L_{t+1} | L_t, A_t)$, TGC constructs a nonparametric approximation at each time point and treatment level.
   • It extracts an empirical transition library from observed data pairs $(L_t, L_{t+1})$ where $A_t = a$.
   • For a query state, it identifies local nearest neighbors to form a subsample, restricting estimation to the observed region to mitigate high-dimensional bias.
   • It solves an entropy-regularized optimal transport problem on this local subsample to generate a "soft-matching" coupling between source and target states. This preserves the multimodality and stochastic structure of the empirical distribution.
   • Gaussian kernel weights are applied to the local source states, and the OT coupling is row-normalized to form a local conditional distribution.
3. Recursive Simulation: Starting from initial state samples, the algorithm recursively simulates trajectories under static interventions by sampling from the estimated local conditional distributions at each time step.
4. Effect Estimation: The Wasserstein distance is computed between the resulting empirical terminal distributions of the two intervention arms.
5. Positivity Handling: The method explicitly checks for sample sufficiency in local neighborhoods. If the effective sample size falls below a threshold, the algorithm flags a positivity violation rather than extrapolating, thereby avoiding spurious estimates in non-identifiable regions.

Key Contributions

1. Distributional Causal Target: The paper introduces the Wasserstein distance as a formal estimand for longitudinal causal effects, capturing shifts in location, spread, and shape, rather than focusing solely on means.
2. Novel Framework: It proposes TGC, which combines local kernel smoothing with optimal transport coupling to estimate terminal distributions under static interventions without global parametric assumptions.
3. Theoretical Analysis: The authors provide identification proofs for the WCE under standard causal assumptions and analyze the approximation properties of the local transition estimator. They establish that the terminal error is bounded by the sum of local estimation errors, highlighting the importance of stability at each recursive step.
4. Diagnostic Capability: The method includes an explicit mechanism to diagnose and report positivity violations, preventing uncontrolled extrapolation.

Results
The performance of TGC was evaluated against parametric g-computation across four simulation scenarios (100 replications each):

• Linear Gaussian Scenario: Parametric g-computation outperformed TGC (RMSE 0.075 vs. 0.659), indicating that when the true model is linear and correctly specified, the parametric approach is superior, and TGC introduces unnecessary smoothing bias.
• Gaussian Mixture Scenario: Parametric g-computation remained more accurate (RMSE 0.080 vs. 0.252). TGC did not demonstrate the anticipated advantage in capturing multimodal structures under these specific settings, suggesting sensitivity to hyperparameters like neighborhood size and regularization strength.
• Confounded Feedback Scenario: TGC showed its strongest advantage, significantly reducing RMSE from 1.821 (parametric) to 0.486. In this complex setting with dynamic feedback, the local OT approach effectively mitigated the systematic bias caused by model misspecification in the parametric baseline.
• Nonlinear Heteroscedastic Scenario: Both methods exhibited substantial errors (RMSE ~3.65 for both), indicating that neither approach fully captured the complex distributional structures in this specific setting.
• Computational Cost: TGC was substantially slower, requiring approximately 136–235 seconds per replication compared to ~1.1 seconds for the parametric method, representing an overhead of roughly two orders of magnitude.
• Stability: Both methods achieved a 0% failure rate across all scenarios, with TGC successfully avoiding positivity collapse through its explicit diagnostics.

Significance and Claims
The paper concludes that TGC is not a universally superior replacement for parametric g-computation. Its primary value lies in scenarios characterized by dynamic feedback confounding and complex distributional transitions where parametric models are prone to misspecification. In such settings, TGC offers a more robust, flexible, and interpretable approach to distributional causal inference.

However, the authors modestly acknowledge that in simple, correctly specified linear settings, traditional parametric methods remain preferable due to lower bias and variance. They also note that TGC currently entails significant computational costs and requires careful hyperparameter tuning. The method is positioned as a specialized tool for observational medical research where the focus is on distributional changes and where feedback structures complicate standard modeling, rather than as a default solution for all longitudinal causal problems. Future work is identified as necessary to improve efficiency, hyperparameter selection, and practical implementation for large-scale data.