Generative Model과 Diffusion Model, 그리고 Denoising Diffusion Probabilistic Model

 

Generative Model

 

Generative Model이란?

 

이에 관한 자세한 내용은 아래 글의 'Generative Model' section을 참고하면 된다.

생성 모델(Generative Model)과 VAE, 그리고 GAN

 

글 작성 시점 기준으로는 diffusion model이 큰 각광을 받고 있다. 이번 글에서는 diffusion model이 무엇이고, 어떠한 이유에서 기존의 생성 모델보다 더 주목을 받고 있는지를 확인해 보고자 한다. 이 모델의 배경이 되는 논문이 바로 DDPM(Denoising Diffusion PRobabilistic Models)인데, 이 글에서는 논문의 내용을 상세히 살펴보기 보다는 핵심적으로 알아두어야 할 내용 위주로 노트 필기와 함께 정리해보고자 한다.

 

 

 

DDPM(Denoising Diffusion Probabilistic Models)

 

DDPM 모델의 개요

 

Diffusion을 이해하려면 그전에 먼저 VAE에 관한 선수 지식이 유용하다. 앞서 설명한 VAE에서는 어떠한 데이터가 들어왔을 때 그 데이터가 어떠한 확률 분포에서 샘플링되었을 확률이 높은지를 encoder를 통해 latent space로 보내고, 그 latent space에서 샘플링한 샘플을 decoder에 통과한다고 했다. 그 encoder의 역할을 입력 데이터에 noise를 추가하는 과정으로 생각하고, decoder를 noise가 추가된 데이터에서 원래 데이터로 복원하는 과정이라고 생각해 보자. 그러면 VAE에서는 극단적으로 큰 noise를 주고 이를 다시 복원하는 과정으로 이해할 수 있다. Diffusion model은 이러한 noise를 추가하고 denoising 하는 과정을 넓은 시간대인 1부터 $T$까지의 과정으로 수행하는 모델이다. 즉, 작은 noise를 매 time step별로 추가해 가고, 이를 역과정으로 매 time step을 거치면서 noise를 제거하는 것이다. 이는 각 time step별로 VAE에서의 encoder와 decoder로 구성되어 있다고 이해해도 된다.

 

[출처] Figure 2, Jonathan Ho, Denoising Diffusion Probabilistic Models

 

즉, 1부터 $T$까지의 시간동안 각 time step 마다 Gaussian noise를 추가하는데, noise를 추가하는 과정에서는 직전 time step의 데이터에만 직접적인 영향을 받는 first order Markov model로 가정한다. 그 noise를 적용하는 과정을 forward process $q$라고 정의한다. 이를 시간 $T$까지 수행했을 때의 결과는 완전한 표준정규분포에 가까운 Gaussian noise가 되도록 한다. 시간 $T$에서 1까지 다시 noise를 단계별로 제거하는 과정은 reverse process $p$라고 정의한다. 이때는 forward process의 역순이므로 first order Markov process를 적용할 때 시간상 이후 time step의 데이터에 영향을 받는다. 이러한 denoising process를 매 time step 마다 거침으로써 원래 데이터가 나오도록 한다. 참고로 forward process와 reverse process는 모두 Gaussian distribution임을 전제로 한다. 그러면 이 모델이 궁극적으로 어떤 걸 학습하는지 묻는다면, 어떤 time step에서의 reverse process가 그 time step의 forward process의 과정을 참고한 posterior를 예측하도록 하여 noise를 제거할 수 있도록 그 forward process의 역과정에 근사하도록 하는 것이다. 이렇게만 보면 이해가 선뜻 되지 않을 수 있다. 그전에 forward process와 reverse process가 어떻게 모델링되는지 살펴보자.

 

 

Forward Process

 

 

