NVIDIA Kumo Tabular 为表格预测树立精度与效率新前沿

亮点（TL;DR）

NVIDIA Kumo Tabular 是 NVIDIA Kumo Structured 模型系列的一部分，是一款面向表格数据的开放基础模型，现已在 Hugging Face 上提供。给定一张带有已标注行的表，它能在单次前向传播中预测新行的标签，无需训练、无需调优、也无需特征工程，分类与回归皆可。它仅在人工数据上预训练，提供三种规模（28M 至 215M 参数），通过我们的开源库运行，并依据 OpenMDW-1.1 许可证发布，可供商业使用。它在 TabArena、BeyondArena、TALENT 和 ScoringBench 四项基准上排名第一。

模型代码： https://github.com/NVIDIA/structured-data-models
模型权重： https://huggingface.co/nvidia/Kumo-Tabular

向表格基础模型的转变

表格数据是企业机器学习的支柱。客户记录、交易、传感器日志、理赔和订单都以表格形式存在，而据此预测流失、违约、需求或价格，是业界最常见的机器学习任务之一。二十年来，这项工作一直借助梯度提升树完成，而且效果良好。但围绕这些模型的生命周期几乎没有变化。每一个新问题都意味着要收集标签、进行特征工程、搜索超参数、加以验证，并部署一个对表格整体一无所知、每个任务都从头学起的模型。

大语言模型展示了应对新任务的另一种方式。只要在提示中给出若干示例，预训练模型就能解决该任务，而无需更新任何一个权重。这就是上下文学习，它适用于表格，也同样适用于文本：一个在数百万张表上预训练的模型，可以将一张带标签的表作为上下文读取，并直接预测新行的标签。

今天，我们发布 NVIDIA Kumo Tabular（GitHub、HuggingFace），一款面向表格分类与回归的开放基础模型。给定一张含有已标注行的表，以及你希望预测的那些行，Kumo Tabular 会在单次前向传播中返回类别概率或数值预测。

Kumo Tabular 的工作原理

Kumo Tabular 是一个围绕表格结构构建的 Transformer，采用了 TabICL 和 TabPFN 中提出的列注意力、行注意力和上下文注意力。要预测一个标签，它必须完成三件事：(1) 理解每个值在其所属列中的含义，(2) 理解一行中各列如何相互作用，以及 (3) 将带有已有标签的上下文行与标签未知的查询行关联起来。Kumo Tabular 的实现方式如下：

单元格嵌入：一组单元格成为一个 token。数值和类别值会经过傅里叶特征，即所学频率的正弦与余弦，且每种类型使用单独的权重。缺失值无需插补，会被特殊处理。最后，上下文中的每个 token 都会获得一个标签嵌入。

行嵌入：随后，我们通过多次交替两种注意力，将每一行转换为嵌入。列注意力沿单列向下查看，并通过诱导自注意力学习某个值在其所在列分布中的含义，例如，某个 42 是典型还是极端。因此，其开销随行数线性增长。行注意力查看单行中的各个 token，学习特征如何相互作用，并用旋转位置来区分各列。四个可学习的 [CLS] token 加入每一行，并作为该行的最终读出。经过这种行压缩后，最终阶段的开销不再取决于列数。

上下文学习：最终的 Transformer 作用于行嵌入。上下文行彼此互相关注，而查询行只关注上下文行。因此，每次预测仅取决于上下文和该行本身，而不取决于与它一同被评分的其他行。由于上下文从不查看查询，其键和值只需计算一次，并可复用于后续预测。查询行使用 Test-GQA，从而缩小每次预测所读取的缓存。一个头将每个查询行转换为分类的类别概率和回归的 999 个分位数，并由此得到点预测和不确定性估计。

长度感知注意力温度：Softmax 注意力会随着键数量的增加而扩散。在几百行上很尖锐的注意力，在数万行上可能会消散，而这正是推理时的表格远大于典型训练表格时的情况。因此，Kumo Tabular 用一个随键数量的对数增长的温度来缩放每个查询，并且每个注意力头单独学习一个系数。其结果是，当表格变长或变宽时，注意力仍能保持尖锐。

Kumo Tabular 是如何构建的

Kumo Tabular 完全在人工表格上预训练。每个训练表格都按下方所示的六个步骤，从一个结构因果模型（SCM）中采样：

