Flow matching: генерация как поток

Зачем это нужно

Урок о диффузии закончился фразой: большие генераторы всё чаще учат через flow matching — модель предсказывает не шум, а скорость движения по прямой от шума к данным. Так обучены Stable Diffusion 3 и FLUX, так устроены Wan и HunyuanVideo, Qwen-Image, синтезаторы речи CosyVoice и F5-TTS. Если вы откроете код генерации FLUX, то найдёте в сердце цикла одну строку — шаг метода Эйлера из школьной физики:

# black-forest-labs/flux, src/flux/sampling.py (сокращено)
for t_curr, t_prev in zip(timesteps[:-1], timesteps[1:]):    # у FLUX время идёт от 1 (шум) к 0 (картинка)
    pred = model(img, txt, t_vec, guidance_vec)               # упрощённо: часть входов опущена
    img = img + (t_prev - t_curr) * pred                      # новое положение = старое + шаг × скорость

К концу урока вы поймёте каждый символ здесь. Flow matching проще DDPM: нет расписания $\beta_t$ и тысячи шагов, loss — MSE по скорости, генерация — несколько шагов Эйлера. А латент, по которому «течёт» генерация, — это латент VAE из прошлого урока.

Генерация как поток

Точки едут по стрелкам

Представьте облако из миллиона точек, сэмплированных из $\mathcal N(0, I)$, и облако картинок (каждая картинка — точка в пространстве латентов). Генерация — это перевезти первое облако во второе. Опишем перевозку полем скоростейполе скоростей (векторное поле)EN velocity fieldстрелка v(x, t) в каждой точке и в каждый момент — куда и как быстро двигаться $v(x, t)$: в каждом месте $x$ и в каждый момент $t \in [0, 1]$ задана стрелка — куда и как быстро двигаться. Точка, стартовав из шума при $t = 0$, послушно едет по стрелкам и к $t = 1$ оказывается картинкой. Как лист в ручье: течение в каждой точке своё, лист просто плывёт.

Записывается это как обыкновенное дифференциальное уравнениеобыкновенное дифференциальное уравнение (ОДУ)EN ordinary differential equation (ODE)dx/dt = v(x, t): скорость изменения величины задана через её текущее значение и время (ОДУ):

$$\frac{dx}{dt} = v(x, t), \qquad x(0) \sim \mathcal N(0, I)$$

Словами: скорость изменения положения точки (производная из урока 1.2.1) равна стрелке поля в том месте, где точка сейчас. Решить ОДУ — значит по старту найти всю траекторию. Нейросеть будет играть роль поля: $v_\theta(x, t)$.

Метод Эйлера

Точно решить ОДУ с нейросетью внутри нельзя, но можно приближённо — маленькими шагами. Метод Эйлераметод ЭйлераEN Euler methodшаг x ← x + h·v(x, t): проехать по текущей стрелке чуть-чуть и посмотреть снова: посмотрели на стрелку, проехали по ней чуть-чуть, посмотрели снова.

$$x_{t+h} = x_t + h \cdot v(x_t, t), \qquad h = \frac{1}{N}$$

Как машина с навигатором, который обновляет направление раз в $h$: если дорога прямая, можно ехать с закрытыми глазами и обновлять редко; на серпантине редкие обновления выкинут в кювет. Проверим на поле, где точный ответ известен: $dx/dt = -x$, $x(0) = 1$. Решение — $x(t) = e^{-t}$, то есть $x(1) = e^{-1} \approx 0.3679$ (урок 1.0.2).

import math

for n in [1, 2, 4, 10, 100]:
    x, h = 1.0, 1 / n
    for i in range(n):
        x = x + h * (-x)                    # шаг Эйлера: скорость в этой точке равна −x
    print(f"{n:3d} шаг(ов): x(1) ≈ {x:.4f}, ошибка {abs(x - math.exp(-1)):.4f}")

  1 шаг(ов): x(1) ≈ 0.0000, ошибка 0.3679
  2 шаг(ов): x(1) ≈ 0.2500, ошибка 0.1179
  4 шаг(ов): x(1) ≈ 0.3164, ошибка 0.0515
 10 шаг(ов): x(1) ≈ 0.3487, ошибка 0.0192
100 шаг(ов): x(1) ≈ 0.3660, ошибка 0.0018

Ошибка падает примерно пропорционально $h$: в 10 раз больше шагов — в 10 раз точнее. Каждый шаг генерации — это прогон огромной нейросети, поэтому главный вопрос всей области — как обойтись малым числом шагов. Ответ подсказывает частный случай: если поле постоянно вдоль траектории (точка едет по прямой с постоянной скоростью), Эйлер точен за один шаг. Отсюда и мечта о прямых траекториях.

Что посмотреть параллельно

Курс MIT 6.S184 «Generative AI with Stochastic Differential Equations» (Питер Хольдеррит и Эзра Эриве): на сайте курса — видео лекций, лабораторные и конспект. Первые лекции идут ровно по пути этого урока: ОДУ и векторные поля → flow matching → связь с диффузией и score. На английском, с формулами чуть строже, чем здесь.

