Die Verteilungsrückverfolgung schafft eine schnellere Konvergenztrajektorie für die einstufige Diffusionsdestillation.

Distribution Backtracking Builds A Faster Convergence Trajectory for One-step Diffusion Distillation

August 28, 2024

Autoren: Shengyuan Zhang, Ling Yang, Zejian Li, An Zhao, Chenye Meng, Changyuan Yang, Guang Yang, Zhiyuan Yang, Lingyun Sun

cs.AI

Zusammenfassung

Die Beschleunigung der Abtastgeschwindigkeit von Diffusionsmodellen bleibt eine bedeutende Herausforderung. Aktuelle Score-Destillationsmethoden destillieren ein umfangreiches Lehrermodell in einen Ein-Schritt-Studentengenerator, der optimiert wird, indem der Unterschied zwischen den beiden Score-Funktionen auf den vom Studentenmodell generierten Proben berechnet wird. Es besteht jedoch ein Problem des Score-Mismatches im frühen Stadium des Destillationsprozesses, da bestehende Methoden hauptsächlich darauf abzielen, das Endpunkt von vorab trainierten Diffusionsmodellen als Lehrermodelle zu verwenden und die Bedeutung der Konvergenztrajektorie zwischen dem Studentengenerator und dem Lehrermodell zu übersehen. Um dieses Problem zu lösen, erweitern wir den Score-Destillationsprozess, indem wir die gesamte Konvergenztrajektorie der Lehrermodelle einführen und Distribution Backtracking Distillation (DisBack) zur Destillation von Studentengeneratoren vorschlagen. DisBack besteht aus zwei Phasen: Degradation Recording und Distribution Backtracking. Degradation Recording ist darauf ausgelegt, die Konvergenztrajektorie der Lehrermodelle zu erhalten, indem der Degradationspfad vom trainierten Lehrermodell zum untrainierten anfänglichen Studentengenerator aufgezeichnet wird. Der Degradationspfad stellt implizit die Zwischenverteilungen der Lehrermodelle dar. Anschließend trainiert Distribution Backtracking einen Studentengenerator, um die Zwischenverteilungen rückgängig zu machen und die Konvergenztrajektorie der Lehrermodelle zu approximieren. Umfangreiche Experimente zeigen, dass DisBack eine schnellere und bessere Konvergenz als die bestehende Destillationsmethode erreicht und eine vergleichbare Generierungsleistung erbringt. Insbesondere ist DisBack einfach zu implementieren und kann auf bestehende Destillationsmethoden verallgemeinert werden, um die Leistung zu steigern. Unser Code ist öffentlich verfügbar unter https://github.com/SYZhang0805/DisBack.

English

Accelerating the sampling speed of diffusion models remains a significant challenge. Recent score distillation methods distill a heavy teacher model into an one-step student generator, which is optimized by calculating the difference between the two score functions on the samples generated by the student model. However, there is a score mismatch issue in the early stage of the distillation process, because existing methods mainly focus on using the endpoint of pre-trained diffusion models as teacher models, overlooking the importance of the convergence trajectory between the student generator and the teacher model. To address this issue, we extend the score distillation process by introducing the entire convergence trajectory of teacher models and propose Distribution Backtracking Distillation (DisBack) for distilling student generators. DisBask is composed of two stages: Degradation Recording and Distribution Backtracking. Degradation Recording is designed to obtain the convergence trajectory of teacher models, which records the degradation path from the trained teacher model to the untrained initial student generator. The degradation path implicitly represents the intermediate distributions of teacher models. Then Distribution Backtracking trains a student generator to backtrack the intermediate distributions for approximating the convergence trajectory of teacher models. Extensive experiments show that DisBack achieves faster and better convergence than the existing distillation method and accomplishes comparable generation performance. Notably, DisBack is easy to implement and can be generalized to existing distillation methods to boost performance. Our code is publicly available on https://github.com/SYZhang0805/DisBack.