[22’ NIPS] CG: Diffusion Models Beat GANs on Image Synthesis

0. Abstract

1. Diffusion model architecture 부분에서 몇 가지를 개선하여 성능을 향상시켰다.
2. Classifier guidance를 도입하여 conditional image synthesis를 수행하였으며, diversity와 fidelity 간 trade-off를 조절할 수 있었다.

1. Introduction

Diffusion model이 GAN을 이기지 못하는 이유에 대해 저자는 두 가지 이유를 들었다. 첫째는 GAN은 이미 너무 많이 연구되어 model architecture가 충분히 정제되어 있다는 점이고, 둘째는 GAN은 fidelity를 높이는 대가로 diversity를 낮게 하였기 때문이다. 따라서 본 논문에서는 모델의 구조를 개선하고, classifier guidance를 통해 diversity와 fidelity 간 trade-off를 조절하였다.

들어가기에 앞서 몇 가지 디테일이 있다. 본 논문에서는 IDDPM의 $L_{\text{hybrid}}$를 사용하여 variance learning을 수행하였고, sampling step이 50보다 적을 때에는 DDIM sampling을 사용하였다. 참고로 IDDPM의 저자인 OpenAI의 Nichol과 Dhariwal이 본 논문의 저자이기도 하다.

2. Architecture Improvements

본 논문에서 제시한 모델 구조 개선 사항은 다음과 같다. 아래의 두 가지 개선사항은 “Score-based generative modeling through stochastic differential equations (2020)” 논문을 참고한 것이다.

• Model size는 비슷하게 한 채 width에 비해 depth를 늘렸다.
• Attention head의 수를 늘렸다. (실제로는 시간 대비 가장 효율이 좋았던 64개를 사용했다.)
• $16 \times 16$ 뿐 아니라 $32 \times 32$, $16 \times 16$, $8 \times 8$에서도 attention을 적용했다.
• BigGAN의 residual block 구조를 사용했다.
• Residual connection 시 $1/\sqrt{2}$를 곱해주어 rescaling하였다.

그러나 실제로 실험해보니 아래와 같이 마지막 rescaling 방법은 성능의 개선이 없어 제거하였다.

그리고 AdaGN(Adaptive Group Normalization)을 사용하여 성능을 향상시켰다. 이는 이미 제안되었던 AdaIN과 유사한 방법으로 다음과 같은 식으로 정의된다. 이때 $y_s$와 $y_b$는 timestep 및 class embedding의 linear projection으로부터 얻는다. 실험 결과 단순 Addition + GN보다 AdaGN의 성능이 높았다.

\[\text{AdaGN}(h, y) = y_s \text{GroupNorm}(h) + y_b\]

3. Classifier Guidance

물론 모델 구조를 바꾸어 성능을 향상시켰다는 점도 중요하지만 본 논문의 핵심은 CG(Classifier Guidance)이라 할 수 있다. CG에서는 classifier $p_\phi (y \vert x_t, t)$를 noisy image $x_t$에 대해 훈련한 후, gradient $\nabla_ {x_t} \log p_\phi (y \vert x_t, t)$를 사용하여 diffusion sampling process에 class label $y$에 대한 guidance를 제공한다. 대강 어떤 뜻인지는 이해되지만 단순히 classifier의 gradient만 있으면 guidance가 된다니, 처음에 이를 읽고 직관적으로 받아들여지지는 않았다. 논문에서는 수식을 이용하여 설명을 이어간다. 참고로 앞으로는 $(x_t, t)$ 대신 brevity를 위해 $x_t$로 표기하겠다.

3.1. Conditional Reverse Noising Process

Conditional reverse noising process를 다음과 같이 나타낼 수 있다. 증명은 본 논문의 Appendix에 있다. $Z$는 normalizing constant이다.

\[p_{\theta, \phi}(x_t \vert x_{t+1}, y) = Z p_\theta (x_t \vert x_{t+1}) p_ \phi (y \vert x_t)\]

한편 $p_\theta (x_t \vert x_{t+1})$는 Gaussian $\mathcal{N} (\mu, \Sigma)$이므로 다음과 같이 나타낼 수 있다.

