Flow Matching

Сопоставление потоков — это метод обучения генеративных моделей, которые преобразуют случайный шум в данные, следуя выученному векторному полю скоростей. Нейронная сеть учится определять направление и скорость движения любой точки в любой момент времени от 0 до 1, а для генерации образца нужно взять гауссов шум и проинтегрировать это векторное поле с помощью решателя обыкновенного дифференциального уравнения (ODE) до времени 1. Ярон Липман и его коллеги из Meta AI и Института науки Вейцмана представили этот метод в работе Flow Matching for Generative Modeling на ICLR 2023. Две параллельные работы пришли почти к той же целевой функции, назвав её выпрямленным потоком и стохастическими интерполянтами. Сопоставление потоков позволяет обучать непрерывные нормализующие потоки без симуляции ODE во время обучения, охватывает диффузионные модели как частный случай и лежит в основе обучения Stable Diffusion 3, FLUX.1, Movie Gen от Meta и многих современных моделей речи и робототехники.

Как работает сопоставление потоков#

Во время обучения ничего не генерируется. На каждом шаге выбираются образец данных x1, образец шума x0 из стандартного гауссова распределения и момент времени t между 0 и 1; затем на прямой между ними задаётся точка: xt = (1 - t) * x0 + t * x1. Точка на этой прямой движется с постоянной скоростью x1 - x0, поэтому сеть v(xt, t) обучается предсказывать x1 - x0 с помощью обычной среднеквадратичной ошибки. Каждая пара «шум — данные» задаёт лишь один путь, но через любую заданную точку проходит множество путей, а скорость полного потока от шума к данным в этой точке равна их средней скорости. В статье доказывается, что регрессия по простым целевым значениям для отдельных пар даёт те же градиенты, что и регрессия по этому трудно вычислимому среднему. Поэтому обучение сводится к обычной регрессионной функции потерь, не требующей симуляции ODE.

При генерации интегрируется dx/dt = v(x, t) от t = 0 до t = 1, обычно с помощью шагов решателя Эйлера или решателя более высокого порядка. Этот код на PyTorch обучает небольшую сеть скоростей на двух кластерах двумерных точек, а затем генерирует образцы за 1, 2 и 8 шагов Эйлера:

import torch
from torch import nn

torch.manual_seed(0)
# Velocity network v(x, t) for 2D points; the data are two clusters at (2, 2) and (-2, -2)
net = nn.Sequential(nn.Linear(3, 128), nn.SiLU(), nn.Linear(128, 128), nn.SiLU(), nn.Linear(128, 2))
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
centers = torch.tensor([[2.0, 2.0], [-2.0, -2.0]])

for step in range(2000):
    x1 = centers[torch.randint(0, 2, (256,))] + 0.2 * torch.randn(256, 2)  # data samples
    x0 = torch.randn(256, 2)  # noise samples
    t = torch.rand(256, 1)
    xt = (1 - t) * x0 + t * x1  # point on the straight path from noise to data
    loss = ((net(torch.cat([xt, t], dim=1)) - (x1 - x0)) ** 2).mean()  # regress the path's velocity
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Generate: start from noise and follow the learned velocity field with Euler steps
for steps in (1, 2, 8):
    x = torch.randn(1000, 2)
    with torch.no_grad():
        for i in range(steps):
            x = x + net(torch.cat([x, torch.full((1000, 1), i / steps)], dim=1)) / steps
    near = (torch.cdist(x, centers).min(dim=1).values < 0.6).float().mean()
    print(f"{steps}-step sampling: {near:.0%} of samples land near a cluster")

Сначала выводится 1-step sampling: 0% of samples land near a cluster, затем 47% для 2 шагов и 94% для 8 шагов. При t = 0 для каждой точки шума ожидаемое место назначения одинаково — среднее значение данных, поэтому за один шаг все образцы попадают в пустое пространство между двумя кластерами. Каждый путь обучения прямой, но пути, начинающиеся рядом друг с другом, заканчиваются в разных кластерах, поэтому усредняющий их выученный поток всё равно изгибается. Поэтому генераторы с небольшим числом шагов используют дополнительные методы, например повторное течение (reflow), при котором модель повторно обучается на собственных парах «шум — образец», чтобы выпрямить поток, или дистилляцию в модели согласованности.

Сопоставление потоков и родственные методы#

• Диффузионные модели: При обучении диффузионных моделей данные зашумляются по расписанию, а сеть учится предсказывать добавленный шум; затем образцы генерируются обратным ходом этого процесса. В статье о сопоставлении потоков показано, что целевая функция допускает любой гауссов путь распределения вероятностей, в том числе пути, используемые диффузионными моделями, и что такой способ обучения более стабилен. Прямой путь — отдельный, более простой вариант. Обе группы методов используют общие архитектуры, такие как диффузионный трансформер (DiT), классификатор-свободное управление и решатели ODE.
• Выпрямленный поток: Использует ту же линейную интерполяцию и функцию потерь скорости, поэтому для линейного пути эти названия в основном взаимозаменяемы. В статье о выпрямленном потоке также предложена процедура reflow для дополнительного выпрямления траекторий.
• Нормализующие потоки: Дискретные нормализующие потоки объединяют обратимые слои и обучаются с помощью точного правдоподобия. Непрерывные нормализующие потоки, представляющие собой разновидность нейронной ODE, задают преобразование с помощью ODE, но обучаются методом максимального правдоподобия, который требует решать ODE на каждом шаге обучения. Сопоставление потоков обучает ту же модель ODE регрессией, без такого решения.
• Оптический поток: Несмотря на похожее название, это не связано с сопоставлением потоков. Оптический поток — это видимое движение пикселей между кадрами видео, тогда как сопоставление потоков перемещает образцы в пространстве возможных данных.

Где применяется сопоставление потоков#

• Генерация изображений: Stability AI обучила Stable Diffusion 3 с помощью трансформера выпрямленного потока, выбирая шаги обучения из логистически-нормального распределения, которое отдаёт предпочтение промежуточным уровням шума. Модель FLUX.1 от Black Forest Labs — это трансформер выпрямленного потока с 12 миллиардами параметров. Изображения, созданные этими моделями преобразования текста в изображение, служат синтетическими данными, которые добавляют в наборы данных редкие объекты, погодные условия и ракурсы для обучения моделей обнаружения YOLO26.
• Генерация видео: Movie Gen от Meta — это трансформер с 30 миллиардами параметров, обученный сопоставлению потоков. Он генерирует видео с разрешением 1080p длительностью до 16 секунд и частотой 16 кадров в секунду.
• Речь: Voicebox от Meta обучалась на более чем 50 000 часов речи как авторегрессионно-независимая модель сопоставления потоков. При генерации речи по тексту в режиме zero-shot она достигла доли ошибок в словах 1,9% против 5,9% у VALL-E и работала до 20 раз быстрее. Позднее в F5-TTS сопоставление потоков объединили с диффузионным трансформером.
• Управление роботом: В модель π0 компании Physical Intelligence добавлен «эксперт по действиям» на основе сопоставления потоков с 300 миллионами параметров, подключённый к модели «зрение — язык» с 3 миллиардами параметров. Модель генерирует блоки из 50 будущих действий за 10 шагов интегрирования и управляет роботами с частотой до 50 Гц.