NVIDIA发布Kumo Relational结构化数据基础模型

NVIDIA发布Kumo Relational结构化数据基础模型。该模型无需重新训练即可执行预测,在41项任务中超越监督式机器学习模型,支持超5000亿行数据,赋能AI代理应用。

核心摘要

NVIDIA 近日发布 Kumo Relational,这是一款专为关系型数据设计的预训练基础模型。该模型无需针对特定任务进行重新训练或特征工程,即可直接对企业数据库执行预测分析。通过收集多达 10,000 个历史样本作为上下文,Kumo Relational 能在单次前向传播中输出关于客户流失、需求预测或欺诈检测等场景的结果。

支持5000亿行数据规模

作为首个在常见关系型基准测试中超越监督式机器学习模型的少样本(few-shot)基础模型,Kumo Relational 在涵盖 15 个数据库的 41 项任务中表现优于现有的监督式关系深度学习模型及专家设计的梯度提升树流水线。在 SAP SALT 企业基准测试中,其平均倒数排名(MRR)比 AutoGluon 集成模型高出 8%。该模型支持扩展至拥有超过 5000 亿行数据的数据库,既面向人类用户,也适配 AI 代理应用,可通过 NVIDIA 开源的 GPU 原生结构化数据模型库使用,或以 NVIDIA NIM 服务形式部署。

从专用模型转向结构化数据基础模型

大语言模型(LLM)已重塑了自然语言处理和计算机视觉领域的开发范式,取代了以往依赖专门代币化、词形还原和 n-gram 管道的任务专用模型。如今,这一转变正延伸至结构化数据领域。尽管关系数据库支撑着全球关键的业务决策,但传统的预测建模往往需要数周的特征工程、特定任务训练以及复杂的管道调整。

Kumo Relational 旨在解决这些痛点,其优势主要体现在三个方面:

  • 更高的精度:在 RelBench V1 分类任务中,Kumo Relational 相比基于专家设计特征的 XGBoost 模型,平均 AUROC 提升了近 4 个百分点。在 SAP SALT 基准测试中,它在工程特征上比 XGBoost 高出 7 个 MRR 点,在微调后更是领先 13 个 MRR 点。这些数据表明,用 Kumo Relational 替代生产环境中的 XGBoost 管道具有显著的性能优势。
  • 简化与灵活性:单一预训练模型取代了为每个新任务构建独立管道的繁琐流程。新的预测任务转化为查询操作,而非完整的项目开发,使得机器学习团队能够更轻松地处理传统方法难以覆盖的用例。
  • 赋能 AI 代理:如果说 LLM 赋予了代理阅读、写作和推理的能力,那么结构化数据基础模型则赋予它们预测未来的能力,例如识别哪些客户即将流失。其预测查询语言(PQL)提供了简洁的接口,用于提出具体的预测问题。

技术架构与工作原理

大多数现有的结构化数据基础模型仅能处理扁平化的单表数据,而现实世界的业务数据通常存在于复杂的关系数据库中,包含用户、订单、产品等多表关联及随时间演变的动态结构。手动聚合(如计算“过去 30 天的平均订单价值”)往往会丢失细粒度的结构信号,且耗费大量人力。

Kumo Relational 直接运行于原始数据库之上,尊重表结构、外键关系和时间维度。模型在合成关系数据上进行预训练,并通过上下文学习适应新任务。在推断阶段,它读取多达 10,000 个上下文示例,每个示例将实体周围的局部子图与其真实目标值(如“该产品已售出 142 台”)相结合。

为了完成预测,模型执行以下三个关键步骤:

  1. 自动上下文生成:任务通过声明式的预测查询语言(PQL)定义。对于涉及时间的任务,模型会重构历史数据库状态:采样截至特定时间点的以实体为中心的子图,并计算该时间点之后的事件作为目标,同时确保模型仅访问之前的数据以防止信息泄露。
  2. 表内推理:轻量级网络独立处理每个表,交替关注列间的相关值和行间的相关特征。模型在此阶段识别对当前预测有意义的列和行,并在信息传播前过滤噪声。
  3. 跨表推理:一个更大的网络通过对主键/外键关系的图表注意力机制,以及对上下文示例的交叉样本注意力机制,在关系结构中分配行嵌入。这是实现关系推理的核心环节,例如将用户个人资料与其交易历史关联起来。由于该阶段作用于行嵌入而非每个单元格,因此避免了二次方计算成本,从而支持扩展到更大的上下文和更深层次的子图。

