Mitigating Label Shift in Tabular In-Context Learning via Test-Time Posterior Adjustment

TL;DR

This paper introduces DistPFN, a test-time posterior adjustment method for tabular foundation models like TabPFN, to mitigate label shift by rescaling class probabilities without retraining.

Contribution

The authors propose DistPFN and DistPFN-T, novel methods that improve classification under label shift for tabular models without architectural changes or additional training.

Findings

DistPFN significantly improves classification accuracy under label shift.

DistPFN maintains strong performance in standard settings without label shift.

Code implementation is publicly available at the provided GitHub repository.

Abstract

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.