[평범한 학부생이 하는 논문 리뷰] Denoising Diffusion Probabilistic Models (DDPM)

이번에 리뷰할 논문은 그 유명한! DDPM! Diffusion의 기초 논문들은 확실하게 이해하고 넘어가는 것이 좋다는 멘토님의 조언에 따라 이번 ddpm도 시간은 오래 걸리겠지만 최대한 꼼꼼하게 읽어서 자세히 리뷰하는 것을 목표로 포스팅을 할 것이다. 이전 NCSN논문도 꼼꼼하게 했다고 생각하지만 아직 100프로 이해했다고 자부할 수 없었다. 그래서 ddpm도 100프로는 아니지만 70,80프로를 넘어 90프로는 이해하자는 목표로 공부했다.

https://arxiv.org/pdf/2006.11239.pdf

그럼 이제 논문 리뷰를 시작하겠다.

1. Introduction

이 논문에서는 diffusion probablilistic model (diffusion model)의 진전을 소개한다고 한다. Diffusion model은 유한한 시간 후에 데이터와 일치하는 샘플을 생성하기 위해 variational inference을 사용하여 훈련된 매개변수화된 마르코프 체인이다. 이 체인의 transitions은 신호가 파괴될 때까지 샘플링의 반대 방향으로 데이터에 점차적으로 노이즈를 추가하는 마르코프 체인인 diffusion process를 reverse하도록 학습된다.

이전까지 diffusion모델이 고품질의 샘플을 생성할 수 있다고 입증된적 없지만 이 논문에서는 diffusion model이 실제로 고품질 샘플을 생성할 수 있으며, 때로는 다른 유형의 생성 모델보다 더 나은 결과를 생성할 수 있음을 보여준다. 또한 diffusion model의 특정 parameterization가 훈련 중 여러 노이즈 레벨에 걸친 denoising score matching과 샘플링 중 annealed Langevin dynamics와 동등하다는 것을 보여준다. 이 parameterization으로 최상의 샘플 quality를 얻었으므로 이 동등성이 주요 기여 중 하나라고 한다.

논문에서 모델의 lossless codelength 대부분이 감지할 수 없는 이미지 detail을 설명하는 데 사용된다는 것을 발견했다고 한다.

2. Background

2.1 Forward(Diffusion) Process

Diffusion model은 위 이미지에서 볼 수 있듯이 두가지 process가 존재한다. 먼저 두 process중에서 forward(diffusion) process부터 알아보자. 이 process는 말 그대로 이미지에 노이즈를 주입해주는 process이다. 이미지 $x_0$에 가우시안 노이즈를 점점 추가해줌으로써 최종적으로 완전한 가우시안 노이즈 $x_T \sim \mathcal{N}(x_T;0,I)$를 얻게 된다.

$$q(\mathbf{x}_{1:T} |\mathbf{x}_0) := \underset{t=1}{\overset{T}{\Pi}}q(\mathbf{x}_t|\mathbf{x}_{t-1}),\\ q(\mathbf{x}_t|\mathbf{x}_{t-1}) := \mathcal{N} (\mathbf{x}_t; \sqrt{1 - β_t}\mathbf{x}_{t-1}, β_tI), \quad (2)$$

여기서 $t-1$시점에서 노이즈를 추가해서 $t$시점의 분포 $q(\mathbf{x}_t|\mathbf{x}_{t-1})$를 얻기에 Markov Chain이다.

또한 $q(\mathbf{x}_t|\mathbf{x}_{t-1}) := \mathcal{N} (\mathbf{x}_t; \sqrt{1 - β_t}\mathbf{x}_{t-1}, β_tI)$로 두는 이유는 $x_t = \sqrt{1-\beta_t}x_{t-1} + \sqrt{\beta_t}z, \ where \ z\sim\mathcal{N}(0,I)$에서 $x_t$의 분산이 너무 커지는 것을 방지하기 위함이라고 한다.