预训练策略与数据来源

鉴于高质量、多样化的公开关系数据库稀缺,Kumo Relational 采用完全合成的关系数据进行预训练。这些数据模拟了反映现实世界结构的多种预测查询场景,从单表到复杂的关系结构逐步增加难度。

数据生成器基于 PluRel(Kothapalli 等,ICML 2026)的方法论,该方法由 Kumo 团队成员参与开发,旨在将单表数据生成器扩展至完整的关系数据库。PluRel 分三个阶段生成数据库:首先创建模式,然后建立与真实数据分布一致的外键链接,最后利用结构因果模型生成具有时间趋势和跨表依赖性的列值。研究表明,随着合成数据库规模的扩大,预训练损失遵循幂律下降,且这种收益能有效迁移至真实数据库场景。

通过这种预训练,模型学会了跨越表格进行关联和信息汇总,不再依赖固定的聚合规则。当追求最高精度时,Kumo Relational 也支持针对特定数据集和任务进行微调。

性能评估

论文显示,Kumo Relational 在 RelBench V1 和 V2 的 15 个关系数据库(涵盖电子商务、医疗、企业、社会及学术领域)的 41 个基准测试中进行了评估。值得注意的是,这些数据集均未出现在预训练过程中。在分类和回归任务中,Kumo Relational 的表现均优于包括关系深度学习模型和专家设计表格集成在内的最强监督模型。

在 RelArena-α 排行榜上,Kumo Relational 位于准确率-效率帕累托前沿。数据显示,其 Elo 分数约为 1,842,每 1,000 个样本的中位推理时间约为 17 秒。相比之下,得分最高的条目 rt-plurel 虽然达到 Elo 1,878,但耗时约 7,300 秒,速度慢了大约 430 倍。其他所有基础模型、特征工程系统和监督学习模型的得分均低于 Kumo Relational,且运行速度更慢。

在 SAP SALT 企业基准测试中,Kumo Relational 的平均相互等级(MRR)超过了大型 AutoGluon 组合约 8%,比最近的表格基础模型高出约 25%。若经过微调,其对 AutoGluon 的优势扩大到 16%。此外,该模型在极端条件下表现稳健:仅使用 0.2% 的训练数据作为上下文时仍有效;即使删除多达 75% 的外键链接,性能保持稳定;添加大量杂乱列时,性能几乎不退化。

局限性与适用场景

Kumo Relational 最适合具有明确关系结构和时间动态的数据,例如随时间推移的用户行为分析。在拥有极大训练集的任务中,特定任务的监督模型可能在某些情况下超越基于上下文的预测,目前基准测试最多使用了 10,000 个上下文示例。针对此类高精度需求场景,建议对基础模型进行微调。此外,预测信号必须能在有限数量的跳跃(hop)内获取,因此建议从几个相关表的集中查询开始,并根据需要逐步扩展图表范围。

部署与代码示例

Kumo Relational 支持内存映射 I/O,可处理超过 5000 亿行的数据,速度高达 5GB/s。开发者可通过 NVIDIA 新发布的 GPU 原生结构化数据模型库使用该模型,首次使用时权重将从 Hugging Face Hub 自动下载。

以下是使用 Python SDK 进行本地推理的代码示例:


import sdm

model = sdm.models.KumoRelational(device="cuda")
pred = model(
    x_context=...,
    y_context=...,
    x_query=...,
    related_context_tables=...,
    related_query_tables=...,
)

若需将其作为独立服务部署,可使用 NVIDIA NIM,并通过开源的结构化数据模型客户端进行调用。客户端负责收集上下文示例和相关表,然后发送请求:


import sdm_client

client = sdm_client.SDMClient(url=...)  
query = sdm_client.PredictiveQuery(
    "PREDICT COUNT(orders.*, 0, 30, days) FOR EACH users.user_id",
    {"users": ..., "orders": ...},
)
pred = client.predict(
    query.materialize(indices=[1, 2, 3]),
    model=sdm_client.KumoRelational(),
)

资源链接

评论 0

0/500

评论需审核后展示,请文明发言

💬
还没有评论,来说两句

相关阅读