Distillazione Migliorata per l'Allineamento della Distribuzione nella Sintesi Rapida di Immagini

Improved Distribution Matching Distillation for Fast Image Synthesis

May 23, 2024

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

cs.AI

Abstract

Approcci recenti hanno mostrato promettenti risultati nel distillare modelli di diffusione in generatori efficienti a un singolo passo. Tra questi, la Distillazione per Matching di Distribuzione (DMD) produce generatori a un passo che corrispondono alla distribuzione del loro insegnante, senza imporre una corrispondenza uno-a-uno con le traiettorie di campionamento dei loro insegnanti. Tuttavia, per garantire un addestramento stabile, DMD richiede una perdita aggiuntiva di regressione calcolata utilizzando un ampio insieme di coppie rumore-immagine generate dall'insegnante con molti passi di un campionatore deterministico. Questo è costoso per la sintesi su larga scala di immagini da testo e limita la qualità dello studente, legandola troppo strettamente ai percorsi di campionamento originali dell'insegnante. Introduciamo DMD2, un insieme di tecniche che superano questa limitazione e migliorano l'addestramento DMD. In primo luogo, eliminiamo la perdita di regressione e la necessità di una costosa costruzione del dataset. Mostriamo che l'instabilità risultante è dovuta al fatto che il critico fittizio non stima accuratamente la distribuzione dei campioni generati e proponiamo una regola di aggiornamento a due scale temporali come rimedio. In secondo luogo, integriamo una perdita GAN nel processo di distillazione, discriminando tra campioni generati e immagini reali. Questo ci permette di addestrare il modello studente su dati reali, mitigando l'imperfetta stima del punteggio reale del modello insegnante e migliorando la qualità. Infine, modifichiamo la procedura di addestramento per abilitare il campionamento a più passi. Identifichiamo e affrontiamo il problema di disallineamento tra input di addestramento e inferenza in questo contesto, simulando i campioni del generatore durante il tempo di inferenza nel tempo di addestramento. Nel complesso, i nostri miglioramenti stabiliscono nuovi benchmark nella generazione di immagini a un singolo passo, con punteggi FID di 1.28 su ImageNet-64x64 e 8.35 su COCO 2014 zero-shot, superando l'insegnante originale nonostante una riduzione di 500X nel costo di inferenza. Inoltre, mostriamo che il nostro approccio può generare immagini megapixel distillando SDXL, dimostrando una qualità visiva eccezionale tra i metodi a pochi passi.

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.