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

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

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

Формула выше означает следующее: мы смотрим на все пары \(x_0,x_1\), для которых промежуточная точка \(x_t\) оказалась равна \(x\), берём их скорости \(x_1-x_0\) и усредняем. Получившееся среднее направление и есть правильное векторное поле в точке \(x\) и времени \(t\). Его мы и будем использовать для семплирования.

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

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

\[\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\). Что мы и хотели доказать.

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

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

Итоговая скорость в момент времени \(t\) для точки \(x\) — среднее среди скоростей \(x_1 - x_0\), для которых \(x_t = x\) (для удобства восприятия векторы мы ужали по длине).

Представим, что в момент времени \(t\) мы смотрим на некоторую точку \(x_t\). Она могла появиться из разных пар \((x_0, x_1)\). Для каждой такой пары есть своя скорость \(v = x_1 - x_0.\) Можно сказать, что она получена применением некоторой силы, толчка. Поэтому в одной и той же точке \(x_t\) может быть не одна «правильная сила», а целое распределение возможных сил.

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

\[p(x_1 - x_0 \mid x_t,t)\]

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

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

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

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

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

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

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

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

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

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

\[x_t = (1-t)x_0 + t x_1.\]

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

\[x_1 - x_0.\]

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

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

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

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

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

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

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

Итак, давайте более формально с помощью пробных функций покажем, почему 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\) и двигаясь по \(u^*(x_t, t)\), мы получим то же самое распределение у \(x_t\), как если бы мы просто семплировали его как \(x_t = (1-t)x_0 + t x_1\).

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

\[x_t = (1-t)x_0 + t x_1,\]

Где:

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

Для каждой конкретной пары \((x_0,x_1)\) скорость вдоль прямой траектории равна:

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

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

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

\[\varphi : \mathbb{R}^d \to \mathbb{R}\]

Тогда для фиксированного \(t\):

\[\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).\]

Значит:

\[\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[ \nabla \varphi(x_t)\cdot (x_1-x_0) \right].\]

Теперь используем условное математическое ожидание по текущей точке \(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) \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)\) зависит только от \(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].\]

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

\[\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\), следовательно, мы получаем уравнение непрерывности для плотности \(p_t\):

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

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

\[x_t = (1-t)x_0 + t x_1,\]

Удовлетворяют уравнению непрерывности с полем \(u^*(x,t)\).

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

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

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

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

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

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

Теперь отдельно покажем, что эта плотность \(q_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].\]

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

\[\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\), получаем уравнение непрерывности для плотности \(q_t\):

\[\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}\) по текущим \(p_t\) и \(u^*\). Значит, если его решение единственно, то:

\[q_t = p_t\]

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

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

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

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

\[p_t(x)u^*(x,t).\]

Поэтому Flow Matching учит макроскопическое поле скоростей \(u^*(x,t)\), а не отдельные траектории между конкретными парами \((x_0,x_1)\). Двигаясь по этому среднему полю из шума, мы получаем ту же эволюцию распределений, что и при движении всех обучающих пар по прямым траекториям.

Если вспомогательная фейк-модель \(u_\psi\) идеально выучила поле генератора — минимаксная цель действительно превращается в «выучивание»:

\[u^\phi(x,t) \approx u^*(x,t).\]

Зафиксируем точку \(\hat{x}_t\) и время \(t\). Обозначим скорость траектории генератора через:

\[v = \hat{x}_1 - x_0.\]

Тогда поле генератора — условное среднее этой скорости:

\[u^\phi(\hat{x}_t,t) = \mathbb{E}_{x_0,z} \left[ v \mid \hat{x}_t,t \right].\]

Для любого вектора \(a(\hat{x}_t,t)\) верно стандартное разложение MSE (мы уже делали это, когда показывали, почему при использовании MSE мы выучиваем условное среднее):

\[\mathbb{E}_{x_0,z} \left[ \left\| a(\hat{x}_t,t)-v \right\|_2^2 \mid \hat{x}_t,t \right] = \left\| a(\hat{x}_t,t) - u^\phi(\hat{x}_t,t) \right\|_2^2 + \operatorname{tr} \operatorname{Cov}_{x_0,z} \left( v \mid \hat{x}_t,t \right).\]

Теперь подставим вместо \(a\) поле учителя \(u^*\):

