Flow Matching: обучение и дистилляция

Flow Matching — современный подход генеративного моделирования. Основная идея — научить модель постепенно превращать простое распределение, например, случайный шум, в сложные данные — изображения, видео, звук, 3D-сцены или действия робота.

В рамках статьи разберём устройство Flow Matching’а: как задаётся путь от шума к данным, чему учится модель и как после обучения из случайного шума появляется новый объект. Обсудим главный недостаток подхода — медленная генерация (чтобы получить один объект, нужно выполнить много последовательных шагов).

Один из способов его решения — дистилляция. Мы заменим многошаговый процесс генератором, который за одно вычисление нейронной сети произведёт нам нужные картинки.

Однако такую модель нельзя просто обучить регрессией по принципу: «Вот — вход, вот — правильный ответ». Мы разберём, почему это так и как из задачи дистилляции естественным образом возникает минимаксная постановка: две модели начинают играть друг с другом, а результатом игры становится качественная генеративная модель!

Сферы применения

Flow Matching сегодня — один из ключевых подходов в современных генеративных моделях, «равный» диффузионным моделям (в каком-то смысле считается, что это одно и то же).

Известные примеры применения — Stable Diffusion 3.5, FLUX.2 и Kandinsky от Сбера. Эти модели способны превращать текстовое описание в новое изображение — от реалистичной фотографии до сложной фантастической сцены.

Но изображениями применение Flow Matching’а не ограничивается. С его помощью можно генерировать и более сложные объекты — например, видео вместе со звуком. MiniMax H3 получает текстовое описание сцены и постепенно превращает случайный шум не просто в одну картинку, а в целую последовательность кадров и соответствующий ей звук.

Но генеративные модели идут ещё дальше: сегодня они умеют создавать уже не отдельные изображения или видео, а целые интерактивные миры, с которыми можно взаимодействовать в реальном времени. Например, семейство Genie от Google DeepMind позволяет сгенерировать виртуальное окружение из изображения или текстового описания, а затем исследовать его, управляя персонажем. Получается, можно придумать игру и почти сразу оказаться внутри неё!

Определение

Flow Matching — один из современных подходов генерации изображений, видео, звука. Его основная идея очень похожа на интуитивное «расшумление»: мы начинаем со случайного шума и постепенно превращаем его в осмысленный объект — например, в чёткое изображение.

Формально у нас есть настоящие данные $x_\text{data} \sim p_{\text{data}}$ . Это могут быть, например, реальные изображения из интернета. Мы хотим научиться порождать новые объекты, которые похожи на данные из этого распределения, но не являются простым копированием обучающих примеров. То есть мы хотим выучить распределение картинок по набору примеров из него.

С другой стороны, нам нужно распределение, из которого легко начинать генерацию. Обычно для этого берут стандартное нормальное распределение: $x_\text{noise} \sim \mathcal{N}(0, I).$ Из него легко семплировать случайные точки, поэтому оно обычно выступает начальным шумом.

Как происходит переход от шума к данным с помощью непрерывного времени $t \in [0,1]$ ? В момент мы находимся в шуме: $x_0 = x_\text{noise} \sim \mathcal{N}(0, I),$ а в момент хотим оказаться в данных: $x_1 = x_\text{data} \sim p_{\text{data}}.$

Можно представить, что каждая точка движется во времени: сначала она выглядит как случайный шум, а затем становится настоящим изображением.

На рисунке выше мы работаем в одномерном пространстве (по вертикали — координата, по горизонтали — время), либо в двумерном пространстве (шум — слева, картинки — справа, и мы движемся в этом двумерном пространстве во времени).

Но для двумерного пространства траектории могут быть и другие, например, петли:

Самый простой способ соединить шум и объект во времени — провести между ними прямую линию:

$x_t = (1-t)x_0 + t x_1, \qquad t\in[0,1].$

При эта формула даёт , то есть шум. При она определяет , то есть объект из данных. А при промежуточных значениях времени точка находится где-то между шумом и изображением.

Скорость движения вдоль такой прямой траектории постоянна и равна:

$\frac{d x_t}{dt} = x_1 - x_0.$

То есть здесь направление движения (вектор скорости) — вектор от начального шума к конечному объекту.