2.2 Reverse Process

Forward Process를 통해 이미지를 완전한 노이즈로 만들었으면 우리는 그 반대 방향인 노이즈를 지우는 reverse process를 알아내야한다. 우리가 알고싶은 것은 $q(\mathbf{x}_{t-1}|\mathbf{x}_t)$이다. 하지만 이를 알기 힘드므로 approximation인 $p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t)$를 구하는 것이 목적이다. Reverse process는 아래와 같이 정의됐다.

$$p_θ(\mathbf{x}_{0:T} ) := p(\mathbf{x}_T )\underset{t=1}{\overset{T}{\Pi}} p_θ(\mathbf{x}_{t-1}|\mathbf{x}_t),\\ p_θ(\mathbf{x}_{t-1}|\mathbf{x}_t) : = \mathcal{N} (\mathbf{x}_{t-1}; µ_θ(\mathbf{x}_t, t), Σ_θ(\mathbf{x}_t, t)) , (1)$$

Reverse process또한 forward process와 마찬가지의 이유로 Markov Chain이다. 또한 완전한 노이즈 $p(\mathbf{x}_T) = \mathcal{N}(\mathbf{x}_T;0,I)$에서 시작한다. $\beta_t$가 작으면 $q(\mathbf{x}_t|\mathbf{x}_{t-1})$와 $q(\mathbf{x}_{t-1}|\mathbf{x}_t)$가 같은 형태를 띈다고 증명이 됐기 때문에 $p_{\theta}$가 가우시안 분포이라고 할 수 있다.

2.3 Objective Function

Diffusion Model은 $p_θ(x_0) := \int p_θ(x_{0:T}) dx_{1:T}​$로 정의됐고, 학습은 $p_θ(x_0)$의 negative log likelihood(NLL)의 variational bound를 이용한다. 아래와 같이 diffusion model의 objective funcion이 유도되는 것을 볼 수 있다.

앞서 우리는 forward process $q(x_{t}|x_{t-1})$의 반대인 $q(x_{t-1}|x_t)$를 알고 싶다고 했고 그를 approximate하는 $p_{\theta}(x_{t-1}|x_t)$를 찾고자했다. 하지만 $q(x_{t-1}|x_t)$는 intractable하기 때문에 이와 직접 비교해서 KL divergence를 구할 수는 없다. 하지만 $q(x_{t-1}|x_t)$대신 $q(x_{t-1}|x_t,x_0)$는 tractable하기 때문에 이와 KL divergence를 줄이는 방식으로 학습이 가능하다.

Forward process는 $q(\mathbf{x}_t|\mathbf{x}_{t-1}) := \mathcal{N} (\mathbf{x}_t; \sqrt{1 - β_t}\mathbf{x}_{t-1}, β_tI)$로 정의됐다. 야래와 같은 방법으로 $x_0$에서 $x_t$로 바로 건너 뛸수가 있다. 이 성질을 이용해서 $x_0$에 condition된 backward process posterior $q(x_{t-1}|x_t,x_0)$를 아래와 같이 구할 수 있다.

위의 결론으로 아래의 식이 나온다. 아래의 식(6)을 논문에서는 forward process posterior라고 한다.

$q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0) = \mathcal{N} (\mathbf{x}_{t-1}; \tilde{\mathbf{µ}}_t(\mathbf{x}_t, \mathbf{x}_0), \tilde{β}_t\mathbf{I}), \quad(6)$

여기서 $\tilde{µ}_t(\mathbf{x}_t, \mathbf{x}_0) : = \frac{\sqrt{\bar{α}_{t-1}}β_t} {1 - \bar{\alpha}_t} \mathbf{x}_0 + \frac{\sqrt{α_t}(1 - \bar{α}_{t-1})}{ 1 - \bar{α}_t} \mathbf{x}_t$ 및 $\tilde{β}_t :=\frac{ 1 - \bar{α}_{t-1}}{ 1 - \bar{α}_t} β_t,\quad (7)$

