[논문리뷰] One-step Diffusion with Distribution Matching Distillation(DMD)

One-step Diffusion with Distribution Matching Distillation

[Arxiv]

One-step Diffusion with Distribution Matching Distillation

 

요약

이미지 품질에 영향을 거의 미치지 않으면서 diffusion model을 one-step image generator로 변환하는 Distribution Matching Distillation(DMD)를 제안

one-step image generator가 Diffusion. model과 분포 수준에서 일치하도록 하며, gradient를 두 score function의 차이로 표현할 수 있는 approximate KL divergence를 최소화하였다. score function은 각각의 분포에 대해 별도로 학습된 두 Diffusion model로 parameter화된다. 

Stable Diffusion과 비슷한 성능을 보이면서도 inference 속도는 수십배 더 빠르다. 

 

기존 모델의 한계점

빠른 샘플링을 위해 이전의 연구들은 diffusion sampler와 distillation을 통해 샘플링 속도를 높이고자 했다. 그러나 여전히 성능이 기존 모델에 미치지 못했다. 

 

방법론

pretrained Diffusion model을 one-step image generator로 distillation

 

1. Pretrained based model and One-step generator

• Pretrained Base Model: 사전 훈련된 Diffusion 모델 $\mu_{base}$ 를 사용하여 Gaussian noise를 점진적으로 줄여 목표 데이터 분포의 샘플을 생선한다. 주로 10에서 100 steps을 거쳐 현실적인 이미지를 생성한다.
• One-step Generator: $\mu_{base}$와 동일한 구조를 갖는 one-step Generator $G_\theta = \mu_{base}(z, T-1), \forall z$를 초기화한다. time condition 없이 초기화하며, $\mu_{base}$의 파라미터를 사용해 시작한다. 

 

2. Distribution Matching Loss

• 빠른 Generator가 실제 이미지와 구별되지 않는 샘플을 생성하도록 한다. 이를 위해 KL divergence를 최소화한다. 

• Gradient Update using Approximate Scores
  • KL divergence의 gradient를 구하여 $G_\theta$를 학습한다. 이 gradient는 두 score function의 차이로 표현된다. 
   

• Diffusion Noise
  • Gaussian noise를 추가하여 데이터를 교란시킨 다음, 두 diffused 분포의 score를 모델링한다. 
  • $s_{real}(x_t,t)와 $s_{fake}(x_t, t)$를 사용하여 gradient를 업데이트한다. 

 

3. Regression Loss and Final Objective

• Regression Loss: 생성된 이미지와 다단계 Diffusion model의 출력 간의 거리를 측정한다. 

• Paired dataset $D = {z, y}$를 사용하여 $G_\theta (z)$ 와 $\mu_{base}$ 출력 간의 LPIPS 거리를 계산한다. 
• Final Objective
  • 최종 목표는 KL divergence와 regression loss의 합으로 정의된다. 
  • $\nabla_\theta D_{KL}$는 두 score function의 차이로 계산되고, $\nabla_\theta L_{reg}$는 자동 미분을 통해 계산된다. 

 

4. Distllation with classifier-free guidance

• Classifier-free guidance: Text to image diffusion model의 품질을 향상시키기 위해 사용
  • Paired dataset을 생성하여 regression loss $L_{reg}$을 계산
  • Distribution matching gradient $\nabla_\theta D_{KL}$를 계산할 때, real score를 guided 모델의 mean prediction으로 대체한다. 

 

실험 결과

• CIFAR10, ImageNet 64x64, MS COCO 512x512에서 평가
• 기존 one-step generator보다 우수한 성능