Милена Газдиева
Несбалансированный оптимальный транспорт помог сделать поиск барицентров распределений устойчивым
В машинном обучении часто возникает задача агрегации — объединения данных, полученных из разных источников. Например, это могут быть МРТ‑снимки мозга, сделанные на разных приборах. При этом важно, чтобы информация, усвоенная моделью, не зависела от типа источника.
В решении задачи агрегации хорошо зарекомендовали себя методы на основе нахождения барицентра по оптимальному транспорту (Optimal Transport, OT). Под оптимальным транспортом понимается задача эффективного перехода от одного вероятностного распределения к другому. Понятие барицентра возникает в тот момент, когда необходимо вычислить геометрически среднее по нескольким распределениям.
За последние несколько лет наметился значительный прогресс в разработке методов вычисления ОТ барицентров, что позволило адаптировать их к многомерным данным, таким как изображения. Однако существующие методы, как правило, игнорируют типичные и вместе с тем важные особенности реальных данных — шумы и выбросы.
На этот недостаток обратила внимание команда исследователей из AIRI и Сколтеха совместно с коллегами из Технологического института Джорджии и Университета Сонгюнгван. Чтобы найти робастный подход на основе OT барицентра, который бы наследовал все достижения в этой области, они обратились к несбалансированному оптимальному транспорту, который естественным образом подходит к работе с неидеальными данными.
В классическом оптимальном транспорте маргиналы (частные распределения) OT плана совпадают с исходным и целевым распределениями. Идея несбалансированного OT опирается на отказ от такого соответствия, что позволяет гибче работать с шумами и выбросами. Авторы исследования остановили свой выбор на полунесбалансировнном оптимальном транспорте (semi‑unbalanced optimal transport, SUOT), в котором маргинал не совпадает лишь с начальным распределением.
Суть новой работы заключается в разработке подхода для оценки барицентра непрерывных распределений с точки зрения SUOT расстояния. Исследователи теоретически показали, что эта задача может быть сведена к минимаксной задаче поиска седловой точки и предложили использовать для её решения нейросети.
Авторы провели серию экспериментов с несколькими популярными датасетами, включая датасет MNIST с цветными цифрами. В последнем случае учёные работали в латентном пространстве StyleGAN модели, которая была предобучена на датасете красных и зелёных цифр «2» и «3». Они добавили небольшой процент белого цвета в качестве выброса, а также рассмотрели датасет цифр («2», «3»), добавив небольшой процент цифр «7».
Проценты допуска (acceptance rate) для разных цифр и цветов, полученные в ходе эксперимента. Справа — результаты для выбросов.
Исследователи показали, что SUOT подход поиска барицентра эффективно игнорирует выбросы на этапе инференса. По их словам, предложенный метод открывает новые возможности для масштабируемых вычислений и решения широкого спектра задач, где требуется надежное усреднение сложных вероятностных распределений.
Статья с описанием метода и экспериментов была принята на конференцию ICLR 2025 и опубликована в её сборнике.
