Diffusion Models Beat GANs on Image Synthesis (NeurIPS 2021)

https://arxiv.org/abs/2105.05233

Diffusion Models Beat GANs on Image Synthesis

참고 영상 : https://www.youtube.com/watch?v=gN1FQhQsUTE 

512x512 ImageNet

 

### GAN vs Diffusion

 image generation task에서 사용하는 metric인 FID,IS, Precision은 image diversity 를 제대로 반영하지 못한다는 단점이 있음. 

GAN은 likelihood based model에 비해 diversity 가 떨어지는 치명적인 단점이 존재한다. 이에 따라 GAN의 전형적인 단점으로 손꼽히는 mode collapse 문제가 발생하기 쉽다.

반면에 diffusion model은 likelihood based model 중 하나로 signal 로 부터 noise를 점진적으로 제거하는 방식으로 새로운 sample을 생성하는 방식을 취한다. 해당 논문 전엔 CIFAR 10 에서는 SOTA를 달성했으나 LSUN이나 ImageNet에선 GAN에게 밀리는 성능을 보여왔다. 

따라서 본 논문의 저자들은 GAN은 fidelity를 위해 diversity를 희생한 것이라고 생각하였다. 그래서 이를 극복하기 위해 diversity를 일부 희생하며 fidelity를 높일 수 있는 scheme를 개발하였다.

여기서 fidelity는 precision과 비슷한 개념으로 generative image 들이 얼마나 real image distribution에 속하는를 판단하는 것이다. 반대로 recall 과 비슷한 개념인 diversity는 generative image들의 분포가 얼마나 넓게 퍼져 있는지를 측정하는 것이다.

 

### DDPM

DDPM training process

DDPM은 init image \(x_0 \)에서 T step 만큼 반복하여  noise를 생성하는 forward 방식과 이를 반대로 noise를 제거하는 backward 방식으로 image 를 reconstruction 한다. 그 후, inference 단계에선 diffusion에선 이를 sampling 이라 하는데 이는 임의의 noise를 입력으로 주고 backward로만 하여 새로운 이미지를 생성해나가는 것이다.

간단한 수식을 보면, forward 에서 이미지의 noise를 추가하는 것은 다음과 같다.

$$ x_t = \sqrt{1 - \beta_{t} x_{t-1} + \sqrt{\beta_t} \epsilon} \ \ \ \ \epsilon \sim \mathcal{N} (0 , I)$$

\( x_{t-1} \)에서 \( x_{t} \)로 noise를 입힐 때, 정규분포를 따르는 noise를 추가하는 과정이다. 여기서 \( \beta \) 의 계산은 normalization 하는 과정이라고 생각하면 된다. 

따라서 이를 수식적으로 다시 정리하자면 다음과 같은 guassian distribution을 따르게 정리한다.

$$ q(x_t | x_{t-1}) := \mathcal{N} (x_{t} ; \sqrt{ 1- \beta_t } x_{t-1}, \beta_t I) $$

하지만 이는 너무 많은 반복을 하게 되므로 computation time 이 길게 소요된다. 따라서, init image \( x_0 \) 에서 t 시점의 이미지 \( x_t \) 로 바로 forward 하는 수식이 제안되었다. \( \alpha_t := 1 - \beta_t \) 라고 정의하고, \( \bar{\alpha}_t := \Pi^{t}_{s=0} \alpha_s \) 라고 할 때

$$ q(x_t | x_0) = \mathcal{N} (x_t ; \sqrt{\bar{\alpha}_t} x_0 , (1- \bar{\alpha}_t) I ) = \sqrt{\bar{\alpha}_t}x_0 + \epsilon \sqrt{1-\bar{\alpha}_t} , \ \ \epsilon \sim \mathcal{N}(0,I) $$

로 표현이 된다. 따라서 skip time step 이 가능해진다. 

다음으로 reverse 과정을 수식적으로 표현해보자.

backward 과정은 forward 와 반대로 \( x_t \) 라는 시점의 이미지가 존재할 때, \( x_{t-1}\) 의 이미지를 구하는 것으로 \( x_t \) 가 조건으로 들어간다. 하지만 이 때, init image 에 대한 guide 가 있어야지 backward 가 효과적으로 되므로, backward distribution은 \(q ( x_{t-1} | x_t , x_0 ) \)가 된다. gaussian distribution의 mean 에 대한 분포를 \( \tilde{\mu}_t (x_t, x_0) \) 과 variance \(\tilde{\beta_t}\) 라고 할 때, 