Траектории точек (или частиц) могут быть совершенно разными. Вместо предварительного знания конечной точки или точки на траектории (что кажется очень сложным, потому что нужно сопоставлять точки из шума с реальными картинками) попробуем описать само движение. Для этого удобно использовать векторное поле.

Векторное поле говорит нам, куда и с какой скоростью нужно двигаться из точки в момент времени . Например, если сейчас мы находимся в точке , то поле показывает направление следующего шага.

На этой и следующей визуализациях векторы короче, чем они есть, потому что удобно брать далёкие и . В таком случае будет большим и сложным для отображения. Тогда мы будем «домножать» векторы на некоторое $d \tau$ .

Если сделать небольшой шаг по времени — новое положение точки можно записать так:

$x_{t+dt} = x_t + dt \cdot u(x_t, t).$

Следовательно, мы берём текущую точку , смотрим на скорость, которую предсказывает поле, и немного передвигаемся в этом направлении.

Итак, задачу генерации можно сформулировать так: необходимо обучить нейросеть $u_{\theta}(x,t)$ , которая приближает правильное векторное поле. После этого мы можем взять случайный шум $x_0 \sim \mathcal{N}(0,I)$ и постепенно двигать его по этому векторному полю от до . Если поле обучено хорошо, то в конце этого движения мы получим реалистичное изображение.

В математике для формального обозначения движения по вектору скорости используют обыкновенные дифференциальные уравнения (ordinary differential equations – ODE). Мы говорим, что хотим найти , зная $x_0 \sim \mathcal{N}(0,I)$ и закон $\frac{d x_t}{d t} = u(x_t, t)$ . Имеем начальную точку и закон скорости, значит, можем получить конечную точку, просто проинтегрировав $x_s = x_0 + \int_{0}^s u(x_t, t) dt$ и $x_1 = x_0 + \int_{0}^1 u(x_t, t) dt$ , где известен. Это мы и аппроксимируем через маленькие шаги .

Если мы для генерации траектории каждый раз будем вычислять — мы не получим траектории, которые в определённый момент времени пересекутся и отойдут друг от друга. Дело в том, что при пересечении мы получим одинаковый , и далее обе траектории пойдут по , то есть уже не разъединятся.

Таким образом, мы избавились от необходимости описывать, какие начальные и конечные точки принадлежат траекториям. Достаточно знать скорость в каждой точке.

Обучение

Вернёмся к простой конструкции с прямыми траекториями. Возьмём пару объектов , где $x_0 \sim \mathcal{N}(0,I)$ — случайный шум, а $x_1 \sim p_{\text{data}}$ — реальное изображение. Между ними можно провести прямую траекторию:

Мы можем соединить все $x_0 \sim p(x_0)$ со всеми $x_1 \sim p(x_1)$ прямыми траекториями (каждый с каждым). В таком случае для текущего бесконечное число пар задаёт прямую траекторию, проходящую через . То есть промежуточная точка появляется из разных пар: шумы могли двигаться к различным изображениям и в какой-то момент оказаться примерно в одном и том же месте, имея скорости .

Итоговую скорость можно получить, если усреднить все скорости, которые проходят через точку в момент времени . То есть мы можем выразить глобальную скорость через среднее локальных скоростей по всем парам .

Итак, Flow Matching показывает, что нужно предсказывать среднюю скорость всех частиц, которые проходят через эту точку в данный момент времени, чтобы получить скорость для искомых траекторий. Формально это записывается через условное матожидание:

$u^*(x,t) = \mathbb{E}_{x_0,x_1} \left[ x_1 - x_0 \mid x_t = x \right].$

Условное математическое ожидание

Математическое ожидание — обобщение понятия среднего элементов. Но элементов может быть бесконечно много. Как и в нашем случае :)

Условное математическое ожидание — среднее элементов, полученных при каких-то условиях. К примеру, если мы хотим найти среднюю продолжительность сна перед сессией, то сессия — это условие.

Формула выше означает следующее: мы смотрим на все пары , для которых промежуточная точка оказалась равна , берём их скорости и усредняем. Получившееся среднее направление и есть правильное векторное поле в точке и времени . Его мы и будем использовать для семплирования.

Здесь время фиксировано: мы усредняем по парам , которые в этот момент проходят через точку , а не по разным моментам времени.

Прямые траектории при соединении с и конечные траектории при семплировании — не одно и то же. Просто нам достаточно учить не все скорости, проходящие через текущую точку, а их среднее.

