Rectified Flow

Rectified flow — метод генеративного моделирования, который обучается полю скоростей, переносящему выборки из простого распределения, обычно гауссовского шума, в распределение данных по как можно более прямым траекториям. Обучение представляет собой обычную регрессию: берется пара из шумовой и реальной выборки, выбирается точка на соединяющем их отрезке, а нейронную сеть обучают предсказывать направление этого отрезка. При генерации изученную скорость интегрирует решатель обыкновенных дифференциальных уравнений (ODE) — тот же механизм, что лежит в основе нейронных ODE; прямые траектории можно пройти за несколько крупных шагов. Xingchao Liu, Chengyue Gong и Qiang Liu из Техасского университета в Остине представили этот метод в сентябре 2022 года в работе Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, одновременно с близким методом сопоставления потоков, предложенным Yaron Lipman и коллегами из Meta.

Обучение rectified flow#

Для шума x0 и данных x1 точка x_t = (1 - t) * x0 + t * x1 перемещается от шума при t = 0 к данным при t = 1 с постоянной скоростью x1 - x0. Сеть v(x_t, t) обучается предсказывать эту скорость с помощью функции потерь на основе квадратичной ошибки, вычисляемой для случайных пар в случайные моменты времени. Через любую точку проходит множество линий, поэтому сеть усваивает их среднее направление, а заданные ею траектории ODE никогда не пересекаются, но изгибаются там, где сходятся линии, ведущие к разным точкам данных. Одномерный выпрямленный поток в PyTorch учится преобразовывать гауссовский шум в два кластера с центрами в -2 и +2:

import torch
from torch import nn

torch.manual_seed(0)
velocity = nn.Sequential(nn.Linear(2, 64), nn.SiLU(), nn.Linear(64, 64), nn.SiLU(), nn.Linear(64, 1))
optimizer = torch.optim.Adam(velocity.parameters(), lr=1e-3)

for step in range(3000):
    x0 = torch.randn(256, 1)  # noise
    x1 = 2.0 * torch.sign(torch.randn(256, 1)) + 0.1 * torch.randn(256, 1)  # data: two modes at -2 and +2
    t = torch.rand(256, 1)
    xt = (1 - t) * x0 + t * x1  # a point on the straight line from noise to data
    loss = ((velocity(torch.cat([xt, t], dim=1)) - (x1 - x0)) ** 2).mean()  # regress the line's direction
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Generate by Euler integration of dx/dt = v(x, t), with different step counts
for steps in [1, 2, 4, 32]:
    x = torch.randn(10000, 1)
    with torch.no_grad():
        for i in range(steps):
            x = x + velocity(torch.cat([x, torch.full_like(x, i / steps)], dim=1)) / steps
    near_mode = ((x.abs() - 2).abs() < 0.3).float().mean()
    print(f"{steps:>2} Euler steps: {near_mode:.0%} of samples within 0.3 of a mode")

При одном шаге Эйлера ни одна выборка не оказывается рядом с модой: при t = 0 среднее направление от каждой точки шума ведет к среднему данных, равному 0 и лежащему посередине между кластерами. При двух шагах в пределах 0,3 от моды оказываются 39% выборок, при четырех — 81%, а при 32 — 96%, поскольку изученные траектории изгибаются в сторону кластеров, и для следования изгибу нужны небольшие шаги.

Reflow: выпрямление траекторий#

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

Та же команда применила reflow и дистилляцию знаний к модели Stable Diffusion 1.5 в InstaFlow (2023). Одношаговая модель показала FID 23,3 на MS COCO 2017-5k против 37,2 у прогрессивной дистилляции, генерировала изображение за 0,09 секунды, а ее обучение стоило 199 дней на GPU A100.

Rectified flow в моделях преобразования текста в изображение#

В статье Stability AI о Stable Diffusion 3, «Масштабирование трансформеров rectified flow для синтеза изображений высокого разрешения» (март 2024 года), сравнивались формулировки обучения; авторы выбрали rectified flow со временем обучения, взятым из логит-нормального распределения, которое концентрирует обучение в середине траектории, где сложнее всего предсказывать скорость. Эта целевая функция сочетается с мультимодальным диффузионным трансформером (DiT), работающим в латентном пространстве. Black Forest Labs, основанная бывшими исследователями Stability AI, описывает свою модель FLUX.1 dev как трансформер rectified flow с 12 миллиардами параметров; дистиллированный вариант FLUX.1 [schnell] генерирует изображения за один — четыре шага. Эти модели преобразования текста в изображение часто служат источником синтетических данных для обучения моделей компьютерного зрения. Платформа Ultralytics включает генерацию изображений по тексту в процесс подготовки наборов данных благодаря функции «Создать похожие изображения»: для одного изображения из набора она создает новые описания сцен, генерирует до 16 вариантов и добавляет выбранные тобой варианты в набор без меток, готовыми к аннотированию.

Rectified flow и связанные термины#

• Сопоставление потоков: Общая схема обучения генераторов на основе ODE регрессией скоростей вдоль выбранных вероятностных траекторий. Прямая линия от независимого шума к данным — наиболее широко используемая траектория в rectified flow; эти два названия часто используют как взаимозаменяемые.
• Диффузионные модели: Добавляют шум по заданному расписанию, при этом вероятностные траектории потока искривляются, а детерминированным сэмплерам обычно требуется от 20 до 50 шагов.
• Нормализующие потоки: Преобразуют шум в данные через фиксированную последовательность обратимых слоев, обученных методом точного максимального правдоподобия.
• Модели согласованности: Обучаются напрямую переходить из любой точки искривленной траектории в ее конец, поэтому для получения выборки может хватить одного вызова сети.