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

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

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

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

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

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

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

基于这些原则，构建了一个混合先验，包括结构因果模型（结合变量间因果依赖关系图和描述每个变量值变化对其因变量影响的概率方程）以及流行的基于树的方法（如梯度提升、随机森林和决策树）。这些先验共同使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