Wasserstein барицентры

Взять два рукописных изображения цифры «2» и усреднить попиксельно - результат мутное серое пятно. Нейросеть на таких данных учится плохо. Усреднить те же изображения в пространстве Вассерштейна - получить морфинг: промежуточную чёткую форму, сохраняющую контур. Agueh & Carlier (2011) показали: это Фреше-среднее в $(\mathcal{P}_2, W_2)$ - барицентр. Применений у него от federated learning до уравнений диффузии.

• **Shape interpolation**: геодезические в W2 между формами органов (МРТ сердца) - промежуточные формы сохраняют топологию. Атласы форм в медицинской визуализации
• **Color transfer и grading**: Sinkhorn-барицентр RGB гистограмм нескольких кадров выравнивает цветовую палитру. DaVinci Resolve использует схожий подход для color grading
• **Federated learning**: Wasserstein-агрегация весов клиентских моделей устойчивее FedAvg при гетерогенных данных (разные распределения классов у разных клиентов)
• **Уравнения диффузии**: JKO-схема (Jordan-Kinderlehrer-Otto 1998) - каждый шаг Fokker-Planck уравнения это проксимальный шаг в W2. Молекулярная динамика и ланжевеновская оптимизация нейросетей

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

• Wasserstein-расстояние $W_2$ и его смысл как стоимость оптимального транспорта
• Алгоритм Sinkhorn: энтропийная регуляризация и матричные итерации

• Wasserstein-расстояния
• Алгоритм Sinkhorn

Фреше-среднее в метрическом пространстве

Взять два рукописных изображения цифры «2» и усреднить попиксельно - получится мутное серое пятно. Пиксели двух изображений геометрически не совмещены, среднее арифметическое смешивает их как краски. Если усреднять те же изображения в пространстве Вассерштейна, происходит другое: масса каждой буквы плавно перетекает в промежуточную форму, сохраняя контур. Это не blur - это морфинг.

Формализация через Фреше (1948): в $\mathbb{R}^n$ среднее нескольких точек - это $\arg\min_x \sum_i \|x - x_i\|^2$. В произвольном метрическом пространстве $(M, d)$ то же определение: $\arg\min_{p \in M} \sum_i d^2(p, p_i)$. Это **Фреше-среднее**. В пространстве вероятностных мер $(\mathcal{P}_2, W_2)$ получаем задачу Вассерштейна барицентра - Agueh & Carlier, 2011:

Веса $\lambda_i$ управляют «притяжением» к каждой мере: $\lambda_1 = 1$ даёт $\bar{\mu} = \mu_1$, $\lambda_i = 1/n$ - равновзвешенное среднее. Для мер на прямой $\mathbb{R}^1$ барицентр вычисляется аналитически через квантильные функции $Q_i$:

Для гауссиан $\mu_i = \mathcal{N}(m_i, \sigma_i^2)$ на прямой барицентр тоже гауссиан: $\bar{\mu} = \mathcal{N}(\bar{m}, \bar{\sigma}^2)$ где $\bar{m} = \sum \lambda_i m_i$ и $\bar{\sigma} = \sum \lambda_i \sigma_i$. Среднее берётся по $\sigma$, не по $\sigma^2$. Это отличает Вассерштейн-геометрию от евклидовой.

**Применение в генеративных моделях.** Латентное пространство VAE - это пространство мер. Интерполяция двух изображений через Wasserstein barycenter ($\lambda_1 = t$, $\lambda_2 = 1-t$) даёт гладкий морфинг по $t \in [0,1]$. Линейная интерполяция латентных векторов - приближение; правильная интерполяция - Фреше-среднее в $(\mathcal{P}_2, W_2)$. FID (Fréchet Inception Distance) неявно использует ту же геометрию через $W_2$ между гауссианами активаций.

Теорема Agueh-Carlier: мультимаргинальный OT

Agueh & Carlier (2011) доказали фундаментальный результат: задача Вассерштейна барицентра эквивалентна **мультимаргинальному** транспортному плану. Вместо $n$ попарных OT-задач $W_2(\mu, \mu_i)$ существует единый план $\gamma$ на произведении $\mu_1 \times \mu_2 \times \cdots \times \mu_n$.

