Исследователи представили новую нейросетевую архитектуру для работы с табличными данными

Лаборатория исследований искусственного интеллекта Yandex Research представила новую нейросетевую архитектуру для работы с табличными данными — TabM. Она позволяет быстро обрабатывать большие массивы данных и строить высокоточные прогнозы, что востребовано в бизнесе, исследованиях и медицине. Модели для работы с табличными данными помогают оптимизировать поставки, прогнозировать энергопотребление, классифицировать пациентов по риску заболеваний и решать многие другие задачи. Архитектура уже доступна разработчикам и исследователям на GitHub, а научная статья — на arXiv. Также работа принята на одну из топовых конференций по ИИ в мире: ICLR.

Разработку использовали на Kaggle — платформе международных соревнований по анализу данных и машинному обучению от Google. В частности, новую архитектуру применяли для предсказания выживаемости пациентов после трансплантации костного мозга. За решение этой и других задач с помощью TabM призёры и победители Kaggle получили в совокупности 60 тысяч долларов.

TabM (от англ. Tabular DL model that makes Multiple predictions) — это эффективная реализация так называемого ансамбля моделей, когда каждая модель проводит свой анализ, после чего прогноз усредняется. Архитектура TabM позволяет добиться оптимального соотношения точности прогноза и необходимых вычислительных мощностей.

По результатам тестирования на 46 наборах данных, TabM превзошла другие решения не только по занимаемому в среднем месту (1,7 у TabM против 2,9 у ближайшего конкурента), но и по стабильности работы, что важно для практического применения. Благодаря способности объединять усилия нескольких подмоделей и эффективному использованию вычислительных ресурсов, TabM успешно конкурирует с классическими моделями градиентного бустинга — CatBoost, XGBoost, LightGBM, — которые долгое время считались лучшим решением для табличных данных.