논문에서는 $q$를 위와 같이 모델링하고 있다. Noise를 적용하는 임의의 시간 $t$ step에서 샘플링할 때 Gaussian distribution을 따르는데, 그 Gaussian distribution의 평균이 이전 time step의 데이터 $x_{t-1}$에 $\sqrt{1 - \beta_t}$를 곱한 분포를 따르는 것이다. 그리고 분산에는 $\beta_t$를 곱한 값을 적용한다. 왜 굳이 $\sqrt{1 - \beta_t}$와 $\beta_t$를 평균과 분산에 각각 곱하는지 궁금할 수 있는데, 이는 noise를 적용하면서 데이터의 분산이 매우 커지지 않고 1이 되는 걸 맞춰주기 위함이다. 자세한 식은 위의 그림의 오른쪽에 설명해 놓았다.

이를 처음 input 데이터인 $x_0$부터 시작해서 각각의 모든 step의 noise 적용 결과인 $x_{1:T}$를 구하는 과정은 위의 그림에서 production으로 정리할 수 있다. 왜 저렇게 정리되는지는 chain rule을 적용하면 알 수 있는데, first order Markov chain이 전제이므로 이전 time step에만 영향을 받는다는 점을 고려하면 $\prod_{t=1}^{T} q(x_t | x_{t-1})$로 정리된다는 걸 알 수 있다.

여기서 $\beta_t$는 매 time step 마다 다른 값을 지니도록 사전에 hyperparameter로 설정했다고 논문에서는 언급하고 있다. 즉, forward process $q$는 데이터만 정해지면 fixed 상태임을 유추할 수 있다.

 

그런데 매번 처음부터 임의의 시간 $t$까지 step을 거쳐서 확률 분포를 모델링한다는 건 번거로운 일이다. 그래서 논문에서는 처음 시간부터 임의의 time $t$까지 바로 확률 분포를 모델링할 수 있도록 했는데, 그 결과는 위의 그림과 같다. 저 $\alpha$가 앞으로 자주 등장하므로 이후 내용을 따라가려면 $\alpha$가 무엇을 의미하는지를 기억해 놓는 것이 좋다. 어떻게 위의 식처럼 도출되었는지는 아래 증명을 참고하면 된다.

 

 

$\beta_t$는 사용자가 직접 정하는 매우 작은 값의 hyperparameter라고 했고, Gaussian 분포로 점차 샘플링을 할수록 $1- \beta_t$가 곱해지는 $\bar{\alpha}_t$로 정의되므로 $t$가 무한히 커지면 $\sqrt{\bar{\alpha}_t}$는 0에 수렴하고, $1 - \bar{\alpha}_t$는 1에 수렴해서 standard Gaussian distribution(표준정규분포)에 가깝게 된다.

 

 

Reverse Process(Denosing Process)

 

 

Reverse process는 위와 같이 정리할 수 있는데, 왼쪽 식을 해석하면 denoising 과정을 통해 $x_T$부터 $x_0$까지의 매 time step의 output을 내는 likelihood를 의미하는 것이다. 즉, likelihood를 구하는 것이므로 $\theta$가 붙은 것을 유의한다. 오른쪽 식은 $x_t$가 주어졌을 때 denoising을 적용하여 $x_{t-1}$를 output으로 내게 하는 조건부 likelihood를 의미한다. 여기서 주목할 점은 time $t-1$에서의 Gaussian distribution의 평균과 분산이 $x_t$와 time $t$에 의존한다는 것이다. 그래서 정리를 하면 임의의 time step $t-1$에서의 reverse process는 Gaussian distribution을 따르는데, 그 분포의 평균과 분산이 $x_t$와 $t$에 의존한다. 그러면 우리가 원하고자 하는 바는 이 평균과 분산이 time step $t-1$에서의 forward process의 과정을 참고하여 잘 모델링하는 것이고, 그것이 denoising diffusion probabilistic model의 loss function을 정의하는 데 필요한 기본적인 아이디어이다. 

 

 

 

Loss Function

 

 

