StablePFN: Stable Prediction with Causal-Aware Tabular Foundation Model
Zhengkang Guan, Yikang Chen, Haoyuan Qian, Kairong Han, Peng Cui, Fei Wu, Kun Kuang
Abstract
Pre-trained tabular prediction models based on Prior-Data Fitted Networks (PFNs), such as TabPFN and LimiX, have achieved remarkable progress in supervised learning, demonstrating immense potential across real-world scenarios and diverse downstream tasks. However, a critical question remains systematically unexplored: Does pre-training on data generated via causal mechanisms truly endow models with the ability to comprehend underlying causal structures? Furthermore, can these models leverage such causal information to achieve stable prediction across environments? To address these fundamental questions, we propose StablePFN, a novel tabular foundation model that integrates explicit causal awareness with stable predictive modeling. Leveraging a key yet largely overlooked advantage of the PFN paradigm, the availability of ground-truth causal structure during synthetic data generation, we train StablePFN to jointly identify the Markov Boundary (MB) of the target variable and perform the primary prediction task. We introduce an end-to-end ''Decouple-Discover-Predict'' architecture that utilizes sample decoupling weights to guide MB discovery, and employs a hard attention masking mechanism during inference to incorporate causal structural knowledge. Extensive experiments on both synthetic and real-world benchmarks demonstrate that StablePFN significantly outperforms state-of-the-art baselines in cross-environment prediction settings, particularly in challenging high-bias scenarios.
Ask about this paper
Ask your agent about it.
Lune has read the top-tier papers around this one, so every answer names the papers it rests on.
Your agent calls
Lunesearch_papers
Free to start. No credit card required.
Terminal
Install the CLIlune papers get dbbe8d16-2c3a-4877-98b6-d0325e1f5315Related papers
- Do-PFN: In-Context Learning for Causal Effect EstimationJake Robertson, Arik Reuter, Siyuan Guo, Noah Hollmann et al.NeurIPS 2025 · 58 citations
- GraphPFN: A Prior-Data Fitted Graph Foundation ModelDmitry Eremeev, Oleg Platonov, Gleb Bazhenov, Artem Babenko et al.ICML 2026 · 15 citations
- Relational In-Context Learning via Synthetic Pre-training with Structural PriorYanbo Wang, Jiaxuan You, Chuan Shi, Muhan ZhangICML 2026 · 8 citations
- Foundation Models for Causal Inference via Prior-Data Fitted NetworksYuchen Ma, Dennis Frauen, Emil Javurek, Stefan FeuerriegelICLR 2026 · 37 citations
- FairPFN: A Tabular Foundation Model for Causal FairnessJake Robertson, Noah Hollmann, Samuel Müller, Noor H. Awad et al.ICML 2025
