Метод

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

Можно попробовать генерировать сначала в низком разрешении, а потом постепенно повышать интерполяцией. Но как 🤔?

Исходная модель может быть не адаптирована под малогашовую генерацию (SDXL в 256x256 выдает безумный поп-арт на любой запрос), да и процедура интерполяции зашумленного латента плохо определена. Наивный upsampling латентов приводит к сильные дефектам. Несколько лучше работает upsampling x0 оценки из предсказанного латента с последующим зашумлением назад, но все равно не здорово.

Потому дообучение неизбежно. При этом дистиллированные модели выполняют две роли одновременно - уменьшение количества шагов сэмплирования и super-resolution в латентном пространстве.

В качестве основы для процедуры дистилляции берется DMD2 как все еще state-of-the-art метод дистилляции. Дополнительно к нему навешивается предложенный Patch Distribution Matching (PDM) лосс, который стремится уравнять прогнанные через исходную диффузионную модель генерации модели-учителя и студента на уровне отдельных токенов.

Итоговый лосс содержит в себе:
🎯обратную KL-дивергенцию
🎯GAN лосс
🎯PDM лосс

В качестве данных для обучения используется синтетика, сгенерированная моделью-учителем (примерно 500к сэмплов).