Distilación EM para Modelos de Difusión en un Solo Paso

EM Distillation for One-step Diffusion Models

May 27, 2024

Autores: Sirui Xie, Zhisheng Xiao, Diederik P Kingma, Tingbo Hou, Ying Nian Wu, Kevin Patrick Murphy, Tim Salimans, Ben Poole, Ruiqi Gao

cs.AI

Resumen

Si bien los modelos de difusión pueden aprender distribuciones complejas, el muestreo requiere un proceso iterativo computacionalmente costoso. Los métodos de destilación existentes permiten un muestreo eficiente, pero tienen limitaciones notables, como la degradación del rendimiento con muy pocos pasos de muestreo, la dependencia del acceso a datos de entrenamiento o la optimización centrada en modos que puede fallar en capturar la distribución completa. Proponemos EM Distillation (EMD), un enfoque basado en máxima verosimilitud que destila un modelo de difusión a un modelo generador de un solo paso con una pérdida mínima de calidad perceptual. Nuestro enfoque se deriva a través de la lente de Expectation-Maximization (EM), donde los parámetros del generador se actualizan utilizando muestras de la distribución conjunta del modelo de difusión maestro y los latentes inferidos del generador. Desarrollamos un esquema de muestreo reparametrizado y una técnica de cancelación de ruido que juntos estabilizan el proceso de destilación. Además, revelamos una conexión interesante de nuestro método con métodos existentes que minimizan la divergencia KL centrada en modos. EMD supera a los métodos generativos de un solo paso existentes en términos de puntuaciones FID en ImageNet-64 e ImageNet-128, y se compara favorablemente con trabajos previos sobre la destilación de modelos de difusión de texto a imagen.

English

While diffusion models can learn complex distributions, sampling requires a computationally expensive iterative process. Existing distillation methods enable efficient sampling, but have notable limitations, such as performance degradation with very few sampling steps, reliance on training data access, or mode-seeking optimization that may fail to capture the full distribution. We propose EM Distillation (EMD), a maximum likelihood-based approach that distills a diffusion model to a one-step generator model with minimal loss of perceptual quality. Our approach is derived through the lens of Expectation-Maximization (EM), where the generator parameters are updated using samples from the joint distribution of the diffusion teacher prior and inferred generator latents. We develop a reparametrized sampling scheme and a noise cancellation technique that together stabilizes the distillation process. We further reveal an interesting connection of our method with existing methods that minimize mode-seeking KL. EMD outperforms existing one-step generative methods in terms of FID scores on ImageNet-64 and ImageNet-128, and compares favorably with prior work on distilling text-to-image diffusion models.