Flow Straight and Fast: Обучение генерации и переносу данных с помощью Rectified Flow

Background & Academic Lineage

Происхождение и академическая преемственность

Задача переноса одной вероятностной меры в другую — так называемая «задача транспортного отображения» (transport mapping problem) — является фундаментальной проблемой машинного обучения и статистики. Исторически эта задача восходит к области Optimal Transport (OT), целью которой является поиск наиболее эффективного способа перемещения массы между распределениями. Несмотря на то что OT предоставляет строгий математический аппарат, его решение в пространствах высокой размерности, характерных для современных задач генерации изображений или переноса доменов (domain transfer), сопряжено с серьезными трудностями.

Предыдущие подходы, в частности генеративные модели, такие как Generative Adversarial Networks (GANs) и Variational Autoencoders (VAEs), пытались решить эту проблему путем обучения отображений между данными и Latent Space. Однако эти модели часто сталкиваются с существенными проблемами: GAN страдают от численной нестабильности и mode collapse, в то время как VAE и другие модели, основанные на правдоподобии, требуют сложных и вычислительно затратных процедур вывода. В последнее время популярность приобрели модели непрерывного времени, такие как диффузионные модели и нейронные Ordinary Differential Equations (ODEs). Будучи мощными инструментами, эти модели по своей сути являются «бесконечношаговыми» процессами; они требуют решения сложных дифференциальных уравнений путем многократного обращения к дорогостоящей нейронной сети, что делает их применение в реальном времени или быстрый вывод практически невозможными. Авторы данной работы установили, что ключевым ограничением этих моделей непрерывного времени является их зависимость от криволинейных, непрямых траекторий, требующих большого количества шагов дискретизации для точной симуляции.

Интуитивно понятные термины предметной области

• Rectified Flow: Это можно представить как «выпрямление магистрали». Вместо того чтобы позволять частицам данных двигаться по извилистым, неэффективным путям между двумя распределениями, этот метод принуждает их следовать кратчайшему прямолинейному пути, что делает процесс значительно быстрее и проще для вычислений.
• Reflow: Представьте курьера, который в первый день выбирает извилистый маршрут. Проанализировав трафик, он «перестраивает» (reflow) свой путь в идеально прямую линию. Путем итеративного обучения на путях, сгенерированных предыдущей моделью, система «выпрямляет» свои собственные траектории, обеспечивая высокое качество результатов при значительно меньшем количестве шагов.
• Coupling: Это «план сопряжения». Если у вас есть куча песка (распределение $\pi_0$) и вы хотите придать ей определенную форму (распределение $\pi_1$), coupling — это набор инструкций, указывающих каждой отдельной песчинке, куда именно она должна переместиться.
• Drift Force: В контексте ODE это «рулевое управление» модели. Это нейронная сеть, которая указывает точкам данных, в каком направлении двигаться в любой момент времени $t$, чтобы они гарантированно достигли пункта назначения.
• Discretization Step: Это «частота кадров» видео. Чтобы симулировать непрерывное движение, мы разбиваем его на малые отрезки. Большое количество шагов означает плавный, но медленный процесс; авторы стремятся достичь высокого качества при очень малом количестве шагов (вплоть до одного).

Таблица обозначений

Постановка задачи и ограничения

Формулировка основной проблемы и дилемма

В статье рассматривается фундаментальная задача обучения транспортного отображения между двумя эмпирически наблюдаемыми распределениями данных, $\pi_0$ и $\pi_1$, в пространствах высокой размерности. Это критически важная задача для различных приложений машинного обучения, включая генеративное моделирование (например, отображение гауссовского шума в изображения) и перенос доменов (например, перевод изображений из одного стиля в другой).

Входные данные / Текущее состояние: Отправной точкой является наличие эмпирических наблюдений (выборок) из двух распределений, $\pi_0$ и $\pi_1$, обычно в $\mathbb{R}^d$. Критическим аспектом этой проблемы является отсутствие парных входных/выходных данных. То есть для каждой выборки $X_0 \sim \pi_0$ не существует соответствующей $X_1 \sim \pi_1$, которая была бы известна как «правильный» целевой объект для генерации или трансляции. Вместо этого у нас есть только независимые наборы выборок из каждого распределения.