그런데 문제는 reverse process의 평균과 분산이 $x_t$와 $t$에 의존적이므로 intractable 하다는 것이다. 다행히 forward process의 $q$는 앞서 살펴 본 바처럼 hyperparameter $\beta_t$에 의해 초기 데이터만 정해지면 fixed 되어서 tractable 하다. 그래서 사실상 forward process에서는 임의의 time step $t$까지 확률분포를 한번에 모델링했지만, reverse process에서는 매 time step 마다 평균과 분산이 다르게 모델링되므로 임의의 시점까지 한번에 모델링하는 과정이 어렵다.

 

 

그러면 $q$를 잘 이용해서 $p$를 최적화하는 것이 필요한데, 문제는 $q$의 방향이 $p$의 방향과 반대라는 것이다. 즉, $p(x_t | x_{t-1})$를 예측하기 위해 $q(x_t | x_{t-1})$을 사용해야 하는데, 둘의 진행 방향이 반대여서 $p$를 $q$에 근사시킬 수 없다. 그러나 다행히 $q(x_{t-1} | x_t, x_0)$를 알 수 있다. 즉, Bayes' rule를 사용하여 $x_0$가 given인 posterior 분포를 유도할 수 있다.

 

 

결론은 negative log likelihood인 $-\log p_{\theta} (x_0)$를 최소화하는 것이다. 이 의미는 $x_0$이 주어지고 forward process를 끝까지 해서 나오는 $X_T$의 분포에 관하여 다시 reverse process를 끝까지 해서 $x_0$가 복원되어 나오게 하는 likelihood를 최대화 하기 위함이고, $\log$와 음수를 붙여서 negative log likelihood를 최소화하는 것으로 바꾼 것이다.

여기서부터는 수식의 향연인데, 노트 필기로 정리한 내용을 차근차근 따라가면 된다.

 

 

위의 식을 잘 정리하면 결국 맨 마지막 식처럼 세 가지 term으로 도출할 수 있다. 맨 앞의 term은 forward process에서 마지막 time step $T$에 관한 loss인 $L_T$인데, 어차피 time $T$에서는 $N(0, I)$의 표쥰정규분포이므로 $L_T$는 0과 다름이 없어서 무시하면 된다. $L_{t-1}$이 중요한데, time step $t$에서의 forward process의 역과정에 해당되는 분포 $q$dhk 그 time step에서의 reverse process의 분포 $p$가 유사해지도록 하는 것이다. 이게 앞서 우리가 잠깐 언급하고 갔던 denoising diffusion probabilistic model의 학습 목표를 내재하고 있는 것이다. 그러면 여기서 $q(x_{t-1} | x_t, x_0)$를 어떻게 구하는지가 관건인데, 이는 bayes rule로 간단히 도출할 수 있다. 이 posterior의 의미를 잘 해석해보면 $x_t$에서 $x_{t-1}$로 denoising을 수행하는 과정에서 $x_0$의 내용을 참고하여 noise를 제거하겠다는 것이며, 결론적으로는 $x_t$에서 $x_0$로 한번에 가기 위해 제거하는 noise를 구해서 이를 $x_{t-1}$로 denoising 하는 과정에 반영하는 것으로 볼 수 있다.

 

 

 

여기도 수식의 향연이라 머리가 좀 아찔해지기는 하지만, 지금 하고자 하는 목적은 $q(x_{t-1} | x_t, x_0)$를 구하는 걸 기억해야 한다. 그런데 $q$와 $p$는 Gaussian distribtuion이라고 앞에서 전제했으므로 $q(x_{t-1} | x_t, x_0)$를 모델링하는 정규분포의 평균과 분산을 구하는 것이다. (실제로 분산은 구하지 않고, 평균만을 예측하는 것으로 한다. 그 이유는 후술할 예정이다.) 정규분포의 확률밀도함수를 사용하여 $q(x_{t-1} | x_t, x_0)$의 평균인 $\tilde{\mu}(x_t, x_0)$을 구하는 식을 정리하는 과정이 위의 노트 필기이다.

 

 