\[\log p_{\theta}(x_t \vert x_{t+1}) = -\frac{1}{2} (x_t - \mu)^T \Sigma^{-1} (x_t - \mu) + C\]

저자들은 $\log p_\phi (y \vert x_t)$의 curvature가 $\Sigma ^{-1}$에 비해 작다고 가정했다. 다음과 같은 Taylor expansion을 위해서 이계도함수를 무시한 것이다.

\[\begin{aligned} \log p_\phi(y \vert x_t) &\approx \log p_\phi(y \vert x_t) \vert_{x_t = \mu} + (x_t - \mu) \nabla_ {x_t} \log p_\phi(y \vert x_t) \vert_{x_t = \mu} \\ &= (x_t - \mu) g + C_1 \quad \text{where } g = \nabla_ {x_t} \log p_\phi(y \vert x_t) \vert_{x_t = \mu} \end{aligned}\]

이로부터 다음과 같은 식을 얻을 수 있다.

따라서 다음과 같은 샘플링 과정을 거친다. 즉, 원래 평균에서 $\Sigma g$만큼 이동한 곳에서 샘플링을 수행한다. 이때 optional scale factor $s$를 도입하였다. 결국 이러한 샘플링은 conditional reverse process를 도입한 결과라고 볼 수 있다.

3.2. Conditional Sampling for DDIM

위 과정은 stochastic DDPM에만 적용되고, deterministic DDIM에는 적용되지 않는다. 따라서 DDIM에도 conditional reverse process를 도입하기 위해 score-based conditioning trick을 사용한다. DDPM의 결과로 score function을 다음과 같이 나타낼 수 있다.

\[\nabla_ {x_t} \log p_\theta (x_t) = - \frac{1}{\sqrt{1- \bar{\alpha}_ t}} \epsilon_ \theta (x_t)\]

따라서 conditional reverse process를 다음과 같이 나타낼 수 있다.

\[\begin{aligned} \nabla_ {x_t} \log (p_\theta (x_t) p_ \phi (y \vert x_t)) &= \nabla_ {x_t} \log p_\theta (x_t) + \nabla_ {x_t} \log p_ \phi (y \vert x_t) \\ &= - \frac{1}{\sqrt{1- \bar{\alpha}_ t}} \epsilon_ \theta (x_t) + \nabla_ {x_t} \log p_ \phi (y \vert x_t) \end{aligned}\]

따라서 대응 관계에 따라 새로운 epsilon prediction $\hat{\epsilon} (x_t)$를 다음과 같이 정의하자.

\[\hat{\epsilon} (x_t) = \epsilon_ \theta (x_t) - \sqrt{1- \bar{\alpha}_ t} \nabla_ {x_t} \log p_ \phi (y \vert x_t)\]

실제 DDIM sampling algorithm은 다음과 같다. 이제 DDPM이든, DDIM이든 classifier-guided sampling을 수행할 수 있다. 이제 classifier graident로 guidance를 제공한다는 것이 어떠한 의미인지 와닿는다.

3.3. Scaling Classifier Gradients

DDPM sampling에서 다음과 같은 식을 얻었고, 실제로는 $\mathcal{N}(\mu + sg, \Sigma)$로부터 샘플링을 수행한다고 하였다.

\[\log (p_\theta (x_t \vert x_{t+1}) p_ \phi (y \vert x_t)) = \log p(z) + C_4, \quad z \sim \mathcal{N}(\mu + \Sigma g, \Sigma)\]

그런데 실제로 scaling factor $s=1$로 두고 sampling을 진행하면 충분한 guidance가 이루어지지 않았다. 따라서 classifier gradient의 크기를 조절하여 sampling을 수행하였다.

저자들은 이러한 효과를 sharpening으로 설명한다. $s$가 커질수록 distribution이 sharp해지면서 classifier guidance가 더 강해진다는 것이다. 이는 다음 결과로 확인할 수 있다. 즉, $p(y \vert x)$대신 $p(y \vert x) ^ s$를 사용하여 guidance를 제공하는 것이다. $s$가 증가할수록 fidelity는 증가하지만 diversity는 감소할 것이다.

