Distillation EM pour les modèles de diffusion en une étape

EM Distillation for One-step Diffusion Models

May 27, 2024

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

cs.AI

Résumé

Bien que les modèles de diffusion puissent apprendre des distributions complexes, l'échantillonnage nécessite un processus itératif coûteux en calcul. Les méthodes de distillation existantes permettent un échantillonnage efficace, mais présentent des limitations notables, telles qu'une dégradation des performances avec très peu d'étapes d'échantillonnage, une dépendance à l'accès aux données d'entraînement, ou une optimisation axée sur les modes qui peut échouer à capturer la distribution complète. Nous proposons EM Distillation (EMD), une approche basée sur le maximum de vraisemblance qui distille un modèle de diffusion en un modèle générateur en une seule étape avec une perte minimale de qualité perceptuelle. Notre approche est dérivée à travers le prisme de l'algorithme Expectation-Maximization (EM), où les paramètres du générateur sont mis à jour en utilisant des échantillons issus de la distribution conjointe du modèle de diffusion enseignant et des latents inférés du générateur. Nous développons un schéma d'échantillonnage reparamétré et une technique d'annulation du bruit qui stabilisent ensemble le processus de distillation. Nous révélons également une connexion intéressante de notre méthode avec les méthodes existantes qui minimisent le KL axé sur les modes. EMD surpasse les méthodes génératives en une étape existantes en termes de scores FID sur ImageNet-64 et ImageNet-128, et se compare favorablement aux travaux antérieurs sur la distillation de modèles de diffusion texte-image.

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.