Technical Summary: Use What You Know: Causal Foundation Models with Partial Graphs

Problem Statement

Causal inference traditionally relies on a two-step process: first, making identification assumptions (often via a causal graph), and second, selecting a bespoke estimator tailored to those assumptions. While recent Causal Foundation Models (CFMs) based on Prior-Data-Fitted Networks (PFNs) offer a unified approach by amortizing causal discovery and inference, they currently suffer from a critical limitation: they cannot incorporate domain knowledge at test time.

Existing CFMs either:

1. Restrict the training prior to specific causal structures (e.g., only back-door adjustment cases), requiring expensive retraining for new assumptions.
2. Marginalize over all possible structures, which leads to unnecessarily conservative estimates and imprecise predictions when partial domain knowledge is available but ignored.

The central challenge addressed is how to enable a general-purpose CFM to condition on partial causal information (such as a partially known causal graph or ancestral relationships) during inference without sacrificing the flexibility of the foundation model paradigm.

Methodology

1. Representation of Causal Knowledge: Partial Ancestral Matrices (PAMs)

The authors propose using Partially Known Ancestral Matrices (PAMs) as the conditioning signal. Unlike adjacency matrices which require knowledge of direct causal edges ($A_{ij}=1$), PAMs encode ancestral relationships ($T_{ij}=1$ if $z_i$ is an ancestor of $z_j$).

• Advantages: Ancestral knowledge is often more robust and easier to obtain than direct causal mechanisms (e.g., knowing smoking causes cancer is easier than specifying the exact biological pathway).
• Partiality: The matrix $\tilde{T} \in \{-1, 0, 1\}^{K \times K}$ allows for:
  • $1$: Known ancestor.
  • $-1$: Known non-ancestor (e.g., ruled out by temporal order).
  • $0$: Unknown relationship.
    This representation subsumes cases of no knowledge, partial knowledge, and complete knowledge.

2. Architectural Conditioning Mechanisms

The paper introduces a Transformer-based CFM architecture modified to ingest PAMs via two complementary mechanisms:

• Structural Attention Biasing:
  • Soft Bias: Learnable scalar biases ($\beta_{anc}, \beta_{non-anc}$) are added to the attention logits. Positive bias encourages attention to known ancestors; negative bias discourages attention to known non-ancestors. This allows the model to prioritize causal parents while remaining robust to incomplete information.
  • Hard Masking: An alternative where attention to known non-ancestors is strictly set to $-\infty$.
• Structural Feature Modulation (GCN):
  • A Graph Convolutional Network (GCN) encoder processes the PAM to generate structure-aware node embeddings.
  • These embeddings are injected into every Transformer block via Adaptive Layer Normalization (AdaLN), modulating the hidden representations of variables based on their global structural position in the causal graph.

3. Training and Data Generation

• Prior: The authors constructed a Tabular Foundation Model (TFM) prior that is natively causal. It generates synthetic datasets by sampling Structural Causal Models (SCMs) with diverse mechanisms (MLPs, XGBoost, various non-linearities) and noise distributions.
• Amortized Training: The model is trained to handle varying degrees of missing information. During training, entries in the PAM are randomly hidden (set to 0), forcing the model to learn a single set of weights that can adapt to full, partial, or no ancestral information.

Key Results

1. Effectiveness of Conditioning

Experiments on synthetic linear-Gaussian and complex non-linear priors demonstrate that:

• Soft Attention + GCN is the most effective conditioning strategy, consistently outperforming baselines without graph conditioning, hard masking, or GCN-only approaches.
• Conditioning on ancestral matrices performs comparably to conditioning on full adjacency matrices, validating the use of ancestral knowledge as a practical proxy.

2. Robustness to Partial Information

• Amortization Efficiency: A single model trained with amortized (randomly hidden) PAMs performs nearly as well as specialized models trained exclusively on full or no information.
• Ignorance vs. Misspecification: A critical finding is that hiding uncertain edges (marking them as "unknown") is far less harmful than misspecifying them (flipping the direction). Performance degrades gradually as edges are hidden but drops precipitously when edges are incorrectly specified. This highlights the practical value of partial matrices.

3. Semi-Synthetic Benchmarks

On the RealCause benchmark (IHDP, ACIC, CPS, PSID):

• Providing partial ancestral information encoding unconfoundedness significantly improves performance (measured by $\sqrt{PEHE}$ and relative ATE error) compared to models without graph conditioning.
• The graph-conditioned model outperforms specialized baselines like Do-PFN and CausalFM across most datasets.
• While CausalPFN achieves the best overall performance on RealCause, the authors attribute this to a prior alignment (CausalPFN's training data mimics the RealCause data-generating process) rather than graph conditioning. Once graph conditioning is introduced, their model surpasses CausalFM and Do-PFN.

4. Comparison to Specialized Models

When the back-door criterion is violated (a setting where specialized models like CausalPFN and CausalFM rely on misspecified assumptions), the graph-conditioned CFM shows a substantial advantage, demonstrating its ability to handle non-identifiable cases by leveraging available partial structure.

Significance and Claims

The paper claims to address a central hurdle in the development of "all-in-one" causal foundation models: the capability to answer causal queries in a data-driven manner while effectively leveraging any amount of domain expertise.

• Unified Framework: It demonstrates that a single CFM can flexibly condition on observational data and varying levels of causal knowledge (none, partial, or full) without retraining.
• Practical Utility: By utilizing Partial Ancestral Matrices, the approach bridges the gap between the rigidity of traditional causal estimators (which require precise graphs) and the "black box" nature of current foundation models (which ignore domain knowledge).
• Modest Scope: The authors acknowledge that while the absolute performance gains on some metrics are small (attributed to the small causal effects in their priors), the method establishes a feasible and effective foundation for future robust causal effect estimation models. They note that the current work focuses on fully observed variables and does not yet address hidden confounders, leaving that for future "all-in-one" developments.