\[\begin{aligned} & \mathbb{E}_{x_0,z} \left[ \left\| u^*(\hat{x}_t,t)-v \right\|_2^2 \mid \hat{x}_t,t \right] \\ &= \left\| u^*(\hat{x}_t,t) - u^\phi(\hat{x}_t,t) \right\|_2^2 + \operatorname{tr} \operatorname{Cov}_{x_0,z} \left( v \mid \hat{x}_t,t \right). \end{aligned}\]

А теперь — вместо \(a\) идеальное вспомогательное поле. Если \(u_\psi\) обучено идеально:

\[u_\psi(\hat{x}_t,t) = u^\phi(\hat{x}_t,t).\]

Поэтому:

\[\begin{aligned} & \mathbb{E}_{x_0,z} \left[ \left\| u_\psi(\hat{x}_t,t)-v \right\|_2^2 \mid \hat{x}_t,t \right] \\ &= \operatorname{tr} \operatorname{Cov}_{x_0,z} \left( v \mid \hat{x}_t,t \right). \end{aligned}\]

Вычтем второе равенство из первого. Дисперсионные члены сократятся:

\[\begin{aligned} & \mathbb{E}_{x_0,z} \left[ \left\| u^*(\hat{x}_t,t)-v \right\|_2^2 - \left\| u_\psi(\hat{x}_t,t)-v \right\|_2^2 \mid \hat{x}_t,t \right] \\ &= \left\| u^*(\hat{x}_t,t) - u^\phi(\hat{x}_t,t) \right\|_2^2. \end{aligned}\]

Теперь возьмём ожидание по всем \(t\), \(x_0\) и \(z\). Получим:

\[\begin{aligned} & \mathbb{E}_{t,\hat{x}_t} \mathbb{E}_{x_0,z | t, \hat{x}_t} \left[ \left\| u^*(\hat{x}_t,t)-(\hat{x}_1-x_0) \right\|_2^2 - \left\| u_\psi(\hat{x}_t,t)-(\hat{x}_1-x_0) \right\|_2^2 \right] \\ &= \mathbb{E}_{t,\hat{x}_t} \mathbb{E}_{x_0,z | t, \hat{x}_t} \left[ \left\| u^*(\hat{x}_t,t) - u^\phi(\hat{x}_t,t) \right\|_2^2 \right] \\ &= \mathbb{E}_{t,\,x_0,z} \left[ \left\| u^*(\hat{x}_t,t) - u^\phi(\hat{x}_t,t) \right\|_2^2 \right]. \end{aligned}\]

Поэтому такая разность двух MSE-ошибок является удобным способом обучать генератор. Она заставляет каждую отдельную скорость \(\hat{x}_1-x_0\) совпадать не \(u^*\), а условные средние поля:

\[u^\phi(x,t) \approx u^*(x,t).\]

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

Теперь покажем другой способ получить тот же минимаксный лосс. Этот вывод основан на приёме линеаризации (linearization trick), который используется в нашей статье RealUID.

Напомним, что идеальная цель дистилляции состоит в согласовании двух средних полей:

\[\min_{G_\phi} \mathbb{E}_{t,\,x_0\sim\mathcal{N}(0,I_d),\,z\sim\mathcal{N}(0,I_m)} \left[ \left\| u^*(\hat{x}_t,t) - u^\phi(\hat{x}_t,t) \right\|_2^2 \right],\]

Где:

\[\hat{x}_1 = G_\phi(z), \hat{x}_t = (1-t)x_0 + t\hat{x}_1,\]

А поле генератора равно:

\[u^\phi(x,t) = \mathbb{E}_{x_0\sim\mathcal{N}(0,I_d),\,z\sim\mathcal{N}(0,I_m)} \left[ \hat{x}_1-x_0 \mid \hat{x}_t=x \right].\]

Проблема в том, что здесь стоит разность с условным математическим ожиданием \(u^\phi\). Напрямую считать и дифференцировать такое выражение нельзя — нужно бесконечное число раз семплировать.

Запишем разность полей в другой форме. Так как \(u^*(\hat{x}_t,t)\) при фиксированных \(\hat{x}_t\) и \(t\) является константой, имеем:

\[\begin{aligned} u^*(\hat{x}_t,t)-u^\phi(\hat{x}_t,t) &= u^*(\hat{x}_t,t) - \mathbb{E}_{x_0,z} \left[ \hat{x}_1-x_0 \mid \hat{x}_t,t \right] \\ &= \mathbb{E}_{x_0,z} \left[ u^*(\hat{x}_t,t)-(\hat{x}_1-x_0) \mid \hat{x}_t,t \right]. \end{aligned}\]

