💻 computer science

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

This paper addresses the challenge of class imbalance in Prior-data Fitted Networks (PFNs) for tabular classification by adapting classical mitigation techniques, finding that thresholding leverages PFN calibration for superior performance while downsampling offers comparable results with reduced inference costs.

Original authors: Samuel McDowell, Nathan Stromberg, Lalitha Sankar

Published 2026-05-22

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Samuel McDowell, Nathan Stromberg, Lalitha Sankar

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a super-smart, pre-trained detective named PFN. This detective has read millions of made-up stories during their training and is incredibly good at solving new cases just by looking at a few clues (a technique called "In-Context Learning"). They don't need to retrain their brain for every new job; they just look at the evidence you give them and make a guess.

However, this detective has a major blind spot: Class Imbalance.

The Problem: The "Majority Vote" Bias

Imagine you ask the detective to find a rare, specific type of coin in a pile of 1,000 coins.

950 coins are common pennies (the Majority Class).
50 coins are rare gold coins (the Minority Class).

Because the detective has seen so many pennies in their training, they get lazy. When they see a new coin, they just guess "Penny" every time.

Result: They are right 95% of the time (because most coins are pennies).
The Catch: They miss every single gold coin. If you are looking for that rare gold coin (like a rare disease or a cyberattack), the detective is useless.

The paper asks: How do we fix this detective without making them study for a new exam (retraining), which they can't do?

The Solution: Three Ways to Tweak the Detective

The researchers tried three different "tricks" to make the detective pay attention to the rare coins.

1. The "Lower the Bar" Trick (Thresholding)

Normally, the detective says, "I'm 51% sure this is a gold coin, so I'll call it gold." But because they are biased toward pennies, they rarely reach that 51% confidence for gold coins.

The researchers realized the detective is actually very honest about their confidence levels (they are "well-calibrated"). They just need a nudge.

The Fix: Instead of waiting for 51% confidence, we tell the detective: "If you are even 10% sure it's a gold coin, call it gold!"
The Analogy: It's like lowering the passing grade on a test. If the student usually needs 90% to pass, but the test is hard, you might lower the bar to 60% so they don't fail everyone.
Result: This worked amazingly well. The detective started catching the rare gold coins without losing too many pennies. It was the simplest and most effective fix.

2. The "Hide the Pennies" Trick (Downsampling)

Instead of changing the rules, we change the evidence. We take the pile of 1,000 coins and throw away 900 of the pennies, leaving only 50 pennies and 50 gold coins.

The Fix: Now the pile is balanced. The detective can't just guess "penny" and be right most of the time; they have to actually look at the clues.
The Bonus: Because there are fewer coins to look at, the detective solves the case faster (less computer power needed).
Result: This also worked very well, almost as good as the "Lower the Bar" trick.

3. The "Fake Coins" Trick (Oversampling & Synthetic Upsampling)

This is where the other methods tried to help.

Oversampling: We take the 50 gold coins and photocopy them until we have 500 gold coins.
Synthetic Upsampling: We use a robot to invent new fake gold coins that look like the real ones and add them to the pile.
The Result: These methods failed.
- Photocopying the coins (Oversampling) confused the detective because the pile looked weird and "spiky."
- The fake coins (Synthetic) weren't good enough to teach the detective anything new. In fact, they made the detective worse at spotting the real rare coins.

The Big Takeaway

The paper found that you don't need complex AI magic to fix this problem.

Don't generate fake data. It just makes things messy.
Do adjust the decision rule. Simply telling the model, "Be more willing to guess the rare thing," works better than almost anything else.
Do remove the common data. If you have too much of the common stuff, throw some away. It makes the model faster and fairer.

In short: The super-smart detective was just being too cautious about the rare items. By simply telling them to be a little less cautious (lowering the threshold) or giving them a smaller, fairer pile of evidence (downsampling), they became excellent at finding the rare things they were missing.

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

