Алгоритм Sinkhorn

Классический OT-солвер: $O(n^3 \log n)$. Sinkhorn: $O(n^2 / \varepsilon^2)$ и полностью параллелизуется на GPU. Cuturi 2013 ускорил задачу в тысячи раз. Один paper - и оптимальный транспорт стал практичным инструментом для машинного обучения, биоинформатики и компьютерного зрения.

• **Wasserstein GAN (Arjovsky 2017)**: Sinkhorn как дифференцируемый proxy для расстояния Вассерштейна при обучении генеративных моделей
• **Single-cell genomics (OTT, Scanpy)**: матчинг клеток до и после perturbation - Sinkhorn на матрицах $10^4 \times 10^4$ за секунды на GPU
• **Domain adaptation**: выравнивание исходного и целевого распределений через Sinkhorn-план для transfer learning
• **JAX/OTT-jax**: дифференцируемый GPU-параллельный Sinkhorn с автоматическим вычислением градиентов по маргиналям
• **Flow matching**: непрерывный OT через Sinkhorn как инициализация планов переноса для генеративных моделей

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

• Метрики Вассерштейна
• KL как Bregman divergence

Энтропийная регуляризация: от LP до матриц

Классический оптимальный транспорт - задача линейного программирования: $O(n^3 \log n)$ по времени и $O(n^2)$ по памяти. Для $n = 1000$ это уже секунды на CPU. GPU не помогает: LP-солверы последовательны по своей природе.

2013 год. Марко Кутюри. NeurIPS. Постер в углу зала. Идея умещается в две строки: добавь энтропийный штраф к функционалу OT - и задача станет строго выпуклой, а решение - аналитически итерируемым. Через 10 лет эту статью цитируют 3000+ раз.

Исходная задача Канторовича: $\min_{\gamma \in U(a,b)} \langle C, \gamma \rangle$ где $U(a,b)$ - транспортный политоп. Добавим энтропийный регуляризатор:

При $\varepsilon \to 0$ решение приближается к оригинальному OT. При $\varepsilon \to \infty$ план вырождается в независимое произведение $ab^\top$.

**Почему это помогает**. Задача с энтропией строго выпуклая с единственным решением. Лагранжиан показывает: оптимальный $\gamma^*$ имеет вид $\gamma^*_{ij} = u_i \cdot K_{ij} \cdot v_j$, где $K_{ij} = e^{-C_{ij}/\varepsilon}$ - ядерная матрица, а $u, v$ - диагональные масштабирующие векторы. Это не LP - это матричное произведение.

Итерации Sinkhorn: поочерёдная нормализация

Оптимальное решение имеет вид $\gamma^* = \text{diag}(u) \cdot K \cdot \text{diag}(v)$. Нужно найти $u, v$ такие, что маргинали совпадают с $a$ и $b$. Это нелинейная система - но она решается элементарной итерацией:

Каждый шаг - два матрично-векторных умножения и поэлементное деление. Никаких симплекс-итераций, никаких pivot-шагов. Весь алгоритм - это $\text{matmul} + \text{div}$ в цикле.

**Геометрическая интерпретация**: каждая итерация Sinkhorn - это Bregman-проекция на одно из маргинальных ограничений. Поочерёдная проекция на два выпуклых множества сходится к их пересечению - классический результат теории выпуклой оптимизации.

Сложность: $O(n^2 \cdot T)$ где $T$ - число итераций. Для фиксированной точности $T = O(1/\varepsilon^2 \cdot \log(1/\varepsilon_\text{err}))$. Классический LP: $O(n^3 \log n)$. Разрыв на порядки - особенно при GPU-параллелизации.

Компромисс epsilon: точность против скорости

Параметр $\varepsilon$ - главная ручка алгоритма. Меньше $\varepsilon$ - точнее аппроксимация OT, но медленнее сходимость и хуже числовая стабильность. Больше $\varepsilon$ - быстро и стабильно, но план размытый.

На практике малые $\varepsilon$ требуют log-domain реализации: вместо $u, v$ хранить $\log u, \log v$ и работать в логарифмическом пространстве, чтобы избежать underflow в матрице $K$. JAX и PyTorch реализации (OTT-jax, geomloss) делают это автоматически.

**Wasserstein GAN и Sinkhorn**: в WGAN (Arjovsky 2017) расстояние Вассерштейна вычисляется через двойственную формулировку. Практические реализации используют Sinkhorn как дифференцируемую аппроксимацию - он даёт градиенты по $a$ и $b$ через автодифференцирование. Single-cell genomics: матчинг клеток между условиями - Sinkhorn на матрицах $10^4 \times 10^4$ на GPU за секунды.

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

• **Энтропийная регуляризация**: $\min \langle C, \gamma \rangle - \varepsilon H(\gamma)$ превращает LP в строго выпуклую задачу с аналитически итерируемым решением
• **Оптимальный план**: $\gamma^* = \text{diag}(u) \cdot K \cdot \text{diag}(v)$, $K_{ij} = e^{-C_{ij}/\varepsilon}$ - всё сводится к матричному произведению
• **Итерации**: $u \leftarrow a / (Kv)$, $v \leftarrow b / (K^\top u)$ - поочерёдная нормализация строк и столбцов
• **Bregman-интерпретация**: каждая итерация - проекция на маргинальное ограничение в KL-метрике
• **Компромисс $\varepsilon$**: малое $\varepsilon$ - точнее, но медленнее и нестабильнее; log-domain реализация решает проблему underflow
• **Сложность**: $O(n^2 T)$ против $O(n^3 \log n)$ классического LP - разрыв в тысячи раз на GPU

Что дальше

Sinkhorn открывает дорогу к практическим приложениям OT:

• Wasserstein GAN — Sinkhorn как дифференцируемый критерий расстояния при обучении генератора
• Flow matching — Непрерывный OT через Sinkhorn-пары как основа современных генеративных моделей
• Domain adaptation через OT — Применение Sinkhorn для выравнивания распределений в transfer learning

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

• Почему энтропийная регуляризация делает задачу OT строго выпуклой, тогда как исходная LP-формулировка может иметь бесконечно много оптимумов?
• Sinkhorn - это поочерёдные Bregman-проекции на два маргинальных ограничения. Что гарантирует сходимость этой процедуры?
• При $\varepsilon \to 0$ регуляризованный план приближается к OT-плану, но всегда имеет полную поддержку. Как это влияет на применение в задачах, где нужна разреженность плана?
• В JAX можно дифференцировать через итерации Sinkhorn по маргиналям $a, b$ и матрице стоимости $C$. Как это используется в flow matching?

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

• ot-03-wasserstein — Базовая постановка OT без регуляризации
• ot-07-wgan — WGAN использует Sinkhorn как дифференцируемый proxy
• ot-11-flow-matching — Flow matching опирается на эффективный OT через Sinkhorn
• ig-04-kl-bregman — Sinkhorn - Bregman-проекция с KL дивергенцией
• it-03 — KL дивергенция как энтропийный регуляризатор
• calc-01-sequences