Mitigating Label Shift in Tabular In-Context Learning via Test-Time Posterior Adjustment
For practitioners using tabular foundation models, this work addresses a critical failure mode (label shift) with a simple, training-free adjustment, improving robustness in real-world classification tasks.
TabPFN, a foundation model for tabular data, is vulnerable to label shift, overfitting to majority classes. The authors propose DistPFN, a test-time posterior adjustment method that rescales predicted probabilities to mitigate this, achieving substantial improvements on over 250 OpenML datasets under label shift while maintaining performance in standard settings.
TabPFN has recently gained attention as a foundation model for tabular datasets, achieving strong performance by leveraging in-context learning on synthetic data. However, we find that TabPFN is vulnerable to label shift, often overfitting to the majority class in the training dataset. To address this limitation, we propose DistPFN, the first test-time posterior adjustment method designed for tabular foundation models. DistPFN rescales predicted class probabilities by downweighting the influence of the training prior (i.e., the class distribution of the context) and emphasizing the contribution of the model's predicted posterior, without architectural modification or additional training. We further introduce DistPFN-T, which incorporates temperature scaling to adaptively control the adjustment strength based on the discrepancy between prior and posterior. We evaluate our methods on over 250 OpenML datasets, demonstrating substantial improvements for various TabPFN-based models in classification tasks under label shift, while maintaining strong performance in standard settings without label shift. Code is available at this repository: https://github.com/seunghan96/DistPFN.