Verbeterde Distributie Matching Distillatie voor Snelle Beeldsynthese

Improved Distribution Matching Distillation for Fast Image Synthesis

May 23, 2024

Auteurs: Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, William T. Freeman

cs.AI

Samenvatting

Recente benaderingen hebben belofte getoond in het destilleren van diffusiemodellen naar efficiënte één-stap generatoren. Onder hen produceert Distribution Matching Distillation (DMD) één-stap generatoren die in distributie overeenkomen met hun leraarmodel, zonder een één-op-één correspondentie af te dwingen met de bemonsteringspaden van hun leraren. Om echter een stabiele training te garanderen, vereist DMD een extra regressieverlies berekend met een grote set ruis-beeldparen gegenereerd door het leraarmodel met veel stappen van een deterministische sampler. Dit is kostbaar voor grootschalige tekst-naar-beeld synthese en beperkt de kwaliteit van het studentmodel, waardoor het te nauw verbonden blijft aan de oorspronkelijke bemonsteringspaden van het leraarmodel. Wij introduceren DMD2, een set technieken die deze beperking opheffen en de DMD-training verbeteren. Ten eerste elimineren we het regressieverlies en de noodzaak voor kostbare datasetconstructie. We tonen aan dat de resulterende instabiliteit wordt veroorzaakt doordat de nep-criticus de distributie van gegenereerde samples niet nauwkeurig schat en stellen een update-regel met twee tijdschalen voor als remedie. Ten tweede integreren we een GAN-verlies in het destillatieproces, waarbij onderscheid wordt gemaakt tussen gegenereerde samples en echte beelden. Hierdoor kunnen we het studentmodel trainen op echte data, wat de onvolmaakte schatting van de echte score door het leraarmodel vermindert en de kwaliteit verbetert. Tot slot passen we het trainingsproces aan om multi-step sampling mogelijk te maken. We identificeren en adresseren het probleem van input-mismatch tussen training en inferentie in deze setting door inferentie-tijd generator samples te simuleren tijdens de training. Samen genomen stellen onze verbeteringen nieuwe benchmarks in één-stap beeldgeneratie, met FID-scores van 1.28 op ImageNet-64x64 en 8.35 op zero-shot COCO 2014, wat het oorspronkelijke leraarmodel overtreft ondanks een 500X reductie in inferentiekosten. Verder tonen we aan dat onze aanpak megapixelbeelden kan genereren door SDXL te destilleren, wat een uitzonderlijke visuele kwaliteit laat zien onder methoden met weinig stappen.

English

Recent approaches have shown promises distilling diffusion models into efficient one-step generators. Among them, Distribution Matching Distillation (DMD) produces one-step generators that match their teacher in distribution, without enforcing a one-to-one correspondence with the sampling trajectories of their teachers. However, to ensure stable training, DMD requires an additional regression loss computed using a large set of noise-image pairs generated by the teacher with many steps of a deterministic sampler. This is costly for large-scale text-to-image synthesis and limits the student's quality, tying it too closely to the teacher's original sampling paths. We introduce DMD2, a set of techniques that lift this limitation and improve DMD training. First, we eliminate the regression loss and the need for expensive dataset construction. We show that the resulting instability is due to the fake critic not estimating the distribution of generated samples accurately and propose a two time-scale update rule as a remedy. Second, we integrate a GAN loss into the distillation procedure, discriminating between generated samples and real images. This lets us train the student model on real data, mitigating the imperfect real score estimation from the teacher model, and enhancing quality. Lastly, we modify the training procedure to enable multi-step sampling. We identify and address the training-inference input mismatch problem in this setting, by simulating inference-time generator samples during training time. Taken together, our improvements set new benchmarks in one-step image generation, with FID scores of 1.28 on ImageNet-64x64 and 8.35 on zero-shot COCO 2014, surpassing the original teacher despite a 500X reduction in inference cost. Further, we show our approach can generate megapixel images by distilling SDXL, demonstrating exceptional visual quality among few-step methods.