Как подключить PEEK к диффузионному сэмплеру и настроить интерполяционный график без переобучения

PEEK — это training-free детерминированный плагин для сэмплера диффузионной модели. Он исправляет exposure bias: разрыв между тем, на каких данных обучался денойзер, и тем, какие зашумлённые состояния он реально получает во время генерации. На каждом подходящем шаге PEEK выполняет пробное одношаговое обновление, оценивает денойзер в состоянии look-ahead и интерполирует текущее чистое предсказание с предсказанием на шаг позже, формируя следующее зашумлённое состояние. Это описание взято из статьи авторов на arXiv.

Что понадобится до начала

Для воспроизведения нужен работающий диффузионный сэмплер с доступом к функции денойзера и к предсказанию чистого изображения на каждом шаге. Конкретный API зависит от того, какой сэмплер вы используете: в статье нет привязки к определённой библиотеке, поэтому подойдёт любой pixel-space или latent-space сэмплер, где можно вызвать денойзер и получить чистое предсказание. Авторы проверяли PEEK на CIFAR-10, LSUN Bedroom, ImageNet и MS-COCO в unconditional, label-guided и text-guided режимах. Переобучение модели не требуется: плагин работает как надстройка поверх существующего сэмплера.

Установка PEEK поверх существующего сэмплера

PEEK не заменяет сэмплер целиком, а добавляет к нему интерполяционный слой. Логика одного шага по описанию в статье такова:

1. Получите текущее чистое предсказание x0_hat из вашего сэмплера на шаге t.
2. Выполните пробное одношаговое обновление к следующему зашумлённому состоянию.
3. Оцените денойзер в этом look-ahead состоянии и получите предсказание на шаг позже.
4. Интерполируйте два чистых предсказания с коэффициентом из интерполяционного графика.
5. Сформируйте следующее зашумлённое состояние из интерполированного предсказания.

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

Выбор интерполяционного графика жадным поиском

Интерполяционный график PEEK выбирается жадным поиском, управляемым метрикой. Поиск перебирает позиции интерполяции, оценивает каждую метрикой и выбирает лучший шаг. Ускорение обеспечивают два кэша: PEEK prefix cache и baseline noise-schedule rebasing cache.

1. Задайте набор кандидатов позиций интерполяции для графика.
2. Для каждого кандидата вычислите метрику качества с учётом кэша.
3. Выберите позицию с лучшей метрикой и зафиксируйте её в графике.
4. Повторяйте, пока график не покроет нужные шаги сэмплирования.

Единоразовая калибровка PEEK для одной модели занимает не более 2,2 минуты и является единственной дополнительной стоимостью сверх обычного сэмплирования, по данным экспериментов авторов. Это значение относится к их настройке и оборудованию, а не к гарантированному времени на любой машине.

Ожидаемый результат и ограничения

Авторы сообщают о снижении Fréchet Inception Distance (FID) вплоть до 49,3% по сравнению с базовыми методами при одинаковом числе оценок функции (NFE) на CIFAR-10, LSUN Bedroom, ImageNet и MS-COCO. Это верхняя граница наблюдаемого сокращения FID в их экспериментах, а не единый показатель для всех моделей и датасетов. Собственное воспроизведение мы не выполняли, поэтому конкретные числа для вашей модели могут отличаться. Сравнение корректно только при выравнивании по NFE — если базовый сэмплер делает меньше вызовов денойзера, прямое сопоставление FID будет некорректным.

Типичная ошибка: учёт NFE и как её исправить

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

Как исправить: считайте NFE явно и подбирайте базовый сэмплер с тем же бюджетом вызовов денойзера, что и у PEEK. Только после этого сравнивайте FID. В статье все заявленные улучшения приведены именно для baselines, matched by number of function evaluations.

Что говорит лицензия и откуда брать код

Статья распространяется по лицензии CC BY 4.0. Проектный сайт авторов находится по адресу dachengruoque.github.io/PEEK-Video — там стоит искать актуальную реализацию и дополнительные материалы. Лицензия кода в отдельном репозитории, если он существует, в статье не указана; проверяйте её непосредственно в репозитории перед использованием в своём проекте.