Paper deep dive
Foundation Models for Partial Causal Identification
Alexis Bellot, Anish Dhir
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 96%
Last extracted: 8/24/2026, 5:28:10 AM
Summary
This paper introduces a causal foundation model for partial causal identification, addressing the challenge of bounding intervention and counterfactual effects from observational data under unobserved confounding. The authors define a canonical prior with full support over structural causal models (SCMs) with discrete observables, enabling the translation of bounding counterfactuals into learning distributions over functions via Prior-Data Fitted Networks (PFN). Theoretical results prove that the PFN's predictive distribution converges to the Bayesian posterior induced by the canonical prior, providing asymptotically valid frequentist coverage for the identified set. Experiments demonstrate that this universal solver outperforms traditional Gibbs sampling methods in inference speed while maintaining comparable coverage and interval width.
Entities (10)
Relation Signals (8)
Anish Dhir → authored → Causal Foundation Model
confidence 99% · Alexis Bellot 1 Anish Dhir 2 Abstract This paper investigates the development of causal foundation models
Alexis Bellot → authored → Causal Foundation Model
confidence 99% · Alexis Bellot 1 Anish Dhir 2 Abstract This paper investigates the development of causal foundation models
Causal Foundation Model → solves → Partial Causal Identification
confidence 97% · Cor. 1 makes precise the claim that the proposed foundation model is a universal solver for partial identification.
Canonical Prior → hassupportover → Structural Causal Model (SCM)
confidence 96% · We show that a canonical prior can be defined with full support over the space of structural causal models with discrete observables.
Causal Foundation Model → implements → Prior-Data Fitted Network (PFN)
confidence 96% · To approximate this posterior distribution for arbitrary queries and data distributions we use a Prior-Data Fitted Network (PFN)
Causal Foundation Model → uses → Canonical Prior
confidence 95% · Our distinctive contribution is the definition of a canonical prior with full support over all SCMs with discretely-valued observables.
Causal Foundation Model → handles → Unobserved Confounding
confidence 94% · This extends the promising causal foundational modelling paradigm to the estimation of partially-identifiable causal effects, i.e., under unobserved confounding
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:This paper investigates the development of causal foundation models for bounding the effect of interventions and counterfactuals from observational data. We show that a canonical prior can be defined with full support over the space of structural causal models with discrete observables. With this canonical prior, we translate the problem of bounding counterfactuals into that of learning distributions over functions that map data (and possibly structural assumptions) to a causal query of interest. This extends the promising causal foundational modelling paradigm to the estimation of partially-identifiable causal effects, i.e., under unobserved confounding, where multiple values are equally compatible with the observed data and prior structural assumptions.
Tags
Links
- Source: https://arxiv.org/abs/2608.20841v1
- Canonical: https://arxiv.org/abs/2608.20841v1
Trouble viewing inline? Open PDF directly →
Full Text
46,450 characters extracted from source content.
Expand or collapse full text
Foundation Models for Partial Causal Identification Alexis Bellot 1 Anish Dhir 2 Abstract This paper investigates the development of causal foundation models for bounding the effect of inter- ventions and counterfactuals from observational data. We show that a canonical prior can be de- fined with full support over the space of structural causal models with discrete observables. With this canonical prior, we translate the problem of bounding counterfactuals into that of learning dis- tributions over functions that map data (and pos- sibly structural assumptions) to a causal query of interest. This extends the promising causal foun- dational modelling paradigm to the estimation of partially-identifiable causal effects, i.e., under un- observed confounding, where multiple values are equally compatible with the observed data and prior structural assumptions. 1. Introduction In the background of any investigation into causal inference are concerns about identifiability: roughly, that the causal query may not be uniquely computable from the available data even in the large sample limit. It is commonly ac- knowledged that a relatively detailed understanding of the data-generating process is necessary for identification (Pearl, 2009; Tian & Pearl, 2002; Bareinboim et al., 2022). In practice, it is natural to ask for results and methods that are more “universally” applicable and do not require having access to qualitative knowledge of the system’s structure. In cases where the query is not uniquely computable we might ask: what is the tightest possible bound that contains all values consistent with the observed data and additional assumptions we might pose? This line of research falls under the rubric of partial causal identification (Manski, 1990; Balke & Pearl, 1994; Tian & Pearl, 2000; Finkelstein & Shpitser, 2020; Zhang & Bareinboim, 2021b; Zhang et al., 2022; Sachs et al., 2023; Duarte et al., 2024; Jung & Kang, 1 Independent 2 University College London. Correspondence to: Alexis Bellot <abellotolivan@gmail.com>. Proceedings of the2 nd ICML Workshop on Foundation Models for Structured Data, Seoul, South Korea. 2026. Copyright 2026 by the author(s). 2026). This generalizes the causal estimation problem to the identification of a set and has found applications in econometrics (Tamer, 2010), epidemiology (Mullahy et al., 2021), and AI (Jalaldoust et al., 2024; Li et al., 2025; Joshi et al., 2024; Bellot et al., 2023). Despite considerable progress, the existing toolbox for par- tial identification is highly fragmented. Each instance of the problem has historically required a bespoke derivation: spe- cific causal structures, queries, and types of input datasets (e.g. observational vs interventional), require separate op- timization runs. On the other hand, Prior-Data Fitted Net- works (Garnelo et al., 2018; M ̈ uller et al., 2022) are estab- lishing a powerful paradigm for learning distributions over functions between two general spaces, such as the space of datasets and causal queries. This idea has led to the develop- ment of causal foundation models (Robertson et al., 2025; Balazadeh Meresht et al., 2025; Dhir et al., 2025; Reuter et al., 2026; Ma et al., 2025; Bynum et al., 2025), mapping data to the value of a causal quantity by training on pairs of synthetically-generated examples. Most causal founda- tion models operate exclusively in settings where the query of interest is point-identifiable. Dhir et al. (2025); Reuter et al. (2026) apply this idea to partially identified regimes, producing a posterior distribution that does not necessarily concentrate around a single value. However, one cannot necessarily derive valid bounds from their posteriors in a principled manner because their prior is not guaranteed to give mass to all values of the query consistent with the data. In this paper, we develop a foundation model for partial identification that, given an observational dataset, outputs a posterior distribution over a causal query whose support provably encodes the identified set. Our distinctive contribu- tion is the definition of a canonical prior with full support over all SCMs with discretely-valued observables. This enables us to derive bounds on any causal query from any observational dataset (and possibly other prior knowledge, like a causal diagram) that are guaranteed to be valid and tight asymptotically as the dataset size grows. 1.1. Preliminaries The basic framework of our analysis rests on structural causal models (SCMs) (Pearl, 2009, Def. 7.1.1), which formalize the notion of a data-generating process. In addi- tion to specifying the distribution of data, these models also 1 arXiv:2608.20841v1 [cs.AI] 21 Aug 2026 Foundation Models for Partial Causal Identification specify the generative mechanisms that produce them. For the purpose of causal inference and learning, SCMs provide a broad, fine-grained hypothesis space. An SCMMis a tuple⟨V ,U,F,P⟩whereVis a set of endogenous variables andUis a set of exogenous variables. Fis a set of functions where eachf V ∈ Fdecides val- ues of an endogenous variableV ∈ Vtaking as argument a combination of other variables in the system. That is, V ← f V (Pa V ,U V )where observed parentsPa V ⊆ V and unobserved parentsU V ⊆ U. An interventiondo(x) replaces the causal mechanismf X with the assignment X ← xcorresponding to a sub-modelM x . We will use the counterfactual notationY x (u)to denote the value of the outcomeYinM x givenU = u. Drawing values of exoge- nous variablesUfollowing the distributionP (U )induces a counterfactual variableY x . Specifically, the eventY x = y (for short,y x ) can be read as “Ywould beyhadXbeen x”. With a given SCM M , for subsets Y,...,Z,X,...,W ⊆ V , the distribution over counterfactualsY x ,...,Z w is de- fined as: P M (y x ,...,w z ) =(1) Z 1Y x (u) = y,...,W z (u) = w dP (u). Given an observational distributions, the best bound on a counterfactual query is defined as follows. Definition 1 (Optimal Counterfactual Bound). For a distri- butionP (V ), the optimal bound[l,u]over a counterfactual probability P (y x ,...,z w ) is defined as, respectively, min/ max M∈M P M (y x ,...,z w )(2) s.t. P M (V ) = P (V ), whereP M denotes the distribution induced by the SCMM. The problem formulation is presented in the language of counterfactuals to encompass any causal or counterfactual statement of interest, such as a causal effectP (y x ) 1 . We will writeI Ψ (P M ) or justI for the true identified set. Each SCMMcan also be associated with a causal diagram G, where nodes represent endogenous variablesV, directed edges represent the argumentsPa V of each functionf V , and bi-directed edges represent the presence of a common latent variable in the arguments of two functions. We will leverage a special type of clustering of nodes inGcalled the confounded component (orc-component for short) from (Tian & Pearl, 2002). For a causal graphG, a subsetC ⊆ V 1 This definition can be extended to a collection of observational and experimental distributions, and additional constraints over the space of SCMs M. is ac-component if any pairV i ,V j ∈ Cis connected by a bi-directed path inG. Given a diagramG, we writeC(U ) for the c-component of U . Throughout this paper, we assume that domains of endoge- nous variablesVare discrete and finite; exogenous vari- ablesUcould a priori take values in any (continuous) do- main. The counterfactual probabilityP (Y x ,...,Z w )de- fined in Eq. (1) is thus a categorical distribution. We use supp V to denote the support of the distribution of V ∈ V . 2. Universal Causal Estimation We work towards the development of a universal solver for partial identification. Rather than targeting a specific causal diagram or a specific bound algorithm, a trained model implicitly solves the optimal bounding problem for any query derived from the SCM. The idea is to bound the value of a counterfactual query Ψ := P (y x ,...,w z )given an observational datasetD ∼ P (V ) by examining the pushforward induced by the poste- rior over SCMs. For any measurable set A⊆ [0, 1] define Π(Ψ∈ A|D) = Z 1Ψ M ∈ AdΠ(M |D).(3) Optimal bounds could then be approximated by the posterior support bounds, inf / sup supp Π(Ψ|D).(4) These correspond to the0- and1-quantile limits. We use I Π(·|D) to denote the set of values defined by the posterior support. The shape of the posterior may depend on the prior, butI Π(·|D) should capture the full range of counterfactual values compatible with the observed data if priors have full support as the dataset size grows. Following the paradigm of causal foundation models, we propose to learn a parametric approximation to the posterior directly from synthetic examples sampled from a prior over the space of SCMs Π(M ),M ∈ M. 2.1. Canonical Priors In general systems, the set of functionsf V : V ∈ V between variables and probability measuresP (U )that de- fine an SCMMmay be arbitrary. This makes it difficult to define good priors. For discrete observablesV, however, any set of counter- factuals derived from an SCMMcan be equivalently ob- tained from a canonical SCMNparameterized by categori- cal probability mass functionsP (U ),U ∈ Uwith a finite and fixed cardinality|supp U |,U ∈ U. This parameteriza- tion is known as the family of canonical SCMs proposed in (Zhang et al., 2022, Def. 2.3). 2 Foundation Models for Partial Causal Identification Definition 2 (Canonical SCM). A canonical SCMMis an SCM parameterized by categorical probability mass func- tions P (U ),U ∈ U with a finite and fixed cardinality, |supp U | = Y V∈C(U ) |supp PA V 7→ supp V |.(5) |supp PA V 7→ supp V | denotes the number of functions from the (discretely-valued) parents ofVto the values ofV. C(U ) denotes the c-component of U . For example, letG : X → Y,X ↔ Ywith binary X,Y ∈ 0, 1. An SCM corresponding to this struc- ture parameterized by a latent variableUwith cardinality |supp U | = |supp X |×|supp X 7→ supp Y | = 2× 4 = 8is a canonical SCM. The canonical parameterization is maximally expressive, that is, it can represent any SCMMwith discrete observ- ablesV(Zhang et al., 2022, Thm. 2.4). This suggests a natural family of priors over SCMs that we call canonical priors. Definition 3 (Canonical Prior). Given a set of observables V, a canonical priorΠ 0 is a probability measure on the joint space of causal diagramsGand the parameter space θ = P (u) : u ∈ supp U ,U ∈ Uof the corresponding family of canonical SCMs. We consider a particular instance of the canonical prior that uses the Erd ̋ os-R ́ enyi model for the causal diagram and the Dirichlet distribution for the exogenous parameters, as follows: Π 0 (G) = Erd ̋ osR ́ enyi(G;q D ,q B ),(6) Π 0 (θ |G) = Y U∈U Dir(θ U ;α1). q D ,q B ∈ (0, 1) are the probabilities of including a directed and a bidirected edge, respectively,α > 0is a concen- tration parameter andθ U ∈ ∆ |supp U |−1 is the exogenous distribution for U . Under this prior, every graph receives positive probability from the Bernoulli edge model and every categorical dis- tributionθ U := P (u) : u ∈ supp U ,U ∈ Ureceives positive probability from the Dirichlet distribution. Proposition 1 (Full Support). The canonical priorΠ 0 in Eq. (6) has full support over M: for any SCM M ∈ M and any open neighborhoodU of M in M, Π 0 (U ) > 0. Proofs are given in Sec. B. The key implication of Prop. 1 is that with full support canonical priors we can guarantee that the posteriorΠ 0 (· | D)remains faithful to the data in the sense that no SCM consistent with the observed distribution is assigned zero mass a priori. 2.2. Prior-Data Fitted Networks To approximate this posterior distribution for arbitrary queries and data distributions we use a Prior-Data Fitted Network (PFN) (M ̈ uller et al., 2022). They are a class of neural processes (Garnelo et al., 2018) that meta-learn to approximate the Bayesian posterior predictive distribution with a parametric modelq ω (Ψ|D)directly from samples from the prior. Following (Robertson et al., 2025; Dhir et al., 2025; Reuter et al., 2026), we adopt a transformer-based architecture with input embeddings designed to preserve variable values, and types (e.g., interventions, observations) as well as correla- tions between them. We then stack multiple transformer blocks with attention mechanisms that alternate at each layer between sample-wise attention, which shares information across samples, and feature-wise attention, which models dependencies among variables within each sample. The counterfactual queryΨis then decoded from the resulting representation through a linear head. The modelq ω is then trained to predict a discrete distribu- tion over the possible binned values ofΨ ∈ [0, 1]given the observational datasetD n . The training data is given by batches of examples(D n , Ψ)generated by repeatedly draw- ing canonical SCMsM ∼ Π 0 from the prior using Eq. (6), generating an observational datasetD n ∼ P M by sampling ni.i.d. observations, and computing the ground-truth coun- terfactual valueΨ(M )fromM. Architectural and training details are given in Sec. D. 2.3. Extensions Existing partial identification methods exploit additional qualitative knowledge about the underlying data generating process, for example encoded in a causal diagram (Zhang et al., 2022) or equivalence class (Bellot, 2023). One can readily condition on a partial or full specification of the causal diagramGby training the PFN on tuples(G,D n , Ψ) derived from the SCMM. This has been shown to be tractable by (Reuter et al., 2026). 3. What do we converge to? The following result shows that the PFN’s predictive dis- tribution converges, as the dataset grows, to the posterior induced by the canonical prior conditional on the observa- tional distribution. Theorem 1 (Consistency). LetD n be a dataset ofni.i.d. observations generated from an SCMM. Asn → ∞, the predictive distributionq ω (Ψ|D n )of the trained PFN converges weakly to the conditional prior: q ω (Ψ|D n )→ Π 0 (Ψ| P M ),(7) whereΠ 0 (Ψ| P M )is the prior distribution restricted to the 3 Foundation Models for Partial Causal Identification Figure 1. Predicted posterior distributions for interventional and counterfactual probabilities in a two-variable binary SCM experiment with 100 data samples. Dashed vertical lines indicate the corresponding analytical lower and upper bounds (given byP (x, y)≤ P (y x )≤ P (x, y) + P (x ′ ) and 0≤ P (y x , y ′ x ′ )≤ P (x, y) + P (x ′ , y ′ )). set of SCMs that generate P M . With sufficient PFN capacity, training drives the PFN pre- dictive distribution to the exact Bayesian posterior over the causal query. Asymptotically, all remaining uncertainty is precisely partial-identification uncertainty and the support of the predictive distribution converges to the identified set I Ψ (P M ). Theorem 2 (Frequentist Coverage). LetC q ω (Ψ|D n ),1−α be a(1 − α)Bayesian credible set derived from the PFN q ω (Ψ | D n ). Then, the credible setC q ω (Ψ|D n ),1−α is an asymptotically valid frequentist confidence set for the iden- tified set: lim n→∞ P M I ⊆C q ω (Ψ|D n ),1−α = 1− α.(8) This result extends existing Bayesian-frequentist equiva- lence results for partially identified models. It implies that the Bayesian credible set for the identified set obtained in Thm. 1 achieves asymptotic frequentist coverage. The following corollary summarizes the main contribution of the paper. Corollary 1 (Universality). For any observational dataset D n and counterfactual queryΨ, the100%Bayesian credi- ble interval induced by the PFNq ω (Ψ|D n ), trained on a canonical prior, converges asymptotically to the true identi- fied set: C q ω (Ψ|D n ),1 n→∞ −→I Ψ (P M ).(9) Cor. 1 makes precise the claim that the proposed foundation model is a universal solver for partial identification. 4. Experiments We demonstrate the use of causal foundation models (CFMs) for partial identification on binary two-variable systems for which analytical bounds are known. Specifically, we con- sider the inference ofP (y x )andP (y x , y ′ x ′ )given samples from the observational distributionP (x,y). For illustration, ModelnCov.WidthInf. time (ms) Gibbs101.000.921053.2 501.000.771095.2 1001.000.711146.5 2001.000.671322.2 4001.000.631521.1 CFM100.990.704.2 501.000.677.0 1001.000.678.9 2001.000.6414.5 4001.000.6324.9 Table 1. Coverage, mean interval width, and end-to-end inference time across context sizes for all evaluated models, averaged over 100 runs. we plot in Fig. 1 posterior distributions from a randomly chosenP (x,y), together with the corresponding tight ana- lytical bounds derived by Manski (1990) and Tian & Pearl (2000). In Tab. 1, we systematically compute the coverage, average width of derived intervals for all possible interventional and counterfactual queries, and inference time as a function of the number of context samplesn, for the proposed method (CFM) as well as the Gibbs sampler proposed by (Zhang et al., 2022). Given that the Gibbs sampler requires a causal diagram as an input we take the union of the predicted sets under all possible causal diagrams of two variables. 5. Conclusion We introduced a causal foundation-model approach to par- tial identification, where a Prior-Data Fitted Network is trained under a canonical prior with full support over dis- crete structural causal models. This formulation interprets counterfactual bounding as amortized posterior prediction over causal queries, with the potential to enable a single model to handle broad classes of queries, datasets, and structural assumptions without deriving problem-specific bounds each time. 4 Foundation Models for Partial Causal Identification References Balazadeh Meresht, V., Syrgkanis, V., and Krishnan, R. G. Partial identification of treatment effects with implicit generative models. Advances in Neural Information Pro- cessing Systems, 35:22816–22829, 2022. Balazadeh Meresht, V., Kamkari, H., Thomas, V., Li, B., Ma, J., Cresswell, J. C., and Krishnan, R. G. Causalpfn: Amortized causal effect estimation via in-context learning. arXiv preprint arXiv:2506.07918, 2025. doi: 10.48550/ arXiv.2506.07918. Balke, A. and Pearl, J. Counterfactual probabilities: Compu- tational methods, bounds, and applications. Proceedings of the Tenth Conference on Uncertainty in Artificial Intel- ligence, p. 46–54, 1994. Balke, A. and Pearl, J. Bounds on treatment effects from studies with imperfect compliance. Journal of the Amer- ican Statistical Association, 92(439):1171–1176, 1997. doi: 10.1080/01621459.1997.10474074. Bareinboim, E., Correa, J. D., Ibeling, D., and Icard, T. On pearl’s hierarchy and the foundations of causal inference. In Probabilistic and causal inference: the works of Judea Pearl, p. 507–556. 2022. Bellot, A. Towards bounding causal effects under markov equivalence. arXiv preprint arXiv:2311.07259, 2023. Bellot, A. and Chiappa, S. Towards estimating bounds on the effect of policies under unobserved confounding. Advances in Neural Information Processing Systems, 37: 104556–104594, 2024. Bellot, A., Malek, A., and Chiappa, S. Transportability for bandits with data from different environments. Advances in Neural Information Processing Systems, 36:44356– 44381, 2023. Bynum, L. E., Puli, A. M., Herrero-Quevedo, D., Nguyen, N., Fernandez-Granda, C., Cho, K., and Ranganath, R. Black box causal inference: Effect estimation via meta prediction. arXiv preprint arXiv:2503.05985, 2025. Dhir, A., Diaconu, C., Lungu, V., Requeima, J., Turner, R., and van der Wilk, M. Estimating interventional distribu- tions with uncertain causal graphs through meta-learning. Advances in Neural Information Processing Systems, 38: 140060–140096, 2025. Doob, J. L. Application of the theory of martingales. Le calcul des probabilit ́ es et ses applications, p. 23–27, 1949. Duarte, G., Finkelstein, N., Knox, D., Mummolo, J., and Shpitser, I. An automated approach to causal inference in discrete settings. Journal of the American Statistical Association, 119(547):1778–1793, 2024. Finkelstein, N. and Shpitser, I. Deriving bounds and inequal- ity constraints using logical relations among counterfac- tuals. Proceedings of the Conference on Uncertainty in Artificial Intelligence, p. 1348–1357, 2020. Garnelo, M., Schwarz, J., Rosenbaum, D., Viola, F., Rezende, D. J., Eslami, S. M. A., and Teh, Y. W. Neural processes. arXiv preprint arXiv:1807.01622, 2018. Hu, Y., Wu, Y., Zhang, L., and Wu, X. A generative adver- sarial framework for bounding confounded causal effects. In Proceedings of the AAAI Conference on Artificial In- telligence, volume 35, p. 12104–12112, 2021. Jalaldoust, K., Bellot, A., and Bareinboim, E. Partial trans- portability for domain generalization. Advances in Neu- ral Information Processing Systems, 37:137768–137805, 2024. Joshi, S., Zhang, J., and Bareinboim, E. Towards safe policy learning under partial identifiability: A causal approach. In Proceedings of the AAAI Conference on Artificial In- telligence, volume 38, p. 13004–13012, 2024. Jung, Y. and Kang, B. Information-theoretic causal bounds under unmeasured confounding. In Fifth Conference on Causal Learning and Reasoning, 2026. Kline, B. and Tamer, E. Bayesian inference in a class of partially identified models. Quantitative Economics, 7 (2):329–366, 2016. Li, M., Zhang, J., and Bareinboim, E. Confounding robust deep reinforcement learning: A causal approach. arXiv preprint arXiv:2510.21110, 2025. Ma, Y., Frauen, D., Javurek, E., and Feuerriegel, S. Foun- dation models for causal inference via prior-data fitted networks. arXiv preprint arXiv:2506.10914, 2025. doi: 10.48550/arXiv.2506.10914. Manski, C. F. Nonparametric bounds on treatment effects. The American Economic Review, 80(2):319–323, 1990. Miller, J. W. A detailed treatment of doob’s theorem. arXiv preprint arXiv:1801.03122, 2018. Mullahy, J., Venkataramani, A., Millimet, D. L., and Man- ski, C. F. Embracing uncertainty: the value of partial identification in public health and clinical research. Amer- ican journal of preventive medicine, 61(2):e103–e108, 2021. M ̈ uller, S., Hollmann, N., Pineda Arango, S., Grabocka, J., and Hutter, F. Transformers can do Bayesian inference. In International Conference on Learning Representations, 2022. 5 Foundation Models for Partial Causal Identification Padh, K., Zeitler, J., Watson, D., Kusner, M., Silva, R., and Kilbertus, N. Stochastic causal programming for bounding treatment effects. In Conference on Causal Learning and Reasoning, p. 142–176. PMLR, 2023. Pearl, J. Causality: Models, Reasoning, and Inference. Cambridge University Press, 2nd edition, 2009. Reuter, A., Dhir, A., Diaconu, C., Robertson, J., Ossen, O., Hutter, F., Weller, A., van der Wilk, M., and Sch ̈ olkopf, B. Use what you know: Causal foundation models with partial graphs. arXiv preprint arXiv:2602.14972, 2026. Robertson, J., Reuter, A., Guo, S., Hollmann, N., Hutter, F., and Sch ̈ olkopf, B. Do-pfn: In-context learning for causal effect estimation. arXiv preprint arXiv:2506.06039, 2025. doi: 10.48550/arXiv.2506.06039. Robins, J. M., Rotnitzky, A., and Scharfstein, D. O. Sen- sitivity analysis for selection bias and unmeasured con- founding in missing data and causal inference models. In Statistical models in epidemiology, the environment, and clinical trials, p. 1–94. Springer, 2000. Rosenbaum, P. R. Design sensitivity and efficiency in ob- servational studies. Journal of the American Statistical Association, 105(490):692–702, 2010. Sachs, M. C., Jonzon, G., Sj ̈ olander, A., and Gabriel, E. E. A general method for deriving tight symbolic bounds on causal effects. Journal of Computational and Graphical Statistics, 32(2):567–576, 2023. Sachs, M. C., Gabriel, E. E., Evans, R. J., and Sj ̈ olander, A. Deriving complete constraints in hidden variable models. arXiv preprint arXiv:2601.11242, 2026. Tamer, E. Partial identification in econometrics. Annu. Rev. Econ., 2(1):167–195, 2010. Tan, Z. A distributional approach for causal inference using propensity scores. Journal of the American Statistical Association, 101(476):1619–1637, 2006. Tian, J. and Pearl, J. Probabilities of causation: Bounds and identification. Annals of Mathematics and Artificial Intelligence, 28(1–4):287–313, 2000. doi: 10.1023/A: 1018912507879. Tian, J. and Pearl, J. A general identification condition for causal effects. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 18, p. 567–573, 2002. Zhang, J. Designing optimal dynamic treatment regimes: A causal reinforcement learning approach. In Interna- tional conference on machine learning, p. 11012–11022. PMLR, 2020. Zhang, J. and Bareinboim, E. Bounding causal effects on continuous outcome. In Proceedings of the AAAI Con- ference on Artificial Intelligence, volume 35, p. 12207– 12215, 2021a. Zhang, J. and Bareinboim, E. Bounding causal effects on continuous outcome. Proceedings of the AAAI Con- ference on Artificial Intelligence, 35(14):12207–12215, 2021b. doi: 10.1609/aaai.v35i14.17408. Zhang, J., Tian, J., and Bareinboim, E. Partial counterfac- tual identification from observational and experimental data. In International conference on machine learning, p. 26548–26558. PMLR, 2022. 6 Foundation Models for Partial Causal Identification A. Related Work In their seminal work in the early 1990’s, Manski (1990); Robins et al. (2000) showed that the effect of an intervention may always be bounded in a non-trivial interval (i.e., probabilities strictly contained in[0, 1]), irrespective of unobserved confounders or the causal structure underlying the variables involved. Causal effects are therefore in general said to be partially identifiable, i.e., one can derive bounds that shrink the range of a priori possible values for a given effect and potentially serve as a useful support for decision-making. Starting from this insight, the problem of partial identification has been attracting growing attention in the literature. To improve upon these bounds, one common restriction on the data generating mechanism is to assume knowledge of a causal graph describing the phenomenon of interest. Several results have derived progressively tighter bounds by exploiting the independencies in observational and interventional distributions implied by the causal graph (Balke & Pearl, 1997; Bellot, 2023; Sachs et al., 2026; Zhang, 2020; Zhang & Bareinboim, 2021a; Zhang et al., 2022; Bellot & Chiappa, 2024). For instance, Balke and Pearl (Balke & Pearl, 1997) (and subsequent refinements, e.g., (Zhang et al., 2022)) defined polynomial optimization programs to derive bounds that are provably optimal given a causal diagram. More recent proposals consider parameterizations in the space of linear combinations of a set of fixed basis functions (Padh et al., 2023) and neural networks (Balazadeh Meresht et al., 2022; Hu et al., 2021). Analytical bounds (rather than numerical approximations) have also been derived in selected settings, such as for the ”instrumental variable” graph (Zhang & Bareinboim, 2021a) and discrete systems with general graphs or equivalence classes (Bellot, 2023; Zhang, 2020). These results are all based on the assumption that the causal diagram (or equivalence class) is known exactly. They might be applied to the data-driven (graph-free) problem of partial identification by considering the union of the predicted sets under all possible causal diagrams of two variables (which, however, is difficult to scale to larger systems in practice). In parallel, a number of works have adopted sensitivity assumptions (as an alternative or in combination with a causal diagram) that quantify the degree of unobserved confounding through various data statistics, such as odds ratios, propensity scores, etc. A rich literature on sensitivity assumptions exists, including Tan’s sensitivity model (Tan, 2006) and Rosenbaum’s sensitivity model (Rosenbaum, 2010). These approaches start with estimators that optimize the average or conditional treatment effect bounds subject to constraints implied by the sensitivity model. While some of these works provide general methods for bounding the value of causal effects and counterfactuals, the approach we take in this paper is distinct. Rather than deriving problem-specific bounds via optimization or sensitivity- constrained estimation for each new dataset and query, we amortize partial identification through a Prior-Data Fitted Network trained once on synthetic examples from a canonical prior with full support over all structural causal models with discrete observables (Zhang et al., 2022). This also distinguishes our approach from existing causal foundation models (Robertson et al., 2025; Balazadeh Meresht et al., 2025; Dhir et al., 2025; Reuter et al., 2026; Ma et al., 2025), which target point-identified causal effects or place priors over restricted families of SCMs that need not assign mass to every counterfactual value compatible with the data; by construction, the support of our posterior predictive distribution asymptotically recovers the identified set, yielding bounds that are both valid and tight in the large-sample limit. 7 Foundation Models for Partial Causal Identification B. Proofs Prop. 1 restated. The canonical priorΠ 0 in Eq. (6) has full support overM: for any SCMM ∈ Mand any open neighborhoodU of M in M, Π 0 (U ) > 0. Proof. Fix an arbitrary SCMM ∈ Mand any open neighborhoodUofM. WriteM = (G,θ), whereGis the semi-Markov graph and θ are canonical parameters. Under Eq. (6), the prior factorizes as Π 0 (G,θ) = Π 0 (G)Π 0 (θ |G). First,Π 0 (G) > 0for every semi-Markov graphGonV. Indeed, the directed-edge and bidirected-edge parts are independent Bernoulli draws with probabilities q D ,q B ∈ (0, 1), so any finite edge pattern has strictly positive probability. Second, conditional onG, the parameter prior is Π 0 (θ |G) = Y U∈U Dir(θ U ;α1), α > 0. Each Dirichlet density is strictly positive on the relative interior of its simplex, hence has full support on that simplex; therefore the product prior has full support on Θ G = Y U∈U ∆ |supp U |−1 . So every non-empty open set in Θ G has positive Π 0 (·|G) mass. By (Zhang et al., 2022), canonical SCMs are representation-complete for discrete observables, and the map from canonical parameters to SCM-induced distributions/queries is continuous (polynomial inθ). Hence any open neighborhoodUofM induces a non-empty open parameter neighborhood around θ within the graph componentG, so Π 0 (U |G) > 0. Combining with Π 0 (G) > 0 gives Π 0 (U )≥ Π 0 (G)Π 0 (U |G) > 0. Since M andU were arbitrary, Π 0 has full support on M. Thm. 1 restated. LetD n be a dataset ofni.i.d. observations generated from an SCMM. Asn → ∞, the predictive distributionq ω (Ψ | D n )of the trained PFN converges weakly to the conditional prior distribution of the causal effectΨ given the true observational distribution P M : q ω (Ψ|D n ) w −→ Π 0 (Ψ| P M ),(10) whereΠ 0 (Ψ| P M )is the prior distribution restricted to the set of SCMs that generate the same observational distribution P M . Proof.The training objective for the PFN is the minimization of the expected KL divergence between the true posterior and the approximation: ω ∗ = arg min ω E M∼Π 0 E D n ∼P M KL P (Ψ|D n )∥q ω (Ψ|D n ) .(11) As shown in M ̈ uller et al. (2022) (Corollary 1.2), for a model with universal approximation capacity, the global optimum q ω ∗ satisfiesq ω ∗ (Ψ | D n ) = P (Ψ | D n )for allD n for which the posterior is defined. Thus, the consistency of the PFN reduces to the consistency of the Bayesian posterior itself. Asn→∞, we aim to show that once the empirical distribution of the data converges to the true observational distribution P M , which is point-identified, the remaining uncertainty characterizes the partially-identified bounds for the counterfactual query exactly. We start by applying Doob’s theorem (Doob, 1949) to the point-identified observational distributionμ = φ(M ) = P M . Following the recent treatment presented in (Miller, 2018, Thms. 2.1 and 2.2), we verify a number of regularity conditions: 8 Foundation Models for Partial Causal Identification •Polish parameter space. With the joint prior, the full parameter space is the disjoint union F G G× Θ G , where Θ G = × U∈U ∆ |Ω U |−1 is the simplex product for a fixed diagramG. Sincep = |V|is fixed there are finitely many semi-Markov diagrams onpnodes, eachΘ G is a compact metric space, and a finite disjoint union of compact metric spaces is again compact and separable, hence a complete separable metric space. •Measurability ofM 7→ P M (A). In the canonical SCM,P M (V )is a polynomial function of the latent probabilities P (u) u∈Ω U . Polynomials are continuous, and continuous functions on Polish spaces are Borel measurable. •Identifiability of the observational distribution. The joint distributionμ = P M (V )is point-identified from the empirical frequencies inD n ; the mappingM 7→ P M is many-to-one (multiple SCMs may induce the sameμ), butμitself is identifiable. • L 1 integrability. The counterfactual queryΨ M = P u 1Y x (u) = y x ,...,W z (u) = w z P (u)is bounded in[0, 1], hence Ψ∈ L 1 (Π 0 ). Sinceμis point-identified, Doob’s theorem guarantees that forΠ 0 -almost allM, the posterior over the observational distribution concentrates: P (μ|D n ) w −→ δ P M as n→∞.(12) While the posterior over the observational distribution concentrates, the posterior over the counterfactual query does not. To see this more precisely, we decompose the prior via the Disintegration Theorem. Letφ : M→ ∆(V )map each SCM to its observational distribution. The prior Π 0 decomposes as: dΠ 0 (M ) = dΠ 0 (M | μ)dΠ μ (μ),(13) whereΠ μ is the marginal over observational distributions andΠ 0 (·| μ) μ is the family of conditional priors on the set of SCMsφ −1 (μ). The Disintegration Theorem guarantees that this decomposition is unique and that the mapμ7→ Π 0 (·| μ)is measurable. Let h : R→ R be any bounded continuous test function. By the tower property: E[h(Ψ)|D n ] = Z h(Ψ) dΠ 0 (M |D n )(14) = Z μ Z M∈φ −1 (μ) h(Ψ) dΠ 0 (M | μ,D n ) ! dΠ 0 (μ|D n )(15) = Z μ Z φ −1 (μ) h(Ψ(M ))dΠ 0 (M | μ) ! |z =:H(μ) dP (μ|D n ).(16) The second equality follows from the Disintegration Theorem. For the third equality, in our setup the likelihood factorizes through the observational distributionμ : L(D n | M ) = L(D n | μ). So ifM 1 ,M 2 ∈ φ −1 (μ), thenL(D n | M 1 ) = L(D n | M 2 ) = L(D n | μ) . Bayes’ rule updates relative weights by prior times likelihood. If likelihood is identical for all M ∈ φ −1 (μ), relative weights stay exactly as in the prior: Π 0 (M | μ,D n )∝ Π 0 (M | μ)× L(D n | μ).(17) The factor L(D n | μ) is constant in M , so it cancels in normalization. Hence Π 0 (M | μ,D n ) = Π 0 (M | μ),for all n.(18) Finally, since all observational and counterfactual quantities are polynomial functions ofθ ∈ ΘandΨ(M )is continuous, H(μ)is a continuous function ofμ. Applying the Continuous Mapping Theorem to the observational distribution (here denoted P M ∗ ) concentration result: Z H(μ)dP (μ|D n )−→ H(P M ∗ ) = Z φ −1 (P M ∗ ) h(Ψ(M ))dΠ 0 (M | P M ∗ ).(19) 9 Foundation Models for Partial Causal Identification Since this convergence holds for every bounded continuous h, the Portmanteau theorem implies: P (Ψ|D n ) w −→ Π 0 (Ψ| P M ∗ ).(20) Combining this with the optimality guarantee of the PFN, we conclude q ω (Ψ|D n ) w −→ Π 0 (Ψ| P M ∗ ). Thm. 2 restated. LetC q ω (Ψ|D n ),1−α be a(1− α)Bayesian credible set derived from the PFNq ω (Ψ | D n ). Then, the credible setC q ω (Ψ|D n ),1−α is an asymptotically valid frequentist confidence set for the identified set: lim n→∞ P M I ⊆C q ω (Ψ|D n ),1−α = 1− α.(21) Proof. The proof is an application of Kline & Tamer (2016, Thm. 5) to our set-up. For this we verify the conditions of their theorem: Assumptions 1, 3, 5, 6 in (Kline & Tamer, 2016). Letμ := P M (V )be the so-called reduced-form parameter (cell probabilities of the observational law) in (Kline & Tamer, 2016), and define the identified-set map Γ(μ) :=I(μ) = [ℓ(μ),r(μ)], whereℓ(μ),r(μ)are the optimal values of the optimization program in Eq. (2). Thus our model is exactly a map from a point-identified reduced form μ to an identified set for the partially identified target. The parameter space forμis a simplex (hence a Borel subset of a finite-dimensional Euclidean space), so Assumption 1 in Kline & Tamer (2016) holds. Since observations are i.i.d. categorical, the reduced-form posterior satisfies a Bernstein–von Mises limit and the frequentist CLT for ˆμ n holds: √ n(μ− ˆμ n )|D n ⇒ N (0, Σ 0 ), √ n(ˆμ n − μ ∗ )⇒ N (0, Σ 0 ), which match Assumptions 3 and 6 in Kline & Tamer (2016). For Assumption 5 (their asymptotic independence condition), we impose the standard smooth-interval sufficient condition from Kline & Tamer (2016, Remark 5 and Lemma 1): nearμ ∗ , Γ(μ) is nonempty and the endpoint maps ℓ(μ),r(μ) are regular enough for the (Bayesian and frequentist) delta method. LetI Π 0 (Ψ|D n ),1−α denote a (1− α) Bayesian credible set for the identified set, i.e. Π Γ(μ)⊆I Π 0 (Ψ|D n ),1−α |D n = 1− α. By Kline & Tamer (2016, Thm. 5), P M ∗ Γ(μ ∗ )⊆I Π 0 (Ψ|D n ),1−α → 1− α. SinceΓ(μ ∗ ) =I(P M ∗ ), this is the desired frequentist coverage statement. Finally, because Thm. 1 givesq ω (Ψ|D n )⇒ Π 0 (Ψ | P M ∗ ) whereM ∗ is the true SCM: the PFN-based credible sets asymptotically coincide with the corresponding Bayesian credible sets for the identified set and inherit the same limit coverage. Remark.If the causal effect happens to be point-identified (e.g., under unconfoundedness and known causal diagram), the identified set collapses to a point and Thm. 2 reduces to the standard frequentist coverage guarantee for a point-identified parameter. Remark. Compared with recent amortized causal PFN theory (Balazadeh Meresht et al., 2025; Ma et al., 2025), our guarantee targets a different object: under partial identification, the PFN predictive converges weakly to the conditional prior Π 0 (Ψ| P M )on the collection of SCMs compatible with the true observational law, whereas they study identified causal functionals and aim for point recovery. The proof strategy is similar: all use Doob-type concentration on the observational parameter and PFN optimality (M ̈ uller et al., 2022) to reduce the network to the exact Bayesian predictive; though an important difference is that Balazadeh Meresht et al. and Ma et al. impose identification restrictions on the prior and assume that the mapM 7→ P M is measurable, we explicitly prove and verify that this is true under the canonical SCM prior for discrete systems that we propose. 10 Foundation Models for Partial Causal Identification Figure 2. Predicted posterior for the interventional queryP (Y x=1 = 1)and the counterfactual queryP (Y x=1 = 0, Y x=0 = 0)as a function of the number of context samples n for a randomly drawn SCM. C. Additional Experiments In this section, we consider additional experiments to evaluate the performance of the CFM. We start by analyzing the predicted posterior distributions for the interventional and counterfactual queries as a function of the number of context samplesnfor a randomly drawn SCM. Fig. 2 shows the predicted posterior distributions for the interventional queryP (Y x=1 = 1)and the counterfactual queryP (Y x=1 = 0,Y x=0 = 0)as a function of the number of context samplesn. We see that the posterior distributions approximately concentrate on the identified set as the number of context samples increases. We also evaluate the performance of the CFM on a three variable system with binary variablesX,Y,Zgiven samples from the observational distributionP (x,y,z). To make comparisons with existing partial identification methods, we will restrict the prior to sampling SCMs with a fixed causal diagramZ → X → Y,X ↔ Y: the instrumental variable graph. We generatensamples from a randomly chosenP (x,y,z)and compute the coverage, average width of the derived interval, and inference time as a function of the number of context samplesn. We evaluate 95 percent credible intervals. The results are shown in Tab. 2. ModelnCov.WidthInf. time (ms) Gibbs101.0000.854757822.981473 201.0000.790407839.420002 501.0000.706258842.181667 1001.0000.660208822.117917 2001.0000.623506951.234029 4001.0000.6010551062.213033 CFM100.9000.5493252.481694 200.9500.5356002.994190 500.9500.4717824.435860 1000.9750.4719995.642881 2000.9700.4729659.045406 4000.9850.46742317.795929 Table 2. Coverage, mean interval width, and end-to-end inference time across context sizes for all evaluated models, averaged over 100 runs. 11 Foundation Models for Partial Causal Identification Figure 3. Predicted posterior distributions forΨ := P (Y x = 1)atx = 0(left) andx = 1(right) in the three-variable binary SCM experiment. Dashed vertical lines indicate the corresponding analytical lower and upper bounds (Balke & Pearl, 1994). D. Architecture and Training Details This appendix gives a complete description of the neural network architecture and training procedure used in the experiments. D.1. Architecture The modelq ω (Ψ | D)is implemented as a Neural Process composed of three components: a context/query encoder, an outcome predictor head, and a loss function. A single forward path handles both interventional and counterfactual queries via a unified query specification. Embedding.Each training episode pairs the shared contextDwith a batch ofN q queries. Queryqtargets a counterfactual probabilityΨ q = P M (y x ,...,z w )that is represented as a conjunction of potential-outcome events (sharing the same latent drawu∼ P (u)). A counterfactual probability is encoded as a 2-d array capturing the index of the node that is intervened upon, its value, and the index and value of the outcome node. For example,P (Y X=1 = 1, Y X=0 = 0)in a two variable system is encoded as: eventint. nodeint. valout. nodeout. val 00 (x)11 (y)1 10 (x)01 (y)0 For queries with less than a pre-specified maximum number of events K, the remaining events are padded. Each query is embedded into a token of shape[num nodes, model dimension]. Each node embedding varies depending on the role (intervention or outcome variable), value, and event index, that is assigned to a node in the query. Each type of embedding is learned separately and the final embedding is the sum of the contributions from all the nodes in the query. For example, for the queryP (Y X=1 = 1, Y X=0 = 0)in a two variable system, the embedding of nodeXis given by the sum of the intervention embedding for X = 1 in the first event and the outcome embedding for X = 0 in the second event. The benefit of this approach is that a query nodejattends to context nodej, giving a clear inductive bias for reading evidence about the variables involved in the query, and is efficient as complex queries can be summarized in a single token while keeping track of the events they are involved in. However, it does have the drawback of summing multiple events into the same node slots, which can blur distinctions when events overlap on the same variables. Encoder. The encoder concatenates context and query tokens along the sample axis and processes a batch throughL alternating attention blocks. Each block applies two sub-layers in sequence. Sample attention operates across tokens for each node independently, with a causal mask that enforces query tokens attend only to context tokens (and not to each other); Node attention operates across nodes within each token. Both sub-layers use pre-norm residual connections (LayerNorm → MHA → add) followed by a position-wise feed-forward network with a SwiGLU activation. The feed- forward hidden dimension is4d. Dropout is applied after each activation. AfterLsuch blocks, only the query slice of the output is retained, yielding a representation tensor of shape [B,N q , num nodes, model dimension]. D.2. Loss Two quantities are extracted from the encoder output for each query: an outcome and an intervention summary. The outcome representation is obtained by averaging the representation vectors at the outcome-node positions across the query’sKactive 12 Foundation Models for Partial Causal Identification worlds. The intervention summary is the mean intervention value across events. These are concatenated to form a predictor input which is fed to a two-layer MLP with SwiGLU activations followed by a linear projection toBlogits. TheBbins partition[0, 1]into equal-width intervals so that the predicted probability of each bin defines a histogram approximation to the predictive distribution over Ψ∈ [0, 1]. The model’s point prediction is the expectation under this histogram. Given a training example(D,Ψ q N q q=1 )where eachΨ q ∈ [0, 1]is the ground-truth probability of the corresponding query, the per-query training loss is a cross-entropy over bins: ℓ( ˆp q , Ψ q ) =− log ˆp q,b ∗ q , b ∗ q = min(⌊Ψ q · B⌋, B− 1). This corresponds to assigning the ground-truth probabilityΨ q to its enclosing bin and maximising the predicted mass on that bin. The loss is averaged over all valid queries in the batch. D.3. Training Procedure Each training step draws a fresh batch of episodes. An episode consists of a context set ofN c i.i.d. observations sampled from the observational distributionP M (V)of a randomly drawn canonical SCMM ∼ Π 0 , together withN q structured queries against that shared context, each paired with its ground-truth counterfactual probabilityΨ q (M ) = P M (y x ,...,z w ). Context sizeN c is sampled uniformly at random from[N min ,N max ]at each step, exposing the model to varying amounts of observational evidence during training. Table 3 summarises the hyperparameters used in the experiments. Table 3. Model and training hyperparameters. ParameterValueDescription d64Model (embedding) dimension L3Number of alternating attention blocks H4Number of attention heads B50Number of histogram bins MLP depth2Depth of predictor MLP F dim4d = 256Feed-forward hidden dimension Dropout0.0Dropout probability Batch size16Episodes per gradient step N q 250Queries per episode W max 2Maximum counterfactual atoms (worlds) per query N min 5Minimum context size N max 5000Maximum context size Steps3000Total training steps Learning rate10 −3 Peak learning rate (AdamW) Weight decay0AdamW weight decay Warmup fraction0.2Fraction of steps for linear warmup LR schedulecosineCosine decay after warmup Gradient clip1.0Max gradient norm 13