Барицентр восстанавливается из плана: применить к каждому «атому» $(x_1, \ldots, x_n)$ взвешенное среднее $\sum_i \lambda_i x_i$ как push-forward $\gamma$. Каждый атом плана вносит в барицентр точку $\sum_i \lambda_i x_i$ с весом $\gamma(x_1, \ldots, x_n)$.

**Существование и единственность.** Если хотя бы одна из мер $\mu_i$ абсолютно непрерывна (имеет плотность), то барицентр существует и единственен. Оптимальные карты $T_i: \bar{\mu} \to \mu_i$ являются градиентами выпуклых функций (следствие теоремы Бренье). Без абсолютной непрерывности единственность может нарушаться - дискретный случай существенно сложнее.

Для гауссиан $\mu_i = \mathcal{N}(m_i, \Sigma_i)$ барицентр тоже гауссиан $\mathcal{N}(\bar{m}, S)$ где $\bar{m} = \sum \lambda_i m_i$ и матрица $S$ - фиксированная точка уравнения Бура:

Это нелинейное матричное уравнение, решаемое простыми итерациями. Связь с метрикой Бура: $W_2^2(\mathcal{N}(m_1,\Sigma_1), \mathcal{N}(m_2,\Sigma_2)) = \|m_1 - m_2\|^2 + B^2(\Sigma_1, \Sigma_2)$ где $B^2(\Sigma_1, \Sigma_2) = \text{tr}(\Sigma_1 + \Sigma_2 - 2(\Sigma_1^{1/2} \Sigma_2 \Sigma_1^{1/2})^{1/2})$.

Матричный барицентр Бура применяется в нейровизуализации: каждый субъект fMRI-исследования задаётся ковариационной матрицей активности мозговых регионов. Барицентр группы субъектов в метрике Бура - геометрически корректное групповое среднее, сохраняющее структуру SPD-многообразия.

Sinkhorn барицентры (Cuturi-Doucet 2014)

Точное решение задачи барицентра через мультимаргинальный OT имеет сложность $O(m^n \cdot d)$ где $m$ - число атомов каждой меры и $n$ - число мер. При $n = 10$ распределениях по $m = 64$ пикселей это $64^{10} \approx 10^{18}$ операций - невычислимо. Нужна регуляризация.

Cuturi & Doucet (2014) предложили энтропийно-регуляризованный барицентр: заменить $W_2$ на $W_\varepsilon$ (Sinkhorn-расстояние) в задаче Фреше-среднего. Это превращает задачу в итерационный алгоритм, где каждый шаг - стандартный Sinkhorn.

Алгоритм Sinkhorn-барицентра чередует два шага. Шаг 1: для каждой пары $(\bar{\mu}, \mu_i)$ запустить Sinkhorn и получить двойственные переменные $(u_i, v_i)$. Шаг 2: обновить $\bar{\mu}$ как геометрическое среднее произведений $v_i$ (в логарифмическом пространстве - взвешенное среднее $\log v_i$). Сложность одной итерации: $O(m^2 \cdot n)$ вместо $O(m^n)$.

**Параллелизм.** Для дискретных мер на одной сетке все $n$ Sinkhorn задач разделяют одну матрицу ядра $K = e^{-C/\varepsilon}$. Это позволяет батчить все $n$ транспортных задач в одну матричную операцию. На GPU Sinkhorn-барицентр $n = 100$ гистограмм по $64 \times 64$ вычисляется за секунды.

Color transfer через Sinkhorn-барицентр: взять гистограммы R, G, B трёх кадров (разные условия съёмки), вычислить барицентр - получить компромиссную цветовую палитру. Применяется в кино для color grading: несколько съёмочных дней выравниваются к единому тону через барицентр распределений пикселей.

Применения: shape analysis, federated learning, gradient flows

**Shape interpolation.** Представить 2D фигуру как меру $\mu$ (равномерную на контуре). Тогда геодезическая в $(\mathcal{P}_2, W_2)$ между $\mu_0$ и $\mu_1$ даёт последовательность промежуточных фигур $\mu_t = \bar{\mu}(t, 1-t)$. Масса перетекает кратчайшим путём, не смешиваясь попиксельно. Применение: медицинская визуализация (атласы форм органов), анимация в графике, статистика форм в биологии.