\[s \cdot \nabla_x \log p_\phi (y \vert x) = \nabla_x \frac{1}{Z} p_\phi (y \vert x) ^ s\]

지금까지는 classifier guidance만 설명하였지만 물론 conditional diffusion model $p(x \vert y)$ 자체를 훈련시킬 수도 있다. 즉 diffusion model에 class label 정보도 같이 embedding하여 훈련시키는 것이다. 결과는 둘 다 사용하는 것이 가장 좋았다.

이제 fidelity와 diversity 간의 trade-off를 분석할 차례이다. FID와 IS를 비교할 때 FID는 fidelity 및 diversity를 모두 고려하는 반면, IS는 비교적 fidelity에 더 중점을 둔다고 이해하면 된다. 한편 precision과 recall을 비교할 때는 precision이 fidelity에 해당하고 recall이 diversity에 해당한다. 결과는 아래와 같다. Figure 4를 통해 gradient scale $s$가 증가하면 fidelity가 증가하고 diversity가 감소하는 것을 확인할 수 있다. Figure 5를 통해서는 BigGAN과 비교하여 CG의 trade-off line이 더 좋은 것을 알 수 있다. Precision-recall curve는 plot이 더 바깥쪽에 있을수록 좋고, IS-FID curve는 plot이 더 아래쪽에 있을수록 좋다.

4. Results

4.1. State-of-the-art Image Synthesis

SOTA 모델들과 비교한 결과는 다음과 같다. 여기서 개선된 모델 구조를 사용한 모델을 ADM(Ablated Diffusion Model)이라 이름지었으며, CG를 사용한 모델을 ADM-G라고 하였다.

대부분 ADM을 활용하여 SOTA를 얻었다. DDIM sampling을 이용해 25 steps만으로도 좋은 성능을 보여주었다. 참고로 아래는 BigGAN과 본 모델을 정성적으로 비교한 결과이다.

4.2. Comparison to Upsampling

저자들은 guidance와 two-stage upsampling stack 기법을 비교하였다. Two-stage upsampling stack은 먼저 low-resolution model이 sample을 생성하면 upsampling model이 이를 high-resolution으로 변환하는 방식이다. 이 방식으로 ImageNet $256 \times 256$에서 훈련한 결과는 Table 5에 있다. 성능이 향상되었음에도 BigGAN을 이기지 못한 결과를 보여준다.

따라서 저자들은 다음과 같이 분석하였다. Upsampling 기법은 diversity는 거의 유지하면서 fidelity를 향상시킬 수 있지만, CG는 diversity를 떨어뜨리는 대신 fidelity를 향상시킨다. 따라서 이 두 가지 기법을 모두 사용하여 가장 좋은 성능을 얻을 수 있을 것이라고 생각하였고, 이를 실험한 결과는 Table 6에 있다. FID가 개선되었고 IS도 높아진 모습이다.

💡 Summary

• 몇 가지 모델 구조를 개선하여 성능을 향상시켰고, classifier guidance를 통해 conditional image synthesis를 수행하였다. 이때 CG(Classifier Guidance)를 통해 diversity와 fidelity 간 trade-off를 조절할 수 있었다.
• 다른 모델들과 비교하여 SOTA 성능을 보여주었으며, two-stage upsampling stack을 사용하여 더 좋은 성능을 얻을 수 있음을 보여주었다.
• Opinion: 수학적으로 CG(Classifier Guidance)를 정당화한 것이 중요한 논문이었다고 생각한다. 한편 실제로는 Classifier를 추가로 훈련해야 하며, labeled dataset에서만 사용할 수 있는 방법이라는 점이 문제가 될 것으로 보인다. 이러한 부분은 GAN에서와 같이 CLIP 등으로 해결할 수 있을지도?

📃 Reference

• [22’ NIPS] CG: Diffusion Models Beat GANs on Image Synthesis

Generation 카테고리 내 다른 글 보러가기