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

CVPR 2024 [paper](https://arxiv.org/abs/2311.18828) [page](https://tianweiy.github.io/dmd/)
Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Fredo Durand, William T. Freeman, Taesung Park
Massachusetts Institute of Technology, Adobe Research
30 Nov 2023

위 이미지 중 하나는 Stable Diffusion으로 2.59초 만에 생성한 것이고, 하나는 저자가 제안한 모델을 이용하여 0.09초 만에 생성한 것이다. (약 29배의 속도 차이 발생)

Abstract

기존 Diffusion 모델의 느린 추론 속도를 개선하기 위해 저자는 One-step으로 생성 가능한 Distribution Matching Distillation (DMD)를 제안하였다. 결과적으로 ImageNet 64x64에서 2.62 FID, zero-shot COCO-30k에서 11.49 FID를 달성하며 SD보다 매우 빠르면서도 유사한 성능을 보였다. FP16 환경에서 일반적인 GPU에서 20FPS의 속도를 낼 수 있다고 언급하였다.

Related Work : Diffusion Acceleration

Diffusion 모델을 가속화하는 접근법은 크게 두 가지로 나뉜다. 먼저 하나는 fast diffusion samplers을 사용하는 것이다. 이를 사용하면 1000개의 sampling step 대신 20~50 step을 사용하여 이미지를 생성할 수 있게된다. 하지만 이보다 더 적은 step을 사용하면 이미지 품질은 급격히 나빠질 수 있다. 이에, Diffusion Distillation 방법이 제안되었다. 대게 Knowledge Distillation으로 접근하였으며 student 모델은 teacher 모델의 multi-step을 single-step에 distill하게 된다. 하지만 student 모델의 loss를 한 번 구하기 위해 전체 denoising 궤적을 계산하는 것은 비용이 많이 든다는 문제가 있었다. 이를 해결하기 위해 Progressive Distillation (PD)는 이전 모델로부터 sampling step을 절반씩 줄이는 방식으로 학습을 하는 등 다양한 방식이 제안되고 있다. 저자는 training objective로 분포 매칭을 하는 것으로 비교적 간단한 방식으로 접근한다.

Distribution Matching Distillation

본 연구의 목표는 pre-trained diffusion 모델 $\mu_{base}$을 별도의 sampling 절차 없이 고품질 이미지를 생성하는 one-step image generator $G_\theta$로 증류(distill)하는 것이다. 이때, 동일한 분포에서 샘플을 뽑고싶은 것이지 정확한 매핑을 재현하려고 하는 것은 아니다.

GAN과 유사하게 증류받는 모델을 fake, 증류하는 모델을 _real_로 간주한다. 본 논문에서 제안하는 접근법은 위 이미지에서 확인할 수 있다.

학습에 사용되는 loss function은 총 두 가지이다.

1. gradient update가 두 score function의 차이로 표현되는 distribution matching objective
2. Generator가 Noise-Image pair의 고정된 데이터셋에서 기본 모델의 출력의 대규모 구조와 일치하도록 하는 regression loss

결과적으로, 다양한 크기의 가우시안 노이즈로 방해 받는 real*과 *fake distribution의 score function 모델로 사용하기 위해 두 개의 diffusion denoiser를 사용하였다.

Pretrained base model and One-step generator

pre-trained diffusion 모델 $\mu_{base}$가 사전에 주어진다고 가정한다. diffusion 모델은 실제 데이터 $x_0 ~ P_{real}$ 에 점진적으로 노이즈를 추가하여 $x_{T} ~ N(0,I)$로 변환하는 Diffusion Process를 역으로 하는 방식으로 학습된다($T$ = 1000). 이것을 $mu_{base}(x_t,t)$라 가정한다. 모델은 $X_T$에서 시작하여 timestep $t \in {0, 1, ..., T\ -\ 1}$이 주어졌을 때의 $x_t$의 노이즈를 반복적으로 제거하여 target 데이터 분포의 샘플을 생성한다. 일반적으로 Diffusion 모델은 실제 같은 이미지를 생성하기 위해서는 수십에서 수백 step이 필요하다.

One-step generator

one-step generator $G_{\theta}$는 base diffusion denoiser에서 time-conditioning이 제외된 아키텍쳐로 구성되어 있다. 저자는 parameter $\theta$의 초기값으로 base model을 사용하였다. 즉, $G_{\theta} = \mu_{base}(z, T-1)$.

Distribution Matching Loss

이상적인 것은, one-step generator가 생성한 이미지와 진짜 이미지를 구별할 수 없는 것이다. 저자는 이를 위해 real 분포($p_{real}$와 fake 분포$p_{fake}$의 KL divergence를 최소하하는 objective를 사용하였다. 수식은 다음과 같다.
$$D_{KL}(p_{fake}\ ||\ p_{real}) = E_{x \sim p_{fake}}(log(\frac{p_{fake}(x)}{p_{real}(x)})) = E_{z \sim N(0;I),\ x=G_{\theta}(x)} - (log\ p_{real}(x) - log\ p_{fake}(x)) \tag{1}$$

위 loss를 추정하기 위해 확률 밀도를 계산하는 것은 일반적으로 다루기 어렵지만 gradient descent로 본 generator를 학습하기 위해 $\theta$에 대한 gradient(경사)만 필요하다.

$\theta$에 대한 gradient는 다음 수식으로 표현될 수 있다.
$$\Delta_{theta}D_{KL} = E_{z \sim N(0;I), x=G_{\theta}(z)} [ -(s_{real}(x) - s_{fake}(x)) \Delta_{\theta}G_{\theta}(z) ] \tag{2}$$

• $s_{real}(x) = \Delta_x log\ p_{real}(x)$
• $s_{fake}(x) = \Delta_x log\ p_{fake}(x)$

여기에서 $s_{real}$은 x를 $p_{real}$의 모드로, $s_{fake}$는 두 모드를 분리시킨다. 위와 같은 score 근사치를 계산하는 것 역시 다음 두 가지 이유로 어렵다.

• 낮은 확률의 샘플에 대한 점수 다양성, 특히 fake 샘플의 경우 미리 사라짐
• 점수를 추정하기 위한 의도된 도구, 즉 Diffusion 모델은 Diffusion 분포의 점수만을 제공한다.

Score-SDE는 위 두 문제에 대한 답을 제공한다. 다양한 표준 편차의 가우시안 노이즈로 데이터 분포를 교란함으로써 주변 공간에서 완전히 지원되고, 따라서 중첩되는 "Blurred" 분포를 생성하여 식 2의 gradient가 잘 정의된다. 그 후, Score-SDE는 학습된 Diffusion 모델이 Diffusion 분포의 점수 함수를 근사화한다는 것을 보여준다. 따라서, 저자의 전략은 가우시안 diffusion 후 real과 fake 분포의 점수를 모델링하기 위해 diffusion 분포 쌍을 사용하는 것이라고 한다.

정리하자면, 저자의 전략은 가우시안을 활용한 diffusion process 이후에 real 분포와 fake 분포의 점수를 모델링하기 위해 diffusion denoiser를 pair로 사용하는 것이다. 각각의 점수는 &s_{real}(x_t,t)&, &s_{fake}(x_t,t)&로 표기하였다. diffused 샘플 $x_t\ \sim \ q(x_t|t)$는 $t$ 번의 timestep을 가진 diffusion $G_{\theta}(z)$의 output에 noise를 추가한 것이다.(= diffusion process를 거친 것이다.) 이에 대한 수식은 아래와 같다.

$$q_t(x_t|x)\ \sim \ N(\alpha_t x; \sigma^2_t I) \tag{3}$$

• $\alpha_t, \sigma_t$ : diffusion noise schedule

Real score

real 분포는 base diffusion 모델의 학습 이미지에 해당하는 고정 분포이므로 pre-trained diffusion 모델 $\mu_{base}(x,t)$의 고정 복사본을 사용하여 점수를 모델링한다. 수식은 아래와 같다.
$$s_{real}(x_t,t)\ = \ - \ \frac{x_t\ - \ \alpha_t \mu_{base}(x_t,t)}{\sigma^2_t} \tag{4}$$

Dynamically-learned fake score

위에서 언급한 수식과 마찬가지로 fake score는 다음 수식으로 정의하였다.
$$s_{fake}(x_t,t)\ = \ - \ \frac{x_t\ - \ \alpha_t \mu^{\phi}_{fake}(x_t,t)}{\sigma^2_t} \tag{5}$$

하지만 학습을 이어나가면서 생성되는 샘플들의 분포는 바뀌게 된다. 이를 바로잡기 위해 fake diffusion 모델인 $\mu^\phi_{fake}$를 동적으로 변화시킨다. 일단, fake diffusion 모델을 초기값으로 pre-trained diffusion 모델 $\mu_{base}$의 weight를 사용하고, 다음 간단한 denoising objective를 최소화하는 방향으로 파라미터 $\phi$를 업데이트한다.
$$L^\phi_{denoise}\ = \ \mid\mid\mu^\phi_{fake}(x_t,t)\ - \ x_0 \mid\mid^2_2 \tag{6}$$

여기에서 $L^\phi_{denoise}$는, diffusion timestep $t$에 따라 가중치가 부여되며 일반적인 diffusion 모델이 학습할 때 사용하는 것과 동일한 weighting strategy를 적용한다.

Distribution matching gradient update

최종 근사 분포 매칭 gradient는 식(2)의 정확한 점수를 교란된 샘플 $x_t$에 대한 두 diffusion 모델에 의해 정의된 score와 다시 배치하고 diffusion time step에 대한 기댓값을 적용함으로써 얻어진다.

$$\Delta_\theta D_{KL} \simeq \mathbb{E}_{z,t,x,x_t}[w_t\alpha_t(s_{fake}(x_t,t)- s_{real}(x_t,t))\Delta_\theta G_\theta(z)] \tag{7}$$

• $z\ \sim \ N(0;I)$
• $x\ =\ G_\theta(z)$
• $t\ \sim \ \mathcal{U}(T_{min}, T_{max})$,
• $x_t\ \sim \ q_t(x_t|x)$

$w_t$는 time-dependent scaler weight로써 학습 개선을 위해 사용되었다. 다양한 noise-level에 걸쳐 gradient의 크기를 정규화하도록 weight factor를 설계했으며 입력이미지와 노이즈 이미지 사이에 spatial 및 channel에 걸친 mean absolute error(MAE)를 계산한다.
$$w_t\ = \ \frac{\sigma^2_t}{\alpha_t}\frac{CS}{||\mu_{base}(x_t,t)-x||_1} \tag{8}$$

• $S$ : Spatial location 수(pixel 수)
• $C$ : 채널 수
• $T_{min} : 0.02T$ (수치 근거 : DreamFusion)
• $T_{max} : 0.98T$ (수치 근거 : DreamFusion)

Regression loss and final objective

Fig. 3 Optimizing various objectives starting from the same configuration

 

위 섹션에서는 $t >> 0$ 즉, 생성된 샘플이 많은 양의 노이즈가 포함된 경우에 대해 표현한 것이다. 그러나 적은 양의 노이즈의 경우, $s_{real}(x_t,t)$가 0으로 가기 때문에 종종 신뢰할 수 없게 된다. 또한, $\Delta_x log(p)$는 확률 밀도 함수 $p$의 스케일링에 불변하기 때문에 최적화는 mode collapse/dropping에 취약하며, 여기서 fake 분포는 mode의 하위 집합에 더 높은 전체 밀도를 할당한다. 이를 방지하기 위해 추가적인 regression loss를 사용한다. (그림 3 참고)

이 loss는 같은 노이즈가 입력되었을 대, geenrator와 base diffusion 모델로부터 생성되는 이미지의 pointwise distance를 의미한다. 구체적으로, paired-dataset $D = {z,y}$ ($z$ : 랜덤 가우시안 노이즈 이미지; $y$ : 그에 상응하는 출력이미지로, deterministic ODE solver를 사용한 pre-trained diffusion 모델 $\mu_{base}$을 통해 생성되었음.)

저자는 CIFAR-10의 경우, 1% 미만의 학습 컴퓨팅을 사용하여 생성된 적은 수의 노이즈-이미지 Pair도 효과적인 정규화 역할을 한다는 것을 발견했다고 한다. 본 회귀식은 다음과 같다.
$$L_{reg} = \mathbb{E}_{(z,y) \sim D} \ell (G_\theta(z),y) \tag{9}$$

• function $ell$ : Learned Perceptual Image Patch Similarity (LPIPS)

Final objective

네트워크 $\mu^\phi_{fake}$는 $\Delta_\theta D_{KL}$의 계산을 돕기 위해 $L^\phi_{denoise}$와 함께 학습되었다.
$G_\theta$ 학습을 위한 최종 objective는 $D_{KL}\ + \ \lambda_{reg} L_{reg}$이다.($\lambda_{reg} = 0.25$) gradient $\Delta_\theta D_{KL}$은 '수식 7'에서 확인할 수 있으며 gradient $\Delta_\theta L_{reg}$는 '수식 9'를 통해 자동으로 연산된다.

정리하자면, 저자는 두 개의 loss를 사용하였다.

• unpaired fake 샘플들은 분포 매칭 gradient 사용
• paired example들은 regression loss 사용
  추가적인 내용은 Appendix B를 참고하면 알 수 있다.

Distillation with classifier-free guidance

Classifier-Free Guidance (CFG)는 text-to-image (t2i) diffusion 모델의 이미지 품질 향상을 높이기 위해 두루두루 사용된다. 저자의 접근법도 CFG를 사용한다. 저자는 먼저, 회귀 손실 $L_{reg}$에 필요한 Pair 데이터셋을 구성하기 위해 guided 모델에서 샘플링하여 해당 노이즈-출력 Pair를 생성한다. distribution matching gradient $\Delta_\theta D_{KL}$을 계산할 때에는, real score를 guided 모델의 평균 예측에서 도출된 점수로 대체한다. 한편, fake score를 위한 수식을 수정할 수는 없다. 저자는 one-step generator를 사용하기 위해 InstaFlow와 LCM-LoRA와 마찬가지로 수정된 guidance scale을 사용하였다고 한다.

Experiments

Table 1.   ImageNet-64x64 에서의 품질 비교

위 테이블은 섹션 별로 1) 자주 사용되는 GAN, Diffusion 모델. 2) distillation diffusion 모델. 3) 저자가 Teacher 모델로 사용한 모델 EDM을 가리킨다. 실험 결과, 저자가 제안한 DMD 모델이 1-step에서의 FID가 가장 낮은 것으로 나타났다.