그런데 앞에서까지 구한 $\tilde{\mu}(x_t, x_0)$는 $x_0$ term을 포함하고 있으므로 이를 제거하고 오로지 $x_t$에 관한 항으로 바꾸고 싶다. 그래서 reparameterization trick을 사용하는데, 여기서 reparameterization은 복잡한 분포를 미분 가능한 함수와 간단한 분포로 noise $\epsilon$를 샘플링하는 식으로 나타내는 기법이다. 그래서 결국 우리는 $\tilde{\mu}$를 $x_t$와 $\epsilon$로 정리했다.

 

 

즉, $\mu_{\theta} (x_t, t)$가 $\tilde{\mu}_t$와 가까워지도록 하는 것이 목표이며, $\mu_{\theta} (x_t, t)$도 마찬가지로 reparameterization trick에 의해 $x_t$와 $\epsilon_{\theta}$에 관한 식으로 정리했을 때의 그 $\epsilon_{\theta}(x_t, t)$가 $\epsilon$에 근사하도록 만드는 것이다. 즉, $x_t$에서 $x_0$로 한번에 가기 위해 제거해야 하는 noise가 실제 $x_0$에서 $x_t$로 한번에 가기 위해 추가되어야 하는 noise와 모델이 유사하게 예측하도록 해야 한다.

 

 

앞서 나온 파라미터들을 바탕으로 $q(x_{t-1} | x_t, x_0)$와 $p_{\theta}(x_{t-1} | x_t)$의 Gaussian distribution 간의 KL divergence를 구하는 과정을 정리한 그림이다. 최종적으로는 뒤의 제곱하는 term만 남겨서 loss function으로 정리하는 게 더 낫다고 논문에서 설명하고 있다.

 

 

학습 방법

 

[출처] Jonathan Ho, Denoising Diffusion Probabilistic Models

 

그래서 궁극적으로 어떻게 모델을 학습하는지가 궁금할 수 있다. 이는 논문에서 pseudo code로 잘 설명해주고 있다. 먼저 최적화 시에는 임의의 시점 $t$를 uniform distribution에서 샘플링하고 $\epsilon$을 표준정규분포에서 샘플링한다. 여기서 $\epsilon$은 앞서 설명한 $\epsilon_t$와 같다. 이를 앞서 정의한 loss 함수 $L_t$에 관해 gradient descent를 수행한다. 여기서 $\epsilon_{\theta}$는 denosing autoencoder이며, $x_t$와 $t$에 관한 정보가 주어졌을 때 $x_t$에서 $x_0$로 한번에 가는 noise인 $\epsilon_{\theta}$을 output으로 예측하는 모델이다. 그 출력값이 $x_0$에서 $x_t$로 한번에 가는 noise인 $\epsilon$과 유사해지도록 학습한다. 임의의 데이터를 샘플링하는 과정은 시간 $T$부터 1까지 reverse process를 통해 $x_0$를 출력하는 것이다. 여기서 유추할 수 있는 사실은 inference 시간이 상당히 걸릴 수 있다는 점인데, 이는 reverse process에서는 $T$에서 임의의 시점 $t$까지 한번에 가는 확률 분포를 모델링하기 어려우므로 하나하나씩 step을 거쳐서 denoising을 수행한다. 그런데 최근에는 fast sampling을 통해 빠르게 샘플링하는 방법이 나왔다고 하는데, 관련 논문도 한 번 찾아서 읽어보는 게 유용할 것 같다는 생각이 들었다.

 

실제로 간단히 구현한 DDPM 코드는 아래를 참고하면 된다.

https://github.com/Glanceyes/ML-Paper-Review/blob/main/ComputerVision/Diffusion/DDPM/DDPM.ipynb

GitHub - Glanceyes/ML-Paper-Review: Implementation of ML&DL models in machine learning that I have studied and written source co

 

 

출처
1. https://arxiv.org/abs/2006.11239, Jonathan Ho, Denoising Diffusion Probabilistic Models