Flow matching: регрессия на скорость

Прямая от шума к примеру

Где взять поле, которое перевозит шум в данные? flow matchingflow matching (сопоставление потоков)EN flow matching (Липман и соавторы, Meta, 2022) и rectified flowrectified flow («выпрямленный поток»)EN rectified flowflow matching по прямым от шума к данным; reflow дополнительно выпрямляет траектории. В китайских текстах часто пишут просто Rectified Flow (Лю, Гун, Лю, 2022) предложили одно и то же простое решение. Возьмём шум $x_0$ и пример из данных $x_1$ и соединим их отрезком:

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

Точка на этом отрезке едет с постоянной скоростью $\frac{d x_t}{dt} = x_1 - x_0$. Учим сеть угадывать эту скорость:

$$\mathcal{L} = \mathbb{E}_{x_0 \sim \mathcal N(0, I),\; x_1 \sim \text{данные},\; t \sim U[0, 1]}\Big[\big\lVert v_\theta(x_t, t) - (x_1 - x_0)\big\rVert^2\Big]$$

Словами: показываем сети точку на полпути (или на 30%, или на 90%) и спрашиваем: «куда и с какой скоростью ехать?» Правильный ответ — вектор от шума к картинке. Обычная MSE, обычная регрессия (урок 2.2.1).

Осторожно с обозначениями

В уроке о DDPM $x_0$ было картинкой, а здесь, как в статьях Липмана и Лю, $x_0$ — шум, $x_1$ — картинка. А в Stable Diffusion 3, FLUX и HunyuanVideo время идёт наоборот: $t = 1$ — шум, $t = 0$ — картинка (отсюда t_prev - t_curr в коде FLUX и флаг --flow-reverse в HunyuanVideo). Скорость при этом меняет знак. Читая чужой код, первым делом выясните, где у него шум.

Парадокс усреднения

Тут есть подвох. Через одну и ту же точку $x$ в момент $t$ проходят отрезки от многих разных пар (шум, картинка), и сеть не может знать, какая пара «настоящая». Минимизируя MSE, она выучит среднюю скорость всех отрезков, проходящих через эту точку:

$$v^*(x, t) = \mathbb{E}\big[\,x_1 - x_0 \;\big|\; x_t = x\,\big]$$

Главная теорема flow matching: если ехать по этому усреднённому полю из $\mathcal N(0, I)$, облако шума превращается ровно в распределение данных. Отдельная точка поедет не по «своему» отрезку, а по кривой, но облако целиком окажется где нужно. Проверим на игрушечном «датасете» из двух чисел: $-2$ и $+2$. Для него идеальное поле считается по формуле (без нейросети):

import numpy as np

Y = np.array([-2.0, 2.0])   # «датасет» из двух картинок-чисел: −2 и +2; шум — N(0, 1)

def v_exact(x, t):
    """Идеальное поле скоростей E[x1 − x0 | xt = x] для этого датасета."""
    # насколько каждая «картинка» y правдоподобна как цель для точки x в момент t:
    # при условии x1 = y точка xt распределена как N(t·y, (1 − t)²)
    logw = -(x[:, None] - t * Y[None, :]) ** 2 / (2 * (1 - t) ** 2)
    w = np.exp(logw - logw.max(axis=1, keepdims=True))
    w /= w.sum(axis=1, keepdims=True)                      # softmax по двум «картинкам»
    return (w * (Y[None, :] - x[:, None])).sum(axis=1) / (1 - t)   # средняя скорость «к цели»

def euler(x, n_steps):
    h = 1 / n_steps
    for i in range(n_steps):
        x = x + h * v_exact(x, i * h)                      # шаг Эйлера
    return x

x0 = np.array([-1.0, -0.3, 0.1, 0.5, 1.2])                  # пять стартовых точек из шума
for n in [1, 2, 10, 100]:
    print(f"{n:3d} шаг(ов): {np.round(euler(x0.copy(), n), 2) + 0.0}")

x, h = np.array([0.1]), 1 / 1000                           # одна точка, мелкие шаги: почти точная траектория
for i in range(1000):
    if i in (250, 500, 750, 900):
        t = i * h
        print(f"t = {t:.2f}: x = {x[0]:.2f}   (прямая от 0.1 к 2 дала бы {0.1 + t * 1.9:.2f})")
    x = x + h * v_exact(x, i * h)

  1 шаг(ов): [0. 0. 0. 0. 0.]
  2 шаг(ов): [-1.93 -1.07  0.39  1.52  1.97]
 10 шаг(ов): [-2. -2.  2.  2.  2.]
100 шаг(ов): [-2. -2.  2.  2.  2.]
t = 0.25: x = 0.09   (прямая от 0.1 к 2 дала бы 0.57)
t = 0.50: x = 0.31   (прямая от 0.1 к 2 дала бы 1.05)
t = 0.75: x = 1.15   (прямая от 0.1 к 2 дала бы 1.52)
t = 0.90: x = 1.66   (прямая от 0.1 к 2 дала бы 1.81)