3 Diffusion models and denoising autoencoder

이번 섹션에서 objective function에서 각 term에 대해 알아보자.

$\mathbb{E}_q [\underbrace{D_{KL}(q(\mathbf{x}_T |\mathbf{x}_0) \Vert p(\mathbf{x}_T ))}_{L_T} + \underset{t > 1}{\Sigma} \underbrace{D_{KL}(q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0) \Vert p_θ(\mathbf{x}_{t-1}|\mathbf{x}_t))}_{L_{t-1}} \underbrace{ - log p_θ(\mathbf{x}_0|\mathbf{x}_1)}_{L_0}], \quad (5)$

$L_T$는 VAE의 loss에서 regularization term과 상응하고, $L_{t-1}$은 denoising process, $L_0$은 VAE의 loss에서 reconstruction term과 상응한다.

3.1 Forward process and $L_T$

논문에서 $p(\mathbf{x}_T) = \mathcal{N}(x_T;0,I)$으로 고정되어있고, $q(\mathbf{x}_T|\mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t;\sqrt{\bar{\alpha}_t}\mathbf{x}_0,(1-\bar{\alpha}_t)I)$으로 learnable parameter가 없기 때문에 학습중에는 상수로 고정된다.

3.2 Reverse process and $L_{1:T-1}$

$L_{t-1} = D_{KL}(q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0) \Vert p_θ(\mathbf{x}_{t-1}|\mathbf{x}_t))$이고 $q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \tilde{\mathbf{µ}}_t(\mathbf{x}_t, \mathbf{x}_0), \tilde{β}_t\mathbf{I}),\ \ p_θ(\mathbf{x}_{t-1}|\mathbf{x}_t) = \mathcal{N} (\mathbf{x}_{t-1}; \mathbf{µ}_θ(\mathbf{x}_t, t), \Sigma_{\theta}(\mathbf{x}_t,t))$이다. $p_{\theta}$에서 $\mu_{\theta}$와 $\Sigma_{\theta}$둘다 모수 $\theta$를 포함해서 learnable parameter이다.

 

- $\Sigma_{\theta}$

논문에서 $\Sigma_{\theta}$를 임의의 상수인 $\sigma^2_tI$로 두었다. $σ ^2_ t = β_t$와 $σ^ 2_ t = \tilde{β}_ t = \frac{1-\bar{α}_{t-1}}{ 1-\bar{α}_t} β_t$ 로 모두 사용가능하다고 한다.

 

- $\mathbf{µ}_θ(\mathbf{x}_t, t)$

$p_θ(\mathbf{x}_{t-1}|\mathbf{x}_t) = \mathcal{N} (\mathbf{x}_{t-1}; \mathbf{µ}_θ(\mathbf{x}_t, t), σ^2_t \mathbf{I})$를 사용하면 다음과 같이 쓸 수 있다:(여기서 $\sigma^2_t$를 $\tilde{\beta}_t$로 두면 $p_{\theta}$와 $q$의 분산이 같아진다.)

$L_{t-1} = D_{KL}(q(\mathbf{x}_{t-1})|\mathbf{x}_{t},\mathbf{x}_{0})\Vert p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t}))\\ \quad \quad = \mathbb{E}_q[ \frac{1} {2σ^ 2_t} \Vert \tilde{µ}_t (\mathbf{x}_t, \mathbf{x}_0) - \mathbf{µ}_θ (\mathbf{x}_t, t)\Vert^ 2] + C,\quad (8)$

여기서 $C$는 $θ$에 의존하지 않는 상수이고 $\tilde{\mu}_t$는 $q(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{x}_0)$의 평균이다.

따라서 $\tilde{µ}_t$를 예측하는 모델이라는 것을 알 수 있다.

 

- $\epsilon_{\theta}$