Обозначим:

\[\zeta = u^*(\hat{x}_t,t)-(\hat{x}_1-x_0).\]

Тогда разность средних полей можно записать как:

\[u^*(\hat{x}_t,t)-u^\phi(\hat{x}_t,t) = \mathbb{E}_{x_0,z} \left[ \zeta \mid \hat{x}_t,t \right].\]

Значит, внутри идеальной цели стоит выражение вида:

\[\left\| \mathbb{E}_{x_0,z} \left[ \zeta \mid \hat{x}_t,t \right] \right\|_2^2.\]

Теперь используем простое тождество (приём линеаризации):

\[\|a\|_2^2 = \max_s \left[ -\|s\|_2^2 + 2\langle s,a\rangle \right].\]

Максимум достигается при \(s=a\), поэтому и слева, и справа производная по \(a\) равна \(2a\) (\(2s = 2a\) для выражения справа). Равенство значений и производных по \(a\) показывают, что мы можем эквивалентно подставлять выражение справа вместо выражения слева при подсчете градиентов.

Применим это тождество к:

\[a = \mathbb{E}_{x_0,z} \left[ \zeta \mid \hat{x}_t,t \right].\]

Тогда:

\[\left\| \mathbb{E}_{x_0,z} \left[ \zeta \mid \hat{x}_t,t \right] \right\|_2^2 = \max_s \left[ -\|s(\hat{x}_t,t)\|_2^2 + 2 \left\langle s(\hat{x}_t,t), \mathbb{E}_{x_0,z} \left[ \zeta \mid \hat{x}_t,t \right] \right\rangle \right].\]

Так как \(s(\hat{x}_t,t)\) зависит только от \(\hat{x}_t\) и \(t\), его можно внести внутрь условного ожидания:

\[\begin{aligned} & -\|s(\hat{x}_t,t)\|_2^2 + 2 \left\langle s(\hat{x}_t,t), \mathbb{E}_{x_0,z} \left[ \zeta \mid \hat{x}_t,t \right] \right\rangle \\ &= \mathbb{E}_{x_0,z} \left[ -\|s(\hat{x}_t,t)\|_2^2 + 2 \left\langle s(\hat{x}_t,t), \zeta \right\rangle \mid \hat{x}_t,t \right]. \end{aligned}\]

Теперь в выражении больше нет нормы от условного математического ожидания. Мы заменили её на максимум по вспомогательной функции \(s\), а внутри ожидания осталось линейное выражение по \(\zeta\). Именно поэтому этот шаг называется приёмом линеаризации.

Параметризуем вспомогательную функцию через дополнительную нейросеть \(u_\psi\):

\[s_\psi(x,t) = u^*(x,t)-u_\psi(x,t).\]

Тогда в оптимуме \(s_\psi\) должен приближать:

\[u^*(x,t)-u^\phi(x,t),\]

Значит, \(u_\psi\) должен приближать поле генератора \(u^\phi\).

Подставим:

\[s_\psi(\hat{x}_t,t) = u^*(\hat{x}_t,t)-u_\psi(\hat{x}_t,t)\]

И:

\[\zeta = u^*(\hat{x}_t,t)-(\hat{x}_1-x_0)\]

В максимизируемое выражение. Получаем:

\[\begin{aligned} & -\left\| u^*(\hat{x}_t,t)-u_\psi(\hat{x}_t,t) \right\|_2^2 \\ &\quad + 2 \left\langle u^*(\hat{x}_t,t)-u_\psi(\hat{x}_t,t), u^*(\hat{x}_t,t)-(\hat{x}_1-x_0) \right\rangle. \end{aligned}\]

Это выражение можно упростить алгеброй. Для краткости обозначим:

\[a = u^*(\hat{x}_t,t), \qquad b = u_\psi(\hat{x}_t,t), \qquad v = \hat{x}_1-x_0.\]

Тогда:

\[-\|a-b\|_2^2 + 2\langle a-b,a-v\rangle = \|a-v\|_2^2 - \|b-v\|_2^2.\]

Возвращаемся к исходным обозначениям, получаем:

\[\begin{aligned} & -\left\| u^*(\hat{x}_t,t)-u_\psi(\hat{x}_t,t) \right\|_2^2 \\ &\quad + 2 \left\langle u^*(\hat{x}_t,t)-u_\psi(\hat{x}_t,t), u^*(\hat{x}_t,t)-(\hat{x}_1-x_0) \right\rangle \\ &= \left\| u^*(\hat{x}_t,t)-(\hat{x}_1-x_0) \right\|_2^2 - \left\| u_\psi(\hat{x}_t,t)-(\hat{x}_1-x_0) \right\|_2^2. \end{aligned}\]

В итоге вместо:

\[\min_{G_\phi} \mathbb{E}_{t,\,x_0\sim\mathcal{N}(0,I_d),\,z\sim\mathcal{N}(0,I_m)} \left[ \left\| u^*(\hat{x}_t,t) - u^\phi(\hat{x}_t,t) \right\|_2^2 \right],\]

Получаем минимаксную задачу:

\[\min_{G_\phi}\max_{u_\psi} \mathbb{E}_{t,\,x_0\sim\mathcal{N}(0,I_d),\,z\sim\mathcal{N}(0,I_m)} \left[ \left\| u^*(\hat{x}_t,t)-(\hat{x}_1-x_0) \right\|_2^2 - \left\| u_\psi(\hat{x}_t,t)-(\hat{x}_1-x_0) \right\|_2^2 \right].\]

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

\[\hat{x}_1-x_0.\]

Роль \(u_\psi\) следующая: он выучивает среднее поле текущего генератора. После этого генератор меняется так, чтобы оно приблизилось к полю учителя \(u^*\).

Рассмотрим нюансы при обучении.

Поле учителя \(u^*\) во время дистилляции не обучается. Оно уже знает, как правильно переносить шум в распределение данных. Поэтому мы один раз переводим его в режим eval и замораживаем параметры для скорости:

\[u^*\texttt{.eval()}, u^*\texttt{.requires_grad_(False)}.\]

Однако здесь есть важная тонкость. Заморозить параметры модели — не то же самое, что запретить градиент через её вход. При обновлении генератора параметры \(u^*\) не должны меняться, но градиент должен проходить через значение \(u^*(\hat{x}_t,t)\) к точке \(\hat{x}_t\), потому что \(\hat{x}_t\) зависит от выхода генератора \(G_\phi(z)\).

Поэтому при обновлении генератора нельзя считать \(u^*(\hat{x}_t,t)\) внутри torch.no_grad() и применять .detach()для \(\hat{x}_t\).

В практическом коде величины ниже обычно являются батчами. Например, для изображений \(x_0\), \(\hat{x}_1\), \(\hat{x}_t\) и \(v\) имеют форму \([B, C, H, W]\), а время \(t\) удобно хранить в форме \([B, 1, 1, 1]\), чтобы оно автоматически broadcast’илось по каналам и пикселям.

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

Шаг 1: обновляем вспомогательное поле \(u_\psi\)

Здесь мы хотим обновлять параметры \(u_\psi\). Генератор \(G_\phi\) используется только для получения текущих сгенерированных объектов, но сам генератор на этом шаге не обновляется.

Поэтому:

• \(u_\psi\) переводим в режим train и разрешаем градиенты по его параметрам;
• \(G_\phi\) переводим в режим eval и замораживаем параметры;
• выход генератора \(\hat{x}_1=G_\phi(z)\) считаем со stop-gradient, например, через torch.no_grad() или .detach();
• \(u^*\) на этом шаге не используем.

Семплируем:

\[x_0 \sim \mathcal{N}(0,I_d), z \sim \mathcal{N}(0,I_m), t \sim \mathcal{U}[0,1].\]

Затем строим:

\[\hat{x}_1 = \operatorname{stopgrad}(G_\phi(z)), \hat{x}_t = (1-t)x_0 + t\hat{x}_1.\]

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

\[v = \hat{x}_1 - x_0.\]

Теперь обучаем \(u_\psi\) предсказывать скорость:

\[\mathcal{L}_\psi = \mathbb{E}_{t,\,x_0\sim\mathcal{N}(0,I_d),\,z\sim\mathcal{N}(0,I_m)} \left[ \left\| u_\psi(\hat{x}_t,t) - v \right\|_2^2 \right].\]

Это обычная MSE-регрессия. Поэтому при достаточно хорошем обучении \(u_\psi\) приближает условное среднее поле генератора:

\[u_\psi(x,t) \approx u^\phi(x,t).\]

Шаг 2: обновляем генератор \(G_\phi\)

Здесь мы хотим обновлять только параметры генератора \(G_\phi\). Поля \(u^*\) и \(u_\psi\) используются как фиксированные функции.

