Controllable Generative Sandbox for Causal Inference

TL;DR

Practical utility is demonstrated in a comparative safety study of metastatic castration-resistant prostate cancer treatments, using CausalMix to compare estimators under calibrated data-generating processes, tune hyperparameters, and conduct simulation-based power analyses under targeted treatment effect heterogeneity scenarios.

Abstract

Method validation and study design in causal inference rely on synthetic data with known counterfactuals. Existing simulators trade off distributional realism, the ability to capture mixed-type and multimodal tabular data, against causal controllability, including explicit control over overlap, unmeasured confounding, and treatment effect heterogeneity. We introduce CausalMix, a variational generative framework that closes this gap by coupling a mixture of Gaussian latent priors with data-type-specific decoders for continuous, binary, and categorical variables. The model incorporates explicit causal controls: an overlap regularizer shaping propensity-score distributions, alongside direct parameterizations of confounding strength and effect heterogeneity. This unified objective preserves fidelity to the observed data while enabling factorial manipulation of causal mechanisms, allowing overlap, confounding strength, and treatment effect heterogeneity to be varied independently at design time. Across benchmarks, CausalMix achieves state-of-the-art distributional metrics on mixed-type tables while providing stable, fine-grained causal control. We demonstrate practical utility in a comparative safety study of metastatic castration-resistant prostate cancer treatments, using CausalMix to compare estimators under calibrated data-generating processes, tune hyperparameters, and conduct simulation-based power analyses under targeted treatment effect heterogeneity scenarios.

Figures (15)

• Figure 1: Overview of the CausalMix generative framework.
• Figure 2: 2D joint embedding of synthetic data under BGMM and Gaussian priors compared with the real mCRPC dataset across causal scenarios. Scenario 1: homogeneous effect, no confounding, perfect overlap; Scenario 2: linear heterogeneous effect, mild confounding, constant overlap; Scenario 3: nonlinear heterogeneity with covariate-dependent confounding and overlap
• Figure 3: Comparison of generated and target treatment effects of the mcrpc dataset from CausalMix with a BGMM prior in Scenario 3.
• Figure 4: Grouped CATE estimates by age quintile. For each age group, boxplots summarize individual-level CATE estimates from each estimator in the subgroup with Charlson = 0, no prior CVD, under a calibrated synthetic data-generating process without unmeasured confounding ($\kappa=0$). Black diamonds indicate the true group-level mean CATE.
• Figure 5: PEHE vs. minimum leaf size for Causal Forests with 2,000 trees. Error bars indicate ±1 standard deviation across replications.