A Structural View of Query Misspecification in Causal Foundation Models

Junha Ham, Deokgyu Kim, Doeun Kim, Serjin Kim, Sanghack Lee

Keywords: Amortized Bayesian Inference, Posterior Predictive Inference, Structural Causal Model, Causal Foundation Model, Prior-Data Fitted Network, Post-Treatment Bias

Abstract: Causal Foundation Models (CFMs) pretrain amortized causal estimators on large collections of synthetic datasets sampled from structural causal model (SCM) priors. In optimal-capacity, they recover the corresponding interventional posterior predictive target for queries on the training query surface. We study the failure mode induced when an inference-time query includes a post-treatment covariate. Structurally, we decompose the resulting CATE bias into three components: loss of the natural indirect effect, an interaction penalty, and treatment-differenced selection bias. Distributionally, we prove that conditioning on post-treatment values yields strictly positive KL divergence from the marginal interventional law on a positive-measure subset, and we provide closed-form KL decompositions under linear-Gaussian SCMs. Empirically, removing post-treatment covariates from the query yields substantial reductions in PEHE across graph topologies and CFM models without retraining. We further introduce Treatment-Centric Local Discovery (TC-LD), a lightweight pre-inference filter that flags likely post-treatment variables and recovers most of this improvement on our synthetic benchmark.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 123