Consistent Distribution Matching сжал диффузию до 1 шага с FID 2,04

Диффузионные и flow-модели генерируют качественные картинки, но платят за это числом шагов численного интегрирования. Группа авторов предложила Consistent Distribution Matching (CDM), метод дистилляции, который ускоряет такие модели без датасета, без прогонов учителя и без вспомогательных сетей. На ImageNet 256×256 студент выдаёт FID 2,04 за одну оценку функции (1-NFE) и FID 1,37 за четыре (4-NFE). Обучение занимает 40 эпох. По данным статьи, результат превосходит лучшие базовые методы дистилляции, работающие без данных. Код и модель опубликованы на consistentdmd.github.io.

Что не так с текущей дистилляцией

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

Каждая из этих зависимостей усложняет масштабирование. Датасет нужно хранить и чистить, прогоны учителя съедают вычисления, а лишняя сеть добавляет гиперпараметры и риск нестабильного обучения. Авторы CDM целятся именно в этот слой инженерных издержек.

Как устроен CDM

Ключевой ход: одна сеть-студент выполняет сразу две задачи, генерирует сэмплы и оценивает скор (градиент логарифма плотности). В схемах с отдельной прокси-сетью эти функции обычно разнесены. Здесь они объединены, поэтому в обучении участвуют только две модели: замороженный учитель и обучаемый студент. Оптимизируется одна цель.

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

Метод называется data-free и simulation-free. Обучение не требует реальных изображений и не требует прогонять учителя по всей траектории.

Ключевые метрики

• Бенчмарк: ImageNet 256×256.
• 1-NFE: FID 2,04.
• 4-NFE: FID 1,37.
• Бюджет обучения: 40 эпох.
• Компоненты: замороженный учитель и обучаемый студент, одна функция потерь.
• Данные для дистилляции: не используются.
• Теория: доказана сходимость по Вассерштейну для flow-map pushforward.
• Релиз: код и модель в открытом доступе.

Что проверять на практике

Все цифры в статье относятся к классовой генерации на ImageNet 256×256. Это стандартный, но довольно узкий полигон. Нет оснований переносить результат на text-to-image модели с миллиардами параметров, длинными промптами и высоким разрешением. Подтверждение масштабируемости станет главным вопросом к CDM. Заявление авторов о сравнении с лучшими базовыми методами без данных стоит проверить по таблицам в самой статье и по воспроизведению на открытом коде.

Тем, кто строит продуктовый инференс с генерацией изображений или видео, метод интересен по трём причинам. Во-первых, не нужно готовить датасет под дистилляцию, а для закрытых или лицензионно сложных учителей это снимает отдельную проблему. Во-вторых, пайплайн из двух моделей и одного лосса проще поддерживать и дешевле отлаживать. В-третьих, один шаг с FID 2,04 и четыре шага с FID 1,37 указывают на удобный компромисс: можно выбирать между максимальной скоростью и качеством без переобучения с нуля.

Логичный первый шаг для команды: взять опубликованный код, повторить дистилляцию на небольшой собственной flow-модели и измерить FID при 1 и 4 шагах. Сравнить затраты GPU-часов с текущим рецептом стоит сразу, потому что экономия на данных и прогонах учителя важна не меньше самой метрики.