NVIDIA Kumo Tabular公開、表形式データ予測で精度と効率の新基準

記事のサマリー（TL;DR）

NVIDIAは、表形式データ向けのオープン基盤モデル「NVIDIA Kumo Tabular」をHugging Faceで公開しました。NVIDIA Kumo Structuredモデルコレクションの一部です。
ラベル付きの行からなるテーブルを与えると、学習・チューニング・特徴量エンジニアリングを行わずに、新しい行のラベルを1回の順伝播（single forward pass）で予測します。分類・回帰の両方に対応します。
人工データのみで事前学習され、パラメータ数28Mから215Mまでの3サイズで提供され、商用利用が可能なOpenMDW-1.1ライセンスのもとで公開されています。TabArena、BeyondArena、TALENT、ScoringBenchの4つのベンチマークで首位にランクしたとされています。

モデルコード: https://github.com/NVIDIA/structured-data-models
モデル重み: https://huggingface.co/nvidia/Kumo-Tabular

詳細

表形式基盤モデルへの移行

表形式データは企業の機械学習の基盤であり、顧客記録、取引、センサーログ、保険請求、注文などはすべてテーブルの形で保存されています。これらのデータから解約、デフォルト、需要、価格を予測することは、業界で最も一般的な機械学習タスクの一つです。

この20年間、こうした作業は勾配ブースティング木（gradient-boosted trees）によって行われてきており、うまく機能してきました。しかし、こうしたモデルを取り巻くライフサイクルはほとんど変わっていません。新しい課題が出るたびに、ラベルの収集、特徴量エンジニアリング、ハイパーパラメータの探索、検証、そして表そのものについて何も知らず各タスクをゼロから学習するモデルのデプロイが必要になります。

大規模言語モデル（LLM）は、新しいタスクへの別の取り組み方を示しました。プロンプト内にいくつかの例を与えるだけで、事前学習済みモデルは重みを一切更新せずにタスクを解けます。これが文脈内学習（in-context learning）であり、テキストと同様に表形式データにも適用できます。数百万のテーブルで事前学習されたモデルは、ラベル付きテーブルを文脈として読み取り、新しい行のラベルを直接予測できます。

NVIDIAは今回、表形式データの分類・回帰向けオープン基盤モデルである「NVIDIA Kumo Tabular」（GitHub、HuggingFace）を公開しました。ラベル付き行と予測対象の行を含むテーブルを与えると、Kumo Tabularは1回の順伝播でクラス確率または数値予測を返します。

Kumo Tabularの仕組み

Kumo Tabularは、TabICLおよびTabPFNで導入された列（column）、行（row）、文脈内（in-context）の各アテンションを活用し、テーブルの構造を中心に構築されたTransformerです。ラベルを予測するために、モデルは次の3つを行う必要があります。(1) 各値がその列の中で何を意味するかを理解すること、(2) 行内の列同士がどう相互作用するかを理解すること、(3) 既知のラベルを持つ文脈行と、ラベル不明のクエリ行とを関連付けることです。

Kumo Tabularは以下の方法でこれを実現します。

セル埋め込み（Cell Embedding）: セルのグループが1つのトークンになります。数値・カテゴリ値はフーリエ特徴（学習された周波数のサインとコサイン）を通過し、型ごとに別々の重みが用いられます。欠損値は補完を必要とせず、特別に扱われます。最終的に、文脈内のすべてのトークンにラベル埋め込みが付与されます。

行埋め込み（Row Embedding）: 2種類のアテンションを複数回交互に適用し、各行を埋め込みに変換します。列アテンションは単一の列を縦方向に見て、ある値がその列の分布の中で典型的か極端かなど、誘導自己アテンション（induced self-attention）を通じて学習します。そのコストは行数に対して線形に増加します。行アテンションは単一行内のトークンを横断的に見て、特徴量同士の相互作用を学習し、列を区別するために回転位置エンコーディング（rotary positions）を用います。学習可能な4つの[CLS]トークンが各行に加わり、行の最終的な読み出し（readout）として機能します。この行の圧縮後は、最終段階のコストは列数に依存しなくなります。

文脈内学習（In-context Learning）: 最後のTransformerが行埋め込みに対して動作します。文脈行は互いにアテンションを行い、クエリ行は文脈行のみにアテンションを行います。したがって各予測は、文脈とその行自体にのみ依存し、同時にスコアリングされる他のどの行にも依存しません。文脈はクエリを一切参照しないため、そのキーと値は一度だけ計算され、以降の予測に再利用できます。クエリ行はTest-GQAを利用しており、これにより各予測が読み取るキャッシュを縮小します。ヘッドは各クエリ行を、分類の場合はクラス確率に、回帰の場合は999個の分位点に変換し、そこから点予測と不確実性の推定値が導かれます。

長さ対応アテンション温度（Length-aware Attention Temperature）: ソフトマックスアテンションは、キーの数が増えるにつれて拡散します。数百行では鋭いアテンションも、数万行になると崩れることがあり、これはまさに推論時のテーブルが学習時の一般的なテーブルよりもはるかに大きい場合に起こる状況です。そのためKumo Tabularは、キーの数の対数に応じて増加する温度で各クエリをスケーリングし、その係数はアテンションヘッドごとに個別に学習されます。この結果、テーブルが長くなったり広くなったりしても、アテンションは鋭いまま維持されます。

