Distillation Améliorée de l'Appariement de Distribution pour la Synthèse Rapide d'Images

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

Résumé

Les approches récentes ont montré des promesses dans la distillation de modèles de diffusion en générateurs efficaces en une seule étape. Parmi elles, la Distillation par Correspondance de Distribution (DMD) produit des générateurs en une étape qui correspondent à leur enseignant en termes de distribution, sans imposer une correspondance un-à-un avec les trajectoires d'échantillonnage de leurs enseignants. Cependant, pour assurer un entraînement stable, la DMD nécessite une perte de régression supplémentaire calculée à l'aide d'un grand ensemble de paires bruit-image générées par l'enseignant avec de nombreuses étapes d'un échantillonneur déterministe. Cela est coûteux pour la synthèse d'images à grande échelle à partir de texte et limite la qualité de l'élève, l'enchaînant trop étroitement aux chemins d'échantillonnage originaux de l'enseignant. Nous introduisons DMD2, un ensemble de techniques qui lèvent cette limitation et améliorent l'entraînement de la DMD. Premièrement, nous éliminons la perte de régression et le besoin de construction coûteuse de jeux de données. Nous montrons que l'instabilité résultante est due au fait que le critique artificiel n'estime pas avec précision la distribution des échantillons générés et proposons une règle de mise à jour à deux échelles de temps comme remède. Deuxièmement, nous intégrons une perte GAN dans la procédure de distillation, en discriminant entre les échantillons générés et les images réelles. Cela nous permet d'entraîner le modèle élève sur des données réelles, atténuant l'estimation imparfaite des scores réels du modèle enseignant et améliorant la qualité. Enfin, nous modifions la procédure d'entraînement pour permettre un échantillonnage en plusieurs étapes. Nous identifions et résolvons le problème de discordance des entrées entre l'entraînement et l'inférence dans ce contexte, en simulant les échantillons du générateur au moment de l'inférence pendant l'entraînement. Ensemble, nos améliorations établissent de nouveaux référentiels dans la génération d'images en une étape, avec des scores FID de 1,28 sur ImageNet-64x64 et 8,35 sur COCO 2014 en zéro-shot, surpassant l'enseignant original malgré une réduction de 500X du coût d'inférence. De plus, nous montrons que notre approche peut générer des images en mégapixels en distillant SDXL, démontrant une qualité visuelle exceptionnelle parmi les méthodes à quelques étapes.

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.