Максим Бобрин
Экспектильная регуляризация десятикратно ускорила оптимальный транспорт
Подходы на основе оптимального транспорта (OT) позволили по-новому взглянуть на множество задач машинного обучения, начиная от генеративного моделирования и заканчивая доменной адаптацией. В них вычисления ведутся на языке вероятностных распределений, а также переходов между ними.
Задачей OT обычно является поиск оптимального отображения (оптимального плана) между некоторыми распределениями для данной функции стоимости. На практике оптимум ищут с помощью явного аналитического решения задачи Канторовича, аппроксимируя оптимальное решение через параметрическую функцию. Такой подход, однако, может демонстрировать нестабильность при поиске оптимальных потенциалов, один из которых является двойственным (сопряжённым) к другому по функции стоимости.
Грубая оценка двойственного потенциала приводит к ситуации, когда сумма потенциалов расходится, а точное решение приводит к дополнительной задаче оптимизации. Существующие методы решения этой проблемы требуют обширной настройки гиперпараметров нейросети, либо используют итеративный поиск более точного решения внутренней задачи на каждом шаге оптимизации, а потому вычислительно затратны.
Ученые из Института AIRI и Сколтеха предложили смягчить вышеуказанные проблемы, ограничив класс решений внутренней оптимизации сопряженных потенциалов с помощью специальной экспектильной регуляризации, которая является несимметричным обобщением подхода на основе среднеквадратической ошибки (MSE). Идея этого подхода основана в вынесении внутренней оптимизации (по сопряженному потенциалу) во внешнюю исходную задачу Канторовича. Новая задача оптимизации, обеспечивает стабильную и сбалансированную сходимость как для поиска оптимальных потенциалов, так и для аппроксимации транспортного плана. Новый метод получил название Expectile-Regularized Neural Optimal Transport (ENOT).
Авторы решили оценить эффективность работы ENOT, сравнив его с существующими методами с помощью популярного бенчмарка для построения OT-планов Wasserstein-2, в создании которого также принимали участие исследователи из AIRI. Бенчмарк предлагает моделям построить связь между несколькими синтетическими парными датасетами разных размерностей и наборами изображений, когда оптимальный план известен заранее. Исследователи экспериментировали с разными функционалами стоимости, а также оценили производительность своего метода в более сложных генеративных img2img-задачах.
Пример генерации 3D-объектов из шума с помощью нескольких методов, включая ENOT
Проверка показала, что новый метод достаточно точно восстанавливает транспортный план, но при этом ему требуется существенно меньше вычислительных затрат, времени на подбор параметров и числа итераций, чем бейзлайнам. В ряде случаев применение экспектильной регуляризации даёт 10-кратное ускорение по сравнению с SOTA-подходами.
ENOT реализован на фреймворке JAX и уже имплементирован в одну из самых популярных библиотек для вычисления OT ott-jax, что позволяет легко использовать метод. Авторы представили разработку на конференции NeurIPS 2024, где статью с описанием исследования определили в Spotlight –– специальный трек для работ, особо отмеченных рецензентами. Кратко ознакомится с методом можно по ссылке.
