Учёные из Сколтеха и Института AIRI предложили новый подход к решению обширного класса сложных вычислительных задач на основе оптимального транспорта (ОТ), повсеместно применимых в машинном обучении и математическом моделировании. Метод позволит ускорить обучение моделей от 3 до 10 раз. Исследователи представят разработку в Ванкувере в ходе NeurIPS 2024, одной из крупнейших мировых ИИ-конференций.

Методы оптимального транспорта сегодня всё чаще используются при обучении генеративных моделей для синтеза искусственных данных, например — изображений или текстов. Ещё одно значимое применение — адаптация моделей к данным из новых источников, что особенно актуально в медицине, где работа часто связана с небольшими и разрозненными выборками. Однако существующие методы решения задач ОТ при помощи нейросетей сталкиваются с рядом проблем, таких как высокая нестабильность обучения и необходимость сложных промежуточных преобразований.

Ключевым преимуществом предложенного учёными метода, реализованного на фреймворке JAX и получившего название ENOT, стало внедрение экспектильной регуляризации. Это позволило достичь значительного ускорения вычислений — от 3 до 10 раз — и улучшения целевых метрик качества работы моделей. Первоначально эксперименты проводились на двумерных данных, а позднее метод был протестирован на задачах генерации изображений, переносе стиля и реконструкции трёхмерных объектов, что подтвердило его универсальность.

«Несмотря на теоретический формат статьи, в основе метода лежит интуитивно понятная идея — предлагается «стягивать» генерируемое к ожидаемому. Мне кажется, сегодня в нашей области слишком много «чёрных ящиков». Конечно, десятикратное ускорение –– это весомый аргумент, но я думаю, что рецензентам NeurIPS наш метод понравился именно своей интуитивностью», –– отметил доцент Центра технологий ИИ Сколтеха и директор лаборатории «Сильный ИИ в медицине» Института AIRI Дмитрий Дылов.

«Предложенный нами метод является самым быстрым и точным на сегодняшний день. Что касается практического применения, то задачи оптимального транспорта мультидисциплинарны, поэтому он может применяться в самых разных областях. В частности, мы применяли его для обучения имитации — когда эксперт показывает определённые действия, агент пытается имитировать поведение, и система оценивает, насколько действия агента похожи на действия эксперта. Примером может служить урок танцев, когда преподаватель показывает движение, а ученик пытается его повторить», –– объяснил Назар Бузун, руководитель группы «Обучение представлений» лаборатории «Сильный ИИ в медицине» Института AIRI.

Метод уже нашёл отклик в научном сообществе: принятую на конференцию статью определили в Spotlight –– специальный трек конференции для работ, особо отмеченных рецензентами. Кроме того, авторы одного из основных решений по теме ОТ (ott-jax) достаточно быстро имплементировали метод в свою библиотеку.