Kumo Tabularの構築方法

Kumo Tabularは、人工的なテーブルのみで事前学習されています。各学習用テーブルは、以下の6つのステップで示される構造的因果モデル（Structural Causal Model、SCM）からサンプリングされます。

まず、テーブル全体の構成（サイズ、タスクから、生成メカニズム、欠損パターンまで）を抽出します。次にランダムな因果グラフが隠れ変数同士を結びつけ、各ノードでランダムに選ばれた関数（線形写像、小規模なニューラルネットワーク、木、ガウス過程など）を用いて根から葉まで評価されます。一部のノードは数値列またはカテゴリ列になり、1つが目的変数となり、残りは実データの背後にある未計測の原因のように隠れたままになります。後処理では列のグループ間の相関を作り、外れ値をクリップし、欠損値を注入し、学習可能な信号を持たないテーブルは簡易的な木アンサンブルチェックによって破棄されます。生成器は学習済みモデルではなく手続き型のサンプラーであるため、新しいグラフと新しいメカニズムを持つテーブルを無限に生成できます。

実世界のテーブルは乱雑であるため、生成器にはそうした不完全さをより多く組み込んでいます。値はいくつかのパターンで欠損し、一部の特徴は粗くなって重複行がラベルで食い違うことがあり、一部のカテゴリ列は多くの水準を持ち、回帰対象はヘビーテイル（裾の重い分布）になることがあります。こうした不完全さを持つテーブルを数百万件見たモデルは、事前のクリーンアップなしにこれらを処理できるようになります。

それぞれの人工テーブルにおいて、モデルはほとんどの行をラベル付きの文脈として見て、残りの行のラベルを予測することを学習します。分類には交差エントロピー損失、回帰には分位点損失が用いられます。分類と回帰は別々のモデルとして学習されます。

TabICLv2と同様に、学習は3つの段階で行われます。最初で最長の段階では、1,024行・最大100列のテーブルを使用し、モデルにテーブルがどのようなものかを教えます。第2段階では文脈を400行から10,240行まで変化させ、第3段階ではこれを最大60,000行まで拡張します（列数は引き続き最大100列）。合計で、Kumo Tabular-Small/Medium/Largeはそれぞれ約3,500万/7,100万/1億3,700万件の人工テーブルを見たとされています。学習レシピと人工データ生成器は近日中に公開される予定です。

性能

NVIDIAは、Kumo Tabularの3つのサイズすべてをデフォルト設定で、チューニング済み勾配ブースティング木、AutoGluon、最新の表形式基盤モデルを含むTabArenaのリーダーボード全体に対して実行しました。Kumo Tabularは、単一のRTX 6000 Proを用いた統一的な評価環境下で、LimiX-2より17倍高速に動作しながら、ELO 1950で総合首位にランクしたとしています。3つのモデルサイズすべてにわたって、Kumo Tabularは精度と効率のパレートフロンティア（Pareto front）において新しい最先端水準を確立したとされています。

また、BeyondArena、TALENT、ScoringBenchでもKumo Tabularを評価しました。BeyondArenaでは、Kumo TabularはELO 1418、Improvabilityスコア7.78%を達成し、リーダーボードで首位となりました。TALENTでは、分類精度、分類の対数損失（log-loss）、回帰のRMSEの全体で最上位となり、それぞれの平均順位は6.67、3.98、4.22でした。予測分布のためのベンチマークであるScoringBenchでは、Kumo Tabular-LargeとMediumが平均順位で1位と2位にランクしたとされています。

制限事項

Kumo Tabularは数値列とカテゴリ列にのみ対応しており、テキスト、画像、タイムスタンプは組み込みの前処理レシピを通じて特徴量に変換する必要があります。1回の順伝播では最大10クラスまで対応し、ライブラリは誤り訂正出力符号（error-correcting output codes）によってこれを任意のクラス数に拡張します。学習範囲を大きく超えるテーブルや、クエリ行が文脈行と異なる分布から来ている場合には精度が低下する可能性があるため、他の予測モデルと同様に、本番導入前に自身の保留データ（held-out data）で精度とキャリブレーションを検証する必要があるとされています。

デモ

Kumo Tabularは、NVIDIAが新たに公開した構造化データモデル向けのGPUネイティブなライブラリ（structured-data-models）を通じて動作します。このライブラリは初回使用時にHubから重みをダウンロードし、評価で使用した前処理、アンサンブル、多クラス処理を提供します。以下のコードだけで、pandas.DataFrameから予測結果までを得られます。

import sdm  # structured-data-models

# 表形式データのテンソル化:
table = sdm.TableTensor.from_pandas(pd.load_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"),
)

Kumo Tabularを使い始める

Kumo Tabularは、OpenMDW License Agreement バージョン1.1のもとで公開されています。NVIDIAは、信頼できるAI（Trustworthy AI）は共有された責任であると考えており、幅広いAIアプリケーションの開発を可能にするためのポリシーと実践を定めているとしています。利用規約に従ってダウンロードまたは使用する場合、開発者はサポートするモデルチームと協力し、このモデルが関連する業界およびユースケースの要件を満たし、予期しない製品の悪用に対応していることを確認する必要があるとされています。モデルの品質、リスク、セキュリティ脆弱性、またはNVIDIA AIに関する懸念は、原文記載の窓口に報告するよう案内されています。