Distributie Backtracking Bouwt een Snellere Convergentietrajectorie voor Eénstaps Diffusiedistillatie

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

August 28, 2024

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

cs.AI

Samenvatting

Het versnellen van de bemonsteringssnelheid van diffusiemodellen blijft een belangrijke uitdaging. Recente score-distillatiemethoden distilleren een zwaar leraarmodel naar een eenstaps studentgenerator, die wordt geoptimaliseerd door het verschil te berekenen tussen de twee scorefuncties op de monsters die door het studentmodel worden gegenereerd. Er doet zich echter een score-mismatchprobleem voor in de vroege fase van het distillatieproces, omdat bestaande methoden zich voornamelijk richten op het gebruik van het eindpunt van vooraf getrainde diffusiemodellen als leraarmodellen, waarbij het belang van het convergentietraject tussen de studentgenerator en het leraarmodel over het hoofd wordt gezien. Om dit probleem aan te pakken, breiden we het score-distillatieproces uit door het volledige convergentietraject van leraarmodellen te introduceren en stellen we Distribution Backtracking Distillation (DisBack) voor voor het distilleren van studentgeneratoren. DisBack bestaat uit twee fasen: Degradatie Registratie en Distributie Terugspoelen. Degradatie Registratie is ontworpen om het convergentietraject van leraarmodellen te verkrijgen, waarbij het degradatiepad wordt vastgelegd van het getrainde leraarmodel naar de ongetrainde initiële studentgenerator. Het degradatiepad vertegenwoordigt impliciet de tussenliggende distributies van leraarmodellen. Vervolgens traint Distributie Terugspoelen een studentgenerator om de tussenliggende distributies terug te spoelen om het convergentietraject van leraarmodellen te benaderen. Uitgebreide experimenten tonen aan dat DisBack snellere en betere convergentie bereikt dan de bestaande distillatiemethode en vergelijkbare generatieprestaties behaalt. Opmerkelijk is dat DisBack eenvoudig te implementeren is en kan worden gegeneraliseerd naar bestaande distillatiemethoden om de prestaties te verbeteren. Onze code is openbaar beschikbaar op 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.