CausalWrap: Model-Agnostic Causal Constraint Wrappers for Tabular Synthetic Data

Amir Asiaee; Zhuohui J. Liang; Chao Yan

cs.LGMar 2, 2026

CausalWrap: Model-Agnostic Causal Constraint Wrappers for Tabular Synthetic Data

Authors: Amir Asiaee, Zhuohui J. Liang, Chao Yan

Organizations: Department of Biostatistics, Vanderbilt University Medical Center, Nashville, TN, USA · Department of Biomedical Informatics, Vanderbilt University Medical Center, Nashville, TN, USA

Abstract

Tabular synthetic data generators are typically trained to match observational distributions, which can yield high conventional utility (e.g., column correlations, predictive accuracy) yet poor preservation of structural relations relevant to causal analysis and out-of-distribution (OOD) reasoning. When the downstream use of synthetic data involves causal reasoning -- estimating treatment effects, evaluating policies, or testing mediation pathways -- merely matching the observational distribution is insufficient: structural fidelity and treatment-mechanism preservation become essential. We propose CausalWrap (CW), a model-agnostic wrapper that injects partial causal knowledge (PCK) -- trusted edges, forbidden edges, and qualitative/monotonic constraints -- into any pretrained base generator (GAN, VAE, or diffusion model), without requiring access to its internals. CW learns a lightweight, differentiable post-hoc correction map applied to samples from the base generator, optimized with causal penalty terms under an augmented-Lagrangian schedule. We provide theoretical results connecting penalty-based optimization to constraint satisfaction and relating approximate factorization to joint distributional control. We validate CW on simulated structural causal models (SCMs) with known ground-truth interventions, semi-synthetic causal benchmarks (IHDP and an ACIC-style suite), and a real-world ICU cohort (MIMIC-IV) with expert-elicited partial graphs. CW improves causal fidelity across diverse base generators -- e.g., reducing average treatment effect (ATE) error by up to 63% on ACIC and lifting ATE agreement from 0.00 to 0.38 on the intensive care unit (ICU) cohort -- while largely retaining conventional utility.

Figures & tables

Panel of Figure 1 — Figure 1: Partial causal knowledge $\mathcal{K}=(E^{+},E^{0},\mathcal{M})$ . Green solid : trusted edges ( $E^{+}$ ). Red dashed : forbidden edges ( $E^{0}$ ). Gray dotted: unknown status.

Panel of Figure 2 — Figure 2: Gradient flow through the CI penalty. Forward (black): raw samples pass through $f_{\phi}$ ; frozen edge models $\hat{m}$ produce residuals; HSIC on residual pairs gives the penalty. Backward (red): gradients flow back through $f_{\phi}$ only (edge models are frozen).

Panel of Figure 3 — Figure 3: CausalWrap pipeline. A pretrained base generator (frozen) produces raw samples; a correction map $f_{\phi}$ produces corrected samples, which—together with real data and partial knowledge $\mathcal{K}$ —feed the loss. Black : forward data flow; red : gradient (only $\phi$ is updated).

Generator	LG	NLA	MT	Overall
CTGAN	0.832	0.090	0.070	0.331
$\Delta$ CW	$-$ 9.1%	$+$ 11.6%	$+$ 23.3%	$-$ 4.9%
Gap closed	$+$ 29.1%	–	$-$ 42.3%	$+$ 23.0%
TabDDPM	0.644	0.180	0.042	0.289
$\Delta$ CW	$-$ 1.7%	$+$ 4.6%	$-$ 4.8%	$-$ 0.6%
Gap closed	$+$ 15.4%	–	$+$ 19.0%	$+$ 5.6%

Table 1: Tier 1: CausalWrap improvement (

\Delta

%) in ATE error across 3 SCM types (5 seeds). Green = CW reduces error; red = CW increases error. Gap closed reports the fraction of the oracle gap closed by CW:

(e_{\mathrm{base}}-e_{\mathrm{CW}})/(e_{\mathrm{base}}-e_{\mathrm{oracle}})

, shown as “–” when

e_{\mathrm{oracle}}\geq e_{\mathrm{base}}

or when the oracle headroom

(e_{\mathrm{base}}-e_{\mathrm{oracle}})/e_{\mathrm{base}}

<5\%

(unstable ratio). Overall is the unweighted average across the 3 SCM types. Best partial-knowledge result underlined .

Panel of Figure 4 — Figure 4: Tier 2: IHDP benchmark (10 replications, 5 seeds; log-scale violin plots). CW reduces mean ATE error for CTGAN and TabDDPM but can increase error for TVAE, and may increase tail variance on some replications.

Figure 5: Tier 2: ACIC-style ATE error across 10 DGP settings (2 seeds). CW improves all three bases on most settings (largest gains for TVAE), illustrating base-model sensitivity of post-hoc correction.

Panel of Figure 6 — Figure 6: Tier 3 ICU causal knowledge $\mathcal{K}$ . Green solid : trusted edges $E^{+}$ (demographics, labs, and ventilation cause vasopressor use; treatment and baseline variables cause mortality). $+$ on an arrow: monotonic constraint in $\mathcal{M}$ (lactate $\uparrow$ , creatinine $\uparrow$ $\Rightarrow$ mortality risk $\uparrow$ ). Red dashed : forbidden directions $E^{0}$ (temporal/biological impossibilities).

Figure 7: Tier 3: ICU cohort (MIMIC-IV). Radar plot of seed-averaged metrics (TSTR, ATE agreement, and

1/(1+\mathrm{MMD})

) for each base generator and its CausalWrap variant. All axes are scaled so larger is better.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 8 — Figure 8: Connecting the theoretical guarantees. Theorem 1 ensures penalties vanish; Proposition 1 converts conditional closeness to joint TV control. The intermediate step (dashed red) is a design motivation: residualized HSIC penalties are a proxy for conditional matching, not a formal equivalence.

	Conventional			Causal
