Context Augmentation for Causal Foundation Models with Interventional Data and Domain Knowledge

Causal effect estimation typically requires expert knowledge to select an appropriate estimator. Recently proposed causal foundation models (CFMs) amortize the prediction of interventional distributions from observational contexts, showing promise for automatic causal effect estimation. However, observational data alone may be insufficient to identify interventional distributions: the model's implicit posterior may spread across multiple causal worlds that are observationally equivalent but imply different interventional distributions. In practice, randomized trial data or structural domain knowledge are sometimes available, yet existing CFMs either provide no interface for such information or require architectural changes and retraining. We propose test-time context augmentation, which supplies this information to a pretrained CFM through additional context samples without retraining. We first consider real interventional samples added without labels distinguishing them from observational samples. Under an idealized Bayesian analysis, we provide a posterior-odds condition under which these samples increase posterior mass on the class of structural causal models (SCMs) that induce the true interventional distribution. When only observational data are available, we fit an auxiliary SCM constrained by domain knowledge and use it to generate pseudo-interventional samples. The framework does not require the ground-truth DAG, and partial structural knowledge suffices as long as it identifies the interventional distribution. We instantiate it with a fully specified DAG and a more accessible causal order. Experiments show that context augmentation breaks observational symmetry in a non-identifiable bivariate setting and strengthens CFM representations of the local causal structure around treatment. Context augmentation also enables a graph-free CFM to match its graph-conditioned counterpart, and improves Do-PFN and CausalPFN on synthetic and semi-synthetic benchmarks.