Flow matching

Сменить вопрос

Диффузия строит путь от данных к шуму физическим процессом и затем его обращает. Flow matching задаёт вопрос иначе: какой путь мы хотели бы? — и обучает поле скоростей, которое по нему ведёт.

Простейший разумный ответ: прямая.

$$\htmlData{k=path}{x_t = (1-t)\,x_0 + t\,x_1}, \qquad \htmlData{k=velocity}{v_t = \frac{dx_t}{dt} = x_1 - x_0}$$

Здесь $x_0$ — шум, $x_1$ — данные. Обратите внимание на скорость: она не зависит от $t$. Вдоль всего пути она одна и та же, и это ровно то, что означает «прямая» — постоянная скорость и нулевая кривизна.

Цель обучения

$\mathcal{L} = \mathbb{E}_{t,\, x_0,\, x_1}\big\|v_\theta(x_t, t) - (x_1 - x_0)\big\|^2$

Опять обычный MSE, и снова та же логика, что в уроке 060: сеть учится предсказывать величину, известную только на обучении (пара $x_0, x_1$ известна, потому что мы её сами составили), и в оптимуме выдаёт маргинальное поле скоростей — среднее по всем парам, проходящим через данную точку.

Стоит проговорить, что здесь усредняется. Через точку $x_t$ проходит много пар, и их скорости различны; сеть выдаёт их среднее. Значит маргинальное поле не прямое — прямыми были условные пути, а усреднение их искривляет. Это важная оговорка, и урок 110 покажет, что с ней делают.

Насколько прямее

Сравните два семейства путей. Смотрите на «кривизну» — отношение длины пути к прямому расстоянию — и на ошибку при малом числе шагов.

flow matchingдиффузия

шагов сэмплера

6

длина пути

2.83

прямое расстояние

2.83

кривизна

×1

ошибка ломаной

4.4e-16

Кривизна равна единице: путь — отрезок, и ломаная из любого числа звеньев совпадает с ним точно.

Разница читается прямо: у прямого пути ошибка дискретизации тождественно нулевая, потому что любая точка ломаной лежит на самом пути. У кривого — растёт при уменьшении числа шагов.

Отдельно стоит заметить, что кривизна диффузионного пути зависит от пары концов. Для показанных трёх пар она около $1.07{-}1.11$, а для пары, где шум и данные почти противоположны ($x_0 \approx -\lambda x_1$), путь оказывается почти прямым — кривизна $1.0002$. Причина геометрическая: при таких концах обе координаты пути пропорциональны одному и тому же вектору, и кривая вырождается в отрезок. То есть «диффузионные пути кривые» — утверждение о типичной паре, а не о каждой.

Здесь нужна честная оговорка, и она существенна. Нулевая ошибка относится к условному пути — тому, который мы задали, зная обе точки. Сэмплер идёт по маргинальному полю, а оно, как сказано выше, кривое. Так что «flow matching генерирует за один шаг» — неверно; верно то, что его пути существенно прямее, и это переводится в меньшее число шагов при том же качестве.

Что общего с диффузией

Больше, чем кажется. Оба метода:

• задают путь от шума к данным явной формулой;
• обучают сеть предсказывать нечто, вычислимое по паре $(x_0, x_1)$, обычным MSE;
• полагаются на то, что оптимум MSE — условное среднее;
• на генерации решают ОДУ (или СДУ) численно.

Более того, диффузию можно записать как flow matching с другим путём: вместо $(1-t)x_0 + t x_1$ взять $\sqrt{\bar\alpha_t}\,x_1 + \sqrt{1-\bar\alpha_t}\,x_0$. Это и есть та кривая, которую рисует виджет во втором режиме.

	диффузия	flow matching
путь	задан физическим процессом	выбран нами
коэффициенты	$\sqrt{\bar\alpha}$, $\sqrt{1-\bar\alpha}$	$1-t$, $t$
сумма квадратов коэффициентов	$1$	не единица
что предсказывает сеть	$\varepsilon$ (или $x_0$, или score)	скорость $v$
кривизна условного пути	больше единицы	ровно единица
шагов на генерацию	десятки–сотни	единицы–десятки

Третья строка стоит внимания: у flow matching масштаб вдоль пути не сохраняется — при $t = 0.5$ дисперсия равна $\frac14 + \frac14 = \frac12$ вместо единицы. Это не ошибка, а другое соглашение; сеть просто видит входы разного масштаба, и нормировка внутри неё это разбирает.

Rectified flow

Оговорка про кривое маргинальное поле имеет красивое решение. Обучим модель, сгенерируем ею пары $(\text{шум}, \text{результат})$ и обучим новую модель на этих парах. Теперь пары не случайны: каждая точка шума сопоставлена с тем результатом, к которому её приводит первая модель. Пересечений путей меньше, усреднение искажает слабее, и маргинальное поле выпрямляется.

Процедуру можно повторять. После одного-двух проходов генерация за несколько шагов действительно становится возможной — и это уже не свойство постановки, а результат отдельной работы по спрямлению.

Итог

• Flow matching выбирает путь, а не выводит его: линейная интерполяция даёт постоянную скорость.
• Цель — MSE по скорости, устроенный так же, как MSE по шуму в диффузии.
• Условные пути прямые точно; маргинальное поле — нет, потому что усреднение искривляет.
• Диффузия — тот же метод с другим путём, и это видно по сравнению кривизны в виджете.
• Rectified flow спрямляет маргинальное поле переобучением на собственных парах.