Improving the Generation and Evaluation of Synthetic Data for Downstream Medical Causal Inference

TL;DR

The paper tackles the challenge of generating synthetic medical data suitable for causal inference by formalizing three key distributions to preserve: the covariate distribution $P_{\mathbf{X}}$, the treatment mechanism $P_{W|\mathbf{X}}$, and the outcome mechanism $P_{Y|W,\mathbf{X}}$. It argues that existing synthetic-data evaluation metrics fail to capture downstream causal utility and introduces a triplet of task-aligned metrics, including covariate precision/recall $P_{\alpha,\mathbf{X}}$ and $R_{\beta,\mathbf{X}}$, $\text{JSD}_{\pi}$ for the propensity mechanism, and $U_{PEHE}$ for outcome preservation. The authors then propose STEAM, a model-agnostic generator that explicitly mimics the data-generating process by separately modelling $Q_{\mathbf{X}}$, $Q_{W|\mathbf{X}}$, and $Q_{Y|W,\mathbf{X}}$, to maximise adherence to the desiderata. Extensive experiments on real medical datasets and simulations show STEAM achieves state-of-the-art performance on the proposed metrics, particularly as the true data-generating process becomes more complex or high-dimensional. The work provides a practical framework for producing high-fidelity synthetic data for treatment-effect analyses with clear implications for privacy-preserving data sharing and downstream causal inference research.

Abstract

Causal inference is essential for developing and evaluating medical interventions, yet real-world medical datasets are often difficult to access due to regulatory barriers. This makes synthetic data a potentially valuable asset that enables these medical analyses, along with the development of new inference methods themselves. Generative models can produce synthetic data that closely approximate real data distributions, yet existing methods do not consider the unique challenges that downstream causal inference tasks, and specifically those focused on treatments, pose. We establish a set of desiderata that synthetic data containing treatments should satisfy to maximise downstream utility: preservation of (i) the covariate distribution, (ii) the treatment assignment mechanism, and (iii) the outcome generation mechanism. Based on these desiderata, we propose a set of evaluation metrics to assess such synthetic data. Finally, we present STEAM: a novel method for generating Synthetic data for Treatment Effect Analysis in Medicine that mimics the data-generating process of data containing treatments and optimises for our desiderata. We empirically demonstrate that STEAM achieves state-of-the-art performance across our metrics as compared to existing generative models, particularly as the complexity of the true data-generating process increases.

Figures (8)

• Figure 1: $U_\text{PEHE}$ and $\text{JSD}_\pi$ as $d$ increases. Averaged over 10 runs, with 95% CIs.
• Figure 2: $\text{JSD}_\pi$ as $K$ increases. Averaged over 10 runs, with 95% CIs.
• Figure 3: $U_\text{PEHE}$ as $K$ increases. Averaged over 10 runs, with 95% CIs.
• Figure 4: DGPs for generic generative models (left), real datasets (middle), and STEAM (right).
• Figure 5: $P_{\alpha, \mathbf{X}}$ ($\uparrow$), $R_{\beta, \mathbf{X}}$ ($\uparrow$), $\text{JSD}_\pi$ ($\uparrow$), and $U_\text{PEHE}$ ($\downarrow$) evaluating $\text{STEAM\xspace}_\text{AIM}$ and standard AIM across privacy budgets. Averaged over 5 runs, shaded area represents 95% CIs.

Theorems & Definitions (9)

• Theorem 1
• proof
• Theorem 1
• proof
• Proposition 1
• proof
• Theorem : Post-Processing Theorem
• Theorem : Composition Theorem
• proof