我们首先为整张表抽取一个配置，涵盖其规模与任务，以及其机制与缺失情况。随后，一张随机因果图将隐变量连接起来，并在每个节点通过随机抽取的函数从根到叶进行求值（例如，线性映射、小型神经网络、树或高斯过程）。一些节点成为数值列或类别列，其中一个成为目标，其余则保持隐藏，就像真实数据背后未被测量的原因。后处理会使各组列相关、截断异常值并注入缺失值，而一次快速的树集成检查会丢弃任何没有可学习信号的表格。由于该生成器是程序化采样器而非训练好的模型，它可以源源不断地生成表格，每张表都有新的图和新的机制。

真实世界的表格杂乱无章，因此我们把更多这类不完美之处构建进了生成器。数值会以多种模式缺失，一些特征被粗化，使得重复行的标签可能不一致，一些类别列包含许多水平，回归目标也可能是重尾的。见过数百万张此类表格的模型，无需任何清理就能学会处理这些不完美之处。

在每一张人工表格上，模型都会看到大部分带有标签的行作为上下文，并学习预测其余行的标签；分类使用交叉熵损失，回归使用分位数损失。分类与回归作为各自独立的模型进行训练。与 TabICLv2 类似，训练分三个阶段进行。第一阶段也是最长的阶段使用 1,024 行、最多 100 列的表格，让模型学习表格的形态。第二阶段将上下文从 400 行变化到 10,240 行，第三阶段则扩展到 60,000 行，列数仍最多为 100。总计，Kumo Tabular-Small/Medium/Large 大约见过 35/71/137 百万张人工表格。

我们的训练方案和人工数据生成器即将发布。

性能

我们以默认设置，将全部三种规模的 Kumo Tabular 与完整的 TabArena 排行榜进行了对比，涵盖经过调优的梯度提升树、AutoGluon 以及最新的表格基础模型。Kumo Tabular 以 1950 的 ELO 位列总榜第一，并在统一的单块 RTX 6000 Pro 评测环境下，运行速度比 LimiX-2 快 17。在全部三种模型规模上，Kumo Tabular 在准确率—效率帕累托前沿上确立了新的业界最优水平：

我们还在 BeyondArena、TALENT 和 ScoringBench 上评估了 Kumo Tabular。在 BeyondArena 上，Kumo Tabular 达到 1418 的 ELO，Improvability 得分为 7.78%，位居排行榜第一。在 TALENT 上，它在分类准确率、分类对数损失和回归 RMSE 上取得总体排名第一，平均排名分别为 6.67、3.98 和 4.22。在面向预测分布的基准 ScoringBench 上，Kumo Tabular-Large 与 Medium 按平均排名分列第一和第二。

局限性

Kumo Tabular 仅适用于数值列和类别列，而文本、图像或时间戳可通过内置预处理方案转换为特征。单次前向传播最多覆盖 10 个类别，该库借助纠错输出码将其扩展到任意数量的类别。对于远超训练范围的表格，或当查询行与上下文行来自不同分布时，准确率可能会下降；因此，与任何预测模型一样，请在部署前于你自己的留出数据上验证准确率与校准。

演示

Kumo Tabular 通过 NVIDIA 新发布的、面向 structured-data-models 的 GPU 原生库运行。该库会在首次使用时从 Hub 下载权重，并提供我们评测中所使用的预处理、集成和多类别处理。下面的代码就是从 pandas.DataFrame 得到预测所需的全部内容：

import sdm  # structured-data-models

# Tensorize tabular data:
table = sdm.TableTensor.from_pandas(pd.load_csv(...), device="cuda")
na_mask = table["target"].isnan()

model = sdm.models.KumoTabular(device="cuda")
pred = model(
    # In-context examples (features/targets):
    x_context=table[~na_mask].drop_columns("target"),
    y_context=table[~na_mask, "target"],
    # Prediction examples (features):
    x_query=table[na_mask].drop_column("target"),
)

开始使用 Kumo Tabular 构建

Kumo Tabular 依据 OpenMDW 许可协议，版本 1.1 发布。NVIDIA 认为可信 AI 是一项共同责任，我们已制定政策与实践，以支持广泛 AI 应用的开发。在按照我们的服务条款下载或使用时，开发者应与为其提供支持的模型团队合作，确保该模型满足相关行业和用例的要求，并应对不可预见的产品误用。请在此处报告模型质量、风险、安全漏洞或对 NVIDIA AI 的关切。

模型代码： https://github.com/NVIDIA/structured-data-models
模型权重： https://huggingface.co/nvidia/Kumo-Tabular

致谢

我们感谢 David Holzmüller 为 Kumo Tabular 贡献了重要的想法和消融实验。我们感谢 Vignesh Kothapalli 在实习期间对 Kumo Tabular 的帮助。