Ключевое наблюдение статьи заключается в следующем:
🎯 Стандартная диффузия делает много шагов интегратора и много раз запускает диффузионную модель
🎯 Дистиллированная диффузия делает мало шагов интегратора и мало раз запускает диффузионную модель
А что если делать много шагов интегратора и мало раз запускать диффузионную модель?
Авторы предлагают модифицировать процесс дистилляции так, чтобы модель выдавала некоторую политику (нечто определяющее процесс интегрирования) на всей траектории. Выходы модели необходимо модифицировать так (расширив число выходных каналов), чтобы она выдавала несколько характеристик, соответствующих разному уровня шума, на ODE траектории.
Предлагают 2 политики:
📌 DX - политику. Выглядит как просто предсказание скоростей на разных участках траектории.
📌 GMFlow. Сеть предсказывает факторизованную по шагам расшумления смесь Гауссиан, каждая компонента отвечает своей моде расшумления.
Процедура обучения напоминает Consistency Distillation:
1️⃣ Сэмплируется начальный шаг зашумления.
2️⃣ Предсказывается политика.
3️⃣ С замороженной политикой (no grad) расшумляют немного.
4️⃣ Предсказывают политику из новой точки.
5️⃣ Считают MSE между ее предсказаниями и предсказаниями из исходной точки.
Траектории на обучении можно генерить расшумляя некоторые входные данные, а можно прямо с нуля генерировать.
🧪 Эксперименты
Тестируют все это дело на ImageNet и Flux-1 / Qwen-Image.
GMFlow работает заметно лучше, чем DX и выдает около SOTA FID при 1- и 2- шаговой генерации.
π-Flow дистиллы Flux и Qwen-Image примерно на одном уровне по метрикам (HPS, ClipScore, FID) с другими SOTA дистиллами. Впрочем, без user preference study сложно сделать конкретные выводы.
Далее на конкретных картинках показывают, что конкуренты шакалят текстуры, а они нет.
Скорость инференса от перехода к многошаговой политике практически не меняется.
💡 Выводы
Выглядит как Consistency Distillation на стероидах, и будто бы не хватает сравнения с sCM / rCM. Здесь, правда, не надо париться с JVP.