NVIDIA Kumo Tabular 发布，在表格基准测试中实现最先进的准确率-效率平衡

TL;DR

NVIDIA Kumo Tabular 是一个开源的表格数据基础模型，能够在单次前向传播中无需任何微调即可预测标签，并在 TabArena、BeyondArena、TALENT 和 ScoringBench 基准测试中排名第一，确立了新的准确率-效率前沿。

表格基础模型的转变

表格数据支撑了大多数企业级机器学习工作负载，但传统流水线仍依赖梯度提升树，需要定制特征工程、超参数搜索以及为每个新任务进行完整重训练。受大型语言模型中上下文学习的启发，Kumo Tabular 将带标签的表格视为提示，并直接为新行预测标签，无需任何权重更新。

"给定一个带标签行的表格以及你希望获得预测的行，Kumo Tabular 可在单次前向传播中返回类别概率或数值预测。" – NVIDIA Kumo Tabular 发布公告

Kumo Tabular 的工作原理

Kumo Tabular 通过三种关键机制将 Transformer 架构适配到表格的内在结构：

  1. 单元格嵌入 – 数值和分类值通过傅里叶特征（学习频率的正弦/余弦）转换。缺失值使用专用标记，每个上下文单元格与一个标签嵌入配对。
  2. 行嵌入 – 两个交替的注意力层捕捉列方向分布（列注意力）和行内特征交互（行注意力）。每行使用四个可学习的 [CLS] 标记作为最终行表示。
  3. 上下文学习 – 顶层 Transformer 处理行嵌入：上下文行相互关注，而查询行仅关注上下文。Test-GQA 缓存减少了每个查询所需的计算量。输出头为分类任务生成类别概率，为回归任务生成 999 个分位数，从而提供点预测和不确定性估计。

长度感知注意力温度

为在表格规模增大时保持注意力清晰，Kumo Tabular 使用键数量的对数缩放 softmax 温度，并采用每个头特定的可学习系数。这确保了在预训练期间未见过的、规模大几个数量级的表格仍能保持区分能力。

在人工表格上的预训练

Kumo Tabular 仅在从结构因果模型（SCMs）生成的合成表格上进行训练：

  • 随机因果图定义隐藏变量和目标关系。
  • 节点通过多种函数（线性映射、小型神经网络、决策树、高斯过程）实例化为数值或分类列。
  • 后处理添加真实世界的缺陷——缺失模式、粗粒度特征、重尾目标和高基数类别。
  • 通过快速树集成检查过滤掉缺乏可学习信号的表格。

训练分为三个阶段，上下文大小从 1,024 行（阶段 1）扩展到最多 60,000 行（阶段 3），同时保持列数 ≤ 100。三种模型尺寸——Small（28 M）、Medium（≈ 100 M）、Large（215 M）参数——分别训练了约 35 M、71 M 和 137 M 个合成表格。

基准测试表现

在统一的 RTX 6000 Pro 设置下评估，所有三个 Kumo Tabular 变体均优于现有的表格基础模型和调优后的梯度提升树，在四个主要排行榜上表现卓越：

  • TabArena – Kumo Tabular 达到最高总体 ELO（1950），比之前的最先进模型 LimiX-2 快 17 倍。
  • BeyondArena – ELO 为 1418，改进潜力得分为 7.78%，排名第一。
  • TALENT – 在分类准确率（6.67）、分类对数损失（3.98）和回归 RMSE（4.22）方面均获得平均排名首位。
  • ScoringBench – Kumo Tabular-Large 和 Medium 在平均预测分布排名上分别位列第一和第二。

结果表明，Kumo Tabular 处于新的准确率-效率帕累托前沿，如附带的基准测试图表所示。

局限性

  • 仅支持数值和分类列；其他模态（文本、图像、时间戳）必须预先处理为特征。
  • 单次前向传播最多可直接处理 10 个类别；更大的标签空间需使用库提供的纠错输出码。
  • 在超出训练范围的表格上，或当查询行来自与上下文行不同分布时，准确率可能下降。部署前必须在保留数据上进行验证。

快速演示

NVIDIA 的 structured-data-models 库提供了 GPU 原生接口。以下 Python 片段将 pandas.DataFrame 转换为 TableTensor，提供带标签的上下文行，并获取缺失标签行的预测结果：

import sdm  # structured-data-models
import pandas as pd

# 将 CSV 加载为 DataFrame 并在 GPU 上张量化
table = sdm.TableTensor.from_pandas(pd.read_csv('data.csv'), device='cuda')
na_mask = table['target'].isnan()

model = sdm.models.KumoTabular(device='cuda')
pred = model(
    x_context=table[~na_mask].drop_columns('target'),
    y_context=table[~na_mask, 'target'],
    x_query=table[na_mask].drop_column('target'),
)

该库在首次使用时会自动从 Hugging Face Hub 下载预训练权重，并处理预处理、集成和多分类扩展。

开始使用

Kumo Tabular 采用 OpenMDW-1.1 许可证发布，允许商业使用。开发者应遵循 NVIDIA 的可信 AI 政策，在特定领域数据上验证模型性能，并通过 GitHub 问题追踪器报告任何质量、风险或安全问题。

致谢

作者感谢 David Holzmüller 提供的重要想法和消融实验，以及 Vignesh Kothapalli 在实习期间的贡献。

Sources