[논문리뷰] classifier guidance: Diffusion Models Beat GANs on Image Synthesis

논문링크

Introduction

Diffusion 모델들의 성능이 아직은 GAN의 성능을 따라잡지는 못했다. 하지만 GAN은 훈련하기 까다롭고, 다양한 도메인에 적응하기 어렵다.
Diffusion 모델은 distribution coverage, a stationary training objective, and easy scalability 같은 바람직한 속성을 제공하며, 고품질 이미지를 생성한다.

저자는 diffusion 모델이 GAN을 따라잡지 못하는 이유로 두 가지를 꼽는다.
1. GAN은 이미 충분히 많이 탐구되고 개선되었다.
2. GAN은 fidelity와 diversity 사이의 trade-off에서, diversity를 희생해 고품질의 샘플을 생성한다.
*fidelity: 현실 소리(or 이미지) 재현(reproduce)이 faithful(충실)한지를 나타내는 용도

먼저 확산 모델에 위와 같은 이점을 제공하는 것을 목표로 하며 섹션 2에서는 이전 연구들, 섹션 3에서는 간단한 아키텍처 개선사항, 섹션 4에서는 분류기의 gradient를 사용하는 방법을 설명한다.

Background

Diffusion 모델은 샘플링은, noise를 점차 덜어내는 과정을 거친다. $Noise \ \ x_T$부터 $x_{T-1}, x_{T-2}...$를 거쳐 마지막으로 $x_0$를 샘플링한다.

DDDPM 논문에서는 함수 $\epsilon (x_t,t)$로 모델을 파라미터화 하여 $x_t$의 noise를 예측하도록 하였다.
또한 $p_{\theta}(x_{t-1}|x_t)$를 가우시안 분포로 가정한다. 평균 $\mu (x_t,t)$는 $\epsilon (x_t,t)$로부터 계산할 수 있고, 분산은 상수로 둔다.

Improvements

DDIM은 랜덤 가우시안 노이즈의 표준 편차를 0으로 지정하여 모델을 implicit(결정)하게 만들었다. DDPM에 비해 훨씬 빠르고 고품질의 샘플링이 가능하며, 생성 결과가 일관된다.

Architecture Improvements

DDPM은 UNet을 베이스로 사용한다. 또한 16x16 feature map에 head가 하나인 global attention을 사용하였으며, 각 residual block에 timestep embedding의 projection을 넣어주었다.

저자는 다음과 같은 변경을 적용하였다.

• 모델의 크기를 일정하게 유지하며 깊이/너비 증가. (훈련 시간 증가로 추후 제거)
• attention head 수를 늘림.
• 기존 16x16 뿐 아니라, 32x32, 8x8에도 attention 사용
• BigGAN의 residual block을 upsampling, downsampling에 사용
• Residual connection을 $1 \over \sqrt {2}$ 로 조정

실험 결과, residual connection을 조정하는 것을 제외하곤 모두 성능이 개선되었다.

또한 저자들은 attention 구성의 변경에 대한 실험도 진행했다. head가 많거나 head 당 채널 수가 적은 경우 FID가 향상됨을 보였다. wall-clock-time을 고려하여, 기본값으로 head 당 64개 채널을 사용한다.

Adaptive Group Normalization

이 layer는 group normalization 연산 후 각 residual block에 timestep embedding과 class embdeeing을 AdaIN과 같은 방법으로 결합한다.

실제로 AdaGN이 FIN를 개선했다.

Classfier Guidance

GAN은 클래스 레이블을 많이 사용한다. classifer처럼 동작하도록 설계된 head가 있는 discriminator를 사용하거나ref, class-conditional normaliztion을 사용한다ref.

diffusion model을 컨디셔닝할 필요가 있다. 먼저, AdaGN을 통해 클래스 정보를 준다. 여기에서 diffusion generator를 개선하기 위해 classifier을 활용하고자 한다.

noise가 있는 이미지 $x_t$에 대해 classifier $p_{\phi}(y|x_t,t)$를 훈련한 다음, gradient $\nabla_{x_t}logp_{\phi}(y|x_t,t)$를 사용하여 임의의 class y로 샘플링 프로세스를 guide 할 수 있다.
간결화를 위해 $p_{\phi}(y|x_t,t)=p_{\phi}(y|x_t)$ , $\epsilon_{\theta}(x_t, t) = \epsilon_{\theta}(x_t)$ 로 사용한다.

Conditional Reverse Noising Process

원래는 y를 조건으로 하는 x를 샘플링 해야하지만, noise가 충분히 많으면 정규분포를 따르기 때문에 아래와 같이 쓸수 있다.

Diffusion model은, $x_{t+1}$로부터 $x_t$를 가우시안 분포를 사용해서 예측한다.

$log p_{\phi} (y|x_t)$가 $\sum^{-1}$보다 낮은 곡률을 가진다고 가정할 수 있다. step의 수가 무한에 가까워진다면, $||\sum|| -> 0$ 이 되어 합리적이다.
이 경우, $x_t=\mu$ 근처에서 테일러 전개를 사용한다면, 아래와 같다.

위 식으로 conditional transition operator는 unconditional transition operator와 유사하지만, 평균이 $\sum_g$ 만큼 이동한 정규 분포로 근사될 수 있음을 확인한다.

Conditional Sampling for DDIM

위 섹션의 경우는 stochastic한 경우에만 유효하다. 즉, DDIM 같은 Implicit한 경우에는 적용할 수 없다.
이를 위해 score-based conditioning trick을 사용한다.

추가된 노이즈를 예측하는 모델 $\epsilon_{\theta}(x_t)$가 있는 경우 score 함수를 유도하는데 사용할 수 있다.

위 식을 score function에 대입하면 아래와 같다.

이에 joint distribution에 대응되는 새로운 epsilon을 아래와 같이 정의할 수 있다.

샘플링 과정에서 $\epsilon_{\theta}(x_t)$ 대신 $\hat{\epsilon}(x_t)$를 사용하면 DDIM에서 사용할 수 있다.