Fig. 5-1. qualitative comparison

위 그림은 저자의 방법론에 distribution matching의 사용 여부에 따른 이미지 생성 결과를 비교한 것이다. 한눈에 보기에도 mode collapse가 발생한 것을 볼 수 있다.

Fig. 5-2. Qualitative comparsion

위 그림은 regression loss 사용 여부에 대한 이미지 생성 결과의 차이를 나타낸 것이다. 사용하지 않는 경우에는 mode collapse가 발생하며 생성 결과의 다양성이 떨어지는 것을 볼 수 있다.(우측 이미지의 빨간색 박스 표시) 모든 이미지는 같은 random seed에 의해 생성되었다.

Table 2. Ablation study

위 테이블은 Ablation study에 대해 FID score를 계산한 표이다. 표의 좌측은 loss의 변화, 우측은 다양한 가중치에 대한 실험이다.

Text-to-Image Generation

저자는 MS COCO 데이터셋에 대해 zero-shot 검증을 진행하였다. 학습에 사용된 모델은 SD1.5이며 base SD에서 가장 좋은 FID를 달성하도록 guidance scale 3으로 설정하였으며 데이터셋은 LAION-Aesthetics-6.25+를 사용하였다. 학습 시간은 72개의 A100을 사용하여 36시간이 걸렸다고 한다.

Table 3. Sample quality comparsion on zero-shot text-to-image generation on MS COCO-30k.

위 표는 다양한 모델 별 zero-shot MS COCO-30k에 대한 latency와 FID score를 나타낸 결과이다.

Table 4. FID/CLIP-Score comparison on MS COCO-30K

위 표는 타 기술 및 LCM-LoRA의 다양한 step보다 저자가 제안한 DMD가 latency, FID, CLIP-score 측면에서 모두 우세하다는 결과를 보인다.

Limitations

정량적 평가나 생성 속도를 보았을 때, 본 논문에서 제안한 모델이 충분히 좋은 성능을 보이지만 여전히 one-step diffusion 모델에는 약간의 품질 저하가 있다고 한다. 또한, 본 모델은 본질적으로 Teach 모델의 역량에 제한된다. SD1.5에서 관측된 한계에 따라 본 모델 또한 텍스트나 작은 얼굴, 사람의 상세한 묘사를 표현하는데 어려움을 겪는다고 한다. 따라서, 성능 향상을 위해 SDXL과 같이 보다 큰 모델에 대한 distill(증류) 실험을 해보아야 한다고 이야기한다.