Почему можно брать именно среднее? Интуитивно, нас интересует не судьба одной конкретной частицы, а то, как меняется всё распределение точек. Если через одну область пространства проходит много частиц с разными скоростями, то общее движение плотности определяется их средним потоком. Поэтому для генерации достаточно выучить не каждую отдельную траекторию, а усреднённое векторное поле .

На практике мы обучаем нейросеть $u_{\theta}(x,t)$ приближать поле. Для этого семплируем пару , выбираем случайное время , строим промежуточную точку и просим модель по предсказать скорость

Получается обычная задача регрессии:

$\mathcal{L}(\theta) = \mathbb{E}_{x_0,x_1,t} \left[ \left\| u_\theta(x_t,t) - (x_1-x_0) \right\|^2 \right].$

Минимум этой MSE-задачи при каждом фиксированном времени как раз и равен условному среднему:

$u_{\theta}(x,t) \approx u^*(x,t) = \mathbb{E}_{x_0,x_1} \left[ x_1-x_0 \mid x_t=x \right].$

Условное среднее

Покажем, что минимум MSE-задачи равен условному среднему.

Распишем наш лосс как математическое ожидание от условного математического ожидания (вынесем ):

$\begin{aligned}\mathcal{L}(\theta)&=\mathbb{E}_{x_0,x_1,t}\left[\left\|u_\theta(x_t,t)-(x_1-x_0)\right\|^2\right] \\&=\mathbb{E}_{t,x_t}\left[\mathbb{E}_{x_0,x_1 \mid x_t,t}\left[\left\|u_\theta(x_t,t)-(x_1-x_0)\right\|^2\right]\right].\end{aligned}$

Отдельно обозначим $\mu(x_t, t)$ как условное среднее $\mu(x_t,t)=\mathbb{E}_{x_0,x_1 \mid x_t,t}\left[x_1-x_0\right]$ . Мы хотим показать, что $u_\theta(x_t, t)$ как раз выучит $\mu(x_t, t)$ .

Распишем отдельно то, что стоит под математическим ожиданием:

$\begin{aligned} & \mathbb{E}_{x_0,x_1 \mid x_t,t} \left[ \left\| u_\theta(x_t,t)-(x_1-x_0) \right\|^2 \right] \\ &= \mathbb{E}_{x_0,x_1 \mid x_t,t} \left[ \left\| u_\theta(x_t,t)-\mu(x_t,t) + \mu(x_t,t)-(x_1-x_0) \right\|^2 \right] \\ &= \mathbb{E}_{x_0,x_1 \mid x_t,t} \left[ \left\| u_\theta(x_t,t)-\mu(x_t,t) \right\|^2 \right] \\ &\quad + \mathbb{E}_{x_0,x_1 \mid x_t,t} \left[ \left\| \mu(x_t,t)-(x_1-x_0) \right\|^2 \right] \\ &\quad + 2 \mathbb{E}_{x_0,x_1 \mid x_t,t} \left[ \left\langle u_\theta(x_t,t)-\mu(x_t,t), \mu(x_t,t)-(x_1-x_0) \right\rangle \right] \\ &= \left\| u_\theta(x_t,t)-\mu(x_t,t) \right\|^2 + \mathbb{E}_{x_0,x_1 \mid x_t,t} \left[ \left\| \mu(x_t,t)-(x_1-x_0) \right\|^2 \right] \\ &\quad + 2 \left\langle u_\theta(x_t,t)-\mu(x_t,t), \mathbb{E}_{x_0,x_1 \mid x_t,t} \left[ \mu(x_t,t)-(x_1-x_0) \right] \right\rangle \\ &= \left\| u_\theta(x_t,t)-\mu(x_t,t) \right\|^2 + \mathbb{E}_{x_0,x_1 \mid x_t,t} \left[ \left\| \mu(x_t,t)-(x_1-x_0) \right\|^2 \right] \\ &\quad + 2 \left\langle u_\theta(x_t,t)-\mu(x_t,t), \mu(x_t,t) - \mathbb{E}_{x_0,x_1 \mid x_t,t} \left[x_1-x_0 \right] \right\rangle \\ &= \left\| u_\theta(x_t,t)-\mu(x_t,t) \right\|^2 + \mathbb{E}_{x_0,x_1 \mid x_t,t} \left[ \left\| \mu(x_t,t)-(x_1-x_0) \right\|^2 \right] \\ &\quad + 2 \left\langle u_\theta(x_t,t)-\mu(x_t,t), \mu(x_t,t) - \mu(x_t, t) \right\rangle \\ &= \left\| u_\theta(x_t,t)-\mu(x_t,t) \right\|^2 + \mathbb{E}_{x_0,x_1 \mid x_t,t} \left[ \left\| \mu(x_t,t)-(x_1-x_0) \right\|^2 \right] \\ &= \left\| u_\theta(x_t,t) - \mu(x_t,t) \right\|^2 + \operatorname{tr} \operatorname{Cov}_{x_0,x_1 \mid x_t,t}(x_1-x_0). \end{aligned}$

