Как делать табличные предсказания без обучения с NVIDIA Structured Data Models

В этом руководстве мы реализуем Structured Data Models — новую библиотеку NVIDIA, которая объединяет несколько предварительно обученных фундаментальных моделей для таблиц под единым интерфейсом PyTorch, и используем её для прогнозирования без обучения. Библиотека включает TabICLv2 и KumoTabular для отдельных таблиц, а также KumoRelational для целых реляционных баз данных, причём каждая модель отвечает с помощью обучения в контексте: мы передаём модели набор размеченных строк и строки, для которых нужно сделать прогноз, а она возвращает вероятности или полное прогнозное распределение за один прямой проход. Сначала мы рассмотрим типизированные табличные тензоры библиотеки, сравним две табличные модели с обученными базовыми моделями scikit-learn, повторно используем закодированный контекст для множества запросов, получим калиброванные интервалы прогнозирования для регрессии и масштабируемся до двадцати классов. Затем мы изучим случаи, когда настройки по умолчанию незаметно приводят к неверному результату: рецепт предварительной обработки, удаляющий столбцы с временными метками и текстом, а также объяснитель с двумя скрытыми проблемами. После этого создадим небольшую базу данных магазина из трёх таблиц и напрямую спрогнозируем активность клиентов, включая единственный аргумент, который отличает честную реляционную модель от модели с утечкой данных. Ознакомьтесь с полным кодом здесь.

import os
import sys
import time
import textwrap
import traceback
import subprocess
 
RESULTS = {}
 
 
def banner(title):
    print("\n" + "=" * 78)
    print(title)
    print("=" * 78)
 
 
def section(name):
    def wrap(fn):
        def run(*a, **kw):
            banner(name)
            try:
                out = fn(*a, **kw)
                RESULTS[name] = out if isinstance(out, str) else "ok"
                return out
            except Exception as e:
                RESULTS[name] = f"SKIPPED / FAILED -> {type(e).__name__}: {e}"
                print(f"\n[!] {name} did not complete: {type(e).__name__}: {e}")
                traceback.print_exc(limit=3)
                return None
        return run
    return wrap
 
 
banner("1. Install SDM, and the one wheel URL its own error message gets wrong")
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "git+https://github.com/NVIDIA/structured-data-models.git@b982fbf4188ea8db17249fe81338723106100b63"], check=True)
 
import torch
 
# The relational sampler needs pyg-lib on CPU. Its wheels are published per MINOR torch version
# (torch-2.14.0+cpu.html), but SDM's own ImportError builds the URL from the exact patch version
# (torch-2.14.1+cpu.html), which is an empty page. Build it from MAJOR.MINOR.0 instead.
major, minor = torch.__version__.split("+")[0].split(".")[:2]
flavour = "cpu" if torch.version.cuda is None else "cu" + torch.version.cuda.replace(".", "")
PYG_INDEX = f"https://data.pyg.org/whl/torch-{major}.{minor}.0+{flavour}.html"
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "pyg-lib", "-f", PYG_INDEX], check=True)
 
import numpy as np
import pandas as pd
import pyg_lib
import sdm
import sdm.processing as sp
from sklearn.datasets import load_breast_cancer, load_diabetes, make_classification
from sklearn.ensemble import HistGradientBoostingClassifier, HistGradientBoostingRegressor
from sklearn.linear_model import LogisticRegression, RidgeCV
from sklearn.metrics import roc_auc_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
 
DEVICE = "cpu"   # verified end to end on CPU; on a T4 runtime set "cuda" for the tabular models
print(f"  sdm {sdm.__version__}  |  torch {torch.__version__}  |  pyg-lib {pyg_lib.__version__}"
      f"  |  device {DEVICE}")
print(f"  pyg-lib wheels from {PYG_INDEX}")
print("\n  One interface over several pretrained structured-data foundation models:")
for name, what, weights in [
    ("TabICLv2", "tabular, in-context learning (ICML '26)", "~110 MB per task head, BSD-3"),
    ("KumoTabular", "tabular, small / medium / large", "small ~110 MB, large ~855 MB, OpenMDW"),
    ("KumoRelational", "multi-table, KumoRFM-2", "~120 MB, OpenMDW"),
    ("TabFM, TimesFM 3", "Google tabular / time series", "6.5 GB / 1.3 GB, non-commercial: not used here"),
]:
    print(f"    {name:18s} {what:42s} {weights}")
print("\n  No model below is trained or fine-tuned. Every prediction is a forward pass that reads")
print("  labelled 'context' rows and answers 'query' rows, like a prompt with examples in it.")

