[논문리뷰] Classifier-Free Diffusion Guidance

[Paper]

Introduction

기존 연구들의 한계점

Classifier guidance는 아래의 문제점이 있다.

• 추가적인 classifier 학습이 필요하기 때문에 diffusion 모델의 학습 파이프라인을 복잡하게 만든다.
• 사용하는 classifier는 노이즈가 섞인 데이터 위에서 학습되어야 하므로, 일반적인 사전학습된 classifier를 바로 사용할 수 없다.

• FID, IS 등의 metric은 사전학습된 classifier를 기반으로 평가하는데, 이 방식은 score를 직접적으로 더하면서 클래스 방향으로 샘플링을 유도하기 때문에 classifier가 잘 맞추게끔 이미지에 작은 조작을 반복하는 것으로 볼 수 있다.

  따라서 진짜 좋은 이미지를 만든 것인지, classifier를 잘 속인 결과인지가 불분명하다.

제안하는 방법

분류기를 사용하지 않는 guidance 기법을 제안한다.

• Conditional 디퓨전 모델의 socre 추정값과 unconditional 디퓨전 모델의 score 추정값을 혼합한다.

Methods

GAN이나 flow-based 모델과 같은 생성 모델은 샘플링 때 입력 노이즈의 분산 또는 범위를 줄임으로써 truncated sampling 또는 low temperature sampling을 수행할 수 있다.

이러한 방식은 샘플의 다양성을 줄이지만, 각 샘플의 품질을 높일 수 있다.

그러나 위의 방식을 디퓨전 모델에 적용하기 위해, score를 스케일링하거나 reverse process에서 가우시안 노이즈의 분산을 줄이는 방식을 사용하면 오히려 저품질의 샘플을 생성하게 만든다.

1. Classifier Guidance

디퓨전 모델에서 truncation과 유사한 효과를 얻기 위해 [Diffusion models beat GANs on image synthesis] 논문에서는 classifier guidance 기법을 도입하였다.

이 기법에서는 디퓨전 모델의 score ${\boldsymbol\epsilon}_\theta(\mathbf{z}_\lambda,\mathbf{c})\approx\sigma_\lambda\nabla_{\mathbf{z}_\lambda}\log p_\theta(\mathbf{z}_\lambda\mid\mathbf{c})$에 auxiliary classifier 모델 $p_\theta(\mathbf{c}\mid\mathbf{z}_\lambda)$의 log-likelihood의 gradient를 포함시킨다.

\[\tilde{\boldsymbol\epsilon}_\theta(\mathbf{z}_\lambda,\mathbf{c}) ={\boldsymbol\epsilon}_\theta(\mathbf{z}_\lambda,\mathbf{c})-w\sigma_\lambda\nabla_{\mathbf{z}_\lambda}\log p_\theta(\mathbf{c}\mid\mathbf{z}_\lambda) \approx\sigma_\lambda\nabla_{\mathbf{z}_\lambda}[\log p_\theta(\mathbf{z}_\lambda\mid\mathbf{c})+w\log p_\theta(\mathbf{c}\mid\mathbf{z}_\lambda)]\]

위의 수식에서 $w$는 classifier guidance의 강도를 조절하는 하이퍼파라미터이다.

샘플링 때는 기존 score ${\boldsymbol\epsilon}_\theta(\mathbf{z}_\lambda,\mathbf{c})$ 대신 변경된 score $\tilde{\boldsymbol\epsilon}_\theta(\mathbf{z}_\lambda,\mathbf{c})$가 사용되며, 이는 아래와 같은 분포로부터 샘플링을 하는 것과 동일하다.

\[\tilde{p}_\theta(\mathbf{z}_\lambda \mid \mathbf{c}) \propto p_\theta(\mathbf{z}_\lambda \mid \mathbf{c}) \, p_\theta(\mathbf{c} \mid \mathbf{z}_\lambda)^w\]

이때, 분류기 $p_\theta(\mathbf{c} \mid \mathbf{z}_\lambda)$가 해당하는 클래스 $c$에 대해 높은 likelihood를 부여하도록 가중치를 조정하며, $w>0$으로 설정하면 생성된 샘플의 다양성이 감소하지만 IS가 향상된다.

아래의 관계 때문에 이론적으로, unconditional 모델에 가중치 $w+1$의 classifier guidance를 적용하면, conditional 모델에 $w$의 guidance를 적용한 것과 동일한 결과를 얻을 수 있다.

\[p_\theta(\mathbf{z}_\lambda \mid \mathbf{c}) \, p_\theta(\mathbf{c} \mid \mathbf{z}_\lambda)^w =p_\theta(\mathbf{z}_\lambda)p_\theta(\mathbf{c} \mid \mathbf{z}_\lambda)^{w+1}\]

Score 관점에서는 아래와 같다.

\[\epsilon_\theta(\mathbf{z}_\lambda) - (w + 1) \, \sigma_\lambda \, \nabla_{\mathbf{z}_\lambda} \log p_\theta(\mathbf{c} \mid \mathbf{z}_\lambda) \approx -\sigma_\lambda \, \nabla_{\mathbf{z}_\lambda} \left[ \log p(\mathbf{z}_\lambda) + (w + 1) \log p_\theta(\mathbf{c} \mid \mathbf{z}_\lambda) \right] \\ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ = -\sigma_\lambda \, \nabla_{\mathbf{z}_\lambda} \left[ \log p(\mathbf{z}_\lambda \mid \mathbf{c}) + w \log p_\theta(\mathbf{c} \mid \mathbf{z}_\lambda) \right]\]

하지만 Diffusion models beat GANs on image synthesis의 저자들은 이미 unconditional 모델에 guidance를 적용하는 것보다 conditional로 학습된 모델에 guidance를 적용하는 경우에 성능이 더 우수함을 발견하였고, 따라서 본 연구에서는 conditional 모델에 guidance를 적용하는 설정을 중심으로 논의를 한다.

2. Classifier-Free Guidance

Classifier guidance는 이미지 classifier의 graident에 의존한다는 한계점이 있다.

별도의 classifier를 학습하는 대신, 아래의 두 모델을 함께 학습한다.

• Unconditional denoising 디퓨전 모델 $p_\theta(\mathbf{z})$: 이 모델의 score 추정값은 $\boldsymbol\epsilon_\theta(\mathbf{z}_\lambda)$
• Conditional 디퓨전 모델 $p_\theta(\mathbf{z}\mid\mathbf{c})$: 이 모델의 score 추정값은 $\boldsymbol\epsilon_\theta(\mathbf{z}_\lambda,\mathbf{c})$

파라미터 증가 없이 위 두 모델을 동시에 학습하기 위해 아래의 전략을 사용하였다.

• Unconditional 모델의 경우 conditional 모델의 class identifier $\mathbf{c}$로 null token $\varnothing$을 입력으로 준다. 즉, $\boldsymbol\epsilon_\theta(\mathbf{z}_\lambda)=\boldsymbol\epsilon_\theta(\mathbf{z}_\lambda,\mathbf{c}=\varnothing)$이다.
• 하이퍼파라미터 $p_{\text{uncond}}$를 설정해 $p_{\text{uncond}}$의 확률로 $\mathbf{c}$를 $\varnothing$으로 만들어 모델에 입력한다.

이후 아래와 같이 conditional과 unconditional score 추정값을 선형 결합하여 샘플링한다.

\[\tilde{\boldsymbol\epsilon}_\theta(\mathbf{z}_\lambda,\mathbf{c})= (1+w){\boldsymbol\epsilon}_\theta(\mathbf{z}_\lambda,\mathbf{c})-w{\boldsymbol\epsilon}_\theta(\mathbf{z}_\lambda)\]

Classifier-Free Guidance의 학습과 샘플링 알고리즘은 아래와 같다.

Experiments

Varying the Classifier-Free Guidance Strength

위의 그림에서 왼쪽은 $w=0$ (non-guided)일 때 생성된 샘플, 오른쪽은 $w=3$일 때 생성된 샘플 결과이다.

Classifier-free guidance strenght가 클수록 fidelity는 증가하지만 다양성이 감소한다.

위의 표와 그림은 classifier-free guidance strenght에 따른 FID와 IS를 비교한 결과를 나타낸다.