Source-linked AI summary
When and Why LLM Causal Priors Help: Closed-Loop Prior Selection for Amortized Causal Inference
Haohao Zhou
TL;DR
Amortized causal inference depends heavily on synthetic priors that are usually designed by hand, leaving the value and selection of LLM-distilled causal graphs uncertain. The paper introduces a budget-constrained closed-loop selection framework and reports significant primary- and adjacent-domain gains, while identifying semantic content and an empirical applicability boundary. Its conclusions remain bounded because the three-condition regularity is an empirical induction rather than a theorem, and out-of-support tasks cannot be repaired by prior-level means.
Problem
The paper addresses the lack of systematic evidence about whether LLM-distilled causal priors help, why they help, and when they help in amortized causal inference.
Method
The framework optimizes over candidate priors using cheap post-training screening followed by full training and paired validation, with selection conditioned on the base and evaluation domain.
Results
The selected prior produced a formally significant 2.75× primary-domain gain, improved adjacent-domain generalization, and showed evidence that semantic graph content matters more than structural diversity alone.
Takeaways & Limitations
LLM causal-prior injection becomes an empirically verifiable selection problem rather than a manual trial-and-error decision when evaluated within the identified applicability conditions.
Takeaways & Limitations
The three-condition regularity is an empirical induction rather than a theorem, and prior-level means cannot repair tasks outside the training-prior support.
Abstract
from arXiv · showhide
Causal effect estimation asks how an outcome would change under an intervention, and medicine, economics, and public policy all treat it as a foundational task. Prior-data fitted networks (PFNs) amortize the task: a model trained on large numbers of programmatically generated synthetic causal tasks reads a new problem's observational data into context and returns an interventional-effect estimate in a single forward pass. The capability of such models is largely determined by the synthetic training prior, which is currently designed by hand, a bottleneck acknowledged by both Do-PFN and CausalPFN. Large language models (LLMs) can now ``draw'' plausible causal graphs for a given domain, suggesting that LLM-distilled graphs could serve as prior material. Whether injecting such graphs helps at all, where any gain comes from, and when injection helps. Practice has so far relied on manual trial and error. We propose a \emph{closed-loop prior selection framework} that casts prior injection as a budget-constrained optimization over a candidate prior pool. Candidates undergo cheap post-training and are scored by a composite metric dominated by real-domain generalization; the winner then receives full training and paired statistical validation. On a 7.34M-parameter Do-PFN, the framework's winner attains a formally significant $2.75\times$ gain on the primary evaluation domain, and its error falls below that of the uninjected official base. Generalization on an adjacent monitoring domain improves significantly, and no monitored capability degrades. Mechanism experiments show that the gain depends on the semantic content of the distilled graph rather than its structural diversity alone does not produce it (directional evidence). With this framework and this regularity in hand, the use of LLM causal priors stops being manual trial and error and becomes an empirically verifiable selection problem.
1 Introduction
The paper addresses hand-designed prior limitations in amortized causal inference by selecting LLM-distilled causal priors through a closed-loop, budget-constrained procedure. On a Do-PFN, the selected prior improves primary and adjacent-domain performance, while experiments identify semantic content and applicability conditions as important.
- PFNs estimate interventional effects in one forward pass, but their capability boundary is largely determined by a manually designed synthetic training prior.Do-PFN and CausalPFN identify prior design as an explicit limitation.
- LLM-distilled causal graphs offer domain-semantic prior material, but their usefulness, gain mechanism, collateral costs, and selection criteria had not been systematically validated.The paper frames these as open questions because distilled graphs are not reliable ground truth.
- The framework screens candidate priors with cheap post-training and real-domain-dominated composite scoring, then fully trains and statistically validates the strongest candidates.The base architecture and training algorithm remain unchanged; prior choice becomes an optimizable, verifiable object.
- 2.75× gain on the primary domain was formally significant (n = 5, p = 0.0086), while adjacent-domain generalization improved significantly (p = 0.0109) without observed capability degradation.The comparison to the uninjected official base is descriptive and cross-lineage rather than part of the paired test.
- Mechanism experiments support semantic content over structural diversity alone, while controlled experiments identify base underfit, prior-domain match, and task support as the induced applicability conditions.The paper also reports negative boundaries involving cross-domain transfer and out-of-support tasks.
2 Related Work
Related work spans amortized inference, LLM causal-knowledge generation, synthetic-data intervention routes, closed-loop search, and causal-estimation baselines. The paper distinguishes its contribution as optimizing the task-distribution generator itself and testing external validity against both classical estimators and standard benchmarks.
- Amortized inference replaces per-dataset optimization with inference learned over synthetic task distributions, extending from tabular prediction to causal-effect estimation.PFN-family models condition on context and answer directly after training.
- LLMs can generate domain-consistent causal graphs, but their outputs vary with prompting, decoding temperature, and graph scale, producing structurally biased knowledge.The paper uses these graphs as candidate prior material rather than treating them as ground truth.
- Unlike curriculum, augmentation, adaptation, or data-selection methods, this work optimizes which parameterizable task-distribution generator to use.Its closed loop applies cheap-proxy evaluation to prior selection rather than selecting subsets from a fixed corpus.
- The framework follows the candidate-pool, cheap-proxy, and full-validation paradigm established by neural architecture search and hyperparameter optimization.The paper instantiates this paradigm through low-fidelity pre-screening for prior selection.
- External-validity evaluation compares injected and base configurations with classical per-dataset estimators on IHDP and Lalonde.These benchmarks represent semi-synthetic individual-effect ground truth and real observational data with an experimental reference, respectively.
3 Method: Closed-Loop Prior Selection Framework
The method section formalizes prior injection and presents the framework and theoretical analysis. Its organization moves from preliminaries to budget-constrained selection, framework design, and conceptual grounding of the selection and support boundaries.
- Section 3.1 establishes preliminaries, Section 3.2 formalizes prior injection, Section 3.3 presents the framework, and Section 3.4 gives theoretical analysis.
3.1 Preliminaries: structural causal models and amortized inference
The preliminaries define structural causal models, interventions, potential outcomes, and amortized causal inference. PFNs learn to predict intervention-query outcomes from observational context by minimizing average task loss over a causal task prior.
- An SCM comprises exogenous variables, endogenous variables, structural equations, and an exogenous distribution that jointly define causal dependencies and observational behavior.Intervening replaces a structural equation with a constant and induces a post-intervention distribution.
- Potential outcomes compare the same individual under treatment and no treatment, but individual treatment effects cannot be directly observed; the paper estimates their population average, the ATE.Graph-encoded assumptions such as ignorability connect observational data to intervention effects.
- A causal task consists of an SCM, observational dataset, and intervention query whose ground truth is solved from the corresponding post-intervention model.
- PFN models condition on observational data as context and directly predict the intervention query after learning over a distribution of causal tasks.
- The training objective minimizes average task loss over tasks sampled from the prior, enabling new tasks under that prior to be answered without per-task optimization.
3.2 Problem formulation: the injection operator and amortized prior selection
The paper formalizes graph-prior injection as continued training of a fixed pretrained base on tasks generated from a candidate prior. It then selects the prior minimizing evaluation risk under a training-budget constraint.
- A distilled graph is compiled into an executable structural causal model specification that generates a task distribution for injection.
- Injection continues training pretrained parameters θ0 on tasks sampled from candidate distribution q, producing the injected model.
- Amortized prior selection chooses q∗ from candidate pool Q to minimize evaluation-domain risk after injection within budget B.
- The optimization treats training cost as the constraint and evaluation risk as the objective, with lower risk preferred.
3.3 Framework design
The framework screens candidate priors cheaply, scores them across real-domain transfer and capability-retention dimensions, then validates finalists with full training and statistical tests. Its design includes safeguards against overinterpreting unreliable proxy rankings.
- Inference-time demonstration injection is rejected because every tested configuration worsened primary-domain error relative to the uninjected official base.The official base had primary-domain NMSE 0.031, while demonstration-injection errors ranged from 0.09–0.74.
- Candidates vary across observational regimes, distillation sources, decoding temperatures, and reference graphs to test content-versus-diversity mechanisms.
- Mixed priors use regime ratios as a dosage knob because single-regime training can overfit regime-specific statistical patterns.
- Composite score: The composite score combines primary-domain generalization, diagnostic fit, diagnostic generalization, and capability retention, weighting real-domain generalization above half.
- Statistical safeguards: If reference-configuration rankings violate expectations, screening is demoted to excluding clearly poor candidates rather than claiming global discrimination.The clause was triggered in Section 5.4; comparisons use paired t-tests with significance requiring p < 0.05.
3.4 Theoretical analysis
The theoretical analysis explains why prior selection should rely on aligned real-domain evidence and identifies a non-repair boundary outside the training support. It also qualifies the empirical regularity as induction rather than theorem.
- Correlated errors impose a non-vanishing lower bound on equal-weight ensemble error, whereas uncorrelated errors yield σ2/K mean-squared error.
- Graph errors from multiple distillation sources are typically positively correlated, so ensembling cannot improve under the theorem and quality dilution may explain further degradation.
- Probe rankings can fully reverse on real evaluation when probe and evaluation task structures differ, so diagnostic proxies require alignment with the target domain.
- When evaluation tasks lie outside the training support, training-risk-based criteria cannot upper-bound evaluation risk.
- For IHDP, whose 26-dimensional inputs exceed the prior support [1] [6], domain-matched injection still degraded performance by Δ= +1.96 with p < 0.00001.The passage describes this as structural non-repair, while noting that the dimensional attribution is not separated.
- The three-condition regularity is an empirical induction from M1–M9, not a provable proposition, and the theorems provide conceptual grounding rather than a stronger claim.
4 Setup
The setup uses Do-PFN with a locally retrained base, evaluates primary-domain causal effect estimation by NMSE, and monitors adjacent-domain collateral damage. Nine controlled modules examine mechanisms, generalization, base dependence, and external validity under pre-registered paired testing.
- The base is Do-PFN with 7.34M parameters, and post-training starts from a plastic locally retrained base v2 rather than the official weights.Full post-training uses 20k steps; cheap pre-screening uses 5k steps.
- The primary domain is law_race, where the model estimates interventional effects from observational data using NMSE as the final adjudication metric.The adjacent sales domain monitors collateral damage.
- The uninjected official base reaches primary-domain NMSE 0.0308±0.0083, while probes monitor capability retention rather than proxying generalization.
- Modules M1–M5 form the mechanism chain, M6–M7 test base dependence and cross-domain generalization, and M8–M9 assess external validity.
- Pre-registered paired comparisons use p < 0.05 and |∆| ≥0.02 for M1–M8, while M9 uses |∆| ≥0.05; comparisons with n ≤3 are directional.
5 Results
Across nine controlled modules, the experiments identify when injected causal priors help, validate closed-loop selection, and expose transfer and support boundaries. The main winner improves the primary and adjacent domains, while mechanism and external-validity tests qualify the result.
- 5.1 Where the gain comes from: M1 attributes about 72% of improvement to distribution narrowing and about 28% to adding knowledge content.The random graph improves 0.112 over the general prior, while replacing it with the distilled graph adds 0.043; this comparison is directional at n = 3.
- 5.2 Observational-regime ratio: M2 finds a non-monotone primary-domain optimum at the 50/30/20 mix, while adjacent-domain sales error worsens monotonically.The 50/30/20 mix reaches primary NMSE 0.052 ± 0.016, whereas sales NMSE progresses 0.166 → 0.387 → 0.414.
- 5.4 Closed-loop selection: The closed-loop winner achieves primary-domain NMSE 0.0211 ± 0.0048, a formally significant 2.75× gain over 0.0580 ± 0.0143.Adjacent-domain sales also improves significantly to 0.356 ± 0.020 versus 0.389 ± 0.013, and no diagnostic-panel degradation is observed.
- 5.5 Mechanism: M5 provides directional evidence that semantic content, rather than structural diversity alone, underlies the gain.The random graph shares the stated order of structural diversity but is 2.1× worse than baseline; the temperature response is V-shaped and T = 0.7 degrades.
6 Discussion
The discussion frames prior injection as condition-dependent rather than universally beneficial: gains are associated with underfit bases, domain-matched priors, and tasks within prior support. It also identifies boundaries in standard benchmarks and warns that synthetic diagnostics may diverge from real-domain transfer.
- Applicability regularity: Injection yields significant gains when base underfit, prior-domain match, and task support all hold simultaneously.This is presented as an empirical regularity and working hypothesis, not an if-and-only-if theorem.
- Applicability regularity: The three-condition regularity is limited by single-instance support, non-separable violations, and no quantitative prediction of gain magnitude.The sales-domain result also shows that satisfying the conditions does not guarantee matching or exceeding a strong native base.
- Benchmark boundaries: Out-of-support dimensionality is a leading explanation for weak standard-benchmark performance, but the attribution remains unseparated.IHDP exceeds Do-PFN’s supported [1] [6] dimensions, and injection does not close the gap to CausalPFN.
- Evaluation implications: Synthetic diagnostic metrics and real-domain generalization diverge in four of five experiments unless the prior and real distribution are aligned.The discussion therefore treats synthetic-panel evidence as insufficient by itself for transfer decisions.
- Evaluation implications: Standard evaluation can overstate individual-effect recovery because naive ATE baselines rank near the top on several domains.The paper cautions that error-type metrics may partly measure heterogeneity recovery rather than average-effect estimation.
- Scope and outlook: The study’s evidence is bounded by a limited set of domains, bases, and candidate priors, with n between 3 and 5 for most experiments.The authors state that the regularity and “content beats diversity” claim are controlled evidence chains, not complete proofs.
- Scope and outlook: Future extensions include higher-dimensional base retraining, automatic prior–base matching, and targeted repair of underfit bases.These directions follow directly from the diagnosed support and matching conditions.
7 Conclusion
The paper replaces manual selection of LLM-distilled causal priors with a two-stage, empirically verifiable optimization procedure. Its winner improves the primary and adjacent domains without observed capability degradation, while the evidence supports a conditional applicability regularity with explicit limits.
- Framework: Prior injection becomes a budget-constrained optimization over candidate priors using cheap screening and full paired validation.The framework changes which prior to inject from a manual decision into an evaluable optimization object.
- Results: 2.75× gain on the primary evaluation domain is formally significant, with error below the uninjected official base.The comparison with the official base is descriptive and cross-lineage.
- Results: Adjacent-domain generalization improves significantly, and no capability degradation is observed.The adjacent-domain result has p = 0.0109; the no-degradation statement is descriptive.
- Mechanism: Mechanism experiments indicate that semantic graph content matters, whereas structural diversity alone does not account for the gain.This mechanism conclusion is reported as directional evidence.
- Applicability boundary: Injection yields significant gains under simultaneous base underfit, prior-domain match, and task support, but each condition rests on a single instance.The conditions are not fully separable, so the regularity remains a working hypothesis.
Reproducibility Statement
The reproducibility statement records the documented protocol, implementation settings, evaluation domains, statistical rules, and planned release materials. It distinguishes what is available or specified from details slated for release.
- Protocol: The base is a 7.34M-parameter Do-PFN local retrain, with frozen 20k-step post-training and 5k-step cheap screening settings.All post-training initializes from the weaker but plastic v2 base during platform diagnosis.
- Candidate pool: Candidate priors come from three distillation sources and decoding temperatures 0.3, 0.7, and 1.5.Some source identifiers and the complete AESD definition are slated for released materials.
- Evaluation: Evaluation covers the law_race and sales domains plus the public IHDP and Lalonde benchmarks.The proprietary domains use ground-truth graphs, while benchmark-specific protocols are documented for the public datasets.
- Statistical protocol: Significance requires p < 0.05 and effect-size thresholds of |∆| ≥0.02 for M1–M8 or |∆| ≥0.05 for M9.Paired tests match seeds or splits, and comparisons with n ≤3 are reported as directional.
- Availability: Code, distilled-graph specifications, candidate-pool definitions, and per-seed result JSONs are scheduled for release.Published CausalPFN Table 8 values are used only as an order-of-magnitude reference and excluded from the paper’s tables.