Выходные данные / Целевое состояние: Желаемый результат — обучить транспортное отображение $T: \mathbb{R}^d \to \mathbb{R}^d$ такое, чтобы в пределе бесконечных данных, если $Z_0 \sim \pi_0$, то $Z_1 := T(Z_0) \sim \pi_1$. Более конкретно, авторы стремятся обучить модель нейронного обыкновенного дифференциального уравнения (ODE), $dZ_t = v(Z_t, t)dt$, которая может переносить выборки из $\pi_0$ в $\pi_1$, следуя путями, которые являются максимально «прямыми». Это ODE должно быть пригодно для прямой симуляции с целью генерации новых данных или выполнения переноса доменов.

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

Предыдущие попытки преодолеть этот разрыв сталкивались с рядом проблем:
1. Наивная линейная интерполяция: Простая линейная интерполяция $X_t = tX_1 + (1-t)X_0$ обеспечивает прямые пути, но является «некаузальной (или предвосхищающей)». Она требует знания конечной точки $X_1$ для определения $X_t$, что делает невозможной прямую симуляцию для генерации новых данных.
2. Optimal Transport (OT): Хотя OT предоставляет теоретически обоснованный каркас для поиска отображений, минимизирующих транспортные затраты, он «крайне сложен вычислительно» для непрерывных мер высокой размерности и часто «не представляет прямого интереса» для специфических целей многих задач машинного обучения.
3. Генеративные модели непрерывного времени (ODEs/SDEs): Недавние достижения в таких моделях, как score-based generative models и denoising diffusion probabilistic models (DDPM), показали впечатляющие результаты. Однако эти модели «фактически являются бесконечношаговыми» и влекут за собой «высокие вычислительные затраты во время вывода», поскольку требуют многократного обращения к дорогостоящему нейронному силовому полю для симуляции ODE/SDE.

Авторы пытаются преодолеть этот разрыв, формулируя задачу как прямолинейную оптимизацию методом наименьших квадратов. Они стремятся обучить поле скоростей $v(Z_t, t)$, которое управляет ODE $dZ_t = v(Z_t, t)dt$ так, чтобы оно следовало направлению линейных путей $(X_1 - X_0)$ как можно точнее, где $X_t = tX_1 + (1-t)X_0$ — линейная интерполяция между эмпирически отобранными точками. Это выражается как:
$$ \min_v \mathbb{E} \left[ \int_0^1 \|(X_1 - X_0) - v(X_t, t)\|^2 dt \right] $$
Эта формулировка направлена на «каузализацию» прямых путей линейной интерполяции, делая их пригодными для симуляции.

Ограничения и режимы отказа

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

Физические, вычислительные или обусловленные данными ограничения:
* Непарные данные: Наиболее значимым ограничением, обусловленным данными, является присущее unsupervised learning «отсутствие парных входных/выходных данных». Это означает, что модель не может просто обучиться прямой регрессии от $X_0$ к $X_1$.
* Высокая размерность данных: Реальные данные, особенно изображения, существуют в пространствах очень высокой размерности ($\mathbb{R}^d$, где $d$ может достигать миллионов). Это делает прямые вычисления Optimal Transport невыполнимыми и усугубляет вычислительные затраты численных решателей ODE/SDE.
* Вычислительные затраты решателей ODE/SDE: Существующие модели непрерывного времени требуют «многократного обращения к дорогостоящему нейронному силовому полю» во время вывода. Это приводит к строгим требованиям по задержке в реальном времени во многих приложениях, где генерация изображения за сотни или тысячи шагов является слишком медленной.
* Свойство отсутствия пересечений ODE: Для корректно определенного ODE его решение должно быть единственным, что означает, что разные пути не могут пересекаться. Это фундаментальное математическое ограничение, которому должен удовлетворять любой обученный поток, в отличие от наивных линейных интерполяций, которые могут пересекаться.

Почему именно этот подход

Неизбежность выбора

Авторы установили, что традиционные генеративные модели — в частности GAN и диффузионные модели — упираются в фундаментальную «вычислительную стену» в отношении скорости вывода. GAN, будучи быстрыми, страдают от известной нестабильности обучения и mode collapse. Напротив, диффузионные модели (и их варианты на основе ODE, такие как PF-ODEs) математически надежны, но вычислительно затратны, поскольку требуют решения сложных криволинейных траекторий, нуждающихся в большом количестве шагов дискретизации для поддержания точности. Авторы осознали, что «криволинейная» природа этих траекторий является основным узким местом; если транспортный путь между двумя распределениями можно сделать «прямым», ODE можно решить с минимальной дискретизацией, возможно, даже за один шаг. Это осознание сместило фокус с простого сопоставления распределений на поиск кратчайшего, самого прямого пути между ними.