Поэтому:

• \(G_\phi\) переводим в режим train и разрешаем градиенты по его параметрам;
• \(u^*\) оставляем в режиме eval с замороженными параметрами;
• \(u_\psi\) переводим в режим eval и замораживаем его параметры;
• не используем torch.no_grad() вокруг \(u^*(\hat{x}_t,t)\) и \(u_\psi(\hat{x}_t,t)\);
• не делаем \(\hat{x}_t.\text{detach}()\);
• проводим градиент через \(u^*(\hat{x}_t,t)\), \(u_\psi(\hat{x}_t,t)\), \(\hat{x}_t\), \(\hat{x}_1\) и далее в параметры \(G_\phi\).

Семплируем:

\[x_0 \sim \mathcal{N}(0,I_d), z \sim \mathcal{N}(0,I_m), t \sim \mathcal{U}[0,1].\]

Теперь считаем выход генератора без stop-gradient:

\[\hat{x}_1 = G_\phi(z), \hat{x}_t = (1-t)x_0 + t\hat{x}_1, v = \hat{x}_1 - x_0.\]

После этого минимизируем лосс генератора:

\[\mathcal{L}_G = \mathbb{E}_{t,\,x_0\sim\mathcal{N}(0,I_d),\,z\sim\mathcal{N}(0,I_m)} \left[ \left\| u^*(\hat{x}_t,t) - v \right\|_2^2 - \left\| u_\psi(\hat{x}_t,t) - v \right\|_2^2 \right].\]

Интуитивно этот шаг меняет генератор так, чтобы среднее поле его траекторий стало ближе к полю учителя \(u^*\).

После обучения для генерации больше не нужны ни поле учителя \(u^*\), ни вспомогательное поле \(u_\psi\). Мы просто семплируем латентный шум:

\[z \sim \mathcal{N}(0,I_m)\]

И один раз применяем генератор:

\[\hat{x}_1 = G_\phi(z).\]

Так, многошаговая Flow Matching-генерация заменяется одним шагом.

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

Что применимо и для обучения, и для дистилляции Flow Matching’а:

• EMA

На практике мы не просто учим модель, мы сохраняем EMA-версию и уже используем её.

EMA-версия модели — скользящее среднее по параметрам модели. Из-за того, что обучение происходит нестабильно, мы берём среднее последних весов модели.

Это стандартная практика в связке с Adam/AdamW, которая дат стабильность и сильный прирост к качеству. Для начала можно взять EMA = 0.999-0.9999.

• Gradient clipping

Ещё одна техника — обрезка значений градиентов. Это заметно улучшает стабильность обучения и является стандартом. Норму градиента для начала можно ограничить единицей.

• Warmup

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

Warmup можно использовать для обеих моделей и выравнивать число warmup шагов из-за \(K\), чтобы они закончили его одновременно.

Warmup можно поставить, например, на 5000 шагов обучения. У основной модели базовый lr = 2e-4.

• Dropout

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

Позже мы обсудим, почему это может не сработать при оптимизации \(u_\psi\).

Dropout для начала можно поставить 0.0-0.2.

• Реализация U-Net

Модель U-Net (которую мы использовали в статье RealUID) можно взять из torchcfm. В этом репозитории понятный код, есть веса предобученных моделей для разных задач.

Подробнее про U-Net можно почитать тут.

Специфичные детали для дистилляции:

• Параметризация генератора

Если латентный шум имеет ту же размерность, что и данные, то есть \(m=d\), генератор удобно параметризовать как \(G_\phi(z) = z + g_\phi(z,0)\) или \(G_\phi(z) = z + g_\phi(z,\tau)\), где \(g_\phi\) — нейросеть той же архитектуры, что и поле Flow Matching’а \(u_\theta\), и \(\tau\) может быть любым. В этом случае \(g_\phi\) инициализируем весами предобученной Flow Matching-модели.

Если же \(m \neq d\) — такая остаточная форма \(z+g_\phi(z,0)\) уже не подходит напрямую, потому что \(z\) и \(G_\phi(z)\) живут в пространствах разной размерности. В этом случае генератор — отображение \(G_\phi : \mathbb{R}^m \to \mathbb{R}^d\).

• Инициализация вспомогательного поля

Вспомогательное поле \(u_\psi\) удобно инициализировать весами предобученного поля учителя \(u^*\). В начале обучения генератор ещё «плохой», поэтому такая инициализация может сделать оптимизацию более стабильной.

