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}ϵ}ϵ\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$ 라고 정의하고, ${\overline{\alpha }}_t:={\mathrm{\Pi }}_{s=0}^t{\alpha }_s$ 라고 할 때

$$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)I)=\sqrt{{\overline{\alpha }}_t}{x}_0+ϵ\sqrt{1-{\overline{\alpha }}_t},ϵ\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 $\widetilde{{\beta }_t}$ 라고 할 때, 

$${\tilde{\mu }}_t(x_t,x_0):=\frac{\sqrt{\overline{{\alpha }_{t-1}}}{\beta }_t}{1-{\overline{\alpha }}_t}{x}_0+\frac{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_{t-1})}{1-{\overline{\alpha }}_t}{x}_t$$

$${\tilde{\beta }}_t:=\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\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),{\mathrm{\Sigma }}_{\theta }(x_t,t))$$

따라서 해당 분포의 \(  \mu_{\theta}(x_t , t) \) 와 ${\mathrm{\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),ϵ\sim \mathcal{N}(0,I)}[||ϵ-{ϵ}_{\theta }(x_t,t)|{|}^2]$$

이다.

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

$${\mu }_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-\overline{{\alpha }_t}}}{ϵ}_{\theta }(x_t,t))$$

$ϵ$ 을 통해서 $\mu$를 구했고, 마지막으로 covariance를 구해야한다. 하지만, $\mu$ 값을 통해 convariance를 구할 순 없다. 따라서 DDPM 에서는 앞서 정의한 constant ${\beta }_t$를 통해 convariance를 constant 로 설정하였다.

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

 

### Improved DDPM

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

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

$${\mathrm{\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_{\varphi }(y|x_t)$$

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

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

$$\log (p_{\theta }(x_t|x_{t+1})p_{\varphi }(y|x_t))\approx \log p(z)+C_4\sim \mathcal{N}(\mu +{\mathrm{\Sigma }}_g,\mathrm{\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)$ 와 ${\mathrm{\Sigma }}_{\theta }(x_t)$ 에서 $\mu$ 와 $\mathrm{\Sigma }$ 를 각각 뽑는 과정을 거친다. 그 다음에 conditional 상태이기 때문에 uncoditional 인 상태에선 $\mu$와 $\mathrm{\Sigma }$ 에서 뽑는 거였더라면 condition 에선 $s\mathrm{\Sigma }{\mathrm{\Delta }}_{x_t}\log p_{\varphi }(y|x_t)$ 만큼 mean shift 한 상태로 sampling 한다. $\mathrm{\Sigma }$ 부분은 앞서 말한 g 를 풀어쓴 것과 같다. 그리고 s 는 mean shift 에 대한 scale constant 인데, 이는 hyper parameter로 1로 주었을 땐, 좋은 결과가 나오진 않았지만 10으로 크게 줬을 때 좋은 생성 결과가 나왔음을 ablation study 로 증명하였다.

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