В итоге получим, что оптимум у $u_\theta(x_t, t)$ равен $\mu(x_t, t)$ , так как второе слагаемое не зависит от $u_\theta$ . Что мы и хотели доказать.

Именно поэтому Flow Matching можно обучать просто: мы сами строим промежуточные точки и знаем целевую скорость для каждой обучающей пары.

В следующих трёх скрытых блоках можно посмотреть примеры траекторий Flow Matching’а на задаче перевода одного нормального распределения в два, а также подробнее почитать про интуицию метода.

Перевод одной гауссианы в две

Рассмотрим пример с отображением одного нормального распределения в смесь двух нормальных распределений. Для начала визуализируем семплы из $x_0 \sim p_0$ и $x_1 \sim p_1$ , а также промежуточные семплы $x_t \sim p_t$ , которые были получены как

Для каждой промежуточной точки мы можем найти и , которые её породили, и соответствующую скорость (на рисунке для удобства восприятия мы ужали векторы по длине).

Итоговая скорость в момент времени для точки — среднее среди скоростей , для которых (для удобства восприятия векторы мы ужали по длине).

Появление средней скорости

Представим, что в момент времени мы смотрим на некоторую точку . Она могла появиться из разных пар . Для каждой такой пары есть своя скорость Можно сказать, что она получена применением некоторой силы, толчка. Поэтому в одной и той же точке может быть не одна «правильная сила», а целое распределение возможных сил.

Теперь вместо того, чтобы сразу заменять все эти силы на среднюю, будем прикладывать их по одной. На очень маленьком промежутке времени случайно выберем одну скорость из условного распределения:

$p(x_1 - x_0 \mid x_t,t)$

и сдвинем точку:

$x_{t+dt} = x_t + dt \cdot (x_1 - x_0).$

Почему такая процедура сохраняет правильное распределение? Потому что в исходном процессе происходит ровно то же самое, если смотреть только на точку . Среди всех частиц, которые оказались в в момент времени , скорости распределены по тому же условному закону:

$p(x_1 - x_0 \mid x_t,t).$

Значит, если мы выбираем скорость из этого распределения, то за маленький шаг получаем такой же локальный перенос частиц, как и в исходном процессе с прямыми траекториями (где мы соединяли и и шли линейно от к ).

Теперь разобьём маленький промежуток времени на множество микропромежутков времени $d \tau$ . На каждом микрошаге мы можем выбрать одну из возможных скоростей из $p(x_1 - x_0 \mid x_t,t)$ . Но отдельный сдвиг очень мал, потому что он умножается на $d \tau$ . Поэтому за каждый конкретный микрошаг мы почти не сдвигаемся, $p(x_1 - x_0 \mid x_t,t)$ не меняется, а за множество таких шагов мы в итоге сдвинемся на среднее направление всех этих скоростей, умноженное на .

Иными словами, если в точке на частицу действует много возможных сил, то на бесконечно малом масштабе их можно заменить одной средней силой:

$u^*(x,t) = \mathbb{E}_{x_0,x_1} \left[ x_1 - x_0 \mid x_t = x \right].$

Именно это среднее поле и нужно выучить модели. Оно не пытается восстановить конкретную пару , а описывает средний поток частиц, который переносит всё распределение шума к распределению данных.

Механическая аналогия

Эту идею можно представить с помощью механической аналогии. Пусть каждая пара задаёт один «удар по мячу»: футболист стоит в точке , ворота находятся в точке , и мяч летит от футболиста к воротам по прямой траектории:

