Optimal Transport
OT и генеративные модели: Flow Matching
Цели урока
- Понять Flow Matching как ODE-альтернативу диффузионным моделям с меньшим NFE
- Объяснить роль OT-coupling: почему прямые траектории уменьшают число шагов при инференсе
- Разобраться в Bridge Matching как обобщении для произвольных пар распределений
Предварительные знания
- Диффузионные модели (DDPM, DDIM) - базовые понятия
- ODE и численное интегрирование (Runge-Kutta)
- Энтропийный OT и алгоритм Синкхорна
2022 год. Stable Diffusion - стандарт. 50 шагов денойзинга, 50 forward passes. 2024 год: Stable Diffusion 3 и Flux - 8 шагов. Разница: Flow Matching с OT-coupling. Один концептуальный сдвиг ускорил generation 6x.
- Stable Diffusion 3: Rectified Flow (OT FM) вместо диффузии - 8 шагов против 50, те же качество и разнообразие
- Flux 12B: MMDiT + OT Flow Matching - state-of-the-art text-to-image, SOTA на Human Arena
- I2SB: Schrodinger Bridge для image restoration - deblurring и super-resolution за 5-10 NFE
- AlphaFold 3 Diffusion: OT-guided flow для сэмплирования 3D структур белков и малых молекул
От нормализующих потоков к Flow Matching
Нормализующие потоки (Real NVP, Glow) существуют с 2016 года - но требуют invertible архитектур. Continuous Normalizing Flows (Chen et al. 2018) освободили от этого ограничения, но обучение через trace трудозатратно. Flow Matching (Lipman et al. 2022) и Rectified Flow (Liu et al. 2022, Liu 2023) решили проблему обучения: simulation-free, без numerical ODE во время обучения. OT Flow Matching (Tong et al. 2023) добавил optimal coupling. К 2024 году Flow Matching полностью вытеснил диффузию в top-1 text-to-image моделях.
Flow Matching: потоки вместо диффузии
2022 год. Stable Diffusion стал стандартом генерации изображений. Но у диффузионных моделей проблема: 50-1000 шагов денойзинга. Каждый шаг - forward pass нейросети. Meta AI предложила Flow Matching: одно ODE вместо цепи Маркова. К 2024 году Stable Diffusion 3 генерирует 1024×1024 за 8 шагов вместо 50. Математика: оптимальный транспорт.
Почему прямые траектории лучше? Нейросеть v_θ обучается предсказывать постоянный вектор x_1 - x_0 вдоль прямой. При inference: одна ODE со step size 1 даёт точный результат. Без шума, без итераций.
Stable Diffusion 3 и Flux
Stable Diffusion 3 (Esser et al. 2024) использует Rectified Flow (вариант Flow Matching) с OT-coupled парами. Генерация 1024×1024: 8-28 шагов против 50 у SD 2.x. Flux (Black Forest Labs) - следующее поколение: MMDiT архитектура + Flow Matching. 12B параметров, SOTA на текст-в-изображение. Ключ: OT-прямые траектории требуют меньше NFE при инференсе.
В чём преимущество прямых траекторий Flow Matching перед кривыми траекториями диффузии при инференсе?
Чем прямее траектории - тем проще поле v_t, тем точнее его дискретная аппроксимация. Поле для прямых траекторий близко к константе x_1-x_0; ODE-солвер с 2-4 шагами даёт высокую точность.
OT Flow Matching: оптимальный транспорт для пар
Обычный Flow Matching соединяет случайные пары (x_0, x_1): шум с произвольным изображением. Траектории пересекаются - поле v_t усредняет направления и становится кривым. OT Flow Matching соединяет пары по оптимальному плану: каждый шум идёт к ближайшему изображению. Траектории параллельны, поле - почти постоянное.
Почему OT Flow Matching уменьшает NFE при инференсе по сравнению с обычным FM?
Случайные пары создают пересекающиеся траектории - поле v_t должно усреднять конфликтующие направления, становится кривым. Кривое поле требует больше шагов для точного интегрирования.
Bridge Matching и мост Шрёдингера в генерации
Flow Matching переводит p_0=Гаусс в p_1=данные. Но что если нужно перевести одно изображение в другое? Или генерировать со структурными ограничениями? Bridge Matching обобщает FM: p_0 - произвольное начальное, p_1 - произвольное конечное. Математический движок - мост Шрёдингера между двумя реальными распределениями.
Практический выбор: OT Flow Matching для text-to-image (p_0 = Гаусс, p_1 = данные). Bridge Matching для задач трансляции: super-resolution, deblurring, colorization - где p_0 и p_1 оба осмысленны.
Flux и следующее поколение генерации
Flux (Black Forest Labs, 2024): 12B параметров, MMDiT архитектура + Rectified Flow (OT-FM вариант). Benchmark Arena: превосходит Midjourney v6 и DALL-E 3 по text following. Ключевой вклад OT: при training размер батча 2048 изображений, OT-coupling через Sinkhorn за 100 итераций на каждый батч. Результат: 6-8 NFE при инференсе против 20+ у диффузионных конкурентов.
Чем Bridge Matching отличается от обычного Flow Matching?
Обычный FM: Гаусс → данные, прямые линии. Bridge Matching: любое p_0 → любое p_1 через мост Шрёдингера. I2SB использует это для: деблюрринг (размытое → чёткое), SR (низкое→высокое разрешение).
Куда ведёт тема
Flow Matching - текущий state-of-the-art в генерации изображений. Следующий шаг: мульти-маргинальный OT (ot-26) для задач с несколькими ограничениями одновременно, и несбалансированный OT (ot-29) для работы с распределениями разного размера.
- Optimal Transport — Связанная тема
Итоги
- Flow Matching: обучаем v_θ предсказывать скорость вдоль прямых траекторий phi_t = (1-t)x_0 + t·x_1
- OT-coupling: Sinkhorn на каждом батче для прямых непересекающихся траекторий - меньше NFE при инференсе
- Bridge Matching: обобщение для произвольных p_0 и p_1 через мост Шрёдингера
- Stable Diffusion 3 и Flux: OT Flow Matching снизил NFE с 50 до 8 при том же качестве
Вопросы для размышления
- Почему пересекающиеся траектории требуют больше шагов ODE при инференсе - и что именно вынуждена делать нейросеть v_θ?
- Как выбрать размер батча для mini-batch OT в Flow Matching: что происходит при m=2 vs m=512?
- В чём принципиальное отличие OT Flow Matching от Score Matching (диффузии) с точки зрения обучающего сигнала?
Связанные уроки
- ot-11-flow-matching — детальное рассмотрение Flow Matching архитектур
- ot-07-wgan — WGAN и Flow Matching - разные пути от OT к генерации
- ot-24-entropic-reg — энтропийная регуляризация - инструмент для OT-coupling
- ot-26-multi-marginal — мульти-маргинальный OT обобщает Flow Matching