Проблема: накопление ошибок при инференсе
Основная проблема итеративной выборки (sampling) в диффузионных моделях заключается в exposure bias. Во время обучения денуайзер (denoiser) обучается на реальных изображениях, интерполированных с шумом. Однако на этапе инференса входные данные формируются на основе собственных, imperfect-предсказаний модели с предыдущих шагов. Модель видит распределение данных, с которым никогда не сталкивалась при обучении, что приводит к накоплению ошибок и артефактам.
Решение: Correction Training
Авторы предлагают introduce correction step — дополнительный модуль, который принимает предсказание базовой модели и обучается восстанавливать ground-truth цель. Ключевая особенность: corrector обучается исключительно на предсказаниях базового денуайзера, что делает его специализированным инструментом для исправления именно тех ошибок, которые допускает основная модель.
В процессе генерации, на каждом шаге денуайзинга, предсказание чистого образца итеративно уточняется corrector'ом. Это позволяет исправлять ошибки на промежуточных этапах, не давая ранним артефактам персистировать в финальном результате.
Результаты и бенчмарки
Метод протестирован на архитектурах DiT, SiT и JiT. Correction Training демонстрирует существенное улучшение качества генерации (FID) при сохранении разнообразия (diversity). Особенно примечательны результаты в unconditional generation (без guidance), где ранее качество было низким.
| Метрика/Условие | Результат Correction Training | Сравнение с базовым подходом |
|---|---|---|
| Unconditional Generation (FID) | Значительно улучшено | Конкурентоспособно с CFG, ранее было слабым местом |
| Сочетание с CFG | Улучшено разнообразие | Снижена oversaturation (перенасыщенность), типичная для CFG |
| Вычислительные затраты | Могут быть амортизированы | Через дистилляцию refined-предсказаний обратно в базовую модель |
Практическая ценность
Метод не только повышает качество, но и предлагает путь к оптимизации инференса. Авторы показывают, что дополнительные вычислительные затраты на коррекцию можно компенсировать путем дистилляции улучшенных предсказаний обратно в базовую диффузионную модель. Это позволяет сохранить прирост качества без увеличения накладных расходов на этапе инференса, что критически важно для коммерческого применения.
Источник: arXiv cs.AI ↗