Wasserstein Gradient Flows

Adam, SGD, LBFGS - всё это дискретные траектории в $\mathbb{R}^n$. Wasserstein gradient flow - тот же градиентный спуск, но в пространстве распределений: каждая итерация двигает не точку, а целое распределение по геодезической в метрике Вассерштейна. Diffusion models, SVGD, нормализующие потоки - всё это разные дискретизации одной непрерывной картины. McCann описал её в 1997, JKO - в 1998. Спустя 25 лет она оказалась языком всей современной генерации.

• DDPM/Stable Diffusion: training implicitly выполняет JKO updates на distribution distance
• SVGD (2016): Bayesian neural network posteriors как Wasserstein gradient flow
• Neural ODE / Flow matching: параметризация gradient flow через нейросеть для генерации

Предварительные знания

• Wasserstein-2 метрика и её свойства (ot-03-wasserstein)
• Schrödinger Bridges и стохастические процессы (ot-13-schrodinger)
• Энтропийная регуляризация и Sinkhorn (ot-04-sinkhorn)

• Wasserstein дистанция
• Мосты Шрёдингера
• Sinkhorn

Gradient flows: минимизирующие движения

Классический gradient flow в $\mathbb{R}^n$: $\dot{x} = -\nabla f(x)$. Точка катится по ландшафту функции $f$ в направлении наискорейшего убывания. Wasserstein gradient flow - та же идея, но **точка живёт в пространстве вероятностных распределений**, а метрика - $W_2$. Каждая итерация двигает целое распределение по геодезической в пространстве мер.

Минимизирующие движения (De Giorgi, 1993): дискретно-временная схема $\rho^{n+1} = \arg\min_\rho \left[F(\rho) + \frac{1}{2\tau} W_2^2(\rho, \rho^n)\right]$. На каждом шаге - компромисс между уменьшением функционала $F$ и расстоянием до предыдущего распределения. При $\tau \to 0$ дискретная траектория сходится к непрерывному PDE: $\partial_t \rho = \nabla \cdot (\rho \nabla \frac{\delta F}{\delta \rho})$.

Главная драма: один и тот же PDE можно получить как gradient flow разных функционалов в разных метриках. Heat equation - gradient flow энтропии в $W_2$, но также gradient flow Дирихле в $L^2$. Метрика определяет геометрию процесса.

Канонические примеры: $F(\rho) = \int \rho \log \rho \, dx$ даёт уравнение теплопроводности; $F(\rho) = \int \rho^m \, dx$ - porous medium equation; функционал взаимодействия $\iint W(x-y)\rho(x)\rho(y) dx dy$ - агрегационные модели типа Keller-Segel.

ML hook: Stein Variational Gradient Descent (Liu, Wang 2016) - это sampling алгоритм, формализованный как Wasserstein gradient flow on KL. Применяется для апостериорных распределений в Bayesian neural networks.

JKO: дискретизация по времени для PDE

Jordan-Kinderlehrer-Otto (1998) - культовая работа: уравнение Фоккера-Планка $\partial_t \rho = \nabla \cdot (\rho \nabla V) + \beta^{-1} \Delta \rho$ есть **Wasserstein gradient flow свободной энергии** $F(\rho) = \int V \rho \, dx + \beta^{-1} \int \rho \log \rho \, dx$. Потенциальная часть + энтропия = свободная энергия из статистической физики.

Частный случай $V = 0$: уравнение теплопроводности есть gradient flow чистой энтропии в $W_2$. Это связь между диффузией и максимизацией энтропии: тепло распространяется не случайно, а по геодезическим в пространстве распределений.

Связь с оптимизацией: Langevin sampling $dX_t = -\nabla V(X_t) dt + \sqrt{2\beta^{-1}} dW_t$ есть JKO scheme + injection шума. На уровне распределений это тот же gradient flow свободной энергии.

Численная сторона: внутренний шаг JKO - задача OT с регуляризацией. Sinkhorn внутри JKO даёт масштабируемый алгоритм. Современная альтернатива - neural JKO: параметризовать каждый шаг нейросетью и тренировать end-to-end.

ML hook: training diffusion model (DDPM, score-matching) неявно выполняет JKO updates. Forward process - это Langevin = Wasserstein gradient flow. Reverse process - обратный по времени gradient flow на KL до распределения данных. Stable Diffusion - это JKO в латентном пространстве.

Displacement convexity: выпуклость на путях OT

Интерполяция McCann (1997): между $\rho_0$ и $\rho_1$ через OT-карту $T$ строится путь $\rho_t = ((1-t)\,\text{Id} + tT)_\# \rho_0$. Это **геодезическая в Wasserstein-пространстве** - аналог прямой между двумя точками, только теперь точки - распределения, а прямая - оптимальная транспортная интерполяция.

