Одношаговые диффузионные модели с согласованием распределений на основе f-дивергенции

One-step Diffusion Models with f-Divergence Distribution Matching

February 21, 2025

Авторы: Yilun Xu, Weili Nie, Arash Vahdat

cs.AI

Аннотация

Сэмплирование из диффузионных моделей представляет собой медленный итеративный процесс, что ограничивает их практическое применение, особенно в интерактивных задачах. Для ускорения генерации современные подходы дистиллируют многошаговую диффузионную модель в одношаговый генератор-студент с помощью вариационного дистиллирования оценок, которое согласует распределение сэмплов, генерируемых студентом, с распределением учителя. Однако эти подходы используют обратную дивергенцию Кульбака-Лейблера (KL), которая известна своей склонностью к поиску мод. В данной работе мы обобщаем подход к согласованию распределений, используя новую минимизационную структуру f-дивергенции, названную f-distill, которая охватывает различные дивергенции с различными компромиссами в плане покрытия мод и дисперсии обучения. Мы выводим градиент f-дивергенции между распределениями учителя и студента и показываем, что он выражается как произведение разницы их оценок и весовой функции, определяемой отношением их плотностей. Эта весовая функция естественным образом акцентирует внимание на сэмплах с более высокой плотностью в распределении учителя при использовании менее модо-ориентированной дивергенции. Мы отмечаем, что популярный подход вариационного дистиллирования оценок с использованием обратной KL-дивергенции является частным случаем в нашей структуре. Эмпирически мы демонстрируем, что альтернативные f-дивергенции, такие как прямая KL и дивергенция Йенсена-Шеннона, превосходят современные методы вариационного дистиллирования оценок в задачах генерации изображений. В частности, при использовании дивергенции Йенсена-Шеннона f-distill достигает современных наилучших результатов в одношаговой генерации на ImageNet64 и в задаче генерации изображений по тексту на MS-COCO. Страница проекта: https://research.nvidia.com/labs/genair/f-distill.

English

Sampling from diffusion models involves a slow iterative process that hinders their practical deployment, especially for interactive applications. To accelerate generation speed, recent approaches distill a multi-step diffusion model into a single-step student generator via variational score distillation, which matches the distribution of samples generated by the student to the teacher's distribution. However, these approaches use the reverse Kullback-Leibler (KL) divergence for distribution matching which is known to be mode seeking. In this paper, we generalize the distribution matching approach using a novel f-divergence minimization framework, termed f-distill, that covers different divergences with different trade-offs in terms of mode coverage and training variance. We derive the gradient of the f-divergence between the teacher and student distributions and show that it is expressed as the product of their score differences and a weighting function determined by their density ratio. This weighting function naturally emphasizes samples with higher density in the teacher distribution, when using a less mode-seeking divergence. We observe that the popular variational score distillation approach using the reverse-KL divergence is a special case within our framework. Empirically, we demonstrate that alternative f-divergences, such as forward-KL and Jensen-Shannon divergences, outperform the current best variational score distillation methods across image generation tasks. In particular, when using Jensen-Shannon divergence, f-distill achieves current state-of-the-art one-step generation performance on ImageNet64 and zero-shot text-to-image generation on MS-COCO. Project page: https://research.nvidia.com/labs/genair/f-distill