Technical Summary: Learning Causal Orderings for In-Context Tabular Prediction

1. Problem Statement

In-context learning architectures, such as Tabular Foundation Models (TFMs), have achieved strong predictive performance on tabular data by leveraging large-scale pre-training. However, these models are primarily associative; they rely on correlational structures rather than explicit causal directionality. Consequently, their predictions can be unreliable under distribution shifts or interventions, as they may exploit non-causal statistical associations that do not generalize.

While established causal discovery methods exist to recover Directed Acyclic Graphs (DAGs), they are often dataset-specific, computationally intensive, and decoupled from modern predictive architectures. Conversely, recent amortized causal discovery approaches rely on supervised training over synthetic graphs, requiring ground-truth causal labels. There is a gap in systematically integrating causal structure discovery into predictive architectures in an unsupervised, in-context learning framework that can handle missing data and interventions.

2. Methodology: TABORDER

The authors propose TABORDER, a transformer-based architecture designed to simultaneously infer causal variable orderings and perform robust predictions. The core innovation is the use of causal order-constrained attention.

Architecture Design

TABORDER operates as a dual-transformer system that processes a dataset $D$ with missing entries indicated by a mask $M$:

1. Embedding: Each cell $x_i^{(r)}$ and its missingness indicator are embedded into a vector representation.
2. Order Inference Branch: A standard multi-head self-attention transformer processes the embeddings to infer a scalar score $s_i$ for each variable $X_i$. Lower scores indicate an earlier position in the causal order. These scores are averaged across rows to produce a global ordering $\hat{\pi}$.
3. Order-Constrained Prediction Branch: A second transformer models the conditional distribution of each cell. Crucially, its column-wise attention is constrained by a mask derived from the learned scores $s$. Specifically, variable $X_i$ can only attend to variables $X_j$ where $s_j \leq s_i$. This enforces an ordered factorization of the joint distribution:
   $$p_\theta(x^{(r)} | D, \pi) = \prod_{k=1}^d p_\theta(x_{\pi(k)}^{(r)} | x_{\pi(1:k-1)}^{(r)}, D)$$
   This prevents the model from using "anti-causal" features (future causes) to predict a target, thereby reducing sensitivity to spurious correlations.

Handling Missing Data and Uncertainty

The model explicitly accounts for missingness in causal parents. Under an Additive Noise Model (ANM) assumption ($X_i = f(X_{pa(i)}) + N_i$), the model decodes a conditional mean $\hat{\mu}$ and variance $\hat{\sigma}^2$.

• Variance Increment: If a causal parent is missing, the model increases the predictive variance for the target variable by a learnable amount $\delta_{i,j}$. This captures the increased uncertainty when relevant inputs are unobserved.
• Training Objective: The model is trained end-to-end by maximizing the log-likelihood of masked entries. The objective couples order learning and prediction: orderings that better explain the masked data (by minimizing residual variance) receive higher likelihood.

Theoretical Justification

The authors justify the likelihood-based approach under the Additive Noise Model (ANM) framework. They demonstrate that for causal additive models, the likelihood-optimal order is consistent with the topological ordering of the true underlying DAG. Furthermore, they prove a theorem regarding missingness-aware likelihood asymmetry: accounting for missingness via variance increments amplifies the likelihood gap between the correct causal direction and incorrect directions, facilitating the identification of causal orders even in the presence of missing data.

3. Key Contributions

1. Integration of Causal Discovery and Prediction: TABORDER bridges the gap between causal structure learning and in-context tabular prediction by inferring causal orderings unsupervised within a predictive transformer framework.
2. Order-Constrained Attention: The introduction of a differentiable, learned attention mask that restricts information flow to precede variables in a learned causal order, ensuring predictions rely only on potential causes.
3. Missingness-Aware Causal Learning: The design of a mechanism where missingness in causal parents explicitly increases predictive variance, which theoretically and empirically aids in distinguishing causal directions.
4. Unsupervised Learning: Unlike amortized causal discovery methods, TABORDER does not require ground-truth causal graphs during training; it learns optimal orderings solely through likelihood maximization on masked data.

4. Experimental Results

The authors evaluate TABORDER on synthetic data with known causal structures and real-world benchmarks.

• Causal Order Inference: On synthetic data generated from non-linear Gaussian process ANMs, TABORDER achieves topological divergence ($d_{TOP}$) comparable to specialized causal discovery algorithms (e.g., CAM, SCORE, NOTEARS) but infers orders in near-instant time via a single forward pass. On the real-world Sachs protein-signaling dataset, TABORDER discovers causal orders with low divergence ($d_{TOP} \approx 0.21$) compared to the consensus graph, outperforming other baselines.
• Predictive Performance & Imputation: In missing value imputation tasks (OpenML CC18 and CTR23 benchmarks), TABORDER achieves competitive performance, particularly at high missingness rates ($\geq 40\%$), outperforming standard methods like KNN, MICE, and GAIN.
• Intervention Robustness: In a three-variable chain ($X \to Y \to Z$) with an intervention on the $Y \to Z$ mechanism, TABORDER maintains accurate prediction of the mediator $Y$ given $X$ and $Z$. In contrast, standard TFMs (like TABPFN) and XGBoost degrade significantly, as they rely on the broken non-causal association between $Y$ and $Z$.
• Biological Insights: On the Sachs dataset under various interventional conditions, TABORDER infers orderings that reflect biologically plausible reorganizations. For instance, upstream interventions (e.g., PKC inhibition) cause global changes in the inferred order, while downstream interventions cause localized changes, aligning with known signaling pathway dynamics.

5. Significance and Claims

The paper claims that TABORDER provides a principled approach to making in-context learning for tabular data more robust to distribution shifts and interventions by explicitly modeling causal directionality.

• Robustness: By enforcing predictions to rely only on predecessors in a learned causal order, the model avoids exploiting spurious correlations that fail under intervention.
• Efficiency: Unlike classical causal discovery which solves optimization problems per dataset, TABORDER infers structure efficiently within a single forward pass, making it scalable for downstream tasks.
• Theoretical Grounding: The work establishes that likelihood maximization in the presence of missing data can serve as a proxy for causal discovery under standard ANM assumptions, providing a theoretical basis for unsupervised order learning.

The authors acknowledge limitations, noting that the current implementation focuses on topological orderings rather than full DAGs and assumes real-valued features. They suggest that while the method is robust, the inferred structures are hypotheses requiring domain validation, and future work could explore richer training priors and mixed-type data.