Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Distributional Treatment Effect Estimation across Heterogeneous Sites via Optimal Transport

Borna Bateni, Yubai Yuan, Qi Xu, Annie Qu

TL;DR

The paper addresses cross-site transportability of intervention effects by moving beyond average effects to full distributional treatment outcomes. It introduces a distributional causal inference framework that treats treatment and control as probability measures on joint feature-outcome spaces and models cross-site heterogeneity as a push-forward map $T$, learned via an optimal transport objective that aligns controls using a regularized fused Gromov-Wasserstein loss and a metric-learning step. It then transports the source treatment distribution to the target site to synthesize a target-treatment dataset, with rigorous guarantees of consistency and asymptotic convergence to the true target distribution. Empirical results on synthetic scenarios and a real PDX dataset demonstrate accurate reconstruction of the entire treatment-effect distribution and superiority over ATE-focused baselines, highlighting practical utility for cross-site causal inference.

Abstract

We propose a novel framework for synthesizing counterfactual treatment group data in a target site by integrating full treatment and control group data from a source site with control group data from the target. Departing from conventional average treatment effect estimation, our approach adopts a distributional causal inference perspective by modeling treatment and control as distinct probability measures on the source and target sites. We formalize the cross-site heterogeneity (effect modification) as a push-forward transformation that maps the joint feature-outcome distribution from the source to the target site. This transformation is learned by aligning the control group distributions between sites using an Optimal Transport-based procedure, and subsequently applied to the source treatment group to generate the synthetic target treatment distribution. Under general regularity conditions, we establish theoretical guarantees for the consistency and asymptotic convergence of the synthetic treatment group data to the true target distribution. Simulation studies across multiple data-generating scenarios and a real-world application to patient-derived xenograft data demonstrate that our framework robustly recovers the full distributional properties of treatment effects.

Distributional Treatment Effect Estimation across Heterogeneous Sites via Optimal Transport

TL;DR

The paper addresses cross-site transportability of intervention effects by moving beyond average effects to full distributional treatment outcomes. It introduces a distributional causal inference framework that treats treatment and control as probability measures on joint feature-outcome spaces and models cross-site heterogeneity as a push-forward map , learned via an optimal transport objective that aligns controls using a regularized fused Gromov-Wasserstein loss and a metric-learning step. It then transports the source treatment distribution to the target site to synthesize a target-treatment dataset, with rigorous guarantees of consistency and asymptotic convergence to the true target distribution. Empirical results on synthetic scenarios and a real PDX dataset demonstrate accurate reconstruction of the entire treatment-effect distribution and superiority over ATE-focused baselines, highlighting practical utility for cross-site causal inference.

Abstract

We propose a novel framework for synthesizing counterfactual treatment group data in a target site by integrating full treatment and control group data from a source site with control group data from the target. Departing from conventional average treatment effect estimation, our approach adopts a distributional causal inference perspective by modeling treatment and control as distinct probability measures on the source and target sites. We formalize the cross-site heterogeneity (effect modification) as a push-forward transformation that maps the joint feature-outcome distribution from the source to the target site. This transformation is learned by aligning the control group distributions between sites using an Optimal Transport-based procedure, and subsequently applied to the source treatment group to generate the synthetic target treatment distribution. Under general regularity conditions, we establish theoretical guarantees for the consistency and asymptotic convergence of the synthetic treatment group data to the true target distribution. Simulation studies across multiple data-generating scenarios and a real-world application to patient-derived xenograft data demonstrate that our framework robustly recovers the full distributional properties of treatment effects.

Paper Structure

This paper contains 20 sections, 2 theorems, 14 equations, 2 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background and Notations
  3. Distributional Causal Inference
  4. The Effect Modification Problem
  5. Optimal Transport for Causal Inference
  6. Metric Learning
  7. Methodology
  8. Problem Definition
  9. Assumptions
  10. Proposed Method
  11. Regularized Fused Gromov-Wasserstein Loss between the control groups
  12. Synthetic treatment generation
  13. Theoretical Results
  14. Simulations
  15. Design overview
  16. ...and 5 more sections

Key Result

Lemma 3.1

For each $(n_0,n_0')$, let $(\widehat{\pi}^*_{n_0,n_0'},\phi^*_{n_0,n_0'})$ denote a measurable selection of global minimizers of eq:0.5 built from control samples $(Z_0,Z_0')$ of sizes $(n_0,n_0')$. Then, under Assumptions (as:map--as:kernel), there exists an event $\Omega^\star$ with $\mathbb P(\

Figures (2)

  • Figure 1: Illustration of the effect modification between the measure spaces $( \mathcal{Z}, \mu )$ and $( \mathcal{Z}', \mu' )$. An arrow represents the mapping $\varphi: \mathcal{Z} \to \mathcal{Z}'$ such that $\mu'$ is a push-forward measure of $\mu$ along $\varphi$; that is, for all $A \in \sigma(\mathcal{Z}')$, we have $\mu'(A) = \mu(\varphi^{-1}(A))$.
  • Figure 2: Comparison of synthesized target–site treatments across robustness scenarios (columns) and methods (rows). Orange points show the synthesized $Z_1^{\prime\,(\mathrm{synth.})}$; blue points show the oracle $Z_1'$. For the higher-dimensional ($d=30$) setting, the first two principal components of features $X$ are shown along with response $Y$, to enable 3D visualization.

Theorems & Definitions (2)

  • Lemma 3.1
  • Theorem 3.1