Nums AI випустила Causilo: таблична базова модель, що очолила TabArena серед одиночних моделей

Nums AI випустила Causilo — попередньо навчену фундаментальну модель для табличних даних, призначену для класифікації та регресії. Causilo має інтерфейс scikit-learn, код під ліцензією Apache-2.0 і попередньо навчені ваги на Hugging Face. У TabArena вона має найвищий Elo серед одиночних моделей як для класифікації, так і для регресії.

Чи придатна вона до розгортання? Так, уже сьогодні — для досліджень і оцінювання, на CUDA або CPU. Для комерційного використання, використання у виробничих середовищах і через розміщений API потрібна окрема ліцензія від Nums AI.

Що робить Causilo

Causilo — це модель навчання в контексті. Виклик fit не оновлює попередньо навчені ваги. Вона зберігає навчальні рядки як контекст і прогнозує рядки запиту за один прямий прохід. Згідно з її поданням до TabArena, Nums AI попередньо навчала Causilo лише на синтетичних даних, без використання наборів даних TabArena.

Вхідними даними можуть бути масиви NumPy або DataFrame pandas, зокрема категоріальні ознаки та пропущені значення. Класифікація підтримує до 10 класів. За замовчуванням регресія повертає середні прогнозовані значення. Версія 1.0.1 додає медіанні та квантильні вихідні дані на основі 999 нативних квантилів.

Архітектура: уточнення, стиснення, навчання в контексті

Nums AI розділяє мережу на 3 фази: уточнення, стиснення та навчання в контексті. Випущений код і конфігурації демонструють принцип роботи кожної фази.

Ознаки групуються в набори по 3. Кожне значення вбудовується за допомогою 16 навчених синусоїдальних і косинусоїдальних частот. Для пропущених значень використовується власний навчений вектор.

2 етапи обробки стовпців узагальнюють кожну групу ознак. На кожному з них 128 латентних слотів зчитують лише навчальні рядки та передають це узагальнення кожному рядку. Між 2 етапами обробки стовпців етап обробки рядків дає змогу групам ознак взаємодіяти через 4 латентні токени. Він використовує перехресну увагу замість повної самоуваги, що, за словами Nums AI, зберігає лінійну залежність обчислювальних витрат від кількості ознак.

Потім блок об’єднання стискає кожен рядок у вектор фіксованого розміру — 512 вимірів. До навчальних рядків додаються мітки. Блок прогнозування з 12 шарів дає змогу рядкам запиту звертатися до цих розмічених рядків. Рядки запиту не можуть змінювати навчальний контекст або один одного.

За замовчуванням 8 ансамблевих учасників мають спільні ваги. Кожен із них циклічно використовує нормалізацію none, rank2gaussian, robust або power, із фіксованими перестановками ознак і класів.

Результати TabArena

Nums AI використала офіційний конвеєр TabArena: 51 набір даних і 816 повних розбиттів, із 8 оцінювачами та початковим значенням 42. Супроводжувач TabArena повторно провів повне оцінювання й отримав таке саме загальне значення Elo — 1794.

До цього переліку входять TabFM від Google Research, EXAONE Tabular від LG AI Research і TabPFN-3 від Prior Labs (1636.2 загалом).

Під час прочитання цих чисел варто врахувати кілька моментів:

• Позиції №1 не враховують системні записи. Якщо включити системи, повторний запуск супроводжувача поставив Causilo на 3-тє місце серед 88 моделей загалом.
• Щодо покращуваності TabFM і далі лідирує загалом і в класифікації. Causilo лідирує в регресії.
• Довірчі інтервали Elo для лідерів перетинаються, тому перевага над TabFM і EXAONE Tabular є незначною.
• Nums AI також вказує Xiaomi-TabLDM і Mitra-v2 від Amazon позаду Causilo. Жодна з цих моделей не представлена у файлах бенчмарків у репозиторії Causilo.

Результати ScoringBench

ScoringBench оцінює регресійні моделі за допомогою коректних правил оцінювання, таких як CRPS, а також RMSE і R². Nums AI подала Causilo 1.0.1 на 101 наборі даних, по 5 фолдів кожен, з обмеженням у 3 000 зразків. Nums AI повідомляє, що Causilo посідає 1-ше місце за CRPS, R² і RMSE. Супроводжувач ScoringBench незалежно перевірив результати перед їхнім додаванням.

Швидкість і пам’ять

Nums AI також повторно протестувала 3 моделі на 1 графічному процесорі H100 із 80 ГБ пам’яті, використовуючи 8 ядер CPU на завдання.

У цьому тесті Causilo є найшвидшою як під час навчання, так і під час прогнозування. TabPFN-3 використовує значно менше пам’яті GPU. Параметр use_kv_cache=True переносить обробку контексту на етап навчання, використовуючи більше пам’яті для прискорення повторних прогнозів.

Початок роботи

Causilo потребує Python версій 3.10–3.12 і PyTorch 2.13 або новішої версії. Під час першого навчання контрольна точка завантажується автоматично.

# pip install causilo
from causilo import CausiloClassifier, CausiloRegressor

clf = CausiloClassifier(n_estimators=8, random_state=42)
clf.fit(X_train, y_train)
proba = clf.predict_proba(X_test)

reg = CausiloRegressor()
reg.fit(X_train, y_train)
bands = reg.predict(X_test, output_type="quantiles", quantiles=[0.05, 0.5, 0.95])

Ви також можете спробувати демонстраційний Space на Hugging Face.

Основні висновки

• Causilo має найвище значення Elo у TabArena серед одиночних моделей — загалом і для кожного завдання.
• Якщо врахувати системні ансамблі, повторний запуск супроводжувача ставить її на 3-тє місце серед 88 моделей.
• Змішування рядків відбувається через 4 латентні токени, зберігаючи лінійну залежність обчислювальних витрат від кількості ознак.
• Код поширюється за ліцензією Apache-2.0; ваги призначені лише для досліджень без комерційної ліцензії.
• Версія 1.0.1 додає квантильні вихідні дані, тому регресійні інтервали працюють одразу після встановлення.

Перегляньте репозиторій на GitHub і модель на HF. Усі заслуги належать досліднику цього проєкту. Також підписуйтеся на нас у Twitter і не забудьте приєднатися до нашого субреддіту про машинне навчання зі 150 тис.+ учасників та підписатися на нашу розсилку. Стривайте! Ви є в Telegram? тепер ви також можете приєднатися до нас у Telegram.

Потрібне партнерство з нами для просування вашого репозиторію на GitHub, сторінки на Hugging Face, випуску продукту, вебінару тощо? Зв’яжіться з нами