La retrotrazabilidad de distribución construye una trayectoria de convergencia más rápida para la destilación de difusión de un paso.

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

Resumen

Acelerar la velocidad de muestreo de los modelos de difusión sigue siendo un desafío significativo. Los métodos recientes de destilación de puntuaciones destilan un modelo profesor pesado en un generador de estudiante de un paso, que se optimiza calculando la diferencia entre las dos funciones de puntuación en las muestras generadas por el modelo de estudiante. Sin embargo, existe un problema de desajuste de puntuaciones en la etapa inicial del proceso de destilación, porque los métodos existentes se centran principalmente en utilizar el punto final de los modelos de difusión preentrenados como modelos profesores, pasando por alto la importancia de la trayectoria de convergencia entre el generador de estudiantes y el modelo profesor. Para abordar este problema, ampliamos el proceso de destilación de puntuaciones introduciendo toda la trayectoria de convergencia de los modelos profesores y proponemos la Destilación de Retroceso de Distribución (DisBack) para destilar generadores de estudiantes. DisBack se compone de dos etapas: Registro de Degradación y Retroceso de Distribución. El Registro de Degradación está diseñado para obtener la trayectoria de convergencia de los modelos profesores, que registra la ruta de degradación desde el modelo profesor entrenado hasta el generador de estudiantes inicial no entrenado. La ruta de degradación representa implícitamente las distribuciones intermedias de los modelos profesores. Luego, el Retroceso de Distribución entrena un generador de estudiantes para retroceder las distribuciones intermedias para aproximar la trayectoria de convergencia de los modelos profesores. Experimentos extensos muestran que DisBack logra una convergencia más rápida y mejor que el método de destilación existente y logra un rendimiento de generación comparable. Es importante destacar que DisBack es fácil de implementar y se puede generalizar a los métodos de destilación existentes para mejorar el rendimiento. Nuestro código está disponible públicamente en 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.