Как несбалансированный оптимальный транспорт помог нам сделать поиск барицентров распределений устойчивым

Привет! Меня зовут Милена Газдиева, я являюсь научным сотрудником Института AIRI, а также инженером-исследователем и аспиранткой Сколтеха. Мои научные интересы лежат в области разработки генеративных моделей на основе оптимального транспорта (optimal transport, ОТ) и их приложений к различных задачам. Мы с коллегами добились успехов в повышении устойчивости таких моделей, и одна из наших статей по этой теме была принята на престижную конференцию по искусственному интеллекту ICLR 2025, которая в этом году будет проходить в Сингапуре. Сегодня я расскажу об этой работе, в рамках которой мы разработали метод оценки барицентров (взвешенных средних) распределений, устойчивый к различным выбросам и дисбалансам в данных.

Что это означает и зачем нужно — читайте далее.

Хотя концепция «барицентра» распределений была введена только в 2010-х годах, она тесно связана с теорией оптимального транспорта (ОТ), которая берет начало еще в XVIII веке в работах Гаспара Монжа, а в XX веке обобщается Леонидом Канторовичем. Говоря простым языком, ОТ барицентр формализует задачу поиска геометрически осмысленного «среднего» вероятностных распределений.

Рис.1. Схематическое представление того, как применять методы поиска барицентров для нахождения общего представления МРТ‑снимков одной модальности, полученных с разных сканеров. Картинка взята из нашей недавней статьи, опубликованной на ICML 2024.

Мы часто встречаем такую задачу на практике, когда хотим привести к общему представлению данные, полученные из разных источников. Задача поиска такого представления, не зависящего от источника получения данных — иначе говоря, задача «агрегации» данных, — может возникать в самых разных областях: от геологии, когда необходимо агрегировать данные, полученные с помощью различных симуляторов месторождений, до медицины, где данные поступают из разных госпиталей или от оборудования с различающимися характеристиками.

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

(a) Формулировка Монжа (b) Формулировка Канторовича
Рис. 2. Классические постановки задачи оптимального транспорта (ОТ). Картинка взята из нашей статьи, опубликованной на NeurIPS 2023.

Увы, в такой формулировке оно не всегда существует. Леонид Канторович предложил более общую постановку задачи, в которой ввел в рассмотрение транспортные планы – оптимальные способы перемещения массы с возможностью распределения массы каждой точки исходного распределения по произвольному числу точек целевого распределения. В отличие от ОТ отображения, которое является детерминированной функцией, ОТ план описывается как вероятностная мера (распределение) на декартовом произведении двух пространств, что позволяет учитывать случаи, когда оптимальное соответствие между распределениями не является однозначным. 

Заметим, что маргиналы (частные распределения) и ОТ планов обязаны совпадать с исходным и целевым распределениями соответственно, и далее мы увидим, что такое свойство влечет за собой ряд проблем. В нашей статье мы получали теоретические результаты в контексте планов (но не обычных, а полу‑несбалансированных), а на практике параметризовали условные (conditional) планы как стохастические или детерминистические отображения, поэтому далее, в зависимости от контекста, я буду оперировать обоими понятиями.

Как в задаче Монжа, так и в задаче Канторовича, стоимость эффективного перехода между распределениями называется ОТ расстоянием между ними. Важно понимать, что в зависимости от того, какие функции стоимости мы будем использовать, ОТ расстояния будут получаться разными. Тогда, строго говоря, оценкой ОТ барицентра заданных распределений называется задача нахождения «центрального» распределения, которое является ближайшим (в среднем) ко всем исходным распределениям с точки зрения ОТ расстояния.

Рис. 3. Визуализация задачи нахождения классического OT барицентра.

В последние годы исследователи начали активно разрабатывать методы вычисления ОТ барицентров непрерывных распределений. Мы с коллегами из AIRI и Сколтеха во главе с Александром Коротиным и Евгением Бурнаевым также занимаемся этим направлением и добились значительных успехов — наши статьи были опубликованы на недавних конференциях ICML 2024 (NOTB) и NeurIPS 2024 (EgBary).

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

