Янник Кильхер объяснил математику Flow Matching в Stable Diffusion 3

Yannic Kilcher 101 тыс. 56 мин 4 мин 08.04.2024
Главное

В современном ландшафте генеративного искусственного интеллекта происходит фундаментальный сдвиг от классических диффузионных моделей к более общим и эффективным подходам. Известный ИИ-исследователь Янник Кильхер (Yannic Kilcher) подробно разбирает концепцию Flow Matching (согласование потоков) — математический метод, который лёг в основу таких прорывных технологий, как Stable Diffusion 3. Этот подход позволяет значительно ускорить генерацию контента за счёт выпрямления траекторий движения от случайного шума к целевому распределению данных.

🔄 От диффузии к генеративному потоку 0:00

Традиционно для генерации изображений по текстовому описанию (например, в моделях вроде DALL-E или ранних версиях Stable Diffusion) использовались диффузионные процессы. В отличие от классических подходов на основе GAN (генеративно-состязательных сетей) или VAE (вариационных автокодировщиков), диффузия представляет собой многошаговый процесс. Генерация начинается с полного случайного шума, сэмплированного из стандартного гауссова распределения, после чего нейросеть шаг за шагом вычисляет всё более очищенные («деноизированные») версии изображения, пока не будет достигнут целевой результат. Янник Кильхер иронично замечает, что в день записи ролика празднуется «Международный день худи», поэтому он ведет трансляцию в толстовке, но это не мешает ему углубиться в суровый математический анализ.

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

Концепция Flow Matching радикально обобщает этот подход. По словам Кильхера, вместо жесткой фиксации процесса зашумления исследователи из Meta AI и Института Вейцмана предложили напрямую обучать модель трансформировать (морфировать) исходное легко сэмплируемое распределение в целевое распределение данных.

🧮 Математический фундамент: путевые плотности и векторные поля 10:42

В основе математического аппарата Flow Matching лежат три главных объекта:

По словам ведущего, векторное поле определяет, с какой скоростью и куда должна двигаться каждая точка, чтобы исходное гауссово распределение со временем перетекло в целевое. Векторное поле само по себе динамично и может меняться от шага к шагу.

На практике, для генерации нового объекта пользователь берет случайную точку из исходного распределения и с помощью обычного численного решателя обыкновенных дифференциальных уравнений (ODE solver) перемещает её вперед по времени до $t=1$ вдоль выученного векторного поля. Главная задача метода — корректно обучить эту функцию векторного поля $v_t$, чтобы регрессировать целевой поток.

🧩 Концепция условных потоков и маргинализация 21:54

Поскольку точное целевое распределение реального мира неизвестно, на практике его аппроксимируют с помощью доступных сэмплов из датасета, например, через модель смеси гауссиан (Gaussian Mixture Model). Чтобы сделать вычисления реализуемыми, авторы исследуемой статьи применили математический трюк: переход к условным путям вероятности (conditional probability paths), привязанным к конкретным конечным точкам $x_1$ из датасета.

Для отдельного сэмпла задаются следующие граничные условия:

Объединение (маргинализация) таких индивидуальных траекторий по всему датасету дает итоговый агрегированный путь плотности, который при $t=1$ с высокой точностью воссоздает реальное распределение данных. Аналогичным образом маргинализируются и условные векторные поля.

Как подчеркивает Янник Кильхер, прямое вычисление таких маргинальных интегралов математически невозможно (intractable) для больших размерностей. Решением стал метод Conditional Flow Matching (CFM). Авторы доказали фундаментальную теорему: градиенты функции потерь CFM, вычисленной для отдельных пар точек, в точности равны градиентам оригинальной, невычислимой функции Flow Matching. Это позволяет обучать нейросеть на индивидуальных парах «шум — изображение» без необходимости интегрирования по всему датасету, гарантируя достижение глобального оптимума.

🏎️ Оптимальный транспорт и ускорение генерации 38:43

В рамках предложенного фреймворка можно выбирать различные траектории интерполяции. Авторы статьи продемонстрировали, что классическая диффузия является лишь частным, избыточно усложненным случаем Flow Matching, где траектории движения искривлены, а для полного зашумления требуется бесконечное время ($t \to \infty$). Из-за этого диффузионные модели требуют сложных вычислений на границах временных отрезков, где функции теряют стабильность.

Вместо этого разработчики Flow Matching предложили использовать пути Оптимального Транспорта (Optimal Transport paths), которые обладают уникальными свойствами:

По оценке Кильхера, визуализация траекторий наглядно показывает преимущество оптимального транспорта: геометрия целевого распределения формируется значительно раньше, чем при диффузии. При использовании ODE-решателя для генерации изображений по выученному прямолинейному полю требуется в разы меньше шагов вычисления функций (FEV), что делает генерацию Stable Diffusion 3 и аналогичных моделей принципиально более быстрой и вычислительно экономичной. Будущее генерации ИИ, по мнению сообщества Кильхера, лежит именно в переходе на прямые траектории Оптимального Транспорта.

💬 Цитаты

«Flow Matching обобщает этот процесс: если задать потоки определенным образом, вы в точности вернетесь к диффузии.»

Янник Кильхер 09:47

«Вместо агрегирования данных перед обучением, мы можем обучать нейросеть предсказывать условное векторное поле для отдельной точки.»

Янник Кильхер 32:04
👥 Спикер
🔗 Упомянутые сайты и проекты
📖 Термины
Flow Matching
Генеративный метод машинного обучения, основанный на непрерывных потоках плотностей и нормализации векторных полей.
Оптимальный транспорт
Математическая концепция перемещения распределения масс из одного состояния в другое с минимальными затратами, задающая прямые траектории в Flow Matching.
ODE solver
Численный решатель обыкновенных дифференциальных уравнений, используемый для пошаговой генерации данных вдоль векторного поля.
📊 Цифры
⚖️ Другая сторона
Искусственный интеллект Flow Matching Янник Кильхер Stable Diffusion 3 Оптимальный транспорт