📊 statistics

Distributional Treatment Effect Transportability across Heterogeneous Sites

This paper proposes a framework for recovering the full distribution of treatment effects in a target site using only control data from that site and both treated and control data from a source site, by learning an optimal transport transformation to account for cross-site heterogeneity and applying it to the source treated samples.

Original authors: Borna Bateni, Yubai Yuan, Qi Xu, Annie Qu

Published 2026-09-04

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Borna Bateni, Yubai Yuan, Qi Xu, Annie Qu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of medical research and public policy, scientists often face a frustrating gap: they have a clear picture of how a treatment works in one specific place, but they need to know how it will work in a different place where they have never tested it. Imagine a new drug that has been rigorously studied in a major hospital in California, where researchers have data on both the patients who received the drug and those who did not. Now, a hospital in a different state wants to use that drug, but they have only observed their own patients who did not receive it; they have no data on what happens to their patients if they do get the treatment. The challenge is to predict the full range of outcomes for the new group, not just an average number, but the entire shape of how different patients might respond. This requires more than just matching patient characteristics; it demands a way to translate the complex reality of one medical environment into another, even when the two places measure things differently, use different equipment, or have populations with different underlying health profiles.

Researchers from universities across the United States have developed a new method to solve this specific puzzle, known as the "cross-site, one-armed target" problem. Their approach moves beyond the traditional goal of simply calculating an average treatment effect, which reduces a complex reality to a single number. Instead, they aim to reconstruct the entire distribution of outcomes—the full spectrum of how patients might fare, including the best cases, the worst cases, and everything in between. The team realized that existing methods often fail because they assume that differences between two sites are limited to the mix of patients they see. In reality, sites can differ in how they record data, the scales they use for measurements, and even the hidden factors that shape how a disease progresses. To handle this, the researchers treated the two sites as distinct worlds with their own internal rules for similarity, and they sought a mathematical bridge that could translate the known data from one world into the other.

The core of their solution involves a two-step process that relies on a concept called optimal transport, which can be thought of as finding the most efficient way to move a pile of sand from one shape to another. First, the researchers used data from the untreated patients in both the source site (where they have full data) and the target site (where they only have untreated data) to learn how the two environments differ. They looked for a transformation that could map the untreated patients from the source site onto the untreated patients in the target site, respecting the unique way each site measures similarity between patients. This step allowed them to identify a "shared transformation," a set of rules that explains how the entire landscape of patient data shifts from one location to another. Crucially, they assumed that this same set of rules applies to treated patients as well, meaning that if the environment changes the way untreated patients respond, it changes the way treated patients respond in the same structural manner.

Once this bridge was built using the untreated data, the researchers applied it to the treated patients from the source site. By running the source treated patients through the learned transformation, they generated a synthetic dataset that represented what the treated patients in the target site would look like. This synthetic data was not a guess based on averages; it was a constructed sample that preserved the complex relationships between patient features and outcomes, including how the treatment reshaped the tails of the distribution where extreme outcomes occur. The researchers tested this method through extensive computer simulations where they knew the true answer in advance. In these tests, their method consistently outperformed other statistical approaches, especially when the differences between the two sites were complex and non-linear. It successfully recovered the correct distribution of outcomes, including the mean, the spread, and the behavior of extreme cases, whereas other methods often failed to capture the full picture or produced distorted results.

The team also applied their method to real-world data from patient-derived xenograft studies, a highly controlled resource used in cancer research where human tumor samples are implanted into mice to test treatments. They used data from breast cancer patients as the source and tried to predict outcomes for patients with melanoma, colorectal cancer, lung cancer, and pancreatic cancer. In every case, their method produced a synthetic distribution of tumor growth times that closely matched the actual observed data, which had been held back for verification. The results showed that their approach could accurately recover the full distribution of treatment responses, including the median and the spread of outcomes, better than competing techniques. This success suggests that the method can effectively translate knowledge from one well-studied environment to a new one, even when the two environments differ in how they measure and report data.

However, the researchers are careful to note that the success of this method depends on a specific condition: the differences between the two sites must affect treated and untreated patients in a similar way. If the treatment itself is administered differently in the new location—for example, if the dosage is changed or the follow-up care is different—the method may not work, because the bridge built from the untreated patients would no longer apply to the treated ones. In their simulations, when they deliberately created scenarios where the treatment was handled differently across sites, the method's accuracy dropped, confirming that the assumption of a shared transformation is vital. While the method cannot be verified with the data alone, the researchers argue that it is a plausible approach when the differences between sites are due to broad environmental or measurement factors rather than specific changes to the treatment protocol. By providing a way to recover the full distribution of outcomes rather than just an average, this framework offers a powerful tool for policymakers and scientists who need to understand how interventions will play out in new and diverse populations.

Technical Summary: Distributional Treatment Effect Transportability across Heterogeneous Sites

