Distributional Treatment Effect Estimation across Heterogeneous Sites via Optimal Transport

TL;DR

This paper introduces a distributional causal inference framework using optimal transport to estimate treatment effects across heterogeneous sites, enabling the synthesis of counterfactual data and improving understanding of effect heterogeneity.

Contribution

It presents a novel optimal transport-based method for modeling and aligning treatment and control distributions across sites, advancing distributional treatment effect estimation.

Findings

Method accurately recovers treatment effect distributions in simulations

Theoretical guarantees ensure consistency and convergence

Effective in real-world patient data application

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.