Problem Statement
Prior-Data Fitted Networks (PFNs), particularly TabPFN, have demonstrated state-of-the-art performance on tabular classification tasks by leveraging in-context learning (ICL) without updating model weights. However, like many classifiers, PFNs suffer significantly when applied to class-imbalanced data. In such scenarios, while average accuracy may remain high, performance on minority classes degrades drastically, limiting the utility of these models for critical applications like rare disease or cyberattack detection.

Standard mitigation strategies for class imbalance fall into three categories: loss-based (reweighting), data-based (resampling), and decision-based (thresholding). However, the unique ICL dynamic of PFNs renders loss-based strategies impossible, as model weights are not updated during inference. Furthermore, existing data-based approaches, such as generating synthetic minority samples, have shown limited efficacy and high computational costs in this setting. There is a lack of proven techniques to address imbalance specifically within the PFN framework.

Methodology
The authors investigate the unique calibration characteristics of PFNs to develop and evaluate correction strategies compatible with ICL. The study focuses on four specific techniques:

Thresholding (Decision-Level): Adjusting the decision boundary away from the standard 0.5 threshold. The authors hypothesize that setting the threshold $\tau$ equal to the minority class prior probability ( $\pi_1$ ) can counteract data imbalance, effectively minimizing the expected risk $R$ defined by unequal costs of false negatives and false positives.
Downsampling (Data-Level): Reducing the number of majority class samples in the context set to match the minority class size, thereby equalizing priors ( $\pi_0 = \pi_1$ ) without generating new data.
Oversampling (Data-Level): Repeating minority class samples in the context set to balance priors.
Synthetic Upsampling (Data-Level): Generating artificial minority samples using a generator model to supplement the context set.

The experiments utilize TabPFN-2.5 on 11 binary classification tasks from the OpenML-CC18 benchmark. The datasets are artificially imbalanced with minority priors ( $\pi_1$ ) ranging from 0.05 to 0.3, while the test set remains balanced. Performance is evaluated using per-class accuracy, balanced accuracy, and Worst-Class Accuracy (WCA).

Key Findings and Results
The study yields several critical observations regarding PFN behavior and the efficacy of correction methods:

Calibration Characteristics: Unlike many non-linear classifiers that are overconfident or underconfident, PFNs exhibit a specific bias: they are well-calibrated on balanced data but become biased toward the majority class when data is imbalanced. This bias is consistent and predictable, suggesting that simple thresholding can effectively correct it.
Thresholding Performance: Experimental results confirm that the optimal decision threshold for maximizing balanced accuracy aligns closely with the minority class prior ( $\tau = \pi_1$ ). This method achieves the best overall balanced accuracy, significantly improving minority class performance with only a minimal decrease in majority class accuracy.
Downsampling Efficacy: Downsampling performs comparably to thresholding in terms of balanced accuracy and achieves the highest Worst-Class Accuracy (WCA). A distinct advantage of downsampling is the reduction in inference computational cost, as TabPFN query complexity scales quadratically with context size.
Ineffectiveness of Resampling: Both oversampling and synthetic upsampling (TabPFGen) underperform compared to the base model. Oversampling leads to a "spiky" posterior fit due to repeated samples, while synthetic upsampling fails to improve accuracy because the generated samples originate from the same model class and data distribution, offering no new information.

Significance and Claims
The paper claims that the performance of state-of-the-art PFN models on imbalanced tabular data can be significantly enhanced through simple, principled corrections rather than complex architectural changes or data generation.

The primary contribution is the identification that PFNs' specific calibration bias toward the majority class allows for effective correction via thresholding (setting $\tau = \pi_1$ ) and downsampling. The authors emphasize that these methods are:

Simple to implement: Requiring no retraining or fine-tuning of the pre-trained model.
Effective: Drastically improving minority class detection while maintaining high balanced accuracy.
Computationally efficient: Particularly downsampling, which reduces inference costs.

The authors conclude that while these simple methods address the immediate problem of imbalance, they may leave room for future performance gains through methods specifically designed for the unique strengths of PFNs. They also note that extending these strategies to multi-class classification remains a non-trivial challenge for future work.