Оптимальный транспорт
Метрики Вассерштейна
KL-дивергенция между двумя гауссианами без перекрытия равна бесконечности. Wasserstein $W_2$ равна расстоянию между их центрами. Именно поэтому Arjovsky в 2017 году заменил JS на $W_1$ в GAN - и обучение, которое раньше схлопывалось за несколько итераций, внезапно стало стабильным. Четыре страницы статьи. Одно число в loss-функции. Конец эпохи mode collapse.
- **WGAN (Arjovsky et al., 2017)**: замена JS-дивергенции на $W_1$ стабилизировала обучение GAN. Critic с gradient penalty аппроксимирует оптимальную 1-Lipschitz функцию через двойственность Канторовича-Рубинштейна.
- **Flow matching (Lipman, 2022)**: FLUX.1, Stable Diffusion 3 используют OT-геодезические для прямолинейных путей в пространстве мер. Нейросеть обучается предсказывать постоянную скорость вдоль этих путей - быстрее диффузии.
- **Sinkhorn / GeomLoss**: энтропийная регуляризация $\varepsilon H(\gamma)$ делает $W_p$ дифференцируемым за $O(n^2 \log n)$. Библиотека GeomLoss от Feydy реализует это на GPU - Wasserstein как loss в нейросетях.
- **Domain adaptation**: DANN и аналоги минимизируют $W_p$ между признаковыми пространствами источника и цели. Wasserstein-distance как мера сдвига доменов.
- **Сравнение облаков точек**: PointNet++ и 3D reconstruction используют Chamfer distance - аппроксимацию $W_1$ для облаков точек в робототехнике и LIDAR-обработке.
Предварительные знания
W_p: infimum над транспортными планами
KL-дивергенция между $\mathcal{N}(0,1)$ и $\mathcal{N}(100,1)$ равна $D_{KL} = 5000$. Расстояние Вассерштейна $W_2$ между ними равно ровно 100 - расстоянию между центрами. Эти два числа не просто разные по величине. Они описывают разные геометрии. И для обучения генеративных моделей правильная геометрия - Wasserstein.
**Расстояние Вассерштейна-p** - минимальная стоимость транспортировки по всем допустимым планам: $$W_p(\mu, \nu) = \left(\inf_{\gamma \in \Gamma(\mu, \nu)} \int_{X \times Y} \|x - y\|^p \, d\gamma(x,y)\right)^{1/p}$$ $\Gamma(\mu,\nu)$ - множество совместных мер с маргиналями $\mu$ и $\nu$. При $p=1$ получается earth mover's distance - буквально стоимость перемещения земли. При $p=2$ - квадратичная задача, связанная с теоремой Брение.
Earth mover's distance - интуиция: есть два распределения, которые мысленно превращаются в рельеф из куч земли. $W_1$ - это минимальная работа по превращению одного рельефа в другой. Каждую единицу массы надо перенести на какое-то расстояние, и сумма $\text{масса} \times \text{расстояние}$ по всем переносам и есть $W_1$.
В 1D вычисление $W_p$ сводится к сортировке: если $F_\mu$ и $F_\nu$ - функции распределения, то $W_p^p = \int_0^1 |F_\mu^{-1}(t) - F_\nu^{-1}(t)|^p dt$. Это $O(n \log n)$ - просто сортировать обе выборки и взять среднюю разность. В многомерном случае задача становится полноценной оптимизацией, но в 1D Wasserstein считается как порядковая статистика.
**Sliced Wasserstein** - трюк для многомерных данных: спроецировать оба распределения на случайные 1D-направления, посчитать $W_p$ для каждой проекции (сортировкой) и усреднить. Аппроксимирует полный $W_p$ за $O(Ln\log n)$, где $L$ - число проекций. GeomLoss от Feydy реализует это на GPU.
Какое значение даст $W_1$ между $\delta_0$ (точечная масса в нуле) и $\delta_5$ (точечная масса в пяти)?
Wasserstein vs KL: разные геометрии
2017 год. Arjovsky, Chintala и Bottou публикуют WGAN. Главный аргумент четырёх страниц: JS-дивергенция (и KL) между распределением данных $p_{data}$ и распределением генератора $p_g$ почти везде равна константе, пока носители не пересекаются. Нет сигнала - нет градиента - mode collapse. Замена одного числа в loss на $W_1$ и обучение внезапно стабилизируется.
**Ключевое различие**: KL не является метрикой. Она асимметрична ($D_{KL}(p\|q) \neq D_{KL}(q\|p)$), нарушает неравенство треугольника и равна бесконечности когда $\text{supp}(p) \not\subseteq \text{supp}(q)$. Wasserstein $W_p$ - настоящая метрика на пространстве вероятностных мер с конечным $p$-м моментом. Она метризует слабую сходимость: $W_p(\mu_n, \mu) \to 0$ тогда и только тогда, когда $\mu_n \Rightarrow \mu$ и $p$-е моменты сходятся.
Почему KL взрывается на непересекающихся носителях? $D_{KL}(p \| q) = \int p(x) \log \frac{p(x)}{q(x)} dx$. Если $p(x) > 0$ в точке, где $q(x) = 0$ - интеграл расходится. В начале обучения GAN носители почти наверняка не пересекаются: генератор производит случайный шум, данные - структурированные изображения. Именно в этот критический момент KL отказывает как сигнал. Wasserstein даёт конечное расстояние и ненулевой градиент всегда.
Дополнительное преимущество - метризация слабой сходимости. Последовательность мер $\delta_{1/n}$ при $n \to \infty$ слабо сходится к $\delta_0$. $W_1(\delta_{1/n}, \delta_0) = 1/n \to 0$ - правильно. KL$( \delta_{1/n} \| \delta_0) = +\infty$ при любом $n$ - бесполезно. Для диффузионных моделей, где цепочка шагов шума описывается как сходимость мер, Wasserstein - естественный язык анализа.
Почему замена JS-дивергенции на $W_1$ в WGAN стабилизирует обучение?
W_2-геодезические и McCann-интерполяция
$W_2$ не просто расстояние. Это метрика на бесконечномерном пространстве вероятностных мер, и у неё есть геодезические - кратчайшие пути между двумя мерами. McCann в 1997 году показал: оптимальный путь из $\mu$ в $\nu$ по $W_2$ имеет явный вид. Этот результат стал фундаментом для flow matching и score-based generation 2022 года.
**McCann-интерполяция**: если $T^* = \nabla \phi$ - оптимальная карта Монжа (градиент выпуклой функции $\phi$, теорема Брение), то геодезическая из $\mu$ в $\nu$ параметризуется: $$\mu_t = \left((1-t) \cdot \text{id} + t \cdot \nabla\phi\right)_\# \mu, \quad t \in [0,1]$$ При $t=0$ получается $\mu$, при $t=1$ - $\nu$. Каждая частица $x$ движется по прямой из $x$ в $T^*(x)$ с постоянной скоростью. Это прямолинейное движение в пространстве мер.
Flow matching (Lipman et al., 2022; Liu et al., 2022) - прямое применение McCann-геодезических. Вместо того чтобы обучать score-функцию $\nabla \log p_t$ как в диффузионных моделях, нейросеть обучается предсказывать скоростное поле $v_t(x)$, которое описывает движение частиц вдоль геодезической. При использовании OT-путей (optimal transport conditional flow matching) траектории прямые, скорость постоянная - и обучение значительно быстрее. FLUX.1 и Stable Diffusion 3 реализуют именно это.
Геодезические $W_2$ также связаны с score-matching. Теорема Отто-Виллани описывает пространство мер с метрикой $W_2$ как риманово многообразие. «Танец» между score-based generation и optimal transport - это две языка для одной геометрии. Диффузионные шаги Ланжевена описывают градиентный поток функционала KL в метрике $W_2$. Уравнение Фоккера-Планка - это градиентный спуск по $W_2$. Это не случайность.
Wasserstein - это просто другой способ вычислить расстояние между распределениями, принципиально не отличающийся от KL
Wasserstein и KL - разные математические объекты с разными областями применения. KL измеряет информационное расхождение (нет метрических свойств, взрывается на непересекающихся носителях). $W_p$ - геометрическое расстояние, метрика, учитывающая структуру пространства.
На непересекающихся носителях KL бесполезна для оптимизации: градиент нулевой. $W_p$ даёт осмысленный сигнал. Это практически важно для ранних стадий обучения генеративных моделей, когда носители далеко друг от друга.
McCann-интерполяция $\mu_t = ((1-t)\cdot\text{id} + t \cdot \nabla\phi)_\# \mu$ задаёт геодезическую в каком пространстве?
Ключевые идеи
- **$W_p(\mu,\nu) = \left(\inf_{\gamma \in \Gamma(\mu,\nu)} \int \|x-y\|^p d\gamma\right)^{1/p}$** - earth mover's distance. Настоящая метрика. Конечна там, где KL = $\infty$.
- **KL vs Wasserstein**: KL взрывается на непересекающихся носителях, нарушает аксиомы метрики. $W_p$ даёт осмысленный градиент всегда - именно поэтому WGAN работает там, где классический GAN падает.
- **McCann-геодезические**: путь $\mu_t = ((1-t)\cdot\text{id} + t \cdot \nabla\phi)_\#\mu$ - кратчайший в $W_2$. Частицы движутся прямолинейно. Flow matching обучает нейросеть аппроксимировать эти пути.
- **Диффузионные модели как OT**: градиентный поток KL в метрике $W_2$ - это уравнение Фоккера-Планка. Score matching и optimal transport - два языка одной геометрии.
Связанные темы
Метрики Вассерштейна - ядро топика. Дальше строится теория и практика:
- Sinkhorn-алгоритм — Энтропийная регуляризация делает W_p вычислимым и дифференцируемым
- Wasserstein GAN — Прямое применение W_1 через двойственность Канторовича как loss
- Flow matching — W_2-геодезические McCann как пути обучения FLUX и SD3
- Двойственность Канторовича — Полная теория Lipschitz-потенциалов и теорема Брение
- KL-дивергенция и взаимная информация — Информационный аналог - другая геометрия на пространстве мер
Вопросы для размышления
- WGAN требует, чтобы critic был 1-Lipschitz. Weight clipping обрезает веса в $[-c, c]$. Почему это приближает Lipschitz-ограничение? В чём принципиальная проблема этого подхода по сравнению с gradient penalty?
- Sinkhorn добавляет энтропийный член $\varepsilon H(\gamma)$ к стоимости транспорта. При $\varepsilon \to 0$ получается точный Wasserstein. При $\varepsilon \to \infty$ - план $\gamma$ стремится к независимому произведению $\mu \otimes \nu$. Почему максимум энтропии при фиксированных маргиналях - это именно независимое произведение?
- Flow matching с OT-путями обучает нейросеть предсказывать постоянную скорость $v_t(x) = T^*(x) - x$. Диффузионные модели предсказывают score $\nabla \log p_t(x)$. Оба подхода описывают одну геодезическую в пространстве мер - но разными словами. Как связаны score и скоростное поле через уравнение непрерывности?
Связанные уроки
- ot-07-wgan — WGAN заменяет JS на W_1 - прямое применение метрики
- ot-04-sinkhorn — Sinkhorn делает W_p вычислимым и дифференцируемым
- ot-11-flow-matching — Flow matching использует W_2-геодезические как пути обучения
- ot-05-dual — Двойственность Канторовича - полная теория Lipschitz-потенциалов
- it-03 — KL-дивергенция - информационная альтернатива Wasserstein
- calc-01-sequences