A Distribuição Backtracking Constrói uma Trajetória de Convergência Mais Rápida para Destilação por Difusão de Um Passo

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

August 28, 2024

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

cs.AI

Resumo

Acelerar a velocidade de amostragem dos modelos de difusão continua sendo um desafio significativo. Métodos recentes de destilação de pontuação destilam um modelo professor pesado em um gerador de estudante de um passo, que é otimizado calculando a diferença entre as duas funções de pontuação nas amostras geradas pelo modelo de estudante. No entanto, existe um problema de incompatibilidade de pontuação na fase inicial do processo de destilação, porque os métodos existentes se concentram principalmente em usar o ponto final dos modelos de difusão pré-treinados como modelos professores, negligenciando a importância da trajetória de convergência entre o gerador de estudante e o modelo professor. Para resolver esse problema, estendemos o processo de destilação de pontuação introduzindo toda a trajetória de convergência dos modelos professores e propomos a Destilação de Retrocesso de Distribuição (DisBack) para destilar geradores de estudantes. DisBack é composto por duas etapas: Gravação de Degradação e Retrocesso de Distribuição. A Gravação de Degradação é projetada para obter a trajetória de convergência dos modelos professores, que registra o caminho de degradação do modelo professor treinado para o gerador de estudante inicial não treinado. O caminho de degradação representa implicitamente as distribuições intermediárias dos modelos professores. Em seguida, o Retrocesso de Distribuição treina um gerador de estudante para retroceder nas distribuições intermediárias para aproximar a trajetória de convergência dos modelos professores. Experimentos extensivos mostram que o DisBack atinge uma convergência mais rápida e melhor do que o método de destilação existente e alcança um desempenho de geração comparável. Notavelmente, o DisBack é fácil de implementar e pode ser generalizado para os métodos de destilação existentes para melhorar o desempenho. Nosso código está publicamente disponível em 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.