Оптимальный транспорт

OT и генеративные модели: Flow Matching

Цели урока

  • Понять Flow Matching как ODE-альтернативу диффузионным моделям с меньшим NFE
  • Объяснить роль OT-coupling: почему прямые траектории уменьшают число шагов при инференсе
  • Разобраться в Bridge Matching как обобщении для произвольных пар распределений

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

  • Диффузионные модели (DDPM, DDIM) - базовые понятия
  • ODE и численное интегрирование (Runge-Kutta)
  • Энтропийный OT и алгоритм Синкхорна
  • Энтропийная регуляризация
  • Flow Matching (подробно)

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
OT и генеративные модели: Flow Matching

0

1

Войти