Correcting Class Imbalance in Prior-Data Fitted Networks for Tabular Classification

2026-05-20

AI Summary

This work addresses the performance degradation of Prior-data Fitted Networks (PFNs) on rare classes in tabular classification due to class imbalance. The authors systematically evaluate and adapt several classical class-imbalance correction techniques, focusing specifically on scenarios that do not require modifications to the loss function. Leveraging PFNs’ strong probabilistic calibration, they demonstrate that simple threshold adjustment substantially improves recognition of minority classes. Furthermore, by exploiting PFNs’ powerful few-shot in-context learning capability, down-sampling strategies effectively reduce inference computational cost while preserving overall performance. Extensive experiments across multiple tabular classification benchmarks show that the proposed approaches significantly mitigate the adverse effects of class imbalance, offering an efficient and practical solution for deploying PFNs in real-world settings with imbalanced data distributions.

Abstract

Prior-data fitted networks (PFNs) have achieved exceptional performance on tabular classification tasks. However, like other classifiers, their performance can suffer under the effect of class imbalance, resulting in poor performance for rare classes. Several techniques exist which attempt to mitigate the deleterious effect of class imbalance on classification performance, but the in-context learning (ICL) dynamic of PFNs means that loss-based strategies are impossible, and other techniques are unproven. We have adapted several classical techniques addressing class imbalance and analyzed their performance on PFN classification. We observe that thresholding performs exceptionally well because of the calibration characteristics of PFNs, and downsampling performs comparably because of PFNs exceptional limited-data performance, with the additional benefit of reduced computation cost for inference.

Problem

• class imbalance
• prior-data fitted networks
• tabular classification
• in-context learning
• rare classes

Innovation

• Prior-Data Fitted Networks
• Class Imbalance
• Thresholding
• Downsampling
• In-Context Learning