В середине прошлого века советский математик Леонид Витальевич Канторович разрабатывал идеи оптимальных решений в экономике, включая проблему оптимальной перевозки грузов. Сегодня его наработки лежат в основе моделей оптимального транспорта (optimal transport, OT), в которых на смену распределениям ресурсов пришли вероятностные распределения для генерации изображений в генеративных нейронных сетях.
Мы уже рассказывали о том, каких успехов в этом направлении добилась команда «Обучаемый интеллект» AIRI, руководимая Евгением Бурнаевым. К примеру, не так давно исследователи смогли представить транспортную задачу, как поиск седловой точки некоторого функционала. Это позволило сделать оптимальный транспорт масштабируемым и теоретически обоснованным.
Недавно же группа обратила своё внимание на тот факт, что непрерывные OT-методы, как правило, полагаются на различные евклидовы дистанции, например, l1 или l2, традиционно используемые в задачах оптимизации. Несмотря на свою простоту и распространенность, подобные функции стоимости могут плохо работать в каких-то специфических случаях. Например, если в результате транспорта требуется сохранить класс объектов, евклидова функция может дать лишь субоптимальное решение.
Иллюстрация к транспортной задаче, в которой требуется сохранить класс объектов при переходе от одних данных (элементы одежды) к другим (рукописные цифры). В данном случае данные целевого представления размечены по классам лишь частично.
Чтобы справится с этой трудностью, авторы применили свои предыдущие наработки, заменив функцию стоимости на обобщенный функционал стоимости (general cost functional), который способен сохранять дополнительную информацию, такую, как метки классов. Таким образом, поиск оптимального транспорта с сохранением классов удалось свести к задаче о поиске седловой точки функционала.
Учёные продемонстрировали этот принцип с помощью функционалов двух типов: ориентированных на классы и пары в выборке. Первый алгоритм решал задачу поклассового сопоставления данных разного типа, второй — перевода изображения из одного домена в другой с учителем. Ориентированные на классы функционалы при этом были способны строить стохастические транспортные карты, при которых данные на входе содержат случайный шум.
Авторы проверили работоспособность своего подходя в экспериментах со множеством различных датасетов, начиная от простых синтетических и заканчивая биологическими данными. Значительная часть датасетов включала в себя пары изображений из разных доменов. Сравнение с методами на основе евклидовых функций показало, что использование обобщенного функционала стоимости не уступает, а в ряде случаев даже показывает лучшие метрики.
Сравнение различных методов в задаче транспорта с сохранением класса из FMNIST (изображения одежды) в MNIST (рукописные цифры). Последние две строчки соответствуют новому методу с добавлением шума z и без.
В заключении исследователи отмечают, что, хотя в задачах с сохранением классов новый метод показал себя хорошо, в иных случаях его следует использовать с осторожностью, поскольку построение функционала в этом случае может быть нетривиальным.
Результаты работы были представлены на конференции ICLR 2024, статья с деталями была опубликована в её сборнике.