Скорость такого мяча постоянна и равна:

Причём футболисты не знают, в какие конкретно ворота они должны попадать, поэтому каждый бьёт в каждые ворота.

Теперь посмотрим на некоторую точку в момент времени . Через неё могут пролетать мячи, запущенные из разных в разные . Поэтому в одной и той же точке мы можем увидеть много различных скоростей: один мяч летит чуть левее, другой — чуть правее, третий — почти прямо.

Представим, что в точке находится тяжёлое подвижное заграждение. Когда через него пролетают мячи, они слегка его толкают. Каждый отдельный толчок очень маленький: за короткое время заграждение успевает сдвинуться только на величину порядка . Поэтому один случайный удар почти не меняет его положения.

Но если за это маленькое время через точку проходит много мячей с разными скоростями, то суммарный эффект определяется не одной конкретной скоростью, а средним направлением всех этих микротолчков. То есть заграждение будет двигаться так, как если бы на него действовала средняя скорость всех мячей, пролетающих через точку в момент (если бы «запульнули» все мячи сразу).

Именно эта средняя скорость и задаёт поле Flow Matching’а:

$u^*(x,t) = \mathbb{E}_{x_0,x_1} \left[ x_1 - x_0 \mid x_t = x \right].$

То есть модель не пытается понять, какой именно мяч сейчас пролетел через точку . Вместо этого она учится предсказывать средний толчок от всех возможных мячей, которые могли оказаться в точке в данный момент времени. Такое усреднённое движение и переносит всё распределение шума к распределению данных.

Алгоритм

На каждой итерации мы берём случайный шум , случайный объект из данных и случайное время . Затем строим промежуточную точку на прямой между шумом и объектом и просим модель предсказать её скорость. На практике всё то же самое делается сразу для целого батча объектов, поэтому алгоритм ниже записан в батчевой форме.

⚙️ Алгоритм: Обучение Flow Matching-модели
Вход: Семплы из распределения $p_{\text{data}}$ (возможность семплировать), модель векторного поля $u_\theta$ , оптимизатор для параметров $\theta$ , количество итераций , размер батча
Результат: Обученное поле $u_\theta \approx u^*$
Для до :
- Семплируем батч шумов, объектов и времён: $x_0^{(b)} \sim \mathcal{N}(0,I), \qquad x_1^{(b)} \sim p_{\text{data}}, \qquad t^{(b)} \sim \mathcal{U}[0,1], \qquad b=1,\ldots,B$
- Строим промежуточные точки: $x_t^{(b)} \gets (1-t^{(b)})x_0^{(b)} + t^{(b)}x_1^{(b)}$
- Считаем целевые скорости прямых траекторий: $v^{(b)} \gets x_1^{(b)} - x_0^{(b)}$
- Считаем среднюю ошибку предсказания скорости по батчу: $\mathcal{L} \gets \frac{1}{B} \sum_{b=1}^{B} \left\| u_\theta(x_t^{(b)},t^{(b)}) - v^{(b)} \right\|_2^2$
- Обновляем параметры $\theta$ шагом оптимизатора по $\mathcal{L}$

После обучения у нас есть модель $u_{\theta}(x,t)$ , которая приближает среднее поле скоростей. Теперь мы можем использовать её для генерации новых объектов.

Идея семплирования следующая — мы стартуем из случайного шума:

$x_0 \sim \mathcal{N}(0,I)$

И постепенно двигаем точку по выученному векторному полю от времени до времени . Если разбить отрезок на $K_{\text{ode}}$ маленьких шагов, то один шаг движения можно записать так:

$x_{t+\Delta t} = x_t + \Delta t \cdot u_\theta(x_t,t),$

Где:

$\Delta t = \frac{1}{K_{\text{ode}}}.$

То есть на каждом шаге модель говорит нам, в каком направлении нужно немного сдвинуть текущую точку. После $K_{\text{ode}}$ таких шагов мы получаем финальный объект , который должен быть похож на данные. Если нужно сгенерировать множество объектов, тот же цикл обычно выполняется параллельно для батча начальных шумов.

