Modèles de diffusion en une étape avec appariement de distributions par f-Divergence

One-step Diffusion Models with f-Divergence Distribution Matching

February 21, 2025

Auteurs: Yilun Xu, Weili Nie, Arash Vahdat

cs.AI

Résumé

L'échantillonnage à partir de modèles de diffusion implique un processus itératif lent qui entrave leur déploiement pratique, en particulier pour les applications interactives. Pour accélérer la vitesse de génération, les approches récentes distillent un modèle de diffusion à plusieurs étapes en un générateur étudiant à une seule étape via une distillation variationnelle des scores, qui aligne la distribution des échantillons générés par l'étudiant sur celle de l'enseignant. Cependant, ces approches utilisent la divergence de Kullback-Leibler (KL) inverse pour l'alignement des distributions, qui est connue pour être orientée vers la recherche de modes. Dans cet article, nous généralisons l'approche d'alignement des distributions en utilisant un nouveau cadre de minimisation de f-divergence, appelé f-distill, qui couvre différentes divergences avec des compromis variés en termes de couverture des modes et de variance d'entraînement. Nous dérivons le gradient de la f-divergence entre les distributions de l'enseignant et de l'étudiant et montrons qu'il s'exprime comme le produit de leurs différences de scores et d'une fonction de pondération déterminée par leur rapport de densité. Cette fonction de pondération met naturellement l'accent sur les échantillons ayant une densité plus élevée dans la distribution de l'enseignant, lors de l'utilisation d'une divergence moins orientée vers la recherche de modes. Nous observons que l'approche populaire de distillation variationnelle des scores utilisant la divergence KL inverse est un cas particulier dans notre cadre. Empiriquement, nous démontrons que des f-divergences alternatives, telles que la divergence KL directe et la divergence de Jensen-Shannon, surpassent les meilleures méthodes actuelles de distillation variationnelle des scores dans diverses tâches de génération d'images. En particulier, lors de l'utilisation de la divergence de Jensen-Shannon, f-distill atteint les performances actuelles de pointe en génération en une étape sur ImageNet64 et en génération texte-image zero-shot sur MS-COCO. Page du projet : 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