Три урока из десяти строк вывода:

• За 10 шагов каждая точка приехала ровно в одну из «картинок», знак старта решил, в какую. Поле работает.
• За 1 шаг все точки оказались в нуле — в среднем двух картинок. При $t = 0$ сеть видит чистый шум и может сказать только «едь к среднему данных». Для настоящих картинок это та самая размытая каша, которую выдаёт модель при слишком малом числе шагов.
• Траектория кривая. Точка из $0.1$ полпути почти стоит на месте (при $t = 0.5$ она в $0.31$, а не в $1.05$): пока неясно, к какой картинке ехать, усреднённые стрелки тянут в обе стороны. Решение принимается ближе к концу, и точка резко сворачивает. Прямые отрезки на обучении — кривые траектории на генерации, поэтому шагов нужно больше одного.

flowchart LR
A["пример x₁ из данных"] --> C["xₜ = (1 − t)·x₀ + t·x₁"]
B["шум x₀ ~ N(0, I)"] --> C
T["t ~ U[0, 1]"] --> C
C --> N["v_θ(xₜ, t)"]
N --> L["MSE с целью x₁ − x₀"]
A --> L
B --> L

Один шаг обучения flow matching: никакого расписания шума, никакой симуляции траекторий

Сравнение с DDPM

	DDPM (урок 9.1.2)	Flow matching
Как получаем зашумлённую точку	$\sqrt{\bar\alpha_t}\,x + \sqrt{1 - \bar\alpha_t}\,\varepsilon$ ($x$ — картинка)	$(1 - t)\,x_0 + t\,x_1$ ($x_0$ — шум, $x_1$ — картинка)
Что предсказывает сеть	шум $\varepsilon$	скорость $x_1 - x_0$
Loss	MSE	MSE
Расписание	$\beta_1, \dots, \beta_T$, $T = 1000$	не нужно, $t$ непрерывно
Генерация	шаги назад со свежим шумом	ОДУ: детерминированные шаги Эйлера
Типичное число шагов	1000 (DDIM — около 50)	20–50, после дистилляции 1–4

Одно знание в разных единицах

Если сеть знает скорость $v$ в точке $x_t$, она «знает» и картинку, и шум. Из $x_t = (1 - t)x_0 + t x_1$ и $v = x_1 - x_0$ получаем:

$$\hat x_1 = x_t + (1 - t)\,v, \qquad \hat x_0 = x_t - t\,v$$

Словами: «доедь с этой скоростью до конца» — получишь картинку; «отъедь назад к началу» — получишь шум. Пример: $x_t = 1.05$, $t = 0.5$, $v = 1.9$ → $\hat x_1 = 1.05 + 0.5 \cdot 1.9 = 2.0$, $\hat x_0 = 1.05 - 0.95 = 0.1$. Так что предсказывать шум (DDPM), картинку или скорость — это одно и то же знание в разных единицах. Помните «поле направлений к данным» (score) из врезки урока 9.1.2? Это предсказанный шум со знаком минус, делённый на $1 - t$. Авторы SD3 сравнили формулировки на одной архитектуре и выбрали rectified flow со своей выборкой времени: $t$ берут не равномерно, а чаще из середины отрезка (распределение logit-normal) — там предсказывать труднее всего.

Сколько шагов и как их сократить

Выпрямление: reflow

Траектории кривые, потому что пары (шум, картинка) случайны и отрезки пересекаются. Лю и соавторы предложили reflow: сгенерировать обученной моделью много пар «шум $x_0$ → получившаяся картинка $\hat x_1$» и дообучить модель на этих парах вместо случайных. Сгенерированные траектории не пересекаются (из каждой точки ОДУ ведёт ровно по одному пути), поэтому новые отрезки почти не спорят друг с другом, и выученное поле становится прямее. Следом InstaFlow (2023) применила reflow и дистилляцию к Stable Diffusion и получила генерацию за один шаг.

Сдвиг расписания

Шаги не обязаны быть равными. Для больших картинок авторы SD3 сдвигают время к шумному концу. Во времени, где 1 — чистый шум, формула такая:

$$t' = \frac{s\,t}{1 + (s - 1)\,t}$$

При $s = 3$: $t = 0.25 \to 0.5$, $t = 0.5 \to 0.75$, $t = 0.75 \to 0.9$. Значит, половина шагов тратится, пока шума больше 75%. Зачем: у большой картинки соседние пиксели сильно связаны, и при том же уровне шума композиция ещё угадывается. «Самые шумные» шаги, где решается, что и где нарисовать, нужно растянуть. Этот параметр вы увидите в README: flow_shift = 5.0 для 720P и 3.0 для 480P у Wan2.1, --flow-shift 7.0 у HunyuanVideo; FLUX вычисляет сдвиг по числу токенов картинки.

Дистилляция в несколько шагов (обзорно)

Десятки прогонов трансформера на картинку — дорого. Поэтому готовую модель «сжимают по времени»:

