Flow Maps: новый подход к ускорению генерации данных в диффузионных моделях

Сэмплирование из диффузионных моделей традиционно представляет собой итеративный процесс, где на каждом шаге модель оценивает направление движения в пространстве входных данных. Фактически, алгоритм вычисляет интеграл по уровням шума, постепенно трансформируя случайный шум в целевое распределение данных. В своей новой публикации автор предлагает рассмотреть альтернативный подход: обучение нейронных сетей напрямую предсказывать этот интеграл. Этот метод, известный как flow maps (карты потоков), позволяет значительно ускорить процесс генерации, минуя необходимость в множестве малых шагов.

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

Ключевым для понимания flow maps является взгляд на диффузионные модели как на определение биекции между шумом и данными. Существуют уникальные пути, соединяющие пары выборок из каждого распределения, которые никогда не пересекаются друг с другом. Алгоритмы сэмплирования можно разделить на стохастические и детерминированные. Стохастический подход, используемый в оригинальном алгоритме DDPM, предполагает выборку из условного распределения на каждом шаге. Детерминированный же метод, например DDIM или Flow Matching, использует фиксированное правило обновления без случайных выборок после начальной точки.

Важным следствием существования детерминированных алгоритмов является наличие детерминированного биективного отображения между отдельными выборками шума и данных. Каждая выборка шума связана с одной конкретной выборкой данных, и наоборот. Стандартный денойзер работает «бессознательно» и «кратковидно»: он знает только текущее положение в пространстве и уровень шума, предсказывая направление движения. Он не может заглянуть вперед к конечной цели или использовать информацию о ранее посещенных точках. Это похоже на навигацию методом пеленга, где путь уникален именно потому, что траектории не пересекаются.

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

Flow maps предлагают элегантное решение этой проблемы. Если стандартный денойзер описывается функцией f(x_t, t), то карта потока обозначается как F(x_s, s, t). Она принимает два временных шага: исходный (s) и целевой (t). Идеальная flow map позволяет мгновенно перейти от любой точки на пути к любой другой: F(x_s, s, t) = x_t. В практике это функция аппроксимируется нейронной сетью. При использовании стандартного расписания шума из Flow Matching, где x_t = (1 - t)x_0 + tε, карта потока может быть построена путем интегрирования скорости по временному интервалу.

Особый интерес представляет случай, когда целевой уровень шума t равен нулю. В этом случае можно напрямую перейти от любой точки на пути к конечной точке в пространстве данных: F(x_s, s, 0) = x_0. Если это предсказание достаточно точно, оно позволяет выполнять сэмплирование за один шаг. Именно этот принцип лежит в основе consistency models. Таким образом, flow maps не только ускоряют генерацию, но и обеспечивают более эффективное обучение на основе вознаграждений и улучшенную управляемость процесса сэмплирования.

Литература по этой теме изобилует различными формализмами и терминологией, что может затруднять понимание связи между методами. Автор статьи опирается в основном на таксономию, предложенную Boffi et al., чтобы прояснить эти различия. Хотя определение flow maps относительно просто, существует множество способов их построения и обучения. Понимание этих механизмов требует некоторой знакомства с основами диффузионных моделей и векторным анализом, но даже без глубокого математического бэкграунда статья предлагает ценные инсайты о геометрии генеративных процессов.