Distillazione EM per Modelli di Diffusione a Singolo Passo

EM Distillation for One-step Diffusion Models

May 27, 2024

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

cs.AI

Abstract

Sebbene i modelli di diffusione possano apprendere distribuzioni complesse, il campionamento richiede un processo iterativo computazionalmente costoso. I metodi di distillazione esistenti consentono un campionamento efficiente, ma presentano limitazioni significative, come il degrado delle prestazioni con un numero molto ridotto di passaggi di campionamento, la dipendenza dall'accesso ai dati di training o l'ottimizzazione orientata alla modalità che potrebbe non catturare l'intera distribuzione. Proponiamo EM Distillation (EMD), un approccio basato sulla massima verosimiglianza che distilla un modello di diffusione in un modello generatore a un solo passaggio con una perdita minima della qualità percettiva. Il nostro approccio è derivato attraverso la lente di Expectation-Maximization (EM), in cui i parametri del generatore vengono aggiornati utilizzando campioni dalla distribuzione congiunta del modello di diffusione insegnante e dei latenti inferiti del generatore. Sviluppiamo uno schema di campionamento riparametrizzato e una tecnica di cancellazione del rumore che insieme stabilizzano il processo di distillazione. Inoltre, riveliamo un'interessante connessione del nostro metodo con i metodi esistenti che minimizzano la KL orientata alla modalità. EMD supera i metodi generativi a un solo passaggio esistenti in termini di punteggi FID su ImageNet-64 e ImageNet-128, e si confronta favorevolmente con il lavoro precedente sulla distillazione di modelli di diffusione testo-immagine.

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.