Оптимальный транспорт: дешевле всего перевезти одну кучу в другую

Есть склады с песком и стройки, которым он нужен. Сколько песка с какого склада на какую стройку везти, чтобы возить пришлось как можно меньше? Из этого простого на вид вопроса вырос целый раздел математики, а заодно и способ, которым современный ИИ измеряет расстояние между распределениями.

Задача и её цена

Пусть в источнике (складе) i лежит масса a_i, а стоку (стройке) j нужна масса b_j, и перевозка единицы из i в j стоит C_{ij} (например, квадрат расстояния). Суммарный запас при этом равен суммарной потребности — иначе ограничения ниже просто несовместны. План перевозки — матрица P_{ij}: сколько массы едет из i в j. Нужно минимизировать суммарную стоимость \min_P \sum_{ij} P_{ij} C_{ij}, \qquad \sum_j P_{ij}=a_i,\quad \sum_i P_{ij}=b_j,\quad P_{ij}\ge 0. Это задача линейного программирования (ЛП): и целевая функция, и ограничения линейны. Её минимум — расстояние Вассерштейна между двумя распределениями массы: чем дороже обходится самая дешёвая перевозка, тем эти распределения дальше друг от друга. С одной оговоркой про C: если C — само расстояние, минимум и есть W_1; если C — квадрат расстояния (как в примере выше), минимум даёт уже W_2^2, а само W_2 — корень из него. В общем виде W_p — корень степени p из минимальной стоимости при C, равном расстоянию в степени p.

Задача старая. Гаспар Монж поставил её в 1781 году в работе о выемках и насыпях: как перевезти грунт с наименьшей работой. У него каждая точка источника едет целиком в одну точку стока — такого отображения может и не существовать. Леонид Канторович в 1942-м разрешил делить массу между несколькими стоками: вместо отображения появился план P. Именно эта, более мягкая постановка записана выше, и именно она оказалась линейной программой.

Синкхорн вместо симплекса

Решать такую ЛП симплекс-методом дорого. Но если добавить к задаче чуть энтропии (-\varepsilon H(P)), решение становится гладким и считается совсем просто — алгоритмом Синкхорна: берём матрицу K_{ij}=e^{-C_{ij}/\varepsilon} и поочерёдно нормируем её строки и столбцы, пока суммы не сойдутся к нужным a и b. Каждый шаг — это матричное умножение, и шагов обычно нужно немного.

Матрица K и нормировки взялись не с потолка. С энтропийной добавкой задача становится строго выпуклой, а условие оптимальности заставляет решение иметь вид P_{ij}=u_i\,K_{ij}\,v_j — «диагональ × K × диагональ», где как раз K_{ij}=e^{-C_{ij}/\varepsilon}. Значит, неизвестных остаётся всего два вектора, u и v, и подбирают их ровно условия на суммы строк и столбцов — то есть поочерёдная нормировка.

Параметр \varepsilon управляет резкостью плана. При \varepsilon\to 0 план стремится к жёсткому ЛП-оптимуму: каждый источник почти целиком обслуживает один сток. При больших \varepsilon масса размазывается по многим стокам, а стоимость растёт.

Тяни точки и крути \varepsilon: видно, как план перестраивается, а толщина линий показывает, сколько массы едет по каждому маршруту.

Зачем это ИИ

Расстояние Вассерштейна сравнивает распределения, учитывая геометрию. Поэлементные метрики этого не умеют: два изображения, сдвинутые на пиксель, по среднеквадратичной ошибке «далеки», хотя глазом неотличимы; по Вассерштейну — близки, потому что массу надо лишь чуть подвинуть. Отсюда применения:

  • Wasserstein GAN — обучает генератор, минимизируя расстояние Вассерштейна W_1 между настоящим и сгенерированным распределениями; градиент ведёт себя лучше, чем у обычного GAN.
  • Сравнение и выравнивание — гистограмм цветов, векторных представлений слов разных языков, данных по отдельным клеткам в биологии.
  • Диффузия и flow matching — современные генеративные модели можно понимать как перенос шумового распределения в распределение данных по оптимальным траекториям.

Платим за это вычислениями. Точное решение ЛП стоит порядка O(n^3\log n) операций на n точек, Синкхорн дешевле — O(n^2) на итерацию, но матрицу n\times n всё равно надо держать в памяти. Хуже другое: чтобы честно оценить расстояние по выборке в пространстве большой размерности, точек нужно тем больше, чем выше размерность, и растёт это требование экспоненциально. Поэтому на практике Вассерштейн почти всегда считают приближённо — по мини-батчам, по одномерным срезам или малоранговыми методами.

Связи

Наверх