EM-distillatie voor éénstaps diffusiemodellen

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

Samenvatting

Hoewel diffusiemodellen complexe verdelingen kunnen leren, vereist het bemonsteren een rekenintensief iteratief proces. Bestaande destillatiemethoden maken efficiënt bemonsteren mogelijk, maar hebben aanzienlijke beperkingen, zoals prestatievermindering bij zeer weinig bemonsteringsstappen, afhankelijkheid van toegang tot trainingsgegevens, of mode-zoekende optimalisatie die mogelijk niet de volledige verdeling kan vastleggen. Wij stellen EM Destillatie (EMD) voor, een op maximale waarschijnlijkheid gebaseerde aanpak die een diffusiemodel destilleert naar een één-staps generatormodel met minimaal verlies van perceptuele kwaliteit. Onze aanpak is afgeleid door de lens van Expectation-Maximization (EM), waarbij de generatorparameters worden bijgewerkt met behulp van steekproeven uit de gezamenlijke verdeling van de diffusie-leraarprior en afgeleide generatorlatenten. We ontwikkelen een herparametriseerd bemonsteringsschema en een ruisonderdrukkingstechniek die samen het destillatieproces stabiliseren. We onthullen verder een interessante verbinding van onze methode met bestaande methoden die mode-zoekende KL minimaliseren. EMD overtreft bestaande één-staps generatieve methoden in termen van FID-scores op ImageNet-64 en ImageNet-128, en vergelijkt gunstig met eerder werk over het destilleren van tekst-naar-beeld diffusiemodellen.

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.