Сравнительное преимущество

Rectified flow качественно превосходит аналоги, поскольку трансформирует транспортную задачу в простую, масштабируемую оптимизацию методом наименьших квадратов без ограничений. В отличие от GAN, требующих тонкого баланса minimax, или диффузионных моделей, полагающихся на сложные решатели SDE/ODE, rectified flow использует процедуру «reflow». Эта процедура итеративно выпрямляет траектории потока. Структурно это значительно снижает ошибку дискретизации. В то время как стандартные диффузионные модели могут требовать сотен вычислений функции (NFE) для получения высококачественных изображений, rectified flow — особенно после reflow — может давать сопоставимые или превосходящие результаты за один шаг Эйлера. Это эффективно преодолевает разрыв между одношаговыми моделями (такими как VAE) и моделями непрерывного времени, предлагая высокое качество последних при скорости первых.

Соответствие ограничениям

Ограничения задачи требовали модель, способную работать с данными высокой размерности (такими как изображения) без нестабильности GAN или непомерных затрат на вывод, характерных для диффузионных моделей. Rectified flow соответствует этим ограничениям благодаря своей «каузализации» транспортного пути. Обучая дрейфовую силу $v$ следовать линейной интерполяции $X_t = tX_1 + (1-t)X_0$, модель учится переносить массу миопическим, непересекающимся и детерминированным способом. Этот «союз» каркаса ODE с целью прямолинейного движения гарантирует, что модель является одновременно вычислительно эффективной (благодаря прямым путям) и теоретически обоснованной (поскольку она сохраняет маргинальные распределения и снижает транспортные затраты).

Математический и логический механизм

Основное уравнение

Ключевой механизм Rectified Flow заключается в обучении поля скоростей $v(z, t)$, которое преобразует исходное распределение $\pi_0$ в целевое $\pi_1$, следуя путями по прямой линии. Целевая функция, используемая для обучения этого поля скоростей:

$$\min_{v} \int_{0}^{1} \mathbb{E} \left[ \left\| (X_1 - X_0) - v(X_t, t) \right\|^2 \right] dt, \quad \text{где } X_t = tX_1 + (1 - t)X_0$$

Пошаговый поток

1. Инициализация: Пара $(X_0, X_1)$ выбирается из распределений данных.
2. Интерполяция: Система вычисляет промежуточную точку $X_t$ в случайно выбранный момент времени $t$.
3. Предсказание скорости: Нейронная сеть $v$ принимает текущее состояние $X_t$ и время $t$ в качестве входных данных и выдает предсказанный вектор скорости.
4. Регрессия: Модель сравнивает предсказанную скорость с целевым направлением $(X_1 - X_0)$.
5. Обновление: Параметры сети обновляются с помощью градиентного спуска для минимизации разности.
6. Вывод: Во время сэмплирования модель начинает с $Z_0 \sim \pi_0$ и решает ODE $dZ_t = v(Z_t, t)dt$ с использованием численного решателя (например, метода Эйлера) для достижения $Z_1 \sim \pi_1$.

Динамика оптимизации

Механизм обучается путем «каузализации» линейной интерполяции. В то время как наивный путь $X_t$ требует знания будущего ($X_1$), обученное поле скоростей $v(Z_t, t)$ является функцией только текущего состояния и времени, что делает его валидным, каузальным ODE.

Процедура «reflow» является критической динамикой оптимизации: после обучения начальной модели она используется для генерации новых пар $(Z_0, Z_1)$ путем симуляции обученного потока. Эти новые пары используются для повторного обучения модели. Поскольку поток, сгенерированный первой моделью, уже «прямее», чем исходное сопряжение данных, вторая итерация создает еще более прямые пути. Этот итеративный процесс эффективно «выпрямляет» поток, снижая ошибку дискретизации численных решателей. Как следствие, ландшафт функции потерь становится все более гладким, позволяя модели сходиться к состоянию, в котором высококачественные выборки могут быть сгенерированы за очень малое (или даже одно) количество шагов Эйлера.

Результаты, ограничения и заключение

Экспериментальный дизайн и базовые модели

