On Distillation of Guided Diffusion Models _CVPR_2023논문 리뷰

Meng_On_Distillation_of_Guided_Diffusion_Models_CVPR_2023_paper

 

 

ABSTRACT

문제 정의:

 

Classifier-free guided diffusion model 은 고화질 이미지 생성에 매우 효과적이었고 DALLE2 나 STABLE-DIFFUSION, IMGEN 에 폭 넓게 사용되었다.

하지만 Classifier-free guided diffusion model 은 하나의 샘플을 생성하기 위해 class-conditional model, unconditional model 두 개의 diffusion model 을 100번 정도 평가하기 때문에 추론 시간이 많이 소요된다는 단점이 있다.

 

본 논문의 해결 방법:

 

pre-trained classifier-free guided model이 주어지면, single model 을 combined conditional, unconditional model 의 output 과 일치하도록 학습시킨다. 그리고 이 모델을 훨씬 적은 sampling step 이 요구하는 diffusion model 로 distill 한다.

픽셀 공간에서 학습 시킨 본 논문의 diffusion model 은 ImageNet 64x64 및 CIFAR-10에서 4개의 샘플링 step만 사용하여 원본 모델과 시각적으로 비교 가능한 이미지를 생성할 수 있고 이는 원본 모델과 유사한 FID/IS 점수를 만족하면서도 최대 256배 빠르게 샘플링할 수 있다.

이 접근 방식은 잠재 공간에서 훈련된 확산 모델(예: Stable Diffusion)을 대상으로도 적용된다. ImageNet 256x256 및 LAION 데이터셋에서 기존 방법과 비교하여 최소 10배 이상의 추론 속도를 제공하면서, 1~4개의 denoising step만 거쳐 고품질의 이미지를 생성할 수 있다.

INTRODUCTION

ddpm은 이미지 생성, 오디오 합성, 분자 생성, 우도 측정에서 sota 를 달성했으며, classifier-free guidance 는 디퓨전 모델의 성능을 향상시켰고 대규모 디퓨전 모델에 사용되었는데 GLIDE, stable-diffusion, DALLE2, IMAGEN 가 이 예시이다.

하지만, classifier-free guidance 의 단점은 sampling 효율이 낮다는 것이다. 한 샘플을 생성하기 위해 두 디퓨전 모델을 100 번 정도 평가해야 하기 때문이다.

Diffusion model 에 대한 distillation 접근은 있었으나 classifier-free-guided diffusion model 에 직접적인 적용은 어려웠다. 이 문제를 해결하기 위하여 본 논문은 two-stage distillation 방법을 소개하여 sampling efficiency 를 향상 시킨다.

 

First stage: teacher 의 두 개의 결합된 diffusion model(class-conditional model, unconditional model)의 출력과 일치하도록 단일 student model 을 학습시킨다.

Second stage: 첫 번째 학습시킨 모델을 더 작은 모델로 distill 한다.

위 방법을 쓰면 single distillation model이 다양한 guidance strengths(외부 정보가 model 에게 얼마나 강하게 적용되는지)를 광범위하게 처리할 수 있으므로 샘플 품질과 다양성 간의 균형을 효율적으로 유지할 수 있다. 본 논문의 모델로부터 샘플링하기 위해 deterministic sampler(샘플링 할 때마다 같은 결과) 과 stochastic sampling process(샘플링 할 때마다 다른 결과) 를 고려한다.

 

본 논문의 distillation framework 는 pixel-space 위에서 훈련된 standard diffusion model 에 적용될 뿐만 아니라, auto-encoder 의 latent-space 위에서 훈련된 diffusion model 에도 적용된다.

 

ImageNet 64x64 및 CIFAR-10에서의 실험 결과를 통해 본 논문의 distilled 모델이 4step 만으로도 teacher 모델과 비슷한 샘플을 생성할 수 있음을 보여준다. 또한, 다양한 guidance strengths 에 대해 4~16 step 을 사용하는 teacher 모델과 비슷한 FID/IS 점수를 달성할 수 있음을 보여준다. 이는 모델이 적은 샘플링 step 으로도 높은 성능을 발휘할 수 있다는 것을 의미한다.

encoder의 latent-space에서 학습된 diffusion model의 경우 최소 1~4개의 샘플링 step(base model 보다 10 배 정도 적은 스텝)을 사용하여 base model과 비슷한 시각적 품질을 달성할 수 있으며, 2~4개의 샘플링 step만으로 teacher의 성능과 일치한다.

Progressive distillation

증류를 반복함으로서 diffusion model 의 sampling 속도를 높이는 효과적인 방법이다.

지금까지 이 방법은 'classifier-free guided models'를 정제하는 데 직접 적용되지 않았으며, 'deterministic DDIM sampler' 이외의 샘플러에 대해서도 연구되지 않았다. 본 논문에서는 이러한 단점을 해결한다.

Latent diffusion models (LDMs)

LDM 의 pixel 공간의 낮은 차원을 의미한다. 이 모델은 훈련 속도와 추론 효율성을 향상 시켰다. 그리고 이 모델은 저차원 이미지를 좋은 해상도를 유지하면서 scale up 가능하게 한다.

본 논문에서는 distillation framework 를 pixel 공간에서 학습 시킨 classifier-free guided diffusion model 과 latent 공간에서 학습시킨 classifier-free guided diffusion model 에 본 논문의 distillation 기법을 적용한다.

Guidance 가중치 w 를 통합하기 위해 w 로 conditioning 된 모델을 도입한다. 여기서 w는 student model에 대한 입력으로 공급된다. Feature를 더 잘 캡처하기 위해 푸리에 임베딩을 w에 적용한 다음 timestep이 통합된 방식과 유사한 방식으로 diffusion model backbone에 통합된다. 성능에서 초기화가 중요한 역할을 하기 때문에 w conditioning과 관련하여 새로 도입된 파라미터를 제외하고는 teacher의 조건부 모델과 동일한 파라미터로 student model을 초기화한다. 사용하는 모델 아키텍처는 U-Net 모델이다.

 

Experiments

민트색이 본 논문의 모델인데 같은 FID SCORE 라고 가정했을 때, 본 논문의 Sampling step 수가 적은 것을 알 수 있다. 또한, guidance strength 변수인 w 가 변화함에 따라서도 높은 성능을 보여준다.

 

Distilled 되지 않은 모델과 본 논문의 모델의 같은 4 step 을 거친 샘플을 봤을 때, (2번째 그림, 1번째 그림) 본 논문의 성능이 좋은 것을 볼 수 있다. 심지어 distilled 되지 않고 8 step 을 거친 샘플보다 더 좋은 성능을 내는 것을 볼 수 있다.

 

 

 

 

 

3번의 denoising step 만으로도 의미 있는 output 을 도출하는 것을 볼 수 있다.

style transfer 에서도 distilled 되지 않은 모델보다 더 선명한 이미지를 도출하는 것을 볼 수 있다.