EM-Destillation für Ein-Schritt-Diffusionsmodelle

EM Distillation for One-step Diffusion Models

May 27, 2024

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

cs.AI

Zusammenfassung

Obwohl Diffusionsmodelle komplexe Verteilungen erlernen können, erfordert das Abtasten einen rechenaufwändigen iterativen Prozess. Bestehende Destillationsmethoden ermöglichen effizientes Abtasten, weisen jedoch bemerkenswerte Einschränkungen auf, wie Leistungsverschlechterung bei sehr wenigen Abtastschritten, Abhängigkeit vom Zugriff auf Trainingsdaten oder modussuchende Optimierung, die möglicherweise nicht die gesamte Verteilung erfasst. Wir schlagen EM-Destillation (EMD) vor, einen Maximum-Likelihood-basierten Ansatz, der ein Diffusionsmodell auf ein Ein-Schritt-Generatormodell destilliert, bei minimalen Verlusten an perzeptueller Qualität. Unser Ansatz wird durch die Brille der Erwartungs-Maximierungs-Methode (EM) abgeleitet, bei der die Generatorparameter unter Verwendung von Stichproben aus der gemeinsamen Verteilung der Diffusionslehrer-Priorität und der erschlossenen Generator-Latenzen aktualisiert werden. Wir entwickeln ein reparametrisiertes Abtastschema und eine Rauschunterdrückungstechnik, die zusammen den Destillationsprozess stabilisieren. Darüber hinaus enthüllen wir eine interessante Verbindung unserer Methode mit bestehenden Methoden, die die modussuchende KL minimieren. EMD übertrifft bestehende Ein-Schritt-generative Methoden in Bezug auf FID-Werte bei ImageNet-64 und ImageNet-128 und vergleicht sich günstig mit früheren Arbeiten zur Destillation von Text-zu-Bild-Diffusionsmodellen.

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.