Verbesserte Verteilungsanpassungsdarstellung für schnelle Bildsynthese

Improved Distribution Matching Distillation for Fast Image Synthesis

May 23, 2024

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

cs.AI

Zusammenfassung

Neuere Ansätze haben vielversprechende Ergebnisse gezeigt, indem sie Diffusionsmodelle in effiziente Ein-Schritt-Generatoren destillieren. Unter ihnen produziert die Distribution Matching Distillation (DMD) Ein-Schritt-Generatoren, die sich in ihrer Verteilung mit dem Lehrer angleichen, ohne eine eins-zu-eins-Korrespondenz mit den Abtasttrajektorien ihrer Lehrer durchzusetzen. Allerdings erfordert DMD zur Sicherstellung eines stabilen Trainings eine zusätzliche Regressionsverlustberechnung unter Verwendung eines großen Satzes von Rauschbildpaaren, die vom Lehrer mit vielen Schritten eines deterministischen Samplers generiert wurden. Dies ist kostspielig für die Synthese von Text-zu-Bild auf großer Skala und begrenzt die Qualität des Schülers, da er zu eng an den ursprünglichen Abtastpfaden des Lehrers gebunden ist. Wir stellen DMD2 vor, eine Reihe von Techniken, die diese Einschränkung aufheben und das DMD-Training verbessern. Erstens beseitigen wir den Regressionsverlust und die Notwendigkeit für die aufwendige Datensatzkonstruktion. Wir zeigen, dass die resultierende Instabilität darauf zurückzuführen ist, dass der Fake-Kritiker die Verteilung der generierten Proben nicht genau schätzt, und schlagen eine Zwei-Zeit-Skalen-Aktualisierungsregel als Abhilfe vor. Zweitens integrieren wir einen GAN-Verlust in das Destillationsverfahren, der zwischen generierten Proben und echten Bildern unterscheidet. Dies ermöglicht es uns, das Studentenmodell mit echten Daten zu trainieren, wodurch die unvollkommene reale Punktschätzung des Lehrermodells gemildert und die Qualität verbessert wird. Schließlich modifizieren wir das Schulungsverfahren, um Mehrschritt-Abtastung zu ermöglichen. Wir identifizieren und lösen das Problem des Trainings-Inferenz-Eingabemissverhältnisses in dieser Einstellung, indem wir während der Trainingszeit Inferenzzeit-Generatorproben simulieren. Zusammenfassend setzen unsere Verbesserungen neue Maßstäbe in der Ein-Schritt-Bildgenerierung, mit FID-Punkten von 1,28 auf ImageNet-64x64 und 8,35 auf Zero-Shot COCO 2014, wodurch der ursprüngliche Lehrer trotz einer 500-fachen Reduzierung der Inferenzkosten übertroffen wird. Darüber hinaus zeigen wir, dass unser Ansatz Megapixel-Bilder durch Destillation von SDXL generieren kann, was eine außergewöhnliche visuelle Qualität unter den wenige-Schritte-Methoden demonstriert.

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.