⚙️ Алгоритм: Семплирование с помощью Flow Matching’а
Вход: Обученное поле скоростей $u_\theta$ , количество шагов семплирования $K_{\text{ode}}$
Результат: Сгенерированный объект $\hat{x}$
Семплируем начальный шум: $x \sim \mathcal{N}(0,I)$
Задаём размер шага по времени: $\Delta t \gets \frac{1}{K_{\text{ode}}}$
Для до $K_{\text{ode}}-1$
- Текущее время: $t_k \gets k\Delta t$
- Предсказываем скорость в текущей точке: $v_k \gets u_\theta(x,t_k)$
- Делаем маленький шаг по полю: $x \gets x + \Delta t \cdot v_k$
$\hat{x} \gets x$

Здесь используется самый простой численный метод — метод Эйлера. Он буквально повторяет нашу интуицию: посмотреть на текущую скорость, сделать маленький шаг в этом направлении, затем снова посмотреть на скорость и снова сделать шаг. Чем больше , тем точнее мы следуем выученному полю, но тем медленнее становится генерация.

В следующем скрытом блоке представлен формальный вывод Flow Matching’а через пробные функции, а также разбор уравнения непрерывности. Это позволит перейти от интуиции к математическому пониманию метода.

Формальный вывод

Итак, давайте более формально с помощью пробных функций покажем, почему Flow Matching учит именно поле:

$u^*(x,t) = \mathbb{E}_{x_0\sim\mathcal{N}(0,I),\,x_1\sim p_{\text{data}}} \left[ x_1 - x_0 \mid x_t = x \right],$

И почему, начиная с и двигаясь по , мы получим то же самое распределение у , как если бы мы просто семплировали его как .

Рассмотрим случайный процесс:

Где:

$x_0 \sim \mathcal{N}(0,I), \qquad x_1 \sim p_{\text{data}}.$

Для каждой конкретной пары скорость вдоль прямой траектории равна:

$\frac{d x_t}{dt} = x_1 - x_0.$

Пусть — плотность распределения промежуточных точек в момент времени . Мы хотим понять, как она меняется со временем.

Возьмём произвольную гладкую пробную функцию с компактным носителем:

$\varphi : \mathbb{R}^d \to \mathbb{R}$

Тогда для фиксированного :

$\mathbb{E}_{x_0\sim\mathcal{N}(0,I),\,x_1\sim p_{\text{data}}} \left[ \varphi(x_t) \right] = \int \varphi(x)\,p_t(x)\,dx.$

Продифференцируем это ожидание по времени:

$\frac{d}{dt} \mathbb{E}_{x_0\sim\mathcal{N}(0,I),\,x_1\sim p_{\text{data}}} \left[ \varphi(x_t) \right] = \mathbb{E}_{x_0\sim\mathcal{N}(0,I),\,x_1\sim p_{\text{data}}} \left[ \frac{d}{dt}\varphi(x_t) \right].$

По правилу цепочки:

$\frac{d}{dt}\varphi(x_t) = \nabla \varphi(x_t)\cdot \frac{d x_t}{dt} = \nabla \varphi(x_t)\cdot (x_1-x_0).$

Значит:

Теперь используем условное математическое ожидание по текущей точке :

$\mathbb{E}_{x_0\sim\mathcal{N}(0,I),\,x_1\sim p_{\text{data}}} \left[ \nabla \varphi(x_t)\cdot (x_1-x_0) \right] = \mathbb{E}_{x_t} \left[ \mathbb{E}_{x_0\sim\mathcal{N}(0,I),\,x_1\sim p_{\text{data}}} \left[ \nabla \varphi(x_t)\cdot (x_1-x_0) \mid x_t \right] \right].$

Так как $\nabla\varphi(x_t)$ зависит только от , её можно вынести из внутреннего условного ожидания:

$\mathbb{E}_{x_0\sim\mathcal{N}(0,I),\,x_1\sim p_{\text{data}}} \left[ \nabla \varphi(x_t)\cdot (x_1-x_0) \mid x_t \right] = \nabla \varphi(x_t) \cdot \mathbb{E}_{x_0\sim\mathcal{N}(0,I),\,x_1\sim p_{\text{data}}} \left[ x_1-x_0 \mid x_t \right].$

Введём среднее поле скоростей:

$u^*(x,t) = \mathbb{E}_{x_0\sim\mathcal{N}(0,I),\,x_1\sim p_{\text{data}}} \left[ x_1-x_0 \mid x_t=x \right].$

Тогда:

