Eénstaps Diffusiemodellen met f-Divergentie Distributie Matching

One-step Diffusion Models with f-Divergence Distribution Matching

February 21, 2025

Auteurs: Yilun Xu, Weili Nie, Arash Vahdat

cs.AI

Samenvatting

Het bemonsteren van diffusiemodellen is een traag iteratief proces dat hun praktische inzetbaarheid belemmert, vooral voor interactieve toepassingen. Om de generatiesnelheid te versnellen, distilleren recente benaderingen een meerstaps diffusiemodel in een enkelstaps studentgenerator via variatie-score-distillatie, waarbij de verdeling van monsters gegenereerd door de student wordt afgestemd op de verdeling van de leraar. Deze benaderingen gebruiken echter de omgekeerde Kullback-Leibler (KL)-divergentie voor verdelingsovereenkomst, die bekend staat om zijn moduszoekende karakter. In dit artikel generaliseren we de verdelingsovereenkomstbenadering met behulp van een nieuw f-divergentie-minimalisatiekader, genaamd f-distill, dat verschillende divergenties omvat met verschillende afwegingen in termen van modusdekking en trainingsvariantie. We leiden de gradiënt van de f-divergentie tussen de leraar- en studentverdelingen af en laten zien dat deze wordt uitgedrukt als het product van hun scoreverschillen en een gewichtsfunctie bepaald door hun dichtheidsratio. Deze gewichtsfunctie benadrukt van nature monsters met een hogere dichtheid in de leraarverdeling bij gebruik van een minder moduszoekende divergentie. We merken op dat de populaire variatie-score-distillatiebenadering met de omgekeerde KL-divergentie een speciaal geval is binnen ons kader. Empirisch tonen we aan dat alternatieve f-divergenties, zoals forward-KL en Jensen-Shannon-divergenties, de huidige beste variatie-score-distillatiemethoden overtreffen in beeldgeneratietaken. Met name bij gebruik van Jensen-Shannon-divergentie bereikt f-distill de huidige state-of-the-art prestaties voor enkelstapsgeneratie op ImageNet64 en zero-shot tekst-naar-beeldgeneratie op MS-COCO. Projectpagina: 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