Генератор также можно инициализировать весами \(u^*\).

• Оптимизаторы

Для обеих сетей можно использовать Adam/AdamW. Для вспомогательного поля \(u_\psi\) часто полезно отключить первый момент, то есть взять Adam с \(\beta_1 = 0\). Для генератора \(G_\phi\) — применять Adam с \(\beta_1 \ge 0\). Такая настройка используется в минимаксных задачах, где одна модель играет роль критика или вспомогательного поля.

Интуиция следующая: \(\beta_1\) — momentum, то есть некоторая инерция при оптимизации. И так как внешняя игра очень сильно меняет внутреннюю, то эта инерция будет мешать, то есть модель Flow Matching не успеет за генератором.

Также рекомендуется для начала отключать dropout для вспомогательного поля \(u_\psi\) — это замедляет обучение модели.

• Несколько шагов \(u_\psi\) на один шаг генератора

На практике полезно делать несколько обновлений \(u_\psi\) на один шаг генератора. Рабочее соотношение: \(K = 5\). То есть сначала \(5\) раз обновляем \(u_\psi\), а затем один раз обновляем \(G_\phi\). Это помогает вспомогательному полю лучше отслеживать текущее распределение траекторий генератора.

Также вместе с K можно адаптировать lr у \(u_\psi\). Вместе с \(K\) это позволит \(u_\psi\) успевать учиться за генератором.

• Низкая скорость обучения генератора

Стоит ставить lr у генератора поменьше (например, 3e-5), иначе он будет учиться нестабильно из-за сложности лосса.

При этом lr у фейк-модели можно брать как 3e-5, так и 2e-4. Остальные гиперпараметры можно взять стандартные.

• Заморозка моделей и stop-gradient

Очень важно правильно управлять градиентами. Когда мы обновляем \(u_\psi\), генератор \(G_\phi\) используется только для получения текущих сгенерированных объектов. Поэтому на этом шаге параметры генератора не обновляются: G_phi.eval(), G_phi.requires_grad_(False). Выход генератора можно считать со stop-gradient: \(\hat{x}_1 = \operatorname{stopgrad}(G_\phi(z))\).

Например, в PyTorch это можно сделать через torch.no_grad() или .detach(). Когда мы обновляем генератор \(G_\phi\), наоборот, параметры \(u_\psi\) и \(u^*\) замораживаются: u_psi.eval(), u_psi.requires_grad_(False), u_star.eval(), u_star.requires_grad_(False).

Но здесь нельзя использовать torch.no_grad() вокруг \(u^*(\hat{x}_t,t)\) и \(u_\psi(\hat{x}_t,t)\) и делать x_t_hat.detach(). Причина следующая: параметры \(u^*\) и \(u_\psi\) не должны обновляться, но градиенту нужно проходить через их вход \(\hat{x}_t\), потому что \(\hat{x}_t = (1-t)x_0 + tG_\phi(z)\) зависит от генератора. Если остановить этот градиент, генератор не получит правильный обучающий сигнал.

• Постоянная фиксация учителя

Предобученное поле \(u^*\) во время всей дистилляции остаётся в режиме eval и не обновляется. Оно играет роль фиксированного учителя, который задаёт правильное направление движения от шума к данным.

Сначала по каждому изображению \(x_i\) из предобученной нейросети (например, Inception-v3) извлекаются числовые признаки \(h_i\). Затем мы смотрим на эти признаки как на семплированные из многомерного нормального распределения. При этом, признаки реальных и сгенерированных картинок семплируются из двух разных нормальных распределений: \(\mathcal{N(\mu, \Sigma)}\) и \(\mathcal{N(\hat{\mu}, \hat{\Sigma})}\).

Для двух нормальных распределений мы знаем метрику: \[d_F\left(\mathcal{N}(\mu,\Sigma),\mathcal{N}(\hat{\mu},\hat{\Sigma})\right)^2=\left\|\mu-\hat{\mu}\right\|_2^2+\operatorname{tr}\left(\Sigma+\hat{\Sigma}-2\left(\Sigma\hat{\Sigma}\right)^{\frac{1}{2}}\right).\]

Среднее и матрицу ковариаций распределения можно оценить по семплам: \(\mu\) — среднее семплов, матрица ковариации — через \(\Sigma=\frac{1}{N-1}(X-\mathbf{1}\mu)^T(X-\mathbf{1}\mu)\).