Для начала скажем, что причиной выявленной проблемы является сама постановка задачи ОТ барицентра, которая основывается на понятии классического ОТ расстояния. Действительно, в концепции классического ОТ необходимо найти эффективный способ преобразования всего исходного распределения в целевое (вспомним, что левый маргинал ОТ плана должен совпадать с исходным распределением). Соответственно, шумы и выбросы изменяют ОТ расстояние между распределениями и итоговый барицентр. К счастью, недавно была предложена новая концепция полу‑несбалансированного транспорта (semi‑unbalanced optimal transport, SUOT).

Рис. 4. Задача полу‑несбалансированного транспорта (SUOT).

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

Рис. 5. Визуализация задачи нахождения полу‑несбалансированного (semi‑unbalanced) OT барицентра. Заметим, что задача ставится в контексте SUOT планов , маргиналы которых на самом деле отличаются от исходных распределений . На практике мы параметризовали условные (conditional) планы в качестве стохастических или детерминистических отображений .

В нашей работе мы разработали подход для оценки барицентра непрерывных распределений с точки зрения SUOT расстояния. Ключевая идея нашего подхода состоит в рассмотрении двойственной задачи полу‑несбалансированного транспорта и использовании нейросетей для параметризации транспортных отображений и потенциалов в данной задаче. Здесь потенциалы не несут какого‑то физического смысла, это лишь дополнительная «переменная», возникающая при переходе к двойственной задаче. В итоге, мы получили минимаксную задачу поиска седловой точки, которую решали с использованием алгоритмов градиентного подъема/спуска.

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

Стоит отметить, что мы ввели в SUOT расстояние параметр (в качестве веса дивергенции), который отвечает за близость исходного распределения и того распределения, с которым мы работаем. При маленьких , такое расстояние игнорирует возможные аномалии в данных. В то же время, если мы выкрутим в , что равносильно рассмотрению определенной дивергенции, то получим классическую ОТ задачу и потеряем свойство устойчивости. В предельном случае, наш алгоритм совпадает с методом для вычисления классических ОТ барицентров под названием NOTB, который мы предложили в статье, опубликованной на ICML 2024.

Так же как и NOTB, новый подход (при любых ) является теоретическим обоснованным, масштабируемым на сложные многомерные данные и способен работать с широким классом функций стоимости. Кроме того, его уникальным свойством стала способность игнорировать аномалии в данных! Благодаря свойствам SUOT задачи наш метод является устойчивым к шумам, выбросам и дисбалансам классов в исходных распределениях, что мы и показали экспериментально в нашей работе.

Яркий пример работы нашего метода — решение задачи нахождения барицентра цветов и черно‑белых цифр. Казалось бы, что сложного в нахождении такого барицентра? Но давайте присмотримся к деталям эксперимента.

Рис. 6. Полу‑несбалансированный барицентр распределений цветов () и цифр (), вычисленный с помощью нашего подхода в латентном пространстве StyleGAN модели, которая была предобучена на датасете окрашенных цифр «2» и «3» из датасета MNIST. Наш метод позволяет успешно игнорировать выбросы в исходных распределениях.

Здесь мы работаем в латентном пространстве StyleGAN'a, предобученного на генерацию раскрашенных цифр («2», «3») из датасета MNIST. В данном пространстве закодированы все разноцветные цифры, но нет белых! Зная этот факт, мы рассмотрели датасет цветов («красный» и «зеленый») и добавили небольшой процент «белого» цвета в качестве выброса, а также рассмотрели датасет цифр («2», «3»), добавив небольшой процент цифр «7». Нашей целью было продемонстрировать, что наш подход поиска барицентра будет игнорировать «белый» цвет в распределении цветов на этапе инференса, а также цифру «7» в датасете цифр. Мы показали это, посчитав acceptance rate для всех цветов и цифр. Помимо этого, мы показали свойства нашего метода в ряде синтетических экспериментов.

Рис. 7. Примеры (черно‑белых цифр) и (фото цветов) и их соответствующего барицентра (раскрашенных цифр).

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

Спасибо всем, кто дочитал до конца. Если у вас остались вопросы, с радостью на них отвечу!