NVIDIA выпустила Kumo Tabular: новые языковые модели для таблиц

Компания NVIDIA представила Kumo Tabular — новое семейство фундаментальных моделей для работы с табличными данными, предназначенных для задач классификации и регрессии. Если вы раньше сталкивались с архитектурами вроде TabPFN или TabICL, то здешний подход покажется вам интуитивно знакомым.

Программный инструмент принимает размеченные строки в качестве контекста и прогнозирует новые строки всего за один проход вперед. При этом отпадает необходимость в трудоемком обучении, ручной настройке гиперпараметров и сложной инженерии признаков. Линейка моделей включает версии Small, Medium и Large, охватывающие диапазон от 28 до 215 миллионов параметров. Все это работает через фирменную библиотеку структурных данных с открытым исходным кодом от NVIDIA под названием SDM. Веса распространяются по лицензии OpenMDW-1.1, разрешающей коммерческое использование, в то время как код SDM написан под Apache-2.0 и требует Python 3.11+, PyTorch 2.7+ и наличия CUDA-совместимого графического процессора.

Возможности библиотеки SDM

Интерфейс SDM представляет собой оптимизированную под графические процессоры библиотеку для предобработки и развертывания фундаментальных моделей структурированных данных. Помимо упорядоченной Kumo Tabular, она поставляется с такими решениями, как TabICLv2, Google TabFM и KumoRelational, предназначенными для многотабличных баз данных.

Все эти архитектуры объединены единым интерфейсом обучения в контексте, который базируется на специализированном контейнере TableTensor. Набор инструментов также берет на себя рутинные задачи по предварительной обработке, ансамблированию и прогнозированию с большим количеством классов.

Как функционирует Kumo Tabular

В основе архитектуры лежит Transformer, адаптированный под специфическую структуру таблиц. В нем задействованы механизмы внимания по колонкам, строкам и в контексте, заимствованные из разработок TabICL и TabPFN.

Этапы обработки данных

Конвейер состоит из трех основных стадий. На первой этапе числовые и категориальные значения проходят через заученные признаки Фурье с раздельными весами для каждого типа, причем пропущенные значения не требуют специальной импутации. На втором этапе колоночное внимание использует индуцированное самовнимание, благодаря чему затраты растут линейно относительно числа строк, а позиционное внимание по строкам изучает взаимодействия между признаками с помощью четырех обучаемых токенов [CLS]. На третьем этапе запускается финальный трансформер для обучения в контексте, где контекстные строки взаимодействуют друг с другом, а строки запросов обращаются исключительно к контексту. Поскольку контекст не видит запросы, его ключи и значения вычисляются единожды и переиспользуются.

Особенности масштабирования

Выходной слой модели выдает вероятности классов или 999 квантилей для регрессии, что обеспечивает точечный прогноз вместе с оценкой неопределенности. Чтобы softmax-внимание не размывалось при росте числа ключей, каждый запрос масштабируется с помощью температуры, увеличивающейся пропорционально логарифму количества ключей. Коэффициент обучается для каждой головки внимания отдельно, сохраняя высокую точность на больших массивах.

Обучение на искусственных таблицах

Предварительное обучение Kumo Tabular проводилось исключительно на синтетических таблицах, сгенерированных из структурных причинных моделей. Случайный причинный граф связывает скрытые переменные через линейные отображения, небольшие нейросети, деревья или гауссовские процессы. Генератор также намеренно добавляет реалистичные изъяны: пропуски, высокую кардинальность категорий, тяжелые хвосты целевых переменных и конфликтующие дубликаты.

Процесс обучения состоял из трех стадий, аналогичных подходу TabICLv2, где контекст постепенно увеличивался с 1024 до 60 000 строк при ширине до 100 столбцов. Младшая, средняя и крупная версии увидели от 35 до 137 миллионов искусственных таблиц, а модели для классификации и регрессии тренировались раздельно.

Результаты бенчмарков

При стандартных настройках модель занимает лидирующую позицию в общем зачете TabArena с рейтингом Elo 1950, демонстрируя при этом в 17 раз большую скорость работы по сравнению с LimiX-2 на одном ускорителе RTX 6000 Pro. Все три варианта размера находятся на фронте Парето по точности и времени инференса. В других тестах вроде TALENT и ScoringBench решения также показывают высшие средние ранги по метрикам точности, лог-потерь и среднеквадратичной ошибки.

Сравнение с конкурентами

Главным дифференциатором среди аналогов выступает лицензионная политика. Веса таких систем, как TabPFN-3, LimiX-2 и TabFM, накладывают ограничения на коммерческое использование, в то время как Kumo Tabular и TabICLv2 предлагают более мягкие условия при лидерстве решения от NVIDIA в бенчмарках.

С чего начать работу

Для запуска необходимо установить библиотеку и передать датафрейм через обертку TableTensor следующим образом:

# pip install structured-data-models
from sklearn.datasets import load_breast_cancer
import sdm
df = load_breast_cancer(as_frame=True).frame
table = sdm.TableTensor.from_pandas(df=df, stypes=sdm.infer_stypes(df, overrides={"target": "categorical"}), device="cuda")
model = sdm.models.KumoTabular(task="classification", device="cuda")

Аргумент размера принимает значения small, medium или large, причем по умолчанию используется крупная версия.