$q(\mathbf{x}_t|\mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t;\sqrt{\bar{\alpha}_t}\mathbf{x}_0,(1-\bar{\alpha}_t)\mathbf{I})$를 $\mathbf{x}_t(\mathbf{x}_0,\mathbf{\epsilon} ) = \sqrt{\bar{α}_t}\mathbf{x}_0 + \sqrt{1 - \bar{α}_t}\epsilon$ for $\mathbf{\epsilon} ∼ \mathcal{N} (0,\mathbf{ I})$로 reparameterize하고 forward 프로세스 posteriors 식 (7)을 적용하면 방정식 (8)을 더 확장할 수 있다.

$L_{t-1} -C = \mathbb{E}_q[ \frac{1} {2σ^ 2_t} \Vert \tilde{µ}_t (\mathbf{x}_t, \mathbf{x}_0) - \mathbf{µ}_θ (\mathbf{x}_t, t)\Vert^ 2]$에서 $\mathbf{x}_t \sim q(\mathbf{x}_t|\mathbf{x}_0)$이므로 $\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 +\sqrt{1-\bar{\alpha}_t}\epsilon, \ \rm{where} \ \epsilon \sim \mathcal{N}(0,\mathbf{I})$로 표현할 수 있다. 이를 $\tilde{\mu}$에 $\mathbf{x}_0$에 대입을 하고 정리를 하면 $\tilde{\mu}_t = \frac{1}{\sqrt{\alpha_t}}(\mathbf{x}_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon)$이 된다.

이를 통해 식(8)을 다시 표현하면 아래와 같은 식이 나온다.

$L_{t-1} - C = \mathbb{E}_{\mathbf{x}_0\sim q(\mathbf{x}_0),\epsilon \sim \mathcal{N}(0,\mathbf{I})}[\frac{1}{2\sigma^2_t}\Vert\frac{1}{\sqrt{\alpha_t}}(\mathbf{x}_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon)-\mu_{\theta}(\mathbf{x}_t,t)\Vert^2]\quad (10)$

위 식을 보면 $\mu_{\theta}$가 $\frac{1}{\sqrt{\alpha_t}}(\mathbf{x}_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon)$를 예측하게 되는 것이고,

$\mu_{\theta} = \frac{1}{\sqrt{\alpha_t}}(\mathbf{x}_t - \frac{\beta_t}{1-\bar{\alpha}_t}\epsilon_{\theta}(\mathbf{x}_t,t))\quad (11)$

로 parameterize할 수 있다. 결국 우리의 backward model은 $\mathbf{x}_{t-1} \sim p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1};\mu_{\theta}(\mathbf{x},t),\sigma^2_t\mathbf{I})$이 된다. 그래서 우리는 algorithm 2에서 나온 것 처럼 샘플링을 할 수 있다.

식 (11)을 (10)에 넣고 $q(\mathbf{x}_t|\mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t;\sqrt{\bar{\alpha}_t}\mathbf{x}_0,(1-\bar{\alpha}_t)\mathbf{I}$)를 이용하면 아래와 같이 식을 정리할 수 있다.

$\mathbb{E}_{\mathbf{x}_0 \sim q(\mathbf{x}_0),\epsilon \sim \mathcal{N}(0,\mathbf{I})}[\frac{ β^ 2_t}{ 2σ ^2_t α_t(1 - \bar{α}_t)}\Vert \epsilon - \epsilon_θ( \sqrt{\bar{α}_t}{\mathbf{x}_0} + \sqrt{1 - \bar{α}_t}\epsilon, t)\Vert^ 2],\quad (12)$

Reverse process 평균 함수 근사치 $µ_θ$를 훈련하여 $\tilde{µ}_t$ 를 예측하거나 parameterize를 수정하여 $\epsilon$을 예측하도록 훈련할 수 있다. $\epsilon$ - predict parameterization은 Langevin dynamics와 유사하며, denoising score matching과 유사한 목표에 대한 diffusion 모델의 variational bound를 단순화한다는 것을 보여준다.