Мы устанавливаем Structured Data Models из репозитория GitHub, зафиксировав коммит, с которым проверялся этот ноутбук, поскольку пакет в PyPI является лишь заполнителем. Для реляционного сэмплера также требуется pyg-lib — именно здесь возникает первая проблема. Колёса pyg-lib публикуются для каждой минорной версии PyTorch на страницах с именами вроде torch-2.14.0+cpu, однако собственный ImportError библиотеки формирует URL на основе точной версии патча — torch-2.14.1+cpu. Это пустая страница, поэтому следование этой рекомендации завершается ошибкой. Вместо этого мы формируем URL индекса по мажорной и минорной версиям. Всё выполняется на CPU — именно этот путь здесь проверен от начала до конца. Мы также выводим модели, доступные в библиотеке, вместе с размером их контрольных точек и лицензиями; две модели Google, которые она тоже поддерживает, TabFM и TimesFM 3, имеют большой размер и не допускают коммерческого использования, поэтому мы их не используем. Ознакомьтесь с полным кодом здесь.

@section("2. Typed tables: what infer_stypes decides, and what a round trip changes")
def typed_tables():
    rng = np.random.default_rng(0)
    words = ["late", "refund", "great", "broken", "fast", "charged", "app", "order", "support", "item"]
    df = pd.DataFrame({
        "ticket_id": np.arange(400),
        "amount": rng.gamma(2, 30, 400).round(2),
        "units": pd.array(rng.integers(1, 5, 400), dtype="Int64"),
        "channel": rng.choice(["web", "app", "phone", None], 400),
        "opened_at": pd.Timestamp("2026-01-01") + pd.to_timedelta(rng.integers(0, 86_400 * 90, 400), unit="s"),
        "comment": [" ".join(rng.choice(words, rng.integers(3, 9))) + f" #{i}" for i in range(400)],
    })
    df.loc[rng.choice(400, 25, replace=False), "units"] = pd.NA
 
    print("  infer_stypes(df) with the defaults, then asking it to look for text and ids:")
    for kw in ({}, {"text": "infer", "id": "infer"}):
        stypes = sdm.infer_stypes(df, **kw)
        print(f"    {str(kw or 'defaults'):34s} {', '.join(f'{c}={s}' for c, s in stypes.items())}")
    print("  With the defaults the free-text comment becomes a 400-level categorical and the id column a")
    print("  number. Both would reach the model as features. Override, or ask: text='infer', id='infer'.")
 
    table = sdm.TableTensor.from_pandas(df, sdm.infer_stypes(df, text="infer", id="infer"))
    print("\n" + textwrap.indent(repr(table), "  "))
 
    back = table.to_pandas()
    print("\n  to_pandas() round trip:")
    for col in df.columns:
        a, b = df[col], back[col]
        if col in ("amount", "units"):
            equal_values = bool(np.allclose(a.astype(float), b.astype(float), equal_nan=True))
        else:
            na_a, na_b = a.isna().to_numpy(), b.isna().to_numpy()
            equal_values = bool((na_a == na_b).all() and
                                (a[~na_a].astype(str).to_numpy() == b[~na_b].astype(str).to_numpy()).all())
        print(f"    {col:10s} {str(a.dtype):16s} -> {str(b.dtype):16s} values preserved: {equal_values}")
    print("  Values survive, including the missing ones; dtypes do not always. Numbers come back float32")
    print("  (the numerical block is one float tensor), nullable Int64 becomes float with NaN, and strings")
    print("  marked categorical come back as a pandas category.")
 
    base = pd.to_datetime(["2026-01-01 08:00:00"] * 3).as_unit("ns")
    whole = pd.DataFrame({"t": base})
    nanos = pd.DataFrame({"t": base + pd.to_timedelta([0, 0, 1], unit="ns")})
    print(f"\n  datetime64[ns], whole seconds:      {sdm.TableTensor.from_pandas(whole, {'t': 'datetime'}).size()} OK")
    try:
        sdm.TableTensor.from_pandas(nanos, {"t": "datetime"})
    except Exception as e:
        print(f"  same column, one value 1 ns later:  {type(e).__name__}: {str(e).split(':')[0]}")
    ok = sdm.TableTensor.from_pandas(nanos.assign(t=nanos.t.dt.floor("us")), {"t": "datetime"})
    print(f"  floor('us') first:                  {ok.size()} OK")
    print("  SDM stores datetimes as int64 microseconds and Arrow refuses any cast that drops a digit. The")
    print("  unit is not the problem; a single value with nanosecond digits fails the whole column.")
    return "6 columns -> 5 typed blocks; text and id need opting in"
 
 
typed_tables()

