Neural Optimal Transport

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

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

Классическая формулировка Монжа предлагает детерминированное отображение, где каждому х соответствует один у. Но это плохо работает, когда нужно отображать один объект в несколько возможных, ведь распределения могут быть разной структуры. Канторович в 1958 году предложил вместо отображения брать совместное распределение π. Но за это платим сложностью: всё становится сильно тяжелее и менее удобным для обучения.

Авторы идут дальше. Вместо того чтобы напрямую работать с этим распределением π и условным распределением p(y|x), которое появляется при переходе от сильного коста к слабому, они предлагают параметризовать отображение x шумом: берём x и случайную переменную z, и уже функция T(x, z) выдаёт результат, лежащий в целевом распределении. За счёт этого условное распределение получается само собой: для одного x можно получать разные y просто благодаря разному z.

Чтобы это оптимизировать, переходят к двойственной задаче. Появляется функция f — её можно воспринимать как вспомогательную. В итоге всё сводится к двум объектам: есть T, которое отвечает за транспорт; есть f, которое участвует в оценке качества этого транспорта. Остаётся задать T и f как нейросетевые функции и приступить к оптимизации.

В итоге приходим к задаче min-max оптимизации, чем-то схожей с GAN: одна нейросеть максимизирует полученный функционал, другая минимизирует. По f делается градиентный подъём, по T — спуск. Но по смыслу это не совсем генератор и дискриминатор в привычном виде — здесь это просто способ решить двойственную задачу. Все интегралы просто оцениваются через Монте-Карло без доптрюков.

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

В экспериментах стандартные примеры с переводом изображений сравниваются с GAN-подходами. Картинки выглядят адекватно, метрики тоже улучшаются, а сама модель получается компактнее по количеству компонентов.

Разбор подготовил Иван Папай
CV Time