StablePFN: Stable Prediction with Causal-Aware Tabular Foundation Model

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.