Всё в библиотеке проходит через TableTensor — подкласс тензора PyTorch, который хранит по одному типизированному блоку для каждого семантического типа: числового, категориального, даты и времени, текста и идентификатора. infer_stypes определяет типы, и в настройках по умолчанию неверно обрабатывает два момента в реалистичном наборе заявок в службу поддержки: свободный текст комментария превращается в категориальный признак с четырьмястами уровнями, а идентификатор заявки — в число, поэтому оба столбца попадают в модель как признаки. Передача флагов для определения текста и идентификаторов исправляет ситуацию. Обратное преобразование в pandas сохраняет каждое значение, включая пропущенные, но не каждый тип данных: числа возвращаются как float32, поскольку числовой блок представляет собой единый тензор float, nullable-целые превращаются в числа с плавающей точкой и NaN, а категориальные строки возвращаются как pandas category. В pandas 2.x строковые столбцы выводятся как object, а не str, а временная метка поступает с наносекундным разрешением; в любом случае те же значения возвращаются без изменений. Перед загрузкой настоящего журнала событий стоит знать ещё об одной проблеме: библиотека хранит даты и время как целые микросекунды, а Arrow отказывается выполнять преобразование, которое приводит к потере цифры. Столбец с наносекундным разрешением, содержащий целые секунды, преобразуется корректно, но единственное значение с наносекундными цифрами приводит к сбою всего столбца, поэтому сначала мы округляем временные метки до микросекунд. Ознакомьтесь с полным кодом здесь.

def as_table(df, target):
    return sdm.TableTensor.from_pandas(df, sdm.infer_stypes(df, overrides={target: "categorical"}), device=DEVICE)
 
 
@section("3. Zero-training classification: two foundation models vs. two trained baselines")
def zero_training():
    df = load_breast_cancer(as_frame=True).frame.sample(frac=1, random_state=0).reset_index(drop=True)
    table = as_table(df, "target")
    n_ctx = 300
    x_ctx, y_ctx = table[:n_ctx].drop_columns("target"), table[:n_ctx, "target"]
    x_q, y_q = table[n_ctx:].drop_columns("target"), df.target.to_numpy()[n_ctx:]
    rows = []
    for name, make in [("TabICLv2", lambda: sdm.models.TabICLv2(device=DEVICE)),
                       ("KumoTabular (small)", lambda: sdm.models.KumoTabular(size="small", device=DEVICE))]:
        model = make()
        torch.manual_seed(0)
        proba = model(x_context=x_ctx, y_context=y_ctx, x_query=x_q, num_estimators=4).to_pandas()
        rows.append((name, "none", proba.idxmax(axis=1).astype(int).to_numpy(), proba["1"].to_numpy()))
    X, y = df.drop(columns="target"), df.target
    for name, est in [("LogisticRegression", make_pipeline(StandardScaler(), LogisticRegression(max_iter=2000))),
                      ("HistGradientBoosting", HistGradientBoostingClassifier(random_state=0))]:
        est.fit(X[:n_ctx], y[:n_ctx])
        rows.append((name, "fit on 300 rows", est.predict(X[n_ctx:]), est.predict_proba(X[n_ctx:])[:, 1]))
    print(f"  breast cancer: {n_ctx} labelled rows, {len(y_q)} to predict, 30 numeric features\n")
    print(f"  {'model':22s} {'training':16s} {'accuracy':>9s} {'AUC':>7s}")
    for name, training, pred, score in rows:
        print(f"  {name:22s} {training:16s} {np.mean(pred == y_q):9.3f} {roc_auc_score(y_q, score):7.3f}")
    print("\n  The output is itself a TableTensor: one probability column per class, named by label.")
    best = max(rows[:2], key=lambda r: np.mean(r[2] == y_q))
    return f"{best[0]} {np.mean(best[2] == y_q):.3f} accuracy with no training"
 
 
zero_training()

Табличные модели получают размеченные строки контекста и немаркированные строки запроса и возвращают TableTensor с одним столбцом вероятности для каждого класса, названным по метке, поэтому этап fit и гиперпараметры не требуются. На наборе данных о раке молочной железы при наличии трёхсот размеченных строк и двухсот шестидесяти девяти строк для прогнозирования TabICLv2 и KumoTabular достигают точности 0,970, а их AUC составляют 0,996 и 0,995 соответственно, без какого-либо обучения. Для логистической регрессии показатели равны 0,970 и 0,993, а для градиентного бустинга — 0,967 и 0,992; обе модели обучались на тех же трёхстах строках. num_estimators задаёт количество по-разному предварительно обработанных представлений данных, по которым модель усредняет результат, то есть реализует встроенный ансамблинг библиотеки.