**Federated learning через барицентр.** Стандартный FedAvg агрегирует веса нейросетей как евклидово среднее: $\bar{\theta} = \sum \lambda_i \theta_i$. При гетерогенных данных у клиентов (разные распределения классов) евклидое среднее размывает модель. Wasserstein-агрегация рассматривает распределение весов каждого слоя $p_i(\theta)$ и вычисляет барицентр:

**Слайдовый барицентр (Sliced Wasserstein Barycenter).** Проецировать меры на случайные 1D-направления и считать барицентр аналитически через квантильные функции. Сложность: $O(mn \log m \cdot L)$ где $L$ - число проекций (обычно 100-500). При $L = 200$ ошибка меньше 3% от точного барицентра, скорость в 50-100 раз выше Sinkhorn. Применяется в быстром color transfer и shape averaging.

**Gradient flows в $W_2$ (Jordan-Kinderlehrer-Otto 1998).** Уравнение диффузии $\partial_t \rho = \Delta \rho$ - это gradient flow функционала энтропии $H[\rho] = \int \rho \log \rho$ в метрике $W_2$. Уравнение Фоккера-Планка - gradient flow свободной энергии $F[\rho] = \int V \rho + \int \rho \log \rho$. JKO-схема дискретизирует эти уравнения как последовательность проксимальных шагов в $(\mathcal{P}_2, W_2)$ - каждый шаг близок к Sinkhorn-барицентру двух мер. Применяется для моделирования диффузных процессов в молекулярной динамике и ланжевеновской динамики оптимизации нейросетей.

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

• **Фреше-среднее в $W_2$**: $\bar{\mu} = \arg\min_{\mu} \sum_i \lambda_i W_2^2(\mu, \mu_i)$. Для 1D гауссиан: $\bar{\sigma} = \sum \lambda_i \sigma_i$ - среднее $\sigma$, не $\sigma^2$
• **Agueh-Carlier 2011**: барицентр эквивалентен мультимаргинальному OT. При абсолютной непрерывности - существует и единственен. Оптимальные карты $T_i$ - градиенты выпуклых функций (Бренье)
• **Уравнение Бура**: барицентр гауссиан $\mathcal{N}(m_i, \Sigma_i)$ - фиксированная точка $S = \sum \lambda_i (S^{1/2} \Sigma_i S^{1/2})^{1/2} S^{-1/2}$, итерационно решаемая
• **Sinkhorn-барицентр** (Cuturi-Doucet 2014): $O(m^2 n)$ вместо $O(m^n)$. Чередовать Sinkhorn и геометрическое среднее двойственных переменных
• **Gradient flow в $W_2$**: JKO-схема - диффузия = последовательность проксимальных шагов в $(\mathcal{P}_2, W_2)$. Объединяет барицентры с уравнениями в частных производных

Связанные темы

Wasserstein барицентр связывает геометрию мер с практическими алгоритмами ML:

• Теорема Бренье — Оптимальные карты $T_i$ к барицентру - градиенты выпуклых функций, что обеспечивает единственность
• Domain adaptation через OT — Multi-source DA - барицентр нескольких source-доменов в feature space
• Алгоритм Sinkhorn — Sinkhorn-барицентр строится как цикл из стандартных Sinkhorn итераций

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

• Sinkhorn-барицентр при $\varepsilon \to 0$ сходится к точному W2-барицентру, а при $\varepsilon \to \infty$ - к евклидову среднему. Почему? Что происходит с транспортным планом при большой регуляризации?
• JKO-схема дискретизирует уравнение диффузии как проксимальные шаги в W2 с шагом $\tau$. Как связаны $\tau$ и точность аппроксимации? При каком $\tau$ схема теряет устойчивость?
• Federated learning через Wasserstein-барицентр требует знать распределение весов $p_i(\theta)$ каждого клиента. Как аппроксимировать $p_i$, имея только точечную оценку $\theta_i$?

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

• ot-03-wasserstein — W2-расстояние - метрика, по которой минимизируется барицентр
• ot-04-sinkhorn — Sinkhorn - базовый строительный блок итеративного алгоритма барицентра
• ot-06-brenier — Теорема Бренье: оптимальные карты T_i - градиенты выпуклых функций, что обеспечивает единственность барицентра
• ot-08-domain-adaptation — Multi-source domain adaptation - частный случай барицентра нескольких source-доменов
• calc-01-sequences