$$ \tilde{\mu}_t (x_t , x_0) := \frac{\sqrt{\bar{\alpha_{t-1}}} \beta_t }{1 - \bar{\alpha}_t} x_0 + \frac{\sqrt{\alpha_t} (1-\bar{\alpha}_{t-1} ) }{1 - \bar{\alpha}_t} x_t $$

$$ \tilde{\beta}_t := \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t $$

$$ q(x_{t-1} | x_t , x_0 ) = \mathcal{N} ( x_{t-1} ; \tilde{\mu}(x_t, x_0), \tilde{\beta}_t I) $$

로 정의된다.

따라서 이러한 reverse 과정을 neural network 에 학습을 시켜서 풀고자 하는 것인데, \( q(x_{t-1} | x_t, x_0) \)을 \( p_{\theta} (x_{t-1} | x_t ) \)에 대한 수식으로 정의하여서 mean 과 covariance matrix (\Sigma\)를 예측하고자 하는 것이다.

$$ p_{\theta} (x_{t-1} | x_t ) := \mathcal{N} x_{t-1} ; \mu_{\theta}(x_t , t), \Sigma_{\theta} (x_t, t) ) $$

따라서 해당 분포의 \(  \mu_{\theta}(x_t , t) \) 와 \( \Sigma_{\theta} (x_t, t) \)를 학습을 하고자 하는 것이 DDPM 의 training 과정이다.

따라서 이를 학습시키기 위한 loss term 은 다음과 같다. model의 distribution \( p(x_0) \)는 true data distribution \( q(x_0) \)을 배우는 것이다. 그래서 \( p_{\theta} (x_0 ) \) 를 위한 variational lower bound 는 다음과 같은 식을 따라 

전체 loss : \( L_{vlb} := L_0 + L_1 + \cdots + L_{T-1} + L_{T} \)

init loss : \( L_0 := - \log p_{\theta} (x_0 | x_1) \)

t-1 loss : \( L_{t-1} := D_{KL} (q(x_{t-1} | x_{t},x_0 ) || p_{\theta} (x_{t-1} | x_{t})) \)

T loss : \( L_{T} := D_{KL}(q(x_{T} | x_0 ) || p(x_T )) \)

이러한 loss 를 따른다. 그러나 실제로는 이를 simplify 한 loss 를 사용한다. 그리고 \( L_{vlb} \) 보다 simplifiy 한 loss 를 사용한 것이 더욱 성능이 좋았다는 것을 실험으로 증명했다. 따라서 DDPM 에서 사용하는 loss는

$$ L_{simple} := E_{t \sim [1,T] , x_{0} \sim q(x_{0}) , \epsilon \sim \mathcal{N} (0 , I )} [ || \epsilon - \epsilon_{\theta} (x_t , t) ||^2 ] $$

이다.

그래서 결국 \( \epsilon_{\theta} \)를 학습하는 과정이고, 이러한  \( \epsilon_{\theta} \) 를 이용하여 sampling 할 때, \( \mu \)를 추정할 수 있는데, 이는 다음과 같은 공식을 따른다.

$$ \mu_{\theta} (x_t , t) = \frac{1}{\sqrt{\alpha_t}} ( x_t - \frac{1 - \alpha_t}{ \sqrt{1 - \bar{\alpha_t}}} \epsilon_{\theta}(x_t, t) ) $$

\( \epsilon \) 을 통해서 \( \mu \)를 구했고, 마지막으로 covariance를 구해야한다. 하지만, \( \mu \) 값을 통해 convariance를 구할 순 없다. 따라서 DDPM 에서는 앞서 정의한 constant \( \beta_t \)를 통해 convariance를 constant 로 설정하였다.

그리고 이를 학습하는 neural network 는 U-Net 구조를 따른다.

 

### Improved DDPM

Improved DDPM 은 DDPM 에서 2가지를 고치면서 성능을 높인 방법이다. 

먼저, DDPM 에서의 covariance값을 constant 로 사용하였는데, 이를 개선하여 covariance도 학습을 시키는 방법을 제안하였다.

$$ \Sigma_{\theta} (x_t, t) = \exp (v \log \beta_t + (1-v) \log \tilde{\beta}_t )$$

여기서 v 는 model의 output 값이다.

그리고 기존 DDPM 에서 loss term 을 수정하였는데, 이는 다음과 같다.

$$ L_{simple} + \lambda L_{vlb} $$