@section("4. fit() once, predict() many: the cached context")
def kv_cache():
    df = load_breast_cancer(as_frame=True).frame.sample(frac=1, random_state=0).reset_index(drop=True)
    table = as_table(df, "target")
    x_ctx, y_ctx = table[:400].drop_columns("target"), table[:400, "target"]
    batches = [table[400 + 30 * i: 430 + 30 * i].drop_columns("target") for i in range(5)]
    model = sdm.models.TabICLv2(device=DEVICE)
 
    torch.manual_seed(0)
    t0 = time.time()
    for b in batches:
        model(x_context=x_ctx, y_context=y_ctx, x_query=b, num_estimators=4)
    t_forward = time.time() - t0
 
    torch.manual_seed(0)
    t0 = time.time()
    model.fit(x=x_ctx, y=y_ctx, num_estimators=4)
    t_fit = time.time() - t0
    t0 = time.time()
    for b in batches:
        model.predict(b)
    t_predict = time.time() - t0
    model.clear()
    print("  five batches of 30 query rows against the same 400 context rows:")
    print(f"    forward() x5                      {t_forward:6.2f}s  (context re-encoded every call)")
    print(f"    fit() once + predict() x5         {t_fit + t_predict:6.2f}s  (fit {t_fit:.2f}s, predicts {t_predict:.2f}s)")
    print(f"    per batch: {t_forward / 5:.2f}s vs {t_predict / 5:.2f}s")
    print("\n  fit() encodes the context once into a key/value cache; predict() only runs the queries.")
    print("  clear() frees it. Nothing is learned: fit() here means 'remember these examples'.")
    return f"{t_forward / (t_fit + t_predict):.1f}x faster for five batches"
 
 
kv_cache()

При прямом вызове модели контекст кодируется заново при каждом обращении. Если одни и те же размеченные строки будут использоваться для ответа на множество пакетов запросов, fit кодирует контекст один раз в кэш ключей и значений, а predict обрабатывает только строки запроса; clear освобождает кэш. При этом в fit ничего не обучается: здесь он означает «запомнить эти примеры». Обработка пяти пакетов по тридцать строк относительно одного и того же контекста из четырёхсот строк на нашем CPU заняла в несколько раз меньше времени, а после кэширования контекста каждый следующий пакет требовал лишь небольшой доли времени полного прямого прохода. Точное соотношение зависит от компьютера, поэтому ячейка измеряет его, а не предполагает заранее.

Для регрессии модель возвращает не одно число. Она возвращает 999 квантилей для каждой строки запроса — полное прогнозное распределение, полученное за один прямой проход. На наборе данных о диабете с тремястами размеченных строк медиана этого распределения даёт RMSE 53,6, немного лучше RidgeCV с результатом 54,3 и заметно лучше градиентного бустинга с результатом 64,7; обе базовые модели обучались на тех же строках. Более важным свойством является калибровка: интервал между 10-м и 90-м процентилями содержит 81,0% отложенных целевых значений, а интервал между 5-м и 95-м процентилями — 91,5%, что близко к номинальным 80% и 90%. При этом диапазон 50% оказывается несколько узким и покрывает 43,7% значений.

Предварительная обработка в библиотеке представляет собой Recipe — набор компонуемых процессоров, а стандартный рецепт TabICLv2 выбирает обработку по семантическому типу и содержит ветви только для категориальных и числовых столбцов. Поэтому столбцы с временными метками и текстом никогда не становятся признаками. В заявках в службу поддержки, где эскалация зависит от времени открытия и от таких слов, как refund, стандартный рецепт даёт AUC 0,525 — почти случайный результат. Библиотека действительно выдаёт RuntimeWarning с указанием проигнорированных типов столбцов. Однако это предупреждение выводится только один раз: оно появляется при первом вызове и подавляется до конца процесса, поэтому каждая последующая модель в том же ноутбуке или сервисе молча отбрасывает эти столбцы. Добавление двух процессоров исправляет ситуацию: AddCalendarFields преобразует временную метку в циклические признаки часа и дня недели, повышая AUC до 0,634, а TFIDF добавляет термины из комментария и повышает показатель до 0,789.

