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

📄 arXiv: 2605.04363v1 📥 PDF

作者: Seunghan Lee, Jaehoon Lee, Jun Seo, Sungdong Yoo, Minjae Kim, Tae Yoon Lim, Dongwan Kang, Hwanil Choi, SoonYoung Lee, Wonbin Ahn

分类: cs.LG, cs.AI

发布日期: 2026-05-06

备注: ICML 2026

🔗 代码/项目: GITHUB

💡 一句话要点

提出DistPFN，通过测试时后验调整缓解表格数据上下文学习中的标签偏移问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 表格数据 上下文学习 标签偏移 后验调整 测试时调整

📋 核心要点

  1. TabPFN在表格数据上表现出色，但易受标签偏移影响，导致模型倾向于训练集中的多数类。
  2. DistPFN通过测试时后验调整，降低训练先验的影响，增强模型预测后验的贡献，从而缓解标签偏移。
  3. 实验表明，DistPFN在多种TabPFN模型上显著提升了分类性能，同时保持了在标准设置下的竞争力。

📝 摘要（中文）

TabPFN作为表格数据集的基础模型，通过上下文学习在合成数据上取得了优异的性能。然而，我们发现TabPFN容易受到标签偏移的影响，经常过度拟合训练数据集中的多数类。为了解决这个局限性，我们提出了DistPFN，这是第一个为表格基础模型设计的测试时后验调整方法。DistPFN通过降低训练先验（即上下文的类别分布）的影响并强调模型预测后验的贡献来重新调整预测的类别概率，无需架构修改或额外训练。我们进一步引入了DistPFN-T，它结合了温度缩放，以基于先验和后验之间的差异自适应地控制调整强度。我们在超过250个OpenML数据集上评估了我们的方法，证明了在标签偏移下，各种基于TabPFN的模型的分类任务性能得到了显著提高，同时在没有标签偏移的标准设置中保持了强大的性能。代码可在https://github.com/seunghan96/DistPFN 获取。

🔬 方法详解

问题定义：TabPFN等模型在表格数据上利用上下文学习，但当训练数据存在标签偏移（即类别分布与测试数据不同）时，模型容易过拟合训练数据的类别分布，导致泛化能力下降。现有方法缺乏有效的测试时调整策略，难以适应标签偏移带来的挑战。

核心思路：DistPFN的核心思路是在测试时调整模型的预测后验概率，通过降低训练数据类别分布（先验）的影响，并增强模型自身预测的后验概率的贡献，从而使模型更加关注测试样本的特征，减少对训练数据偏差的依赖。这样可以在不修改模型结构或重新训练的情况下，提升模型在标签偏移环境下的鲁棒性。

技术框架：DistPFN的整体框架包括以下步骤：1) 使用TabPFN等模型进行初始预测，得到类别概率分布；2) 计算训练数据的类别先验分布；3) 使用DistPFN或DistPFN-T方法，基于先验和后验概率，调整预测的类别概率；4) 输出调整后的类别概率作为最终预测结果。DistPFN-T在DistPFN的基础上，引入了温度缩放机制，自适应地调整先验和后验的权重。

关键创新：DistPFN的关键创新在于提出了一个简单有效的测试时后验调整方法，专门用于缓解表格数据上下文学习中的标签偏移问题。与现有方法相比，DistPFN无需额外的训练或模型修改，即可显著提升模型在标签偏移环境下的性能。DistPFN-T进一步引入了温度缩放，使得调整强度可以根据先验和后验的差异自适应地调整。

关键设计：DistPFN的核心公式是调整后的概率 = (后验概率)^α * (先验概率)^(1-α)，其中α是一个超参数，控制后验和先验的权重。DistPFN-T通过温度缩放自适应地调整α，α = sigmoid(T * (KL散度(后验 || 先验) - 阈值))，其中T是温度参数，阈值是一个超参数。KL散度用于衡量后验和先验的差异，温度参数控制sigmoid函数的陡峭程度，阈值用于控制调整的敏感度。这些参数需要根据具体数据集进行调整。

📊 实验亮点

在超过250个OpenML数据集上的实验结果表明，DistPFN和DistPFN-T能够显著提升TabPFN等模型在标签偏移环境下的分类性能。例如，在某些数据集上，DistPFN可以将模型的准确率提升超过10%。同时，DistPFN在没有标签偏移的标准设置中也能保持良好的性能，证明了其通用性和实用性。

🎯 应用场景

DistPFN可应用于各种表格数据分类任务，尤其是在训练数据存在标签偏移的情况下。例如，在医疗诊断、金融风控等领域，由于数据收集的偏差，训练数据可能无法真实反映实际情况，导致模型性能下降。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.