Lune

KDD2026Top-tier venue

StablePFN: Stable Prediction with Causal-Aware Tabular Foundation Model

Zhengkang Guan, Yikang Chen, Haoyuan Qian, Kairong Han, Peng Cui, Fei Wu, Kun Kuang

2026Year

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.

Questions to start from

Your agent calls

Lunesearch_papers

Ask in Lune

Free to start. No credit card required.

lune papers get dbbe8d16-2c3a-4877-98b6-d0325e1f5315

Related papers

Dusk over the sea between two cliffs drawn in fine vertical lines