Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Диффузионную модель научили исправлять собственные промежуточные изображения до того, как ошибка перейдёт в следующий шаг генерации. В работе UC Berkeley и Reve показатель FID у DiT снизился более чем в четыре раза при сопоставимом объёме вычислений, хотя препринт не рецензирован, а числа получили сами авторы. Метод можно совместить с управлением без классификатора (CFG), а затем убрать дополнительные вычисления с помощью дистилляции.
Как корректор останавливает накопление ошибок
Обычный диффузионный генератор учится восстанавливать чистое изображение из смеси реальной картинки и шума. Во время генерации условия меняются: на каждом следующем шаге модель получает результат, который зависит от её предыдущих, неидеальных предсказаний.
Так возникает разрыв между обучением и использованием. Ошибка в форме объекта появляется рано, а следующие шаги не исправляют её, а добавляют поверх текстуры и мелкие детали. В статье показан характерный случай: генератор прорисовывает глаза, зубы и язык, сохраняя неверную структуру объекта.
Correction training добавляет вторую задачу. Базовый удалитель шума сначала предсказывает чистое изображение, после чего корректор получает именно это предсказание и учится восстанавливать ту же исходную картинку. Градиент от корректора не проходит обратно через его вход, поэтому базовая модель не может искусственно приспособить своё предсказание под последующий этап.
При генерации корректор несколько раз применяют к собственному результату на одном уровне шума. Затем исправленное изображение используют для следующего шага удаления шума. Хотя при обучении показывают только один проход коррекции, повторные проходы продолжают улучшать структуру без отдельного обучения для каждой итерации.
Это важное отличие от простого увеличения числа шагов диффузии. Дополнительный шаг удаления шума продолжает траекторию, уже содержащую ошибку, а корректор пересматривает само предсказание чистого изображения до перехода дальше.
Качество выросло при том же бюджете вычислений
На DiT-XL/2 коррекция снизила FID с 9,58 до 2,07. FID измеряет расстояние между распределениями реальных и сгенерированных изображений: чем он ниже, тем ближе наборы. В безусловной генерации по ImageNet показатель улучшился в 4,6 раза, хотя модель не получала ни класс, ни текстовую подсказку.
В генерации по тексту результат GenEval вырос на 21,79 пункта без CFG. Этот тест проверяет, выполняет ли изображение условия запроса: правильно ли передаёт объекты, их количество, цвета, положение и связи. При совместном использовании с CFG прибавка сохранилась, но стала меньше.
Основные сравнения проводили при бюджете около 250 запусков сети на изображение. Когда добавляли коррекцию, число обычных шагов удаления шума сокращали, поэтому выигрыш нельзя объяснить тем, что модель просто получила больше вычислений. Коррекция также сохраняла более высокую полноту — разнообразие охваченных вариантов — при качестве, сопоставимом с CFG, и реже давала перенасыщенные изображения.
Метод проверили на архитектурах DiT, SiT, SiTREPA и JiT. Эксперименты охватывали условную и безусловную генерацию ImageNet, а также генерацию по тексту в NanoGen на смеси JourneyDB и наборов подписей BLIP-3o. Результаты относятся к этим моделям, данным и протоколам, а не ко всем диффузионным системам.
Когда correction training меняет план разработки
Работа меняет план команды, если она обучает или дообучает собственную диффузионную модель и выбирает между дополнительными шагами генерации и CFG. Вместо того чтобы тратить весь бюджет на продолжение удаления шума, можно выделить часть запусков сети на пересмотр промежуточного предсказания. Для SiTREPA дополнительное обучение корректора заняло 1,25% от числа обновлений при предобучении базовой модели.
Коррекция не заменяет CFG во всех режимах. Её можно применять самостоятельно, когда важны разнообразие и естественная насыщенность, либо поставить перед CFG, если нужна максимальная привязка к условию. Глубина коррекции становится отдельной настройкой: дополнительные проходы повышают точность и связность, но постепенно сокращают разнообразие и при избытке могут усиливать неестественные контуры.
До оптимизации корректор увеличивает стоимость генерации, потому что требует дополнительных проходов сети. Авторы предлагают дистиллировать исправленные предсказания обратно в базовую модель: ученик воспроизводит результат повторной коррекции одним проходом на каждом шаге. После этого сохраняются архитектура, число параметров и стоимость инференса базовой модели, а значительная часть выигрыша в качестве остаётся.
Практический план состоит из трёх этапов: проверить коррекцию при фиксированном вычислительном бюджете, подобрать баланс качества и разнообразия, затем дистиллировать результат для рабочего контура. Для продукта, который использует только закрытый API генерации и не управляет обучением модели, работа сама по себе ничего не меняет: метод требует доступа к промежуточным предсказаниям и обучающему контуру.
Источники
Иллюстрация: рисунок из статьи «Correction Training for Diffusion Models», Donghoon Ahn, Runqian Wang, Alexei A. Efros и др., CC BY 4.0
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
