# Learning Causal Orderings for In-Context Tabular Prediction

arXiv: https://arxiv.org/abs/2605.22335 · HTML: https://mlsanity.com/paper/2605.22335/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 43 references that are in the map.
- First posted on arXiv 21 May 2026.

## Builds on (most cited first, 25 of 25)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [DAGs with NO TEARS: Continuous Optimization for Structure Learning](https://mlsanity.com/paper/1803.01422/) (2018)
- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [GAIN: Missing Data Imputation using Generative Adversarial Nets](https://mlsanity.com/paper/1806.02920/) (2018)
- [TabICL: A Tabular Foundation Model for In-Context Learning on Large Data](https://mlsanity.com/paper/2502.05564/) (2025)
- [Gradient-Based Neural DAG Learning](https://mlsanity.com/paper/1906.02226/) (2019)
- [CAM: Causal additive models, high-dimensional order search and penalized regression](https://mlsanity.com/paper/1310.1533/) (2013)
- [OpenML Benchmarking Suites](https://mlsanity.com/paper/1708.03731/) (2017)
- [TabDPT: Scaling Tabular Foundation Models on Real Data](https://mlsanity.com/paper/2410.18164/) (2024)
- [Score matching enables causal discovery of nonlinear additive noise models](https://mlsanity.com/paper/2203.04413/) (2022)
- [Matrix Completion and Low-Rank SVD via Fast Alternating Least Squares](https://mlsanity.com/paper/1410.2596/) (2014)
- [Amortized Inference for Causal Structure Learning](https://mlsanity.com/paper/2205.12934/) (2022)
- [HyperImpute: Generalized Iterative Imputation with Automatic Model Selection](https://mlsanity.com/paper/2206.07769/) (2022)
- [Learning to Induce Causal Structure](https://mlsanity.com/paper/2204.04875/) (2022)
- [MIRACLE: Causally-Aware Imputation via Learning Missing Data Mechanisms](https://mlsanity.com/paper/2111.03187/) (2021)
- [Do-PFN: In-Context Learning for Causal Effect Estimation](https://mlsanity.com/paper/2506.06039/) (2025)
- [CausalPFN: Amortized Causal Effect Estimation via In-Context Learning](https://mlsanity.com/paper/2506.07918/) (2025)
- [Scalable Causal Discovery with Score Matching](https://mlsanity.com/paper/2304.03382/) (2023)
- [Causal Discovery with Score Matching on Additive Models with Arbitrary Noise](https://mlsanity.com/paper/2304.03265/) (2023)
- [A Scale-Invariant Sorting Criterion to Find a Causal Order in Additive Noise Models](https://mlsanity.com/paper/2303.18211/) (2023)
- [Demystifying amortized causal discovery with transformers](https://mlsanity.com/paper/2405.16924/) (2024)
- [Not Another Imputation Method: A Transformer-based Model for Missing Values in Tabular Datasets](https://mlsanity.com/paper/2407.11540/) (2024)
- [Algorithmic causal structure emerging through compression](https://mlsanity.com/paper/2502.04210/) (2025)

## Built on by (most cited first, 1 of 1)

- [DAG-FM: A Foundation Model for Causal Discovery under Heterogeneous Causal Mechanisms](https://mlsanity.com/paper/2607.11510/) (2026)

## Abstract (arXiv)

In-context learning for tabular data sets strong predictive standards in observational settings; it however primarily relies on correlational structure, which becomes unreliable under distribution shift or intervention. While established methods to discover causal structure exist, they are often focused on structure identifiability and decoupled from the predictive architectures that could benefit from them. To bridge these perspectives, we study how to simultaneously infer and enforce causal structure in the form of topological variable orderings into tabular prediction. Unlike standard architectures, our model TabOrder uses causal order-constrained attention, basing predictions only on features that precede a target under a learned causal order. Similar to causal discovery methods, TabOrder learns the optimal variable ordering in an unsupervised manner through a likelihood-based objective. We justify this choice under standard functional model classes and also study how sample missingness, a common challenge in tabular data, interacts with causal direction identification. Empirically, we confirm that TabOrder recovers accurate variable orderings while addressing prediction and imputation tasks, as well as gives insight into real-world biological data under intervention.