[논문 리뷰] On Distillation of Guided Diffusion Models

Chenlin Meng, Rombach, Robin외 5명|arXiv (Cornell University)|2022. 10. 06.

Advanced Neuroimaging Techniques and Applications인용 수 8

한 줄 요약

이 논문은 분류기 없는 지도 학습 기반 확산 모델의 추론 속도를 향상시키기 위해 이중 단계의 디스틸레이션 프레임워크를 제안한다. 이는 최소 1~4단계의 노이즈 제거 단계만으로도 고해상도 이미지 생성을 가능하게 한다. 먼저 조건부 및 비조건부 교사 모델의 병합된 출력을 단일 학생 모델로 디스틸레이션하고, 이후 점진적으로 더 적은 단계로 디스틸레이션함으로써, 표준 확산 모델 대비 최대 256배 빠른 추론 속도를 달성하면서도 ImageNet 및 LAION 벤치마크에서 경쟁적인 FID 및 IS 점수를 유지한다.

ABSTRACT

Classifier-free guided diffusion models have recently been shown to be highly effective at high-resolution image generation, and they have been widely used in large-scale diffusion frameworks including DALLE-2, Stable Diffusion and Imagen. However, a downside of classifier-free guided diffusion models is that they are computationally expensive at inference time since they require evaluating two diffusion models, a class-conditional model and an unconditional model, tens to hundreds of times. To deal with this limitation, we propose an approach to distilling classifier-free guided diffusion models into models that are fast to sample from: Given a pre-trained classifier-free guided model, we first learn a single model to match the output of the combined conditional and unconditional models, and then we progressively distill that model to a diffusion model that requires much fewer sampling steps. For standard diffusion models trained on the pixel-space, our approach is able to generate images visually comparable to that of the original model using as few as 4 sampling steps on ImageNet 64x64 and CIFAR-10, achieving FID/IS scores comparable to that of the original model while being up to 256 times faster to sample from. For diffusion models trained on the latent-space (e.g., Stable Diffusion), our approach is able to generate high-fidelity images using as few as 1 to 4 denoising steps, accelerating inference by at least 10-fold compared to existing methods on ImageNet 256x256 and LAION datasets. We further demonstrate the effectiveness of our approach on text-guided image editing and inpainting, where our distilled model is able to generate high-quality results using as few as 2-4 denoising steps.

연구 동기 및 목표

• 분류기 없는 지도 학습 기반 확산 모델의 높은 계산 비용 문제를 해결하기 위해, 조건부 및 비조건부 두 개의 별도 모델을 추론 시 수십에서 수백 번 이상 평가해야 하는 문제를 해결한다.
• 이미지 품질이나 지도 학습의 유연성을 희생시키지 않고도, 디스틸레이션된 모델에서 빠른 추론가능성을 확보하는 디스틸레이션 방법을 개발한다.
• 기존의 픽셀 공간 및 잠재 공간 확산 모델을 포함한 최신 모델(예: Stable Diffusion)에까지 디스틸레이션 기법을 확장한다.
• 최소한의 추론 단계로 효율적인 텍스트 기반 이미지 생성, 이미지 편집 및 인painting을 가능하게 한다.

제안 방법

• 교사 모델의 조건부 및 비조건부 확산 모델 출력을 병합한 결과를 모방하도록 학습하는 단일 학생 모델을 도입한다.
• 이중 단계 디스틸레이션 프로세스를 사용한다: 첫 번째 단계에서는 학생 모델이 교사의 출력을 모방하도록 학습하고, 두 번째 단계에서는 점진적으로 더 적은 추론 단계를 요구하는 모델로 디스틸레이션한다.
• 디스틸레이션된 모델을 사용한 추론 시 기존의 결정론적 방법 외에 새로운 확률적 추론 프로세스를 도입한다.
• 픽셀 공간 모델(예: CIFAR-10 및 ImageNet 64×64)과 잠재 공간 모델(예: ImageNet 256×256 및 LAION 512×512에서의 Stable Diffusion)에 모두 이 디스틸레이션 프레임워크를 적용한다.
• 다양한 지도 학습 가중치를 지원하도록 디스틸레이션된 모델을 학습시켜, 샘플의 다양성과 품질 사이의 트레이드오프를 가능하게 한다.
• 다양한 지도 학습 강도에서 학생 모델의 예측이 교사의 출력과 일치하도록 가중치가 부여된 MSE 손실을 디스틸레이션 과정에서 사용한다.

실험 결과

연구 질문

• RQ1 분류기 없는 지도 학습 기반 확산 모델은 두 개의 별도 모델에서 추론이 필요하므로, 이에 효과적으로 디스틸레이션을 적용할 수 있는가?
• RQ2 단일 디스틸레이션된 모델이 1~4단계의 노이즈 제거 단계만으로도 높은 품질의 이미지 생성을 가능하게 하면서도 지도 학습의 유연성을 유지할 수 있는가?
• RQ3 제안된 이중 단계 디스틸레이션 접근법이 원본 모델 대비 FID 및 IS 점수를 유지하거나 향상시키는가?
• RQ4 이러한 디스틸레이션 기법이 픽셀 공간 및 잠재 공간 모델을 포함한 다양한 확산 모델 아키텍처에 일반화 가능한가?
• RQ5 이 방법은 최소한의 추론 단계로 텍스트 기반 이미지 생성, 이미지 편집 및 인painting 등의 후행 작업에 성공적으로 적용될 수 있는가?

주요 결과

• 디스틸레이션된 모델은 ImageNet 64×64 및 CIFAR-10에서 단 4단계의 추론만으로도 교사 모델과 유사한 FID 점수를 달성하며, 추론 속도 향상 최대 256배를 기록한다.
• Stable Diffusion와 같은 잠재 공간 모델에서는 추론 단계를 1~4단계로 줄여, ImageNet 256×256 및 LAION 512×512에서 기존 방법 대비 최소 10배 이상의 빠른 추론 속도를 확보한다.
• 다양한 지도 학습 가중치(w = 0에서 w = 4)에 걸쳐 고해상도의 시각적 품질을 유지하며, 다양성과 정확도 사이의 효과적인 트레이드오프를 가능하게 한다.
• 텍스트 기반 이미지 편집 및 인painting 작업에서는 단 2~4단계의 노이즈 제거만으로도 고품질 결과를 생성하여 편집 작업에서 뛰어난 성능을 보였다.
• 특히 낮은 단계 수에서 결정론적 추론 대비 확률적 추론 프로세스가 샘플 품질을 더욱 향상시켰다.
• 이 방법은 픽셀 공간 및 잠재 공간 분류기 없는 지도 학습 기반 확산 모델에 대해 효과적인 디스틸레이션을 처음으로 입증했으며, 다양한 벤치마크에서 일관된 성능 향상을 보였다.