$\frac{d}{dt} \mathbb{E}_{x_0\sim\mathcal{N}(0,I),\,x_1\sim p_{\text{data}}} \left[ \varphi(x_t) \right] = \mathbb{E}_{x_t} \left[ \nabla \varphi(x_t)\cdot u^*(x_t,t) \right].$

Переходим к интегралу по плотности и получаем:

$\frac{d}{dt} \mathbb{E}_{x_0\sim\mathcal{N}(0,I),\,x_1\sim p_{\text{data}}} \left[ \varphi(x_t) \right] = \int \nabla \varphi(x)\cdot u^*(x,t)\,p_t(x)\,dx.$

Теперь проинтегрируем по частям. Так как $\varphi$ имеет компактный носитель, граничные члены исчезают (на бесконечности они равны 0):

$\int \nabla \varphi(x)\cdot u^*(x,t)\,p_t(x)\,dx = - \int \varphi(x)\, \nabla\cdot \left( p_t(x)u^*(x,t) \right) dx.$

С другой стороны:

$\frac{d}{dt} \mathbb{E}_{x_0\sim\mathcal{N}(0,I),\,x_1\sim p_{\text{data}}} \left[ \varphi(x_t) \right] = \frac{d}{dt} \int \varphi(x)p_t(x)\,dx = \int \varphi(x) \frac{\partial p_t(x)}{\partial t} dx.$

Значит:

$\int \varphi(x) \frac{\partial p_t(x)}{\partial t} dx = - \int \varphi(x) \nabla\cdot \left( p_t(x)u^*(x,t) \right) dx.$

Или:

$\int \varphi(x) \left[ \frac{\partial p_t(x)}{\partial t} + \nabla\cdot \left( p_t(x)u^*(x,t) \right) \right] dx = 0.$

Это верно для любой пробной функции $\varphi$ , следовательно, мы получаем уравнение непрерывности для плотности :

$\frac{\partial p_t}{\partial t} + \nabla\cdot \left( p_t u^* \right) = 0.$

Итак, мы показали — плотности из прямых траекторий:

Удовлетворяют уравнению непрерывности с полем .

Но почему можно двигаться по среднему полю?

Рассмотрим новый процесс генерации. Возьмём начальную точку:

$\tilde{x}_0 \sim \mathcal{N}(0,I)$

И будем двигать её по ODE:

$\frac{d \tilde{x}_t}{dt} = u^*(\tilde{x}_t,t).$

Пусть — плотность распределения случайной величины $\tilde{x}_t$ .

Теперь отдельно покажем, что эта плотность тоже удовлетворяет уравнению непрерывности.

Снова возьмем гладкую пробную функцию $\varphi$ :

$\mathbb{E}_{\tilde{x}_0\sim\mathcal{N}(0,I)} \left[ \varphi(\tilde{x}_t) \right] = \int \varphi(x)q_t(x)\,dx.$

Продифференцируем это ожидание по времени:

$\frac{d}{dt} \mathbb{E}_{\tilde{x}_0\sim\mathcal{N}(0,I)} \left[ \varphi(\tilde{x}_t) \right] = \mathbb{E}_{\tilde{x}_0\sim\mathcal{N}(0,I)} \left[ \frac{d}{dt}\varphi(\tilde{x}_t) \right].$

По правилу цепочки:

$\frac{d}{dt}\varphi(\tilde{x}_t) = \nabla\varphi(\tilde{x}_t)\cdot \frac{d \tilde{x}_t}{dt}.$

Так как $\tilde{x}_t$ движется по ODE:

$\frac{d \tilde{x}_t}{dt} = u^*(\tilde{x}_t,t),$

Получаем:

$\frac{d}{dt} \mathbb{E}_{\tilde{x}_0\sim\mathcal{N}(0,I)} \left[ \varphi(\tilde{x}_t) \right] = \mathbb{E}_{\tilde{x}_0\sim\mathcal{N}(0,I)} \left[ \nabla\varphi(\tilde{x}_t)\cdot u^*(\tilde{x}_t,t) \right].$

Теперь перепишем это ожидание как интеграл по плотности :

$\mathbb{E}_{\tilde{x}_0\sim\mathcal{N}(0,I)} \left[ \nabla\varphi(\tilde{x}_t)\cdot u^*(\tilde{x}_t,t) \right] = \int \nabla\varphi(x)\cdot u^*(x,t)q_t(x)\,dx.$

