Modelos de Difusión de un Paso con Emparejamiento de Distribución mediante f-Divergencia

One-step Diffusion Models with f-Divergence Distribution Matching

February 21, 2025

Autores: Yilun Xu, Weili Nie, Arash Vahdat

cs.AI

Resumen

El muestreo en modelos de difusión implica un proceso iterativo lento que dificulta su implementación práctica, especialmente en aplicaciones interactivas. Para acelerar la velocidad de generación, enfoques recientes destilan un modelo de difusión de múltiples pasos en un generador estudiantil de un solo paso mediante la destilación variacional de puntuaciones, que iguala la distribución de muestras generadas por el estudiante con la distribución del profesor. Sin embargo, estos enfoques utilizan la divergencia inversa de Kullback-Leibler (KL), conocida por ser propensa a la búsqueda de modas. En este artículo, generalizamos el enfoque de igualación de distribuciones utilizando un novedoso marco de minimización de f-divergencia, denominado f-distill, que abarca diferentes divergencias con distintos equilibrios en términos de cobertura de modas y varianza en el entrenamiento. Derivamos el gradiente de la f-divergencia entre las distribuciones del profesor y del estudiante y mostramos que se expresa como el producto de las diferencias de sus puntuaciones y una función de ponderación determinada por su razón de densidad. Esta función de ponderación enfatiza naturalmente las muestras con mayor densidad en la distribución del profesor cuando se utiliza una divergencia menos propensa a la búsqueda de modas. Observamos que el popular enfoque de destilación variacional de puntuaciones que utiliza la divergencia inversa de KL es un caso especial dentro de nuestro marco. Empíricamente, demostramos que f-divergencias alternativas, como la divergencia directa de KL y la divergencia de Jensen-Shannon, superan a los mejores métodos actuales de destilación variacional de puntuaciones en tareas de generación de imágenes. En particular, al utilizar la divergencia de Jensen-Shannon, f-distill logra el mejor rendimiento actual en generación de un solo paso en ImageNet64 y en generación de texto a imagen sin entrenamiento previo en MS-COCO. Página del proyecto: 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