GradientExplainer библиотеки возвращает градиент выбранного выхода по отношению к предварительно обработанным входным данным запроса. Мы проверяем его на синтетических данных, где значение имеют только три из восьми признаков, причём их знаки известны. Здесь есть две проблемы, заметные только при анализе того, что получает функция выхода. Во-первых, функция вызывается один раз для каждого оценщика, и каждый вызов перезаписывает предыдущий результат. Поэтому при четырёх оценщиках функция выхода, активная только при первом вызове, даёт объяснение, состоящее из одних нулей, а полученное объяснение относится только к последнему оценщику. Во-вторых, необработанный вывод каждого оценщика имеет собственный порядок классов, поскольку стандартный рецепт перемешивает метки классов для каждого оценщика. При объяснении с одним оценщиком на шести случайных начальных значениях выбор вероятности класса 1 по позиции столбца меняет знак каждого признака в двух случаях, когда порядок меняется. Выбор по имени столбца даёт правильный знак во всех шести случаях, а усреднение по начальным значениям восстанавливает три настоящих признака, тогда как признаки-«шумы» остаются близкими к нулю.

KumoRelational принимает реляционную базу данных вместо отдельной таблицы. Мы создаём небольшой магазин с восемьюстами клиентами, шестьюдесятью товарами и заказами за год. У некоторых клиентов активность резко падает в случайную дату. Затем мы спрашиваем, сделает ли каждый клиент повторный заказ в течение тридцати дней после контрольной даты. RelationalData хранит таблицы и внешние ключи между ними, а его сэмплер проходит по этим ключам от каждой строки задачи, собирая до шестнадцати соседей на каждом шаге. Объявление даты заказа временным столбцом делает этот обход времязависимым: мы проверяем, что последний заказ, возвращённый сэмплером для строки задачи с датой 1 июня, был размещён 31 мая, поэтому строки, определяющие метку, недоступны модели.

Затем мы отвечаем на тот же вопрос тремя способами. TabICLv2, работающая только с таблицей клиентов, где доступны лишь регион и тарифный план, достигает AUC 0,556. TabICLv2 на шести вручную созданных признаках истории — количестве заказов и сумме покупок за последние тридцать и девяносто дней, количестве дней с момента последнего заказа и общем числе заказов, рассчитанных строго до каждой контрольной даты, — достигает 0,828. KumoRelational на исходных таблицах без созданных вручную признаков достигает 0,823. Реляционная модель почти полностью повторяет результат ручного конвейера, читая сами строки заказов через ключи и учитывая момент времени каждой строки задачи.

Наконец, мы запускаем ту же реляционную модель, не объявляя временной столбец. AUC повышается до 0,940. Без временного столбца сэмплер может вернуть заказы, размещённые после контрольной даты, то есть именно те заказы, которые определяют метку. Поэтому модель уже не прогнозирует, а читает ответ, причём никаких ошибок не возникает. В настоящей базе данных реляционная модель, которая подозрительно превосходит ручную базовую модель, — это первое, что следует повторно запустить после объявления всех временных столбцов.

В итоговом выводе печатается однострочный результат, возвращённый каждым разделом, а также три полезные привычки: намеренно объявляйте семантические типы и рецепты, выбирайте классы по имени внутри объяснителя и используйте один оценщик при объяснении, а каждую таблицу с временным столбцом добавляйте в time_columns. Затем приводятся ссылки на дообучение, пример RelBench и выполнение на GPU.

В заключение можно сказать, что Structured Data Models от NVIDIA выполняет своё основное обещание: в задачах классификации, регрессии, двадцатиклассовой классификации и на базе данных из трёх таблиц предварительно обученные модели не уступили или превзошли базовые модели scikit-learn, обученные на тех же строках, сами при этом не обучаясь. Они также бесплатно возвращали калиброванные прогнозные распределения, а в реляционной задаче достигли точности ручного конвейера создания признаков, работая с исходными таблицами. Настройки по умолчанию заслуживают большего внимания, чем сами модели. Свободный текст и идентификаторы не определяются без явного запроса, стандартный рецепт отбрасывает столбцы с временными метками и текстом и предупреждает об этом лишь один раз за процесс, объяснитель сообщает только о последнем оценщике и требует выбирать классы по имени, URL pyg-lib в собственном сообщении об ошибке библиотеки ведёт на пустую страницу, а реляционная модель без временного столбца выдаёт впечатляющее, но бессмысленное число. Для каждой проблемы существует исправление в одну строку, и каждая из них в противном случае проходит незаметно. Практический вывод, который стоит перенести в работу с реальными данными: проверяйте, что именно попадает в модель, прежде чем доверять тому, что из неё выходит.

Ознакомьтесь с полным кодом здесь. Все заслуги принадлежат исследователю этого проекта. Также подписывайтесь на нас в Twitter, не забудьте присоединиться к нашему сабреддиту о машинном обучении с более чем 150 тысячами участников и подписаться на нашу рассылку. Стоп! Вы есть в Telegram? теперь вы также можете присоединиться к нам в Telegram.