• consistency models (OpenAI, 2023) учат сразу отображать любую точку траектории в её конец;
• состязательная дистилляция (SDXL-Turbo, 2023; для SD3 — её латентный вариант) добавляет дискриминатор из прошлого урока — так GAN вернулись в генерацию;
• FLUX.1 [schnell] по умолчанию генерирует за 4 шага (так прописано в коде flux), а «полная» FLUX.1 [dev] — за 50;
• MeanFlow (2025) учит не мгновенную скорость, а среднюю на отрезке времени, и делает хорошие картинки за один шаг.

Общий принцип тот же, что у дистилляции LLM: большая медленная модель-учитель, быстрая модель-ученик.

CFG в форме скоростей

Classifier-free guidance из урока 9.1.2 переносится один в один: модель учат и с условием, и без (условие иногда выбрасывают), а при генерации экстраполируют скорости:

$$v = v_\theta(x_t, t, \varnothing) + w \cdot \big(v_\theta(x_t, t, c) - v_\theta(x_t, t, \varnothing)\big)$$

Пример: $v_\varnothing = 0.2$, $v_c = 0.5$, $w = 4$ → $v = 0.2 + 4 \cdot 0.3 = 1.4$. Цена — два прогона на шаг. FLUX.1 [dev] обходит её дистилляцией guidance: сила $w$ подаётся сети как ещё один вход (guidance в коде выше), и сеть за один прогон выдаёт то, для чего раньше нужны были два.

Где это используется

Модель	Что известно из первоисточника
Stable Diffusion 3 (Stability AI, 2024)	rectified flow + MMDiT, до 8 млрд параметров; выборка $t$ из logit-normal и сдвиг для высокого разрешения (статья)
FLUX.1 (Black Forest Labs, 2024)	rectified flow transformer; [schnell] — 4 шага по умолчанию, Apache-2.0; [dev] — 50 шагов, дистилляция guidance (код и README)
Wan 2.1 (Alibaba, 2025)	«designed using the Flow Matching framework» поверх DiT; flow_shift 3.0 или 5.0 (README)
HunyuanVideo (Tencent, 2024)	планировщик flow matching, --flow-shift 7.0, 50 шагов (README)
Qwen-Image (Alibaba, 2025)	MMDiT на 20 млрд параметров, обучение flow matching (техотчёт)
CosyVoice, F5-TTS	flow matching превращает речевые токены или текст в мел-спектрограмму (урок 9.1.5)

Архитектуры быстро меняются, поэтому про каждую новую модель сверяйтесь с её техотчётом: формулировку (шум или скорость), направление времени, сдвиг и число шагов.

Практика

В браузере: flow matching на «лунах»

Обучите маленький MLP с ручным backprop предсказывать скорость на 2D-датасете «две луны», а потом сгенерируйте точки методом Эйлера за 1, 10 и 50 шагов. Готовы данные, прямой проход, кодирование времени и Adam. Ваши — батч flow matching, обратный проход, цикл обучения и сэмплер. Тест с «точным полем» для гауссовых данных проверит то, что вы видели на двух точках: один шаг стягивает всё в среднее.

Flow matching на двух «лунах»numpy · нужен интернет6 тестов

import numpy as np

# ── Готовые детали (их писать не нужно) ──────────────────────────────────────


