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

arXiv: https://arxiv.org/abs/2605.04363 · HTML: https://mlsanity.com/paper/2605.04363/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 44 references that are in the map.
- First posted on arXiv 6 May 2026.

## Builds on (most cited first, 27 of 27)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [Revisiting Deep Learning Models for Tabular Data](https://mlsanity.com/paper/2106.11959/) (2021)
- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [TabNet: Attentive Interpretable Tabular Learning](https://mlsanity.com/paper/1908.07442/) (2019)
- [Long-tail learning via logit adjustment](https://mlsanity.com/paper/2007.07314/) (2020)
- [Tabular Data: Deep Learning is Not All You Need](https://mlsanity.com/paper/2106.03253/) (2021)
- [TabTransformer: Tabular Data Modeling Using Contextual Embeddings](https://mlsanity.com/paper/2012.06678/) (2020)
- [Balanced Meta-Softmax for Long-Tailed Visual Recognition](https://mlsanity.com/paper/2007.10740/) (2020)
- [Detecting and Correcting for Label Shift with Black Box Predictors](https://mlsanity.com/paper/1802.03916/) (2018)
- [SAINT: Improved Neural Networks for Tabular Data via Row Attention and Contrastive Pre-Training](https://mlsanity.com/paper/2106.01342/) (2021)
- [TabICL: A Tabular Foundation Model for In-Context Learning on Large Data](https://mlsanity.com/paper/2502.05564/) (2025)
- [When Do Neural Nets Outperform Boosted Trees on Tabular Data?](https://mlsanity.com/paper/2305.02997/) (2023)
- [OpenML Benchmarking Suites](https://mlsanity.com/paper/1708.03731/) (2017)
- [Well-tuned Simple Nets Excel on Tabular Datasets](https://mlsanity.com/paper/2106.11189/) (2021)
- [Regularized Learning for Domain Adaptation under Label Shifts](https://mlsanity.com/paper/1903.09734/) (2019)
- [TabM: Advancing Tabular Deep Learning with Parameter-Efficient Ensembling](https://mlsanity.com/paper/2410.24210/) (2024)
- [Better by Default: Strong Pre-Tuned MLPs and Boosted Trees on Tabular Data](https://mlsanity.com/paper/2407.04491/) (2024)
- [Benchmarking Distribution Shift in Tabular Data with TableShift](https://mlsanity.com/paper/2312.07577/) (2023)
- [Retrieval & Fine-Tuning for In-Context Tabular Models](https://mlsanity.com/paper/2406.05207/) (2024)
- [TuneTables: Context Optimization for Scalable Prior-Data Fitted Networks](https://mlsanity.com/paper/2402.11137/) (2024)
- [Mixture of In-Context Prompters for Tabular PFNs](https://mlsanity.com/paper/2405.16156/) (2024)
- [Drift-Resilient TabPFN: In-Context Learning Temporal Distribution Shifts on Tabular Data](https://mlsanity.com/paper/2411.10634/) (2024)
- [TabFlex: Scaling Tabular Learning to Millions with Linear Attention](https://mlsanity.com/paper/2506.05584/) (2025)
- [MambaTab: A Plug-and-Play Model for Learning Tabular Data](https://mlsanity.com/paper/2401.08867/) (2024)
- [On the Efficiency of NLP-Inspired Methods for Tabular Deep Learning](https://mlsanity.com/paper/2411.17207/) (2024)
- [AdapTable: Test-Time Adaptation for Tabular Data via Shift-Aware Uncertainty Calibrator and Label Distribution Handler](https://mlsanity.com/paper/2407.10784/) (2024)

## Built on by (most cited first, 4 of 4)

- [Fusion Anything: A Generalized Multimodal Foundation Model](https://mlsanity.com/paper/2609.22107/) (2026)
- [The Geometry of Logic: Stratification Induces Semantic Structure and Robust Reasoning](https://mlsanity.com/paper/2609.32927/) (2026)
- [Parameter-Efficient Distributionally Robust Adaptation of Tabular Foundation Models under Subpopulation Shift](https://mlsanity.com/paper/2610.01143/) (2026)
- [OneShot: Label-Efficient Detector Weighting for Time Series Anomaly Detection](https://mlsanity.com/paper/2610.16219/) (2026)

## Abstract (arXiv)

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.