Авторы оценивают Rectified Flow преимущественно на безусловной генерации изображений с использованием набора данных CIFAR-10 и наборов данных высокого разрешения (LSUN, CelebA-HQ, AFHQ). Для создания строгого базиса они используют архитектуру U-Net из каркаса DDPM++ (Song et al., 2020b). Экспериментальный дизайн структурирован для проверки эффективности процедуры «reflow» и результирующей «прямолинейности» обученных траекторий ODE.

Что доказывают полученные данные

Представленные доказательства убедительны, особенно в отношении эффекта «выпрямления» процедуры reflow. Авторы демонстрируют, что, хотя начальный (1-rectified) поток эффективен, он не является идеально прямым. Применяя процедуру reflow — где модель переобучается на данных, сгенерированных предыдущим потоком, — траектории становятся все более линейными.

Окончательные доказательства этого механизма двояки:
* Количественные: На CIFAR-10 дистиллированный 2-rectified flow достигает FID 4.85, что значительно превосходит лучшую известную одношаговую генеративную модель (TDPM, FID 8.91). Более того, recall 0.51 превышает показатель StyleGAN2+ADA (0.49), доказывая, что метод поддерживает высокое разнообразие.
* Визуальные/Геометрические: Рисунок 4 и Рисунок 18 предоставляют визуальное доказательство того, что траектории 2-rectified flow являются почти прямыми линиями. Экстраполяция $\hat{z}_1^t = z_t + (1-t)v(z_t, t)$ остается почти постоянной независимо от $t$, что является признаком ODE с прямолинейными путями. Это подтверждает, что модель успешно «каузализировала» транспортный процесс, позволяя проводить точную симуляцию с минимальными шагами дискретизации.

Ограничения и будущие направления

Будущие направления для этого исследования могут включать:
* Теоретическое уточнение: Исследование того, существует ли теоретический предел количества шагов reflow, после которого накопление численной ошибки перевешивает преимущества выпрямления траектории.
* Более широкие приложения: Изучение того, можно ли использовать свойство «выпрямления» в негенеративных задачах, таких как моделирование физических систем или прогнозирование временных рядов.
* Интеграция Optimal Transport: Как упоминают авторы, rectified flow не гарантирует строго $c$-оптимальный транспорт для конкретной функции стоимости $c$. Будущая работа может быть сосредоточена на ограничении поля скоростей $v$ так, чтобы оно было градиентным полем (например, $v = \nabla f$), для явного обеспечения оптимальности.

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

Изоморфизмы с другими областями

Структурный скелет

Механизм, который трансформирует некаузальный, пересекающийся путь интерполяции между двумя вероятностными распределениями в детерминированный, непересекающийся и прямолинейный поток обыкновенных дифференциальных уравнений (ODE).

Дальние родственники

1. Целевая область: Гидродинамика

   • Связь: Процедура «reflow», которая итеративно выпрямляет траектории для минимизации транспортных затрат, является зеркальным отражением задачи поиска ламинарного потока в трубе. Подобно тому, как rectified flow «перепрокладывает» траектории, чтобы избежать пересечений и минимизировать диссипацию энергии, гидродинамика стремится устранить турбулентные вихри (пересечения) для достижения плавных, параллельных линий тока, минимизирующих вязкое сопротивление.

2. Целевая область: Городская транспортная инженерия

   • Связь: Переход от некаузальной линейной интерполяции к rectified flow аналогичен переходу от статичной, сеточной дорожной сети к адаптивной, интеллектуальной системе управления трафиком. В исходной интерполяции пути пересекаются вслепую (как затор на перекрестке). Rectified flow действует как центральный контроллер, который «перепрокладывает» транспортный поток, гарантируя, что частицы (транспортные средства) движутся по наиболее эффективным, не сталкивающимся путям для достижения пункта назначения, эффективно оптимизируя пропускную способность всей системы.

Сценарий «Что, если»

Если бы исследователь в области квантовой теории поля «позаимствовал» это уравнение, он мог бы применить механизм rectified flow к формулировке интеграла по путям (Path Integral). Рассматривая переход между квантовыми состояниями как rectified flow, а не как сумму по всем возможным путям, он мог бы потенциально вывести «выпрямленный» путь наименьшего действия, который вычислительно тривиален для симуляции. Это позволило бы проводить точный расчет амплитуд перехода в квантовых системах высокой размерности без необходимости дорогостоящего сэмплирования методом Монте-Карло, эффективно превращая сложные нелинейные квантовые взаимодействия в серию детерминированных одношаговых «прямых» переходов.

Универсальная библиотека структур

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