[논문 리뷰] BOOT: Data-free Distillation of Denoising Diffusion Models with Bootstrapping

Jiatao Gu, Shuangfei Zhai외 3명|arXiv (Cornell University)|2023. 06. 08.

Advanced Neuroimaging Techniques and Applications인용 수 4

한 줄 요약

BOOT는 실제 데이터 또는 광범위한 오프라인 데이터 생성 없이, 연속적인 샘플링 단계에서의 부트스트랩을 통해 단일 단계 학생 모델을 훈련시키는 데이터 프리 지식 정복 방법을 제안한다. 이는 안정된 확산 모델의 품질에 근접한 이미지 품질을 달성하면서도, 대규모 텍스트-이미지 모델인 Stable Diffusion나 DeepFloyd IF와 같은 모델에서도 수배수에 이르는 빠른 추론을 가능하게 한다.

ABSTRACT

Diffusion models have demonstrated excellent potential for generating diverse images. However, their performance often suffers from slow generation due to iterative denoising. Knowledge distillation has been recently proposed as a remedy that can reduce the number of inference steps to one or a few without significant quality degradation. However, existing distillation methods either require significant amounts of offline computation for generating synthetic training data from the teacher model or need to perform expensive online learning with the help of real data. In this work, we present a novel technique called BOOT, that overcomes these limitations with an efficient data-free distillation algorithm. The core idea is to learn a time-conditioned model that predicts the output of a pre-trained diffusion model teacher given any time step. Such a model can be efficiently trained based on bootstrapping from two consecutive sampled steps. Furthermore, our method can be easily adapted to large-scale text-to-image diffusion models, which are challenging for conventional methods given the fact that the training sets are often large and difficult to access. We demonstrate the effectiveness of our approach on several benchmark datasets in the DDIM setting, achieving comparable generation quality while being orders of magnitude faster than the diffusion teacher. The text-to-image results show that the proposed approach is able to handle highly complex distributions, shedding light on more efficient generative modeling.

연구 동기 및 목표

• 확산 모델의 느린 추론 속도 문제를 해결하기 위해, 반복적인 노이즈 제거를 다수의 단계를 거쳐 수행하는 데 기인한다.
• 기존의 지식 정복 방법이 대규모 합성 데이터 생성이나 실제 훈련 데이터에 접근이 필요로 하는 한계를 극복하기 위해.
• 개인정보나 규모 제약으로 인해 훈련 데이터에 접근이 불가능한 대규모 텍스트-이미지 확산 모델의 효율적인 지식 정복을 가능하게 하기 위해.
• 연속적인 시간 단계 예측에서 유도된 부트스트랩을 활용하는 데이터 프리, 효율적인 훈련 프레임워크를 개발하기 위해.
• 전체 확산 교사 모델과 동등한 품질을 갖춘 고해상도 단일 단계 생성을 달성하면서도 추론 시간을 극적으로 단축시키기 위해.

제안 방법

• 모델은 단일 노이즈 입력을 기반으로, 사전 훈련된 확산 교사 모델의 모든 시간 단계에서의 출력을 예측할 수 있도록 시간 조건부 학생 모델을 훈련시킨다.
• PF-ODE에서 유도된 새로운 신호-ODE를 사용하여 훈련을 안정화하고 시간 단계 간의 일반화 성능을 향상시킨다.
• 동일한 확산 궤적을 따라 결정론적 매핑이 존재하므로, $ t' > t $인 $ \bm{x}_{t'} $의 지식을 기반으로 $ \bm{x}_t $를 예측하도록 모델을 훈련시켜 부트스트랩을 구현한다.
• 샘플링 품질과 다양성을 향상시키기 위해 훈련 중 주기적으로 경계 손실을 적용하며, 효율성을 유지하기 위해 빈도를 낮춘다.
• 분류기 자유 가이던스(CFG)와 호환되어 조건부 및 텍스트-이미지 모델의 효과적인 지식 정복이 가능하다.
• 실제 데이터 없이 교사 모델의 순방향 전파만을 사용하여 엔드 투 엔드로 학생 모델을 훈련시기 때문에, 개인정보에 민감하거나 대규모 모델에 적합하다.

실험 결과

연구 질문

• RQ1 실제 훈련 데이터나 대규모 합성 데이터 생성에 접근할 수 없는 상황에서 단일 단계 학생 모델을 효과적으로 훈련시킬 수 있는가?
• RQ2 연속적인 시간 단계 예측에서의 부트스트랩은 데이터 프리 환경에서 지식 정복의 효율성과 성능을 어떻게 향상시킬 수 있는가?
• RQ3 제안된 방법은 전체 확산 교사 모델과 비교해도 높은 품질의 이미지 생성을 달성하면서도 수배수에 이르는 빠른 추론을 가능하게 할 수 있는가?
• RQ4 Stable Diffusion나 DeepFloyd IF와 같은 복잡하고 고용량의 텍스트-이미지 확산 모델에 대해 이 방법은 얼마나 잘 일반화되는가?
• RQ5 신호-ODE와 경계 손실은 훈련 안정성 향상과 샘플링 다양성 향상에 어떤 역할을 하는가?

주요 결과

• BOOT는 FFHQ, LSUN, ImageNet를 포함한 조건부 및 비조건부 이미지 생성 벤치마크에서 전체 확산 교사 모델과 비교해 유사한 생성 품질을 달성한다.
• 데이터셋과 설정에 따라 교사 모델 대비 50~1000배 빠른 단일 단계 추론 속도 향상을 달성한다.
• 텍스트-이미지 생성 분야에서 BOOT는 DeepFloyd IF 및 Stable Diffusion 모델을 성공적으로 지식 정복하여, 훈련 데이터에 접근하지 못함에도 불구하고 다양한 프롬프트에서 고품질의 이미지를 생성한다.
• 정교한 다양성과 해상도를 유지하며, 동일한 초기 노이즈를 사용한 교사 및 학생 모델 간의 정성적 비교 결과에서 이를 입증한다.
• 매 4개의 훈련 단계마다 적용되는 경계 손실은 훈련 비용을 크게 증가시키지 않으면서도 샘플링 품질과 다양성을 향상시킨다.
• 계산적으로 효율적이며, 1B 파rameter 이하의 모델에 대해 8장의 A100 GPU와 3~7일간의 훈련만으로도 가능하며, 텍스트 프롬프트 외에 데이터 로딩 오버헤드가 없다.