분포 백트래킹은 일단의 확산 증류에 대해 빠른 수렴 궤적을 형성합니다.

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

August 28, 2024

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

cs.AI

초록

확산 모델의 샘플링 속도를 가속화하는 것은 여전히 중요한 도전 과제입니다. 최근의 점수 증류 방법은 무거운 교사 모델을 1단계 학생 생성기로 증류하는데, 이는 학생 모델이 생성한 샘플에서 두 점수 함수 간의 차이를 계산하여 최적화됩니다. 그러나 기존 방법은 주로 사전 훈련된 확산 모델의 끝점을 교사 모델로 사용하며, 학생 생성기와 교사 모델 간의 수렴 궤적의 중요성을 간과합니다. 이 문제를 해결하기 위해 우리는 교사 모델의 전체 수렴 궤적을 도입하여 학생 생성기를 증류하기 위한 분포 역추적 증류(DisBack)를 제안합니다. DisBack는 소멸 기록과 분포 역추적 두 단계로 구성됩니다. 소멸 기록은 교사 모델의 수렴 궤적을 얻기 위해 설계되었으며, 훈련된 교사 모델에서 미훈련 초기 학생 생성기로의 소멸 경로를 기록합니다. 이 소멸 경로는 교사 모델의 중간 분포를 암시적으로 나타냅니다. 그런 다음 분포 역추적은 학생 생성기를 훈련하여 교사 모델의 수렴 궤적을 근사하기 위해 중간 분포를 역추적합니다. 광범위한 실험 결과 DisBack가 기존 증류 방법보다 빠르고 더 나은 수렴을 달성하며 비교 가능한 생성 성능을 보여준다는 것을 보여줍니다. 특히 DisBack는 구현하기 쉽고 성능을 향상시키기 위해 기존 증류 방법에 일반화할 수 있습니다. 우리의 코드는 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.