Generator	MMD $\downarrow$	JSD $\downarrow$	TSTR $\uparrow$	ATE err $\downarrow$	CI Pass $\uparrow$
CTGAN	0.023	0.052	0.744	0.832	0.669
$\Delta$ CW	$-$ 72.7%	$-$ 31.7%	$+$ 1.0%	$-9.1\%$	$+$ 4.7%
TabDDPM	0.011	0.020	0.763	0.644	0.486
$\Delta$ CW	$-$ 39.2%	$-$ 7.6%	$+$ 0.3%	$-$ 1.7%	$+$ 34.1%
TVAE	0.011	0.056	0.761	0.766	0.916

Appendix

Table 2: Tier 1: CausalWrap on the LG SCM (

d=10

n=5{,}000

, 5 seeds). Each pair shows base

\to

+CW.

\Delta

% row: green = improvement, red = degradation. Underline marks the largest ATE error reduction among wrapped bases.

	$E^{0}=\emptyset$			Full $\mathcal{K}$ (with $E^{0}$ )
Base	$\Delta$ MMD	$\Delta$ TSTR	$\Delta$ ATE	$\Delta$ MMD	$\Delta$ TSTR	$\Delta$ ATE
CTGAN	$-$ 69%	$+$ 0.1%	$+$ .19	$-$ 64%	$-$ 1.5%	$-$ .09
TabDDPM	$+$ 72%	$-$ 17%	$+$ .24	$+$ 61%	$-$ 28%	$+$ .38
TVAE	$-$ 6%	$-$ 0.4%	$+$ .31	$+$ 40%	$-$ 0.6%	$+$ .13

Appendix

Table 3: Forbidden-edge ablation on the ICU cohort (5 seeds).

\Delta

CW reports the change from base to +CW within each knowledge configuration.

E^{0}{=}\emptyset

: trusted edges and monotonicity only. Full

\mathcal{K}

: adds 7 temporal forbidden edges.

Panel of Figure 9 — Figure 9: Ablation studies on the Linear-Gaussian SCM with TabDDPM base.

Panel of Figure 10 — Figure 10: Additional ablation studies on the Linear-Gaussian SCM with TabDDPM base.

Synthetic Data Generation

Synthetic data generation creates artificial examples, simulations, and datasets to train, test, or protect machine-learning systems.

See all 606 papers

Generative Synthetic Data for Causal Inference: Pitfalls, Remedies, and Opportunities
Apr 26, 2026Yichen XuCausal Effect Estimation Synthetic Data Generation
University of California, Berkeley
Beyond Distributional Fidelity: Causal-Penalized Diffusion for Synthetic Tabular Data
Oct 8, 2026Lan Tao, Yongxian He, Shirong Xu +2Synthetic Tabular Data Generation Tabular Diffusion Models
Department of Statistics & Data Science, University of California, Los Angeles · Herbert Wertheim School of Public Health and Human Longevity Science, University of California, San Diego · Wang Yanan Institute for Studies in Economics, Xiamen University
Causal-Aware Tabular GANs with Reinforcement Learning
Oct 28, 2025Tu Anh Hoang Nguyen, Dang Nguyen, Tri-Nhan Vo +3Synthetic Data Generation Synthetic Tabular Data Generation
Applied Artificial Intelligence Initiative (A2I2), Deakin University, Australia · School of Computer Science and Information Technology, Adelaide University, Australia

Apr 26, 2026J/K move · Enter open · S save

Generative Synthetic Data for Causal Inference: Pitfalls, Remedies, and Opportunities

Yichen Xu

University of California, Berkeley

Causal Effect Estimation Synthetic Data Generation Synthetic Tabular Data Generation Synthetic Data Evaluation Causal Intervention

Open paper AskPDF arXiv

Synthetic tabular data are often evaluated by distributional similarity, privacy distance, or train-on-synthetic-test-on-real predictive performance, but these criteria do not ensure validity for causal inference. We show that fully generative tabular synthesizers, including GAN- and LLM-based models, can preserve predictive utility while distorting average treatment effect (ATE) estimates. The failure is structural: ATE preservation requires both a realistic covariate law and an accurate treatment-effect contrast, whereas prediction loss penalizes treatment-effect error only through an overlap-weighted term. Thus, under imbalance or limited overlap, a generator may reproduce dominant observed outcomes while underlearning intervention-relevant contrasts. We formalize this mismatch through sensitivity and loss-decomposition results. Motivated by this causal analysis and intuition, we propose a hybrid synthetic-data framework for causal inference that generates covariates while modeling treatment and outcome mechanisms separately. We evaluate the framework in three settings: ATE preservation under fully generative versus hybrid synthesis, augmentation for practical positivity problems, and diagnostic simulation engines for comparing OR, IPW, AIPW, and TMLE before real-data analysis. We also stress-test the hybrid construction across settings that vary overlap, covariate dimension, seed sample size, and treatment-effect complexity, including a logistic outcome-model misspecification check. Across controlled simulation experiments, hybrid synthesis improves causal fidelity relative to fully generative baselines; the ACTG application shows improved predictive fidelity and potential for finite-sample estimator benchmarking. LLM-based hybrid synthesis is often more faithful than CTGAN in settings where causal fidelity can be assessed.

Abstract

Synthetic Data Generation

Explore similar work

Generative Synthetic Data for Causal Inference: Pitfalls, Remedies, and Opportunities

Beyond Distributional Fidelity: Causal-Penalized Diffusion for Synthetic Tabular Data

Causal-Aware Tabular GANs with Reinforcement Learning