테이블 형식 데이터의 컨텍스트 학습에서 테스트 시간 사후 조정(Posterior Adjustment)을 통한 레이블 변화(Label Shift) 완화

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

TabPFN은 최근 테이블 형식 데이터셋을 위한 기초 모델로 주목받으며, 합성 데이터를 활용한 컨텍스트 학습을 통해 뛰어난 성능을 보여줍니다. 하지만, 본 연구에서는 TabPFN이 레이블 변화에 취약하며, 종종 훈련 데이터셋에서 다수 클래스에 과적합되는 현상이 발생한다는 것을 확인했습니다. 이러한 한계를 해결하기 위해, 본 연구에서는 테이블 형식 기초 모델을 위한 최초의 테스트 시간 사후 조정 방법인 DistPFN을 제안합니다. DistPFN은 모델의 예측 사후 분포에 대한 기여도를 강조하고 훈련 사전(context의 클래스 분포)의 영향을 감소시키는 방식으로, 아키텍처 수정이나 추가 훈련 없이 예측 클래스 확률을 재조정합니다. 또한, 사전 분포와 사후 분포 간의 차이에 따라 조정 강도를 적응적으로 제어하기 위해 온도 스케일링을 통합한 DistPFN-T를 소개합니다. 저희는 250개 이상의 OpenML 데이터셋에서 저희 방법들을 평가했으며, 레이블 변화가 있는 환경에서 다양한 TabPFN 기반 모델의 분류 성능을 크게 향상시키는 것을 확인했습니다. 동시에, 레이블 변화가 없는 표준 환경에서는 강력한 성능을 유지합니다. 관련 코드는 다음 저장소에서 확인할 수 있습니다: https://github.com/seunghan96/DistPFN.

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.