Functional $F$ называется **displacement convex**, если он выпуклый вдоль Wasserstein-геодезических: $F(\rho_t) \le (1-t) F(\rho_0) + t F(\rho_1)$. Это эквивалент сильной структурной гарантии: gradient flow по такому $F$ ведёт себя предсказуемо, без бифуркаций и патологий.

Канонические displacement convex функционалы: внутренние энергии $\int U(\rho) dx$ при $U$ выпуклой и неубывающей, потенциальные $\int V \rho \, dx$ при $V$ выпуклой, $L^p$-нормы. Энтропия $\int \rho \log \rho$ - displacement convex в $\mathbb{R}^n$ и логарифмически вогнутых пространствах.

Из displacement convexity выводятся **функциональные неравенства**: log-Sobolev, неравенство Talagrand transportation cost ($W_2^2 \le 2 D_{KL}$ для гауссовских мер), Бакри-Эмери. Эти неравенства - инструменты для оценок скорости сходимости MCMC и Langevin sampling.

ML hook: Wasserstein-2 регуляризатор в WGAN с gradient penalty приводит к displacement convex loss landscape (в идеализированном пределе). Score-based diffusion: KL до целевого распределения displacement convex вдоль $W_2$-геодезических при разумных условиях, что даёт глобальную сходимость training.

Production: SVGD, neural ODE, диффузионные модели

**SVGD** (Stein Variational Gradient Descent, Liu & Wang 2016): kernelized Wasserstein gradient flow on KL. Частицы $\{x_i\}$ движутся $\dot{x}_i = \frac{1}{N}\sum_j [k(x_j, x_i)\nabla \log p(x_j) + \nabla_{x_j} k(x_j, x_i)]$ - первый член тянет в моду $p$, второй разводит частицы. В пределе - аппроксимация Wasserstein gradient flow on KL до целевого $p$.

**Neural ODE / Continuous Normalizing Flows / Flow Matching**: параметризовать векторное поле $v_\theta(x, t)$ нейросетью, обучать так, чтобы траектория $\dot{x}_t = v_\theta(x_t, t)$ переводила базовое распределение в целевое. Flow matching обучает $v_\theta$ имитировать OT-карту - детерминированная версия gradient flow.

Information-geometric связь: Fisher metric на параметрическом семействе - инфинитезимальная версия $W_2$ при определённых условиях. Natural gradient на параметрах ~ Wasserstein flow на распределениях. Это объясняет, почему second-order методы (K-FAC, Shampoo) работают: они приближают geometric flow.

Статистическая физика: равновесная статистическая механика = Wasserstein gradient flow при $t \to \infty$. Скорость сходимости MCMC выводится через curvature bounds (условие Бакри-Эмери: гессиан потенциала $\succeq \lambda I$ влечёт экспоненциальную сходимость с rate $\lambda$).

Распространённое заблуждение: gradient flow ассоциируют с forward diffusion в DDPM. На самом деле наоборот: forward diffusion идёт **против** gradient flow KL-функционала (увеличивает энтропию, отдаляется от данных). Sampling/denoising - это и есть forward gradient flow по KL до распределения данных.

ML production: noising schedule в Stable Diffusion соответствует Wasserstein-геодезической на гауссовом многообразии. Score-based модели тренируются вдоль траекторий Wasserstein gradient flow. Диффузионные модели текста (D3PM, MaskGIT) - дискретный аналог JKO.

Ключевые идеи

• Wasserstein gradient flow - это градиентный спуск в пространстве вероятностных мер с метрикой $W_2$, реализуемый через minimizing movements (De Giorgi)
• JKO scheme (1998): уравнение Фоккера-Планка есть gradient flow свободной энергии $\int V\rho + \beta^{-1}\int \rho \log\rho$ в $W_2$
• Displacement convexity вдоль McCann-интерполяции даёт сильные функциональные неравенства (log-Sobolev, Talagrand) и гарантии сходимости
• SVGD, neural ODE, flow matching и диффузионные модели - разные дискретизации Wasserstein gradient flow для генеративного моделирования

Что это разблокирует

Wasserstein gradient flows ведут в:

• Schrödinger Bridges — SB - стохастический gradient flow
• Wasserstein vs Fisher — Две метрики на пространстве распределений
• Natural gradient — Аналог Wasserstein flow в Fisher геометрии

Вопросы для размышления

• Почему JKO scheme выбирает свободную энергию (потенциал + энтропия), а не один потенциал, как functional для gradient flow Фоккера-Планка?
• Как displacement convexity функционала KL связана со скоростью сходимости Langevin sampling и MCMC?
• В чём принципиальная разница между forward diffusion в DDPM и Wasserstein gradient flow на KL до распределения данных?

Связанные уроки

• ot-13-schrodinger — SB is stochastic version of gradient flow
• ot-03-wasserstein — W2 metric is the geometry for the flow
• ig-07-natural-gradient — Natural gradient = Fisher version of Wasserstein flow
• ig-11-wasserstein-vs-fisher — Compares the two metrics for gradient flows
• ot-11-flow-matching — Flow matching parametrizes deterministic version