def make_moons(n, noise=0.05, seed=0):
    """Две «луны» — классический 2D-датасет. Возвращает массив (n, 2), центрированный около нуля."""
    rng = np.random.default_rng(seed)
    a1 = rng.uniform(0, np.pi, n // 2)
    a2 = rng.uniform(0, np.pi, n - n // 2)
    upper = np.stack([np.cos(a1), np.sin(a1)], axis=1)
    lower = np.stack([1 - np.cos(a2), 0.5 - np.sin(a2)], axis=1)
    X = np.concatenate([upper, lower]) + noise * rng.normal(size=(n, 2))
    return (X - np.array([0.5, 0.25])) * 1.5


def time_features(t):
    """Время t (форма (B, 1)) → 7 признаков: t, sin и cos от πt, 2πt, 4πt (как позиции в трансформере)."""
    k = np.array([1.0, 2.0, 4.0]) * np.pi
    return np.concatenate([t, np.sin(k * t), np.cos(k * t)], axis=1)


def init_mlp(hidden=64, seed=0):
    """MLP: [x, y, 7 признаков времени] → hidden → hidden → 2 (скорость). Активация tanh."""
    rng = np.random.default_rng(seed)
    d_in = 2 + 7
    return {
        "W1": rng.normal(0, 1 / np.sqrt(d_in), (d_in, hidden)), "b1": np.zeros(hidden),
        "W2": rng.normal(0, 1 / np.sqrt(hidden), (hidden, hidden)), "b2": np.zeros(hidden),
        "W3": rng.normal(0, 0.1 / np.sqrt(hidden), (hidden, 2)), "b3": np.zeros(2),
    }


def mlp_forward(params, x, t):
    """Прямой проход. Возвращает (скорость формы (B, 2), cache для обратного прохода)."""
    inp = np.concatenate([x, time_features(t)], axis=1)
    h1 = np.tanh(inp @ params["W1"] + params["b1"])
    h2 = np.tanh(h1 @ params["W2"] + params["b2"])
    out = h2 @ params["W3"] + params["b3"]
    return out, (inp, h1, h2)


def adam_update(params, grads, state, lr):
    """Один шаг Adam (урок 1.4.2). state — словарь, который хранится между шагами (начните с {})."""
    b1, b2, eps = 0.9, 0.999, 1e-8
    state["step"] = state.get("step", 0) + 1
    s = state["step"]
    for k in params:
        m = state.setdefault("m_" + k, np.zeros_like(params[k]))
        v = state.setdefault("v_" + k, np.zeros_like(params[k]))
        m[:] = b1 * m + (1 - b1) * grads[k]
        v[:] = b2 * v + (1 - b2) * grads[k] ** 2
        params[k] -= lr * (m / (1 - b1 ** s)) / (np.sqrt(v / (1 - b2 ** s)) + eps)


# ── Ваша часть ────────────────────────────────────────────────────────────────


def fm_batch(x1, rng):
    """Готовит батч для flow matching из батча данных x1 (форма (B, 2)).

    Порядок случайных чисел важен для воспроизводимости:
      1) x0 = rng.normal(size=x1.shape)       — шум;
      2) t  = rng.uniform(size=(len(x1), 1))  — время в [0, 1);
    затем xt = (1 − t)·x0 + t·x1 (точка на прямой от шума к данным) и target = x1 − x0 (её скорость).
    Возвращает (x0, t, xt, target).
    """
    # TODO
    ...


def mlp_backward(params, cache, grad_out):
    """Обратный проход через MLP. grad_out = dLoss/d(выход), форма (B, 2).

    Возвращает словарь градиентов с теми же ключами и формами, что params.
    Напоминание: d tanh(a)/da = 1 − tanh(a)², а для y = h @ W + b: dW = h.T @ dy, db = dy.sum(0), dh = dy @ W.T.
    """
    inp, h1, h2 = cache  # входы слоёв, сохранённые в mlp_forward
    # TODO: идите от выхода к входу: W3, b3 → через tanh → W2, b2 → через tanh → W1, b1
    ...


def train_flow_matching(data, steps=400, batch=128, lr=1e-2, hidden=64, seed=0):
    """Обучает MLP предсказывать скорость x1 − x0. Возвращает (params, losses).

    rng = np.random.default_rng(seed); params = init_mlp(hidden, seed); state = {}.
    На каждом шаге: x1 = data[rng.integers(0, len(data), batch)] → fm_batch → mlp_forward →
    loss = ((pred − target)²).sum(axis=1).mean() → градиент 2·(pred − target)/batch → mlp_backward → adam_update.
    """
    # TODO
    ...


def euler_sample(velocity_fn, x0, n_steps):
    """Решает dx/dt = v(x, t) от t = 0 до t = 1 методом Эйлера за n_steps равных шагов.

    velocity_fn(x, t) принимает x формы (B, 2) и t формы (B, 1) и возвращает скорость формы (B, 2).
    Шаг: x ← x + h·v(x, t), t ← t + h, где h = 1 / n_steps. Возвращает x в момент t = 1.
    """
    # TODO
    ...


# ── Проверьте себя кнопкой «Запустить» ────────────────────────────────────────


def ascii_scatter(points, width=56, height=20, lim=2.6):
    """Рисует облако 2D-точек символами: чем больше точек в клетке, тем «темнее» символ."""
    grid = np.zeros((height, width))
    for x, y in points:
        c, r = int((x + lim) / (2 * lim) * width), int((lim - y) / (2 * lim) * height)
        if 0 <= c < width and 0 <= r < height:
            grid[r, c] += 1
    chars = " .:-=+*#%@"
    levels = np.ceil(grid / grid.max() * (len(chars) - 1)).astype(int)
    for row in levels:
        print("".join(chars[v] for v in row))


x1 = np.array([[1.0, 1.0]])
print(fm_batch(x1, np.random.default_rng(0)))                       # (x0, t, xt, target)
print(euler_sample(lambda x, t: -x, np.array([[1.0, 0.0]]), 10))    # ожидается [[0.3487, 0]]


def demo():
    data = make_moons(1000)
    params, losses = train_flow_matching(data, steps=1500)
    print(f"loss: {np.mean(losses[:20]):.2f} → {np.mean(losses[-50:]):.2f}")
    x0 = np.random.default_rng(1).normal(size=(2000, 2))
    v = lambda x, t: mlp_forward(params, x, t)[0]  # noqa: E731
    # Траектория облака: как оно выглядит в середине пути (t = 0.5) при 50 шагах
    x = x0.copy()
    for i in range(25):
        x = x + v(x, np.full((len(x), 1), i / 50)) / 50
    print("\n— облако на полпути, t = 0.5 —")
    ascii_scatter(x)
    for n in [1, 10, 50]:
        print(f"\n— t = 1 после {n} шаг(ов) Эйлера —")
        ascii_scatter(euler_sample(v, x0, n))


# demo()   # ← раскомментируйте, когда всё реализуете: обучение займёт несколько секунд

ЗапуститьПроверитьПодсказкаСбросить

Раскомментируйте в конце заготовки demo(): модель обучится и нарисует облака точек символами — на полпути ($t = 0.5$) и в конце. За 1 шаг получится сгусток посередине, за 10–50 — две дуги.

Локально: цифры MNIST с guidance

Тот же алгоритм на PyTorch, но «точка» — картинка из 784 пикселей, а к сети добавлена метка цифры. Метку в 10% примеров заменяем на «без условия» — это позволит включить CFG.

"""Flow matching на MNIST с условием по цифре и classifier-free guidance."""
import gzip
import math
import time
import urllib.request
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"
URL, DATA = "https://ossci-datasets.s3.amazonaws.com/mnist/", Path("data/mnist")

def load(name, offset):
    DATA.mkdir(parents=True, exist_ok=True)
    if not (DATA / name).exists():
        urllib.request.urlretrieve(URL + name, DATA / name)
    with gzip.open(DATA / name, "rb") as f:
        return np.frombuffer(f.read(), dtype=np.uint8, offset=offset)

X = torch.tensor(load("train-images-idx3-ubyte.gz", 16).reshape(-1, 784), dtype=torch.float32)
X = (X / 255.0 * 2 - 1).to(device)                       # пиксели в [−1, 1]: масштаб, как у шума N(0, 1)
Y = torch.tensor(load("train-labels-idx1-ubyte.gz", 8), dtype=torch.long).to(device)

class Velocity(nn.Module):
    """v_θ(x_t, t, цифра): MLP поверх картинки, времени и метки. Метка 10 = «без условия»."""
    def __init__(self, hidden=1024):
        super().__init__()
        self.label = nn.Embedding(11, 128)
        self.net = nn.Sequential(
            nn.Linear(784 + 64 + 128, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, 784),
        )

    def forward(self, x, t, y):
        freqs = torch.exp(-math.log(1000) * torch.arange(32, device=x.device) / 32)
        ang = 1000 * t[:, None] * freqs[None, :]         # синусоидальное кодирование времени, как в 9.1.2
        temb = torch.cat([ang.sin(), ang.cos()], dim=1)
        return self.net(torch.cat([x, temb, self.label(y)], dim=1))

model = Velocity().to(device)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
print(f"параметров: {sum(p.numel() for p in model.parameters()) / 1e6:.1f} млн")

EPOCHS, B = 10, 256
t0 = time.time()
for epoch in range(1, EPOCHS + 1):
    perm, total = torch.randperm(len(X), device=device), 0.0
    for i in range(0, len(X), B):
        x1, y = X[perm[i:i + B]], Y[perm[i:i + B]].clone()
        y[torch.rand(len(y), device=device) < 0.1] = 10     # 10% примеров без условия — для CFG
        x0 = torch.randn_like(x1)                            # шум
        t = torch.rand(len(x1), device=device)               # время ~ U[0, 1]
        xt = (1 - t[:, None]) * x0 + t[:, None] * x1         # точка на прямой шум → картинка
        loss = ((model(xt, t, y) - (x1 - x0)) ** 2).mean()   # регрессия на скорость x1 − x0
        opt.zero_grad(); loss.backward(); opt.step()
        total += loss.item() * len(x1)
    print(f"эпоха {epoch}: loss {total / len(X):.4f}   ({time.time() - t0:.0f} с)")

@torch.no_grad()
def sample(digits, n_steps=20, w=2.0):
    """Эйлер от шума (t = 0) к картинке (t = 1); скорость с classifier-free guidance."""
    y = torch.tensor(digits, device=device)
    none = torch.full_like(y, 10)
    x = torch.randn(len(y), 784, device=device)
    for i in range(n_steps):
        t = torch.full((len(y),), i / n_steps, device=device)
        v_c, v_u = model(x, t, y), model(x, t, none)
        v = v_u + w * (v_c - v_u)                            # CFG в форме скоростей
        x = x + v / n_steps                                  # шаг Эйлера
    return (x.clamp(-1, 1) + 1) / 2

def ascii_row(imgs):
    """Печатает картинки рядом, уменьшив до 14×14; яркость → символ по порогам."""
    small = F.avg_pool2d(imgs.view(-1, 1, 28, 28), 2).squeeze(1).cpu()
    for r in range(14):
        print("  ".join("".join(" .+#"[(v > 0.3) + (v > 0.5) + (v > 0.7)] for v in img[r].tolist()) for img in small))

torch.save(model.state_dict(), "fm_mnist.pt")
for steps, w in [(1, 2.0), (20, 2.0), (20, 1.0)]:
    torch.manual_seed(1)                                     # один и тот же стартовый шум для всех вариантов
    print(f"\n— цифры 0…4: {steps} шаг(ов) Эйлера, w = {w:g} —")
    ascii_row(sample(list(range(5)), n_steps=steps, w=w))

Запуск: python fm_mnist.py. На одном ядре загруженного CPU десять эпох шли около 8.5 минут. Вывод (символы . + # — яркость пикселя по порогам 0.3, 0.5 и 0.7; числа и картинки у вас будут немного другими):

параметров: 3.9 млн
эпоха 1: loss 1.2015   (42 с)
эпоха 2: loss 1.1096   (85 с)
эпоха 3: loss 0.9994   (136 с)
эпоха 4: loss 0.9209   (187 с)
эпоха 5: loss 0.8692   (236 с)
эпоха 6: loss 0.8325   (286 с)
эпоха 7: loss 0.8049   (340 с)
эпоха 8: loss 0.7823   (391 с)
эпоха 9: loss 0.7644   (447 с)
эпоха 10: loss 0.7488   (505 с)

— цифры 0…4: 1 шаг(ов) Эйлера, w = 2 —
            .                                   .           .
                                .    .  . .       .        ..           .
 . .   .+.             . .  .         +#.+           .#++
     ++###+            .  ..         ##. +          ..+###          ..   .+
    #+###...           #              .  #...          +++        .  .. .
    +#+   #+           #+          .     ++      .    .+#.         +...  +
    ##    ++    .     .#.             . +#           +###+ .      .+#.  .#.
   .#.   +#. .        #.    .        .+#.+            #+....      .######+.  .
  ..#    .#+     .    .# .  .      .###+#+##.     .     #+          ++####.
  .+#   +##           .. .        .+##+#++         +    ##+          . ...
   +###+#.   .      ..+.          . #+. ..+.        +. ##.            . .
    +#+## .         .+.             + . ...  .     .#+##+           .  .
     .                 ...                   .  .     . .              ..
    .      .              .        .        .

— цифры 0…4: 20 шаг(ов) Эйлера, w = 2 —
                        .
                       .        .       .                  ..       .    .
 .    .++.             .               #.+           .#..                .
    .#####.         .  .  +.         ##. #.         ..++##          +.  .#.
    #+++#+       .  .  #              .  ##..           +#        . +.  ..
   .++  ++#            #+                ##      .     +#.         #.   ..
    ##. ....           +.               +#           ##++. .      +#.   .#+
    ++.  +#           +..   .        ++#.+.  .       .##+  .      +########+
    #.  ..#.     .    .# .         +###.####           .#+          ..+.+#+.
   .#.. +##            . .        .###.#.+.        +    #+           .  .+.
    ##+ #+.  .      . +. .        ..#+. . .            .#.            . +.
    .#+##+.          ..            .# .. .   .     .+.##+.          +  .
     +.+              ....          .                .#..              .  .
    ..       .                                        .

— цифры 0…4: 20 шаг(ов) Эйлера, w = 1 —
            .           .                  .
                       .        .       .                  ..       .    .
 .    .++.             .               # +            # .
    +#####.         .  .  +.         ##+.#.          .++##          ..  .#.
  ..#...#.       . .+. #             .+  ##..           ++        . +...+.
   ...  #.#        ... ++          . .   ##      .     .#.         #. .  .
    +#+ +. .           ..             . .#           ##.+. .      +#.  ..++
    .+#+.+.           + .  ..        ++#.+.  .       .##+  .      .#+####+#+
    #.  ..+.     .    .# .          ##+ #+##          .+#.          . +..++.
    +.. .++            . .         .## ..+         +    #+           .   +.
    +#. ++.  .      . +. .          #+ ....            .#             . +.
    .#.##+.         .+.             #.##..   .      ..+#+           +  .
     +.+              ....          .  .             .#..              .  .
    ..       .                                        .

Как читать результат:

• Loss падает, но не к нулю — как и на «лунах»: цель $x_1 - x_0$ содержит случайный шум, который угадать нельзя.
• Один шаг дал «среднюю цифру» своего класса — расплывчатую, но узнаваемую. Это ровно предсказание теории: из $t = 0$ сеть может сказать только «едь к среднему данных», а с условием — к среднему среди нулей, единиц и так далее. Без условия получилось бы серое пятно.
• 20 шагов с $w = 2$ — штрихи чётче и цельнее, чем за один шаг; без guidance ($w = 1$) штрихи у тех же цифр более рваные и неуверенные. Guidance усиливает «то, что добавляет условие».
• Фон в точках и неровные штрихи — предел маленького MLP после десяти эпох на CPU. На GPU обучите 30+ эпох или замените MLP на маленькую U-Net или DiT по патчам — цифры станут чистыми.

🖥 На вашей машине (4×GPU)

• Скрипт из урока на GPU. Поставьте EPOCHS = 30: на одной RTX 4080 это минуты. Сравните 1, 2, 4, 8 и 20 шагов и $w$ = 1, 2, 4 — увидите и размытость малого числа шагов, и «пережжённые» цифры при большом $w$.
• Reflow своими руками. Сгенерируйте обученной моделью 60 000 пар (шум, цифра) за 50 шагов — на четырёх картах параллельно, по четверти на каждую (CUDA_VISIBLE_DEVICES=0..3, четыре процесса). Дообучите копию модели на этих парах вместо случайных и сравните генерацию за 1–2 шага до и после.
• FLUX.1 [schnell] в diffusers. Трансформер FLUX — около 12 млрд параметров, в bf16 это ≈24 ГБ только весов, а вместе с текстовым энкодером T5 на одну карту не помещается: включите pipe.enable_model_cpu_offload() или держите энкодеры на второй карте. Сгенерируйте одну картинку за 1, 2 и 4 шага (num_inference_steps) — это ваша дистилляция в действии.

Время и память — ориентиры: зависят от разрешения, версий библиотек и того, что ещё занимает карты. Подробнее о машине — на странице «Ваша лаборатория».

Китайский угол: прямые потоки в китайских моделях

Авторы rectified flow — Лю Синчао, Гун Чэнъюэ и Лю Цян из Техасского университета в Остине. Они же сделали InstaFlow — Stable Diffusion за один шаг через reflow и дистилляцию.

Многие китайские генераторы последних лет учатся на потоках. В README Wan2.1 прямо написано, что модель построена во фреймворке flow matching; у HunyuanVideo в параметрах запуска — --flow-shift и --flow-reverse; Qwen-Image обучен с целью flow matching. DeepSeek в JanusFlow встроила rectified flow прямо в языковую модель вместо VQ-токенов (урок 9.1.3). В речи то же самое: CosyVoice (Alibaba) переводит речевые токены в спектрограмму моделью flow matching, а GLM-4-Voice (Zhipu) построила свой декодер на её архитектуре. F5-TTS (лаборатория X-LANCE Шанхайского университета Цзяотун) синтезирует речь чистым flow matching поверх DiT.

Читать на китайском: Су Цзяньлинь разбирает rectified flow как «просто идти по прямой» в 17-й части серии о диффузии, а в 30-й — переход «от мгновенной скорости к средней», то есть генерацию в один шаг в духе MeanFlow. Среди авторов MeanFlow — Каймин Хэ, автор ResNet (урок 3.4.2).

Проверь себя

Проверь себя6 вопросов·

1 Решаем $dx/dt = -x$ из $x(0) = 1$ методом Эйлера за 4 шага ($h = 0.25$). Чему равно $x(1)$? Округлите до трёх знаков.

2 Шум $x_0 = 0.5$, пример $x_1 = 3$, $t = 0.2$. Какая точка $x_t$ подаётся сети и какая у неё цель?

$x_t = 1.0$, цель $2.5$$x_t = 2.5$, цель $0.5$$x_t = 1.0$, цель $3$$x_t = 0.6$, цель $2.5$

3 Почему генерация за один шаг Эйлера из обычной (не дистиллированной) модели даёт размытую картинку?

Потому что на один шаг не хватает памятиПри $t = 0$ сеть видит чистый шум и предсказывает усреднённую скорость, так что все точки едут к среднему данныхПотому что сеть обучали только на $t > 0.5$Потому что шаг Эйлера всегда добавляет шум

4 Сеть предсказала скорость $v = 2$ в точке $x_t = 1.2$ при $t = 0.6$. Какую картинку $\hat x_1$ она «имеет в виду»?

5 Какие утверждения верны?Можно выбрать несколько ответов

Если поле постоянно вдоль траектории, Эйлер точен за один шагReflow дообучает модель на парах «шум → сгенерированная ей же картинка», чтобы траектории стали прямееFlow matching требует расписания $\beta_1, \dots, \beta_T$, как DDPMCFG для скоростей: $v_\varnothing + w\,(v_c - v_\varnothing)$

6 Сдвиг расписания $t' = s\,t / (1 + (s - 1)\,t)$ при $s = 3$. Куда переходит $t = 0.5$? (1 — чистый шум)

Проверить ответы

Итоги

• Генерация — это решение ОДУ $dx/dt = v(x, t)$ от шума к данным; на практике — несколько шагов Эйлера $x \leftarrow x + h\,v$. Чем прямее траектории, тем меньше шагов нужно.
• Flow matching учит сеть скорости $x_1 - x_0$ на точках отрезка $(1 - t)x_0 + t x_1$. Усреднённое поле переносит облако шума точно в данные, хотя отдельные траектории получаются кривыми, а один шаг даёт «среднюю» картинку.
• Скорость, шум и картинка переводятся друг в друга линейно — DDPM и flow matching говорят на одном языке.
• Шаги экономят reflow, сдвиг расписания и дистилляция (consistency, состязательная, MeanFlow; FLUX.1 [schnell] — 4 шага). CFG работает со скоростями так же, как с шумом.

Дальше: речь — как звуковая волна превращается в спектрограмму и токены, как работают Whisper, нейрокодеки и синтез речи на тех же flow matching и VQ, что вы только что изучили.