混合先验增强表格基础模型技术

Mitra：通过混合合成先验增强表格基础模型

生成多样化的合成先验分布，可产生优于特定任务基线的表格基础模型。

表格数据驱动着医疗、金融、电子商务和科学等领域的核心决策。然而，传统用于表格数据的机器学习方法（如随机森林和XGBoost）通常会产生针对单个数据集的定制模型，跨分布迁移能力有限。

受大语言模型成功的启发，表格基础模型（TFMs）有望改变这一现状：无需为每个任务单独训练模型，单个预训练模型只需通过适量示例进行条件化即可泛化到新任务，这种技术称为上下文学习（ICL）。

作为某中心自动机器学习框架AutoGluon最新版本的一部分，推出了Mitra——一个基于ICL范式训练的表格基础模型。类似于大语言模型（LLMs）在多样化文本语料库上训练的方式，Mitra通过精心设计的先验分布（priors）混合生成的合成数据集进行预训练。

初看之下，在预训练Mitra时未使用真实世界数据可能令人惊讶。但真实世界的表格数据通常有限且异构，具有不同的特征类型、依赖关系和噪声水平。实践证明，模拟覆盖广泛可能数据模式的多样化合成数据集更为实用。

研究发现这些合成先验的质量对模型泛化能力起关键作用。有效的先验往往（1）在真实任务上表现良好；（2）展现多样性，防止过拟合；（3）提供其他先验中未发现的独特模式。

基于这些原则，构建了一个混合先验，包括：结合变量间因果依赖图与描述每个变量值变化对其因变量影响的（概率）方程的结构因果模型；以及流行的基于树的方法，如梯度提升、随机森林和决策树。这些先验共同使Mitra能够学习鲁棒表示，并有效泛化到各种真实世界的表格问题。

框架概述

Mitra框架在混合合成数据先验（包括结构因果模型和基于树的模型）上预训练表格基础模型（TFMs）。每个数据集分为支持集和查询集。Mitra支持跨行和列的二维注意力以及一维行向注意力。在推理时，模型以真实数据集的支持示例为条件，使用上下文学习（ICL）预测查询标签，无需梯度更新。

在选定的先验混合上预训练Mitra。每个合成任务包含支持集和查询集。模型通过学习关注支持集来预测查询集的标签；不需要梯度更新。经过数百万个此类任务，Mitra学会了可推广的推理和适应模式。该架构基于跨行和特征的二维注意力，允许灵活处理不同表格大小和特征交互。

性能评估

在分类和回归任务上评估Mitra，覆盖主要表格基准如TabRepo、TabZilla、AMLB和TabArena。与强大的TFMs（如TabPFNv2和TabICL）以及数据集特定模型（如CatBoost、RealMLP和AutoGluon 1.3最佳质量预设）相比，Mitra展示了最先进的性能。

在二维正弦棋盘数据上，Mitra比TabPFNv2显示出更规则和更少碎片化的决策边界。

未来展望

正如基础模型重塑了计算机视觉和自然语言处理领域，Mitra为表格数据预测提供了更通用有效的方法。随着领域发展，设想更丰富的先验空间和自适应混合策略。Mitra在AutoGluon 1.4版本中开源，可供使用。邀请研究人员和实践者探索这一表格预测的新基础。

相关资源：

• Mitra分类器
• Mitra回归器

致谢：Junming Yin, Nick Erickson, Abdul Fatir Ansari, Boran Han, Shuai Zhang, Leman Akoglu, Christos Faloutsos, Michael W. Mahoney, Cuixiong Hu, Huzefa Rangwala, George Karypis, Bernie Wang