NVIDIA Kumo Tabular задает новый стандарт точности и эффективности для прогнозирования на основе табличных данных

Основные моменты (TL;DR)

NVIDIA Kumo Tabular, входящая в коллекцию моделей NVIDIA Kumo Structured, представляет собой открытую базовую модель для табличных данных, которая теперь доступна на Hugging Face. Имея на входе таблицу с размеченными строками, она прогнозирует метки новых строк за один проход вперед без обучения, настройки и инжиниринга признаков как для задач классификации, так и регрессии. Модель предобучена исключительно на искусственных данных, представлена в трех размерах (от 28M до 215M параметров), работает через нашу библиотеку с открытым исходным кодом и выпускается под лицензией OpenMDW-1.1 license для коммерческого использования. Она занимает первое место в четырех бенчмарках TabArena, BeyondArena, TALENT и ScoringBench.

• Код модели:
• Веса модели: https://huggingface.co/nvidia/Kumo-Tabular

Переход к базовым табличным моделям

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

Большие языковые модели показали иной подход к работе с новыми задачами. Получив несколько примеров в промпте, предобученная модель решает задачу без обновления ни единого веса. Это контекстное обучение, которое применимо к таблицам точно так же, как и к тексту: модель, предобученная на миллионах таблиц, может читать размеченную таблицу в качестве контекста и напрямую прогнозировать метки новых строк.

Сегодня мы выпускаем NVIDIA Kumo Tabular (, HuggingFace), открытую базовую модель для табличной классификации и регрессии. Получая на входе таблицу с размеченными строками и строки, для которых необходимо сделать предсказание, Kumo Tabular возвращает вероятности классов или числовые прогнозы за один прямой проход.

Как работает Kumo Tabular

Kumo Tabular — это трансформер, построенный на основе структуры таблицы, использующий колоночное, строчное и контекстное внимание, как описано в TabICL и TabPFN. Чтобы спрогнозировать метку, модель должна выполнить три действия: (1) понять значение каждого значения в пределах его колонки, (2) понять, как колонки одной строки взаимодействуют друг с другом, и (3) сопоставить контекстные строки с существующими метками со строками запроса, имеющими неизвестные метки. Kumo Tabular достигает этого следующим образом:

Эмбеддинг ячеек (Cell Embedding): Группа ячеек превращается в токен. Числовые и категориальные значения проходят через признаки Фурье, синусы и косинусы выученных частот с раздельными весами для каждого типа. Пропущенные значения не требуют импутации и обрабатываются особым образом. Наконец, каждый токен в контексте получает эмбеддинг метки.

Эмбеддинг строк (Row Embedding): Затем мы превращаем каждую строку в эмбеддинг, многократно чередуя два типа внимания. Колоночное внимание (Column attention) просматривает отдельную колонку и узнает значение величины в распределении ее колонки (например, является ли число 42 типичным или экстремальным) с помощью индуцированного самовнимания (induced self-attention). Стоимость этого процесса растет линейно с количеством строк. Строчное внимание (Row attention) просматривает токена одной строки и изучает взаимодействие признаков с использованием ротационных позиций для различения колонок. К каждой строке присоединяются четыре обучаемых токена [CLS], которые выступают в роли финального считывания строки. После этого сжатия строк стоимость финального этапа больше не зависит от количества колонок.

Контекстное обучение (In-context Learning): Финальный трансформер работает с эмбеддингами строк. Контекстные строки обращаются друг к другу, в то время как строки запросов обращаются только к контекстным строкам. Таким образом, каждое предсказание зависит только от контекста и от самой строки, а не от того, какие другие строки оцениваются вместе с ней. Поскольку контекст никогда не смотрит на запросы, его ключи и значения вычисляются один раз и могут быть переиспользованы для последующих предсказаний. Строки запроса используют Test-GQA, что уменьшает кэш, который читает каждое предсказание. Специальная «голова» превращает каждую строку запроса в вероятности классов для классификации и 999 квантилей для регрессии, на основе которых формируются точечный прогноз и оценка неопределенности.

Температура внимания с учетом длины (Length-aware Attention Temperature): Внимание 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 с настройками по умолчанию против полной таблицы лидеров , включающей настроенные деревья с градиентным бустингом, AutoGluon и новейшие табличные базовые модели. Kumo Tabular занимает первое место в общем зачете с рейтингом ELO 1950, работая при этом в 17 раз быстрее LimiX-2 в условиях единой тестовой среды на одном GPU RTX 6000 Pro. Во всех трех размерах моделей Kumo Tabular устанавливает новый уровень передовых достижений на границе Парето точности и эффективности:

Мы также оценили Kumo Tabular на , и . На BeyondArena модель Kumo Tabular достигает рейтинга ELO 1418 с показателем улучшаемости (Improvability) 7,78%, занимая первое место в таблице лидеров. На TALENT она занимает первое место в общем зачете по точности классификации, логарифмическим потерям классификации и среднеквадратичной ошибке (RMSE) регрессии со средними рангами 6,67, 3,98 и 4,22. На ScoringBench, бенчмарке для предиктивных распределений, версии Kumo Tabular-Large и Medium занимают первое и второе места по среднему рангу.

Ограничения

Kumo Tabular работает только с числовыми и категориальными столбцами, в то время как текст, изображения или временные метки могут быть преобразованы в признаки с помощью встроенных рецептов предварительной обработки. Один прямой проход охватывает до 10 классов, что библиотека расширяет до любого количества классов с помощью помехоустойчивых кодов с исправлением ошибок. Точность может снижаться на таблицах, значительно выходящих за рамки диапазонов обучения, или когда строки запроса происходят из другого распределения, чем контекстные строки; поэтому, как и в случае с любой предиктивной моделью, проверяйте точность и калибровку на собственных отложенных данных перед развертыванием.

Демонстрация

Kumo Tabular работает через недавно выпущенную библиотеку NVIDIA с нативной поддержкой GPU для . Библиотека загружает веса с Hub при первом использовании и обеспечивает предварительную обработку, ансамблирование и работу с большим количеством классов, использовавшиеся в наших оценках. Приведенный ниже код — это все, что требуется для перехода от pandas.DataFrame к предсказанию:

Начните создание с Kumo Tabular

Kumo Tabular выпускается под лицензионным соглашением OpenMDW License Agreement, версия 1.1. NVIDIA считает, что создание надежного искусственного интеллекта — это общая ответственность, и мы установили правила и практики, способствующие разработке для широкого спектра приложений ИИ. При загрузке или использовании в соответствии с нашими условиями обслуживания разработчики должны взаимодействовать со своей командой поддержки моделей, чтобы гарантировать соответствие этой модели требованиям для конкретной отрасли и варианта использования, а также предотвращение непредвиденного неправильного использования продукта. Пожалуйста, сообщайте о качестве модели, рисках, уязвимостях безопасности или проблемах с NVIDIA AI here.

• Код модели:
• Веса модели: https://huggingface.co/nvidia/Kumo-Tabular

Благодарности

Мы благодарим Дэвида Хольцмюллера за значительный вклад в виде идей и абляционных исследований для Kumo Tabular. Мы благодарим Вигнеша Котхапалли за его помощь в работе над Kumo Tabular во время его стажировки.