DDPM(Denoising Diffusion Probabilistic Models) 요약 정리 🧐

들어가기에 앞서...

요즘 Diffusion을 공부하고 있다. 그 시작이 되는 DDPM을 I-DDPM 논문에서 잘 요약, 정리해 놨길래 그에 기반하되, 내가 공부했던 내용을 더해 DDPM을 요약 정리해 보았다. 

+ 사실 이 논문의 원조는 <Deep unsupervised learning using nonequilibrium thermodynamics> 라는 2015년에 출간된 논문이다. 

---

Definitions

Forward Process

Forward Process는 $x_0$이 주어졌을 때, latent vector $x_1$부터 $x_T$를 차례로 가우시안 노이즈를 더하며 생성하는 과정이다. 노이즈를 더해가는 과정을 수식으로 표현하면 다음과 같다.

이 때 $\beta_t$는 0부터 1까지 점점 커져가도록 스케줄링되어있다. DDPM에서는 $10^{-4}$부터 0.02까지 Linear 하게 커지도록 스케줄링하였다. Learnable 한 파라미터로 두는 실험 또한 하였으나, 상수로 두어도 큰 차이가 없는 것을 확인했다고 한다. (I-DDPM 논문에서는 코사인 스케줄링을 사용하였고, 성능 향상에 도움이 되었다고 리포트하였다.) 

우측 식에서 좌측 식을 유도하는 방법은 마르코프 체인에 기반한 것으로 다음과 같이 증명할 수 있다.

위의 수식만 보면 $x_t$를 계산하기 위해서 $x_0$에서 $t$번의 연산을 거쳐야할 것 같지만, 연산을 잘 정리하면 $x_0$에서 $x_t$를 곧바로 계산할 수 있음을 알 수 있다. 

충분히 큰 $T$와 적절한 $\beta_t$ 스케줄링 하에 이 과정을 반복하면, 마지막 latent vector인 $x_T$는  거의 완벽하게 가우시안 분포를 따르게 된다. 

 

Reverse Process

Reverse Process는 $x_T$를 조금씩 denoising해 $x_0$, 즉 이미지를 만드는 과정이다. 우리가 알고 싶은 것은 정확한 reverse 분포인 $q(x_{t-1}|x_t)$이지만, 이 분포를 실제로 계산할 수는 없다. (사실 알 수 있다면, 이런 짓도 안 해도 됨! 🥵) 따라서 이를 근사하는 분포 $p_\theta(x_{t-1}|x_t)$를 계산하는 것을 목표로 한다.

그리고 이 분포의 평균 $\mu_{\theta}(x_t, t)$과 분산 $\sum_{\theta}(x_t, t)$을 예측하는 것이 뉴럴 네트워크의 목표가 된다.

뉴럴 네트워크는 negative log likelihood를 최소화하는 방향으로 학습된다. Variational Inference 방식을 통해 식 (3)에서 식 (5)가 유도된다. (DDPM의 Appendix A 참고)

우선 $L_0$ 항을 제외한 $L_T$, $L_{t-1}$항은 KL Divergence로 해석할 수 있다.  VAE에서 $L_T$항은 Posterior($q(z|x)$)가 Prior (일반적으로 가우시안)을 따르도록 하는 loss에 해당한다. 하지만 diffusion에서는 forward process가 input($x_0$) 을 가우시안 노이즈($x_t$)로 만드는 '연산'이기 때문에 이 값은 항상 0에 가까운, 무시할 수 있는 값이 된다. $L_{t-1}$항은 우리가 알고자 하는 $p_\theta(x_{t-1}|x_t)$와 forward process posterior $q(x_{t-1}|x_t, x_0)$를 비교하는 항이다. 본래는 $q(x_{t-1}|x_t)$에 대해 계산하는 것이 정확하나, 애초에 $q(x_{t-1}|x_t, x_0)$은 우리가 알고자 하는 분포로, 그 값을 알 수 없다. 따라서 $x_0$에 대해 condition을 주어 그 값을 근사했다. 마지막으로 $L_0$항은 latent $x_1$로부터 데이터 $x_0$을 추정하는 확률 모델의 파라미터를 최적화하는 항이다.

 

Training in Practice

$L_{t-1}$항의 $q(x_{t-1}|x_t, x_0)$ 분포는 Bayes rule 등을 이용해 다음과 같이 평균과 분산을 계산할 수 있다.

계산을 해보니 $\tilde{\beta_t}$는 상수이다. 따라서 $p_\theta$ 분포의 분산은 $\tilde {\beta_t}$로 고정하고, $\mu(x_t, t)$만 학습하면 될 듯하다. 그런데 논문에서는 $\beta_t$로 분산을 두었을 때 학습이 더 잘되는 것을 발견, $\tilde {\beta_t}$가 아닌 $\beta_t$를 분산으로 사용하였다고 한다. (추측하건데, 사실 우리가 정답이라고 가정하고 있는 $q(x_{t-1}|x_t, x_0)$이 정확한 정답이 아니기 때문에 발생한 현상이 아닐까 싶다.)

Equation (7)을 바탕으로 $L_{t-1}$항을 정리하면 다음과 같다. Gaussian끼리의 KLD는 이미 Closed form이 알려져 있다.



이제 위 식에 Equation (7)에서 구한 값들을 대입해서 식을 정리하면 된다. 이때, 저자는  Equation (4)를 이용해 $x_0$을 $x_t$로 reparameterizing 하였다. ($\alpha_t\bar {\alpha_{t-1}} = \bar {\alpha_t}$, $\beta_t= 1 - \alpha_t$ 임을 이용하기.)

새롭게 정리한 $\tilde {\mu_t}$와 $\mu_\theta$를 이용해 equation 10을 다시 표현하면 다음과 같다.

이제 문제는 noise $\epsilon$을 예측하는 문제로 바뀌었다. 

최종적으로 사용한 Loss는 Equation (12)에서 weight를 제거한 Equation (14)이다. $\tilde {\mu}$를 예측하던 방식과 달리, $\epsilon$을 예측하는 방식에서는 $L_{simple}(\theta)$이 더 효과적이었다고 한다. 저자들이 예측한 이유는 $t$가 증가할 때 이 weight가 감소하는 경향을 보이는데, Reverse process에서는 $t$가 클 때 노이즈를 더 많이 제거해야 하는 상황에 처하기 때문이다. 

---

Reference

https://sang-yun-lee.notion.site/Denoising-Diffusion-Probabilistic-Models-ade353a9abaa41d98c3ae2c18b52c129

Denoising Diffusion Probabilistic Models

https://developers-shack.tistory.com/8

[논문공부] Denoising Diffusion Probabilistic Models (DDPM) 설명

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

[Paper Review] DDPM: Denoising Diffusion Probabilistic Models 논문 리뷰