Problem Formulation
This paper addresses the challenge of distributional transportability in a "cross-site, one-armed target" design. In this setting, a source site provides observed data for both treated and control units, while a target site provides data only for control units. The objective is not merely to estimate the Average Treatment Effect (ATE), but to recover the full distribution of the treated outcome in the target site ( $\mu'_1$ ).

The authors highlight two primary limitations in existing literature:

Focus on ATE: Most transportability methods focus on scalar summaries (ATE), failing to capture distributional features such as quantiles, tail behavior, and heterogeneity in impacts across subpopulations.
Restrictive Invariance Assumptions: Standard methods rely on "unconfounded location" (conditional invariance), which assumes cross-site differences arise solely from covariate distributions. This fails when sites differ in measurement scales, outcome reporting conventions, latent factors, or the geometry of the joint feature–outcome space.

The proposed framework models cross-site heterogeneity as a transformation acting on the joint feature–outcome distribution, allowing for differences in measurement systems, population composition, and latent contextual factors that cannot be resolved by covariate reweighting alone.

Methodology
The proposed method, OTSynth, operates in two main stages: learning a cross-site transformation from control data and applying it to synthesize the target treated distribution.

Identification Framework:
- The authors replace the conditional invariance assumption with a shared-transformation assumption. They posit that a bijective map $\phi$ exists that transports the source control distribution ( $\mu_0$ ) to the target control distribution ( $\mu'_0$ ) and, crucially, that this same map transports the source treated distribution ( $\mu_1$ ) to the target treated distribution ( $\mu'_1$ ).
- To handle sites with different feature spaces and metrics, the framework introduces an anchor subspace ( $H$ ) containing features directly comparable across sites (e.g., age, gender). An alignment function $K$ ensures the transformation preserves comparability on these anchors.
Learning the Transformation (Optimal Transport):
- The method learns the transformation $\phi$ by comparing observed control samples from the source and target sites.
- It utilizes a Regularized Fused Gromov–Wasserstein (RFGW) criterion. Unlike standard Optimal Transport (OT) which learns a coupling between distributions, this approach is "transformation-aware." It jointly optimizes:
  - Geometric Alignment: Preserving the relational geometry (distances between pairs) of the feature–outcome spaces.
  - Anchor Consistency: Penalizing mismatches on the anchor subspace.
  - Graph Consistency: Encouraging the coupling to concentrate on the graph of a valid transformation map.
- The target metric (similarity structure) is treated as a learnable component (pull-back metric) rather than a fixed prior.
Synthesis of Target Treated Distribution:
- Once the transformation $\hat{\phi}$ is learned from control data, it is applied to the observed source treated samples ( $Z_1$ ).
- For each source treated unit, a synthetic target analogue is constructed by solving an optimization problem that balances preserving the unit's relative geometry to controls, maintaining anchor compatibility, and adhering to the learned transformation.
- The resulting synthetic sample serves as an estimator for the unobserved target treated distribution.

Theoretical Guarantees
The paper establishes weak convergence results. Under the stated assumptions (shared transformation, absolute continuity of treated over control distributions, and regularity conditions on the transformation class), the empirical distribution of the synthetic treated sample converges weakly to the true target treated distribution as sample sizes increase. A sensitivity analysis is also provided, showing that approximate violations of the shared-transformation assumption yield bounded Wasserstein errors.

Empirical Results
The framework was evaluated through extensive simulations and a real-world application:

Simulations: The method was tested against baselines including Two-Way Fixed Effects (TWFE), Matching, Generalized Synthetic Control (GenSynth), and CycleGAN.
- Linear Heterogeneity: OTSynth performed competitively with CycleGAN.
- Nonlinear Heterogeneity: OTSynth significantly outperformed all competitors, particularly in recovering outcome distribution tails, feature–outcome dependence, and Wasserstein distances. Competing methods (especially TWFE and Matching) failed to capture nonlinear shifts.
- Robustness: The method remained robust under non-random treatment assignment and limited cross-site overlap. Performance degraded only when the core shared-transformation assumption was explicitly violated (e.g., when treatment and control units experienced different cross-site shifts).
Real-World Application (PDX Data):
- The method was applied to Patient-Derived Xenograft (PDX) data to predict tumor doubling times (TTD) for various cancer types (NSCLC, Melanoma, CRC, PDAC) using Breast Cancer (BRCA) as the source.
- OTSynth achieved the lowest Wasserstein distance and KL divergence compared to the "oracle" (observed target treated) distribution, accurately reproducing means, medians, and quantiles across all target cancer types.

Key Contributions and Significance
The authors claim five primary contributions:

Conceptual: Formulating the problem as distributional transportability rather than ATE transportability, addressing the need to recover full counterfactual distributions.
Identification: Replacing the restrictive "unconfounded location" assumption with a shared cross-site transformation of joint feature–outcome distributions, accommodating richer heterogeneity (e.g., measurement scale shifts).
Methodology: Developing a transformation-aware Optimal Transport criterion (RFGW) that learns both the cross-site map and the target geometry simultaneously.
Theory: Proving the weak convergence of the synthetic empirical distribution to the true target distribution.
Empirical: Demonstrating superior performance in recovering distributional properties under nonlinear heterogeneity and complex real-world data structures compared to existing benchmarks.

Significance
The paper argues that its framework is significant because it enables the recovery of full distributional treatment effects in settings where traditional covariate reweighting fails due to structural shifts in measurement or outcome reporting. By leveraging optimal transport to learn a geometric transformation between sites, the method provides a robust tool for policy evaluation and clinical trial generalization where only control data is available in the target population. The authors note that, like other transportability frameworks, the validity of the shared-transformation assumption relies on scientific knowledge regarding whether cross-site differences affect treated and control units similarly.

Technical Summary: Distributional Treatment Effect Transportability across Heterogeneous Sites

More like this