Российские учёные создали универсальный бенчмарк, который позволяет сравнивать методы построения мостов Шрёдингера

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

В её основе лежит преобразование одного распределения данных в другое: сложного в простое на этапе обучения и наоборот — на этапе генерации. В пространстве данных (например, изображений) это эквивалентно перемещению из одной точки в другую. Для быстрой работы моделей этот процесс необходимо оптимизировать. Один из способов добиться этого — использование оптимального транспорта между этими точками. Для его поиска часто применяют регуляризацию на основе энтропии, что помогает эффективнее искать такое решение. В таком виде эта проблема называется задачей энтропийного оптимального транспорта (Entropic Optimal Transport, EOT).

На самом деле этой задаче больше лет, чем самому машинному обучению. Ещё в 30-х годах прошлого века ею интересовался Эрвин Шрёдингер. Он известен как один из основателей квантовой механики, но, помимо этого, учёный занимался проблемами статистической физики и броуновского движения частиц. Сегодня мостом Шрёдингера (Schrödinger Bridge, SB) называют наиболее вероятный способ стохастического перехода между распределениями, а задача о его построении эквивалентна задаче EOT.

Схематическое изображение моста Шрёдингера

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

Чтобы помочь коллегам лучше оценивать свои методы, команда специалистов из AIRI и Сколтеха создала первый в мире теоретически обоснованный бенчмарк для методов решения SB- и EOT-задач. Его работа основана на новом способе генерации пар вероятностных распределений, для которых известно аналитическое OT-решение. Бенчмарк оценивает метод, сравнивая решение, найденное методом, с эталонным. Благодаря универсальности методологии, разработка подходит к широкому классу солверов.

Авторы бенчмарка проверили его на нескольких популярных нейронных солверах, обучив их на сгенерированных парах в многомерных пространствах, включая пространство портретов знаменитостей размерами 64 на 64 пикселя. Эффективность метода оценивала специальная метрика, основанная на усредненной и нормализованной метрике Буреса — Васерштейна.

Качественное сравнение эталонных выборок с выборками, полученными солвером [ENOT]. Управление параметром модели увеличивает разнообразие, но снижает точность восстановления изображений