Modelli di Diffusione a Un Passo con Allineamento di Distribuzione basato su f-Divergenza

One-step Diffusion Models with f-Divergence Distribution Matching

February 21, 2025

Autori: Yilun Xu, Weili Nie, Arash Vahdat

cs.AI

Abstract

Il campionamento dai modelli di diffusione comporta un processo iterativo lento che ne ostacola l'impiego pratico, specialmente per applicazioni interattive. Per accelerare la velocità di generazione, approcci recenti distillano un modello di diffusione multi-step in un generatore studente a singolo passo tramite distillazione variazionale del punteggio, che allinea la distribuzione dei campioni generati dallo studente a quella del modello insegnante. Tuttavia, questi approcci utilizzano la divergenza inversa di Kullback-Leibler (KL), nota per essere orientata alla ricerca di modi. In questo articolo, generalizziamo l'approccio di allineamento delle distribuzioni utilizzando un nuovo framework di minimizzazione della f-divergenza, denominato f-distill, che copre diverse divergenze con diversi compromessi in termini di copertura dei modi e varianza dell'addestramento. Deriviamo il gradiente della f-divergenza tra le distribuzioni del modello insegnante e dello studente e mostriamo che è espresso come il prodotto delle differenze dei loro punteggi e di una funzione di ponderazione determinata dal loro rapporto di densità. Questa funzione di ponderazione enfatizza naturalmente i campioni con densità più alta nella distribuzione del modello insegnante quando si utilizza una divergenza meno orientata alla ricerca di modi. Osserviamo che il popolare approccio di distillazione variazionale del punteggio che utilizza la divergenza inversa di KL è un caso speciale all'interno del nostro framework. Empiricamente, dimostriamo che f-divergenze alternative, come la divergenza diretta di KL e la divergenza di Jensen-Shannon, superano i migliori metodi attuali di distillazione variazionale del punteggio in vari compiti di generazione di immagini. In particolare, quando si utilizza la divergenza di Jensen-Shannon, f-distill raggiunge le migliori prestazioni attuali per la generazione a singolo passo su ImageNet64 e la generazione zero-shot da testo a immagine su MS-COCO. Pagina del progetto: 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