Интегрируем по частям:

$\int \nabla\varphi(x)\cdot u^*(x,t)q_t(x)\,dx = - \int \varphi(x) \nabla\cdot \left( q_t(x)u^*(x,t) \right) dx.$

С другой стороны:

$\frac{d}{dt} \mathbb{E}_{\tilde{x}_0\sim\mathcal{N}(0,I)} \left[ \varphi(\tilde{x}_t) \right] = \frac{d}{dt} \int \varphi(x)q_t(x)\,dx = \int \varphi(x) \frac{\partial q_t(x)}{\partial t} dx.$

Значит:

$\int \varphi(x) \frac{\partial q_t(x)}{\partial t} dx = - \int \varphi(x) \nabla\cdot \left( q_t(x)u^*(x,t) \right) dx.$

Или:

$\int \varphi(x) \left[ \frac{\partial q_t(x)}{\partial t} + \nabla\cdot \left( q_t(x)u^*(x,t) \right) \right] dx = 0.$

Так как это верно для любой пробной функции $\varphi$ , получаем уравнение непрерывности для плотности :

$\frac{\partial q_t}{\partial t} + \nabla\cdot \left( q_t u^* \right) = 0.$

Теперь сравним два процесса. Для исходных прямых траекторий мы получили:

$\frac{\partial p_t}{\partial t} + \nabla\cdot \left( p_t u^* \right) = 0, \qquad p_0=\mathcal{N}(0,I).$

А для движения по среднему полю:

$\frac{\partial q_t}{\partial t} + \nabla\cdot \left( q_t u^* \right) = 0, \qquad q_0=\mathcal{N}(0,I).$

Это одно и то же уравнение с одинаковым начальным условием. При этом, в уравнении явно показывается изменение плотности $\frac{\partial p_t}{\partial t}$ по текущим и . Значит, если его решение единственно, то:

для всех $t\in[0,1]$ .

Зачем нужна единственность решения уравнения

Рассмотрим одномерное поле $\frac{dx}{dt}=2\sqrt{|x|}$ . Из точки существуют разные решения. Можно навсегда остаться в нуле , а можно начать двигаться .

Более того, можно некоторое время стоять, а затем начать движение

$x(t)= \begin{cases} 0, & 0 \le t\le \tau,\\ (t-\tau)^2, & t>\tau. \end{cases}, \tau \ge 0$

Все эти траектории имеют одну и ту же начальную точку и удовлетворяют одинаковому полю. Причина в том, что функция $u(x)=2\sqrt{|x|}$ не является локально липшицевой в окрестности точки .

На практике наш сетап удовлетворяет условию единственности (можно вывести из теоремы Коши — Липшица), примем это за факт не будем к нему возвращаться далее в статье 🙂

Именно это и объясняет, почему можно двигаться по среднему полю. Мы не просто «на глаз» заменили разные скорости на среднюю. Мы показали, что исходный процесс с прямыми траекториями и новый процесс, который движется по ODE, порождают одну и ту же эволюцию плотности:

$\frac{d \tilde{x}_t}{dt} = u^*(\tilde{x}_t,t),$

Интуитивно это означает, что в одной точке в момент времени разные частицы могут иметь различные скорости. Одни летят левее, другие — правее. Одни быстрее, другие — медленнее. Но изменение плотности зависит не от индивидуальной скорости каждой частицы, а от суммарного потока массы через эту точку. Он равен:

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

Итак, первая часть про обучение Flow Matching завершена. Мы показали, что суть Flow Matching’а — выучивание средних скоростей. Самое время обновить чай и перейти ко второй части — дистилляции в одношаговый генератор! 😇

Мотивация

Итак, мы разобрали, что такое Flow Matching. Давайте посмотрим, почему его не очень удобно использовать на практике 😢

После обучения Flow Matching-модели генерация требует нескольких шагов. Мы стартуем из шума:

$x_0 \sim \mathcal{N}(0,I)$

И численно движемся по выученному полю $u_{\theta}(x,t)$ от до :

$x_{t+\Delta t} = x_t + \Delta t \cdot u_\theta(x_t,t).$ ВКонтакте Telegram Twitter

📊 Источник: Habr | Оригинал