Дата публикации
15.05.2024
Авторы
Александр Коротин Марат Хамадеев
Поделиться

Энергетическое моделирование сделало мост Шрёдингера лёгким


С начала 00-х годов в машинном обучении набирают популярность модели, называемые энергетическими (energy-based model, EBM). В их основе лежит применение идеи канонического ансамбля, возникшей в статистической физике, при работе с вероятностными распределениями внутри модели. Так там называют систему, которая находится в тепловом равновесии с окружающей средой, из-за чего вероятность найти её в определенной конфигурации зависит от энергии по закону, выраженному через распределение Больцмана.

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

Параллельно с ними активно развивается подход к генерации, основанный на оптимальном транспорте. Под оптимальным транспортом (OT) понимают класс задач об эффективном переходе от одного вероятностного распределения к другому. Больших успехов в этом направлении добилась команда исследователей из AIRI и Сколтеха под руководством Евгения Бурнаева. Так, мы уже рассказывали, как они предложили строгую формулировку для задачи о построении оптимального транспорта с энтропийной регуляризацией с помощью нейросетей (Entropic Neural Optimal Transport, ENOT), а также математически обосновали поиск теоретически наилучшего перевода домена с использованием непарных обучающих выборок.

На этот раз учёные представили новую методологию, которая использует достижения EBM, чтобы улучшить ENOT. Сначала они математически разработали процедуру оптимизации и соответствующий алгоритм, который позволяет неявно восстанавливать планы оптимального транспорта с помощью энергетического представления. Также авторы провели тщательный теоретический анализ метода.

Чтобы продемонстрировать преимущества нового подхода, команда провела серию экспериментов с EBM-аугментированными моделями на датасете с простыми двумерными распределениями, а также решала задачи о трансформации гаусс-в-гаусс и о превращении кошачьих мордочек (а также мордочек других диких зверей) в собачьи, взятых из датасета AFHQ, в непарной постановке. Опыты показали, что предложенное нововведение позволяет достигать значений FID-метрики, сравнимых с бейзлайнами.

«Превращение» кошек в собак с помощью EBM-транспорта, примененного в латентном пространстве модели StyleGAN2-ADA

На этом команда не остановилась. Оказалось, что энергетический подход позволяет построить ещё более эффективное решение задачи о построении моста Шрёдингера, то есть нахождение наиболее вероятного стохастического перехода от одного распределения к другому. Это стало возможным благодаря тому, что задача о мосте Шрёдингера с некоторыми уточнениями эквивалентна построению энтропийного оптимального транспорта. Ранее исследователи уже делились бенчмарком для тестирования солверов, которые строят мосты Шрёдингера с помощью нейросетей.

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

Авторы нового исследования смогли преодолеть эту трудность, создав легковесный солвер моста Шрёдингера, который опирается на две идеи: параметризация шрёдингеровских потенциалов с помощью гауссовских смесей и их использование в качестве энергетической функции. Теоретический анализ показал, что солвер может служить универсальным аппроксиматором для мостов Шрёдингера. При этом гарантируется его сходимость, так как ошибка метода исчезает по мере увеличения размера выборки.

В качестве практической проверки исследователи применили новый метод в серии экспериментов с набором синтетических и реальных данных, включая лица людей и данные, полученные в биологических системах. Оказалось, что скорость построения моста Шрёдингера с помощью солвера занимает всего несколько минут, так как теперь нет необходимости в max-min-оптимизации, симуляции траекторий полных процессов и прочих времязатратных процедурах. 

«Превращение» мужских лиц в женские с помощью нового солвера, примененного в латентном пространстве модели ALAE

Оба результата были приняты на конференцию ICLR 2024, подробности можно прочитать в статьях её сборника трудов (EBM-ENOT и лёгкий солвер). Код исследований доступен на GitHub (EBM-ENOT и лёгкий солвер).


Присоединяйтесь к AIRI в соцсетях