실험을 통해, 아래와 같이 계수를 없앤 단순화한 loss가 더 품질이 우수한 이미지를 생성하여 최종적으로는 아래의 loss 함수를 쓴다고 한다.

$L_{simple}(θ) := \mathbb{E}_{t,\mathbf{x}_0,\epsilon}[\Vert \epsilon - \epsilon_θ( \sqrt{\bar{α}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon, t) \Vert^2\quad(14)$

3.3 Data scaling, reverse process decoder, and $L_0$

Reverse process를 통해 나온 $\mathbf{x}_1$을 실제 데이터 분포인 $q(\mathbf{x}_0)$의 domain으로 보내줘야한다. 이 분포를 independent discrete decoder로 한다.

$$p_θ(\mathbf{x}_0|\mathbf{x}_1) = \underset{i=1}{\overset{D}{\Pi}} \int ^{δ_+(x^ i_0 )}_{ δ_-(x^i_ 0 )}\mathcal{ N} (\mathbf{x}; µ^ i_ θ (\mathbf{x}_1, 1), σ^2_1 ) dx$$

$$δ_+(x) = \begin{cases}∞&{\mathbf{if} \ x = 1} \\ {x + \frac{1} {255}}& {\mathbf{if} \ x < 1}\end{cases} \quad δ_-(x) =\begin{cases} -∞ &{\mathbf{if} \ x = -1} \\ {x - \frac{1}{ 255}}&{\mathbf{if} \ x > -1}\end{cases}\quad (13)$$

여기서 $D$는 데이터 차원이고 $i$ 윗첨자는 한 좌표의 추출을 나타낸다.

VAE loss에서의 reconstruction loss와 같은 부분인데 이미지 $\mathbf{x}_0$에서 노이즈가 조금 들어간 $\mathbf{x}_1$과 $\mathbf{x}_0$의 MSE loss를 최소화 하는 것과 같다고 한다.

위 과정을 통해,

$Loss_{Diffusion} = \mathbb{E}_q [\underbrace{D_{KL}(q(\mathbf{x}_T |\mathbf{x}_0) \Vert p(\mathbf{x}_T ))}_{L_T} + \underset{t > 1}{\Sigma} \underbrace{D_{KL}(q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0) \Vert p_θ(\mathbf{x}_{t-1}|\mathbf{x}_t))}_{L_{t-1}} \underbrace{ - log p_θ(\mathbf{x}_0|\mathbf{x}_1)}_{L_0}]$

이 논문에서는 위 diffusion의 loss함수를 아래와 같이 단순화 했다.

$Loss_{DDPM} = \mathbb{E}_{t,\mathbf{x}_0,\epsilon}[\Vert \epsilon - \epsilon_θ( \sqrt{\bar{α}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon, t) \Vert^2$

이러한 점이 이 논문의 가장 큰 contribution이라고 한다.

4. Experiments

실험 설정은 위와 같고 더 자세한 설명은 부록을 참고하면 된다.

4.1 Sample quality & Reverse process parameterization and training objective ablation

표 2의 결과를 보면 논문에서 제안한 $\epsilon$ 예측은 고정 분산으로 variational bound에 대해 훈련할 때 $\tilde{µ}$를 예측하는 것과 거의 비슷한 성능을 보이지만, 단순화된 목표로 훈련할 때 훨씬 더 나은 성능을 보인다.

 

 

Diffusion 두번 째 논문인 DDPM논문 리뷰를 이렇게 마쳤는데 이번에도 최대한 자세하게 하려고 노력했다. 하지만 역시 모르는게 많았고 스터디를 진행하면서 모르는 것들을 채워야겠다.... 이번 논문도 자세하게 하려다 보니 길이가 길어진 것 같은데 긴 포스팅을 읽어주셔서 너무 감사합니다!

QnA

-> 스터디 진행 후에 추가할 예정!