\( L_{vlb} \) 대신 \(L_{simple} \)을 사용한 DDPM 과 달리 Improved DDPM 은 두 loss 를 모두 사용하여서 이를 통해, sampling  step 을 적게 가져가면서 학습을 시킬 수 있었다.

 

### DDIM

참고

https://happy-jihye.github.io/diffusion/diffusion-2/

[Paper Review] DDIM: Denoising Diffusion Implicit Models 논문 리뷰

 

DDIM 에서는 DDPM과 동일하게 marginal distribution을 구성하지만 inference 과정에서 사용되는 joint distribution 과정을 다르게 distribution을 구성하였다.

DDPM은 \( x_t \)가 바로 이전 step 값 \(x_{t-1} \)에 의해 결정되는 markov chain 이지만, DDIM은 \( x_t \)가 바로 이전 step 값 \( x_{t-1} \) 과 x_{0} 에 의해 결정되는 non-markovian chain 이다. 

최근엔 DDPM 으로 학습을 하고, DDIM 으로 sampling 하는 과정을 주로 사용하기 때문에 sampling 과정만 보겠다.

왜냐하면 DDIM은 sampling 속도가 빠르기 때문이다. 

빠른 이유에 대해서 보면, non-markovian chain 으로 임의의 step에 대해서만 모델을 학습했기 때문에, generative process 에서도 일부만 sampling 할 수 있어서 sampling 을 빠르게 할 수 있다. 그러나 이러한 방식으로 인해, training process 에서는 더 많은 step 에 대해 필요하다. 따라서 주로 DDPM으로 학습시키고 DDIM 으로 sampling 하는 과정을 쓰는 것이다. 

 

### To make diffusion models beat GAN

Classifier Guidance 를 이용

diffusion model에 classifier \( p(y|x_t ) \) 를 같이 학습을 시키면 더욱 좋은 conditional image synthesis를 할 수 있다는 것이다. 여기서 classifier 는 t 시점의 image 를 주었을 때, 해당 이미지에 대한 class 가 무엇인지 맞추게하는 것이다. 

그러면 이걸 어떻게 condition을 줄까? 바로 calssifier 의 gradient 를 전달해주는 것이다. 

conditional reverse process 는 다음과 같다.

$$ q(x_t | x_{t+1},y) \propto p_{\theta}(x_t | x_{t+1}) p_{\phi}(y|x_t ) $$

여기서 \( p_{\theta} \) 는 Nerual Network of reverse process 를 의미하고 \( p_{\phi}\)는 classifier on image 를 의미한다. 여기서 y는 label 을 의미한다.

그러면 이 식을 계산의 용이성을 위해 log를 취하여 풀어쓰면

$$ \log(p_{\theta} (x_t | x_{t+1}) p_{\phi} (y | x_t ) ) \approx \log p(z) + C_{4} \sim \mathcal{N} ( \mu + \Sigma_{g} , \Sigma ) $$

처럼 된다. 여기서 C 는 constant 이고, g는 gradient of classifier를 의미한다. 이런 식으로 classifier의 gradient 를 전달해주는 방식으로 conditional reverse process 가 진행된다. 즉, gradient의 covariance 만큼 mean 을 이동시켜서 condition 을 주게 하는 것이다. 

classifier sampling guide algorithm box

Algorithm 1 은 DDPM method 를 사용하였을 때, classifier guide 를 주는 알고리즘 박스이고, Algorithm 2 는 DDIM 을 사용하였을 때 classifier guide 를 주는 알고리즘 박스이다.

Algorithm 1 에선 \( \mu_{\theta} (x_t ) \) 와 \( \Sigma_{\theta} (x_t ) \) 에서 \( \mu \) 와 \( \Sigma \) 를 각각 뽑는 과정을 거친다. 그 다음에 conditional 상태이기 때문에 uncoditional 인 상태에선 \( \mu \)와 \( \Sigma \) 에서 뽑는 거였더라면 condition 에선 \( s \Sigma \Delta_{x_{t}} \log p_{\phi} (y | x_t ) \) 만큼 mean shift 한 상태로 sampling 한다. \( \Sigma \) 부분은 앞서 말한 g 를 풀어쓴 것과 같다. 그리고 s 는 mean shift 에 대한 scale constant 인데, 이는 hyper parameter로 1로 주었을 땐, 좋은 결과가 나오진 않았지만 10으로 크게 줬을 때 좋은 생성 결과가 나왔음을 ablation study 로 증명하였다.

그리고 Algorithm 2 에서도 마찬가지로 normalization 한 상태로 class condition을 주어서 sampling을 한다.