[22′ NIPS] Classifier Guidance : Diffusion Models Beat GANs on Image Synthesis

1. Introduction

이번 글에서는 2022년 NIPS에 발표된 Diffusion Models Beat GANs on Image Synthesis 논문을 리뷰합니다. 이 논문은 ADM (Ablated Diffusion Model) 또는 Classifier Guidance 라는 별칭으로 불리며, 이번 글에서는 Classifier Guidance 라고 지칭하겠습니다.

이번 글에서는 Diffusion Model 배경 지식에 대한 설명은 생략했습니다. Diffusion Model의 기초와 응용 글을 참고해주세요.

Classifier Guidance 이전에 다양한 Diffusion Model 들이 발표되었는데요. GAN과 비교했을때 이들 방법의 공통적인 한계는 Condition을 반영할 수 없다는 것입니다. 예를들어 2014년 GAN이 처음 제안된 뒤 곧바로 Conditional GAN이 발표되었죠. Conditional GAN은 Class를 Condition으로 입력받아 원하는 Class 이미지만 생성할 수 있었습니다. 반면 지금까지 발표된 Diffusion Model들은 이렇게 원하는 Class 이미지만 생성하는 기능이 없었죠.

따라서 Classifier Guidance 에서는 이렇게 Class 정보를 Condition으로 입력받아 해당 Class 이미지만 생성할 수 있는 Diffusion Model을 제안합니다. 이때 별도의 Classifier를 학습시켜 이로부터 나오는 Gradient를 활용하는데요. 이 때문에 Classifier Guidance 라는 별칭이 붙었습니다.

2. Method

이번 챕터에서는 Classifier Guidance의 제안 방법을 하나씩 살펴보겠습니다. 크게 두 가지면에서 살펴볼건데요.
첫 번째는 Architecture 의 개선점입니다. 기존 Diffusion Model 대비 몇 가지 Architecture에서의 개선을 통해 성능을 향상할 수 있었습니다.
두 번째는 Classifier Guidance 입니다. 이번 논문의 핵심 내용이라고 할 수 있는데요. 어떻게 Classifier를 사용하여 특정 Class 이미지만 생성할 수 있는지 위주로 살펴보겠습니다.

2-1. Architecture Improvements

먼저 Architecture 에서의 개선점을 살펴보겠습니다. 기존 Diffusion Model 들은 전부 UNet 구조를 사용하는데요. Classifier Guidance 에서도 UNet 구조를 동일하게 사용하면서 몇 가지를 개선합니다.

이 부분은 간단하게 살펴보고 넘어가겠습니다. 먼저 Depth를 증대해주고요, Attention Head를 늘려줍니다. 기존에는 Attention 을 8×8 사이즈에만 적용했는데 이를 16×16, 32×32에도 적용해주죠. BigGAN의 Residual Block을 차용하고요 Rescale Residual Connection 기법을 사용합니다. 또한 Adaptive Group Normalization을 적용해주죠.

이렇게 Arhcitecture를 변경하여 얻은 효과는 다음과 같습니다.

위 표를 통해 Architecture 개선을 통해 약간의 FID 성능 향상이 있음을 알 수 있습니다.

2-2. Classifier Guidance

이번 논문의 가장 중요한 포인트는 Classifier Guidance 인데요. 어떻게 Classifier를 사용하여 Class Condition에 대한 Guidance 효과를 냈는지 살펴보겠습니다.

먼저 Classifier를 학습해줍니다. 예를 들어 최종적으로 생성하고자 하는 이미지 데이터셋이 ImageNet이라면, ImageNet Classifier를 학습해주는 거죠. 다만 최종적으로 Diffusion Model에서 활용할 목적이니까 Noisy Image를 입력으로 사용해줍니다.

위 그림을 보면 입력으로 강아지 이미지뿐만 아니라 Noisy 강아지 이미지가 들어가는 모습을 볼 수 있습니다. 이렇게 들어간 Noisy 강아지, 고양이 이미지등에 대해 Class를 구분하도록 학습해줍니다.

Classifier 학습이 완료되면 이제 Classifier는 Freeze하고 Diffusion Model을 학습해줍니다.

위 그림은 Classifier Guidance 방식의 Diffusion Model 학습 과정을 도식화한 그림입니다. Classifier로부터 입력 이미지에 대한 해당 클래스의 Gradient가 나오고 이는 Noise에 영향을 주는 모습을 표현하고 있습니다. 기존 Diffusion Model의 학습 방법과 동일하게 학습하되, Classifier로부터 나오는 Gradient가 예측한 Noise에 추가되는 것이죠.

이러한 방법이 왜 Class Guidance 역할을 할 수 있을까요? 강아지 이미지를 생성하려고 하는 상황을 가정해볼게요. 이때 Diffusion Model이 고양이 이미지를 생성하는 방향으로 진행되면 어떻게 될까요? Classifier는 Noisy Image Classification을 학습했죠. 따라서 Noisy 고양이 이미지에 대해 큰 Gradient를 줄겁니다. 이 Gradient는 Diffusion Model의 Output Noise에 강아지 이미지를 생성하는 방향으로 반영되는거죠.

이를 알고리즘으로 표현하면 다음과 같습니다.

기존 Diffusion Model과 달리 예측한 Noise에 Classifier로부터의 Gradient가 반영되는 모습을 볼 수 있습니다.

3. Results

이렇게 학습한 Classifier Guidance의 Image Generation 성능을 살펴보겠습니다.

위 표는 다양한 모델의 Image Generation 성능을 비교한 표입니다. 제안 방법인 ADM이 대체로 가장 좋은 모습을 볼 수 있습니다.

위 그림은 BigGAN (가장 왼쪽) 과 Classifier Guidance (가운데) 모델이 생성한 클래스별 이미지를 보여주고 있습니다. Classifier Guidance 모델도 Class에 해당하는 이미지만 잘 생성하는 모습을 볼 수 있습니다.

4. Conclusion

지금까지 Classifier Guidance 논문의 주요 내용을 살펴봤습니다. Classifier Guidance는 기존 Diffusion Model에 Class를 Condition으로 입력받을 수 있는 방법을 제안했다는 의의를 갖습니다. 반면 아직도 GAN보다 느린 속도와 Labeled Dataset만 사용할 수 있다는 점은 한계로 남아있죠. 뿐만 아니라 Classifier Guidance 방식을 사용하기 위해서는 별도의 Classifier를 학습해야만 한다는점 또한 한계라고 할 수 있습니다. 이는 필연적으로 Classifier Free Guidance 모델의 필요성을 제기했습니다.