Mitra：混合合成先验增强表格基础模型
生成多样化的合成先验分布，打造出超越任务特定基线的表格基础模型。
表格数据在医疗保健、金融、电子商务和科学等领域支撑着关键决策。然而，传统用于表格数据的机器学习方法（如随机森林和XGBoost）通常会产生针对单个数据集定制的模型，跨不同分布迁移的能力有限。
受大型语言模型成功的启发，表格基础模型有望改变这一现状：无需为每个任务单独训练模型，单个预训练模型只需通过调节适量示例即可泛化到新任务，这种技术被称为上下文学习。
作为某中心自动机器学习框架AutoGluon最新版本的一部分，我们推出Mitra——一个基于上下文学习范式训练的表格基础模型。与大型语言模型在多样化文本语料库上训练类似，Mitra通过精心设计的先验分布混合生成的合成数据集进行预训练。
初看之下，我们在预训练Mitra时未使用任何真实世界数据可能令人惊讶。但真实世界的表格数据通常有限且异构，具有不同的特征类型、依赖关系和噪声水平。事实证明，模拟覆盖广泛可能数据模式的多样化合成数据集更为实用。
我们发现这些合成先验的质量对模型泛化能力起着关键作用。有效的先验倾向于：（1）在真实任务上产生良好性能；（2）展现多样性，防止过拟合；（3）提供其他先验中未发现的独特模式。
基于这些原则，我们构建了一个混合先验，包括结构因果模型（结合变量间因果依赖关系图和描述每个变量值变化对其依赖变量影响的概率方程）以及流行的基于树的方法（如梯度提升、随机森林和决策树）。这些先验共同使Mitra能够学习鲁棒表示，并有效泛化到各种真实世界表格问题。
Mitra框架概述
我们在合成数据先验的混合上预训练表格基础模型，包括结构因果模型和基于树的模型。每个数据集被划分为支持集和查询集。Mitra支持跨行和列的二维注意力以及一维行间注意力。在推理时，模型通过上下文学习调节真实数据集的支持示例来预测查询标签，无需梯度更新。
我们在选定的先验混合上预训练Mitra。每个合成任务包含一个支持集和一个查询集。模型通过学习关注支持集来预测查询集的标签；不需要梯度更新。经过数百万个这样的任务，Mitra学会了可泛化的推理和适应模式。该架构基于跨行和特征的二维注意力，允许灵活处理不同表格大小和特征交互。
性能评估
我们在分类和回归任务上评估了Mitra，涵盖主要表格基准测试，如TabRepo、TabZilla、AMLB和TabArena。与强大的表格基础模型（如TabPFNv2和TabICL）以及数据集特定模型（如CatBoost、RealMLP和AutoGluon 1.3最佳质量预设）相比，Mitra展示了最先进的性能。
评估结果显示，每个评估指标的优胜者和亚军分别以绿色和蓝色显示。缩写+e表示上下文学习中的集成，+f表示微调。Elo的95%置信区间显示在括号中。聚合指标列显示相应指标的均值和标准差（显示在括号中）。
在二维正弦棋盘数据上，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
研究领域
机器学习
标签
表格数据
