Score-Based / Diffusion Model[3] - DDPM

이번 논문의 주인공은 DDPM입니다. Denoising Diffusion Probabilistic Model입니다. Score-based generative model이랑 거의 흡사하지만, 기본 개념이 조금 다릅니다. 따라서 이에 대해서도 한번 리뷰해보고자 합니다.

 

 

 

• 2023.12.11 Experiment 부분 추가

Diffusion model

Diffusion model의 가장 기본적인 아이디어는 stochastic process입니다. 여기서 stochastic process란, time-dependent한 변수 $X_t$ 즉, 시간 t에 따라 변하는 변수 $X_t$를 stochastic process라고 합니다. 그 중에서 diffusion process는 확률 미분방정식의 한 솔루션이며, continuous-time Markov process이고 대표적인 예로 Brownian motion (a.k.a Wiener process) 가 있습니다.

 

생성모델에서의 diffusion model은 위의 figure에서 보이는 것 처럼 Data $ x_0 $에서 Noise $x_T$로 가는 점선 화살표 방향을 forward process, Noise $x_T$에서 Data $x_0$로 가는 실선 화살표 방향을 backward process개의 process가 존재합니다.

 

생성 모델 관점에서 생각해보면 diffusion model은 backward process를 통해서 explict한 distribution에서 sampling한 random noise로부터 denoising을 통하여 데이터를 만들어내는 생성 방식입니다. 그런데, 우리가 흔히 data -> noise로 가는 forward process는 쉽게 정의가 가능하지만, backward process를 정의하기는 쉽지 않습니다. 그래서 그 backward process를 학습하는 모델이 diffusion probabilistic model이라고 합니다. 좀만 더 정확하게 말하자면, backward process에서 forward process의 intermediate distribution을 학습하는 모델입니다.

 

Forward Process

이전 state인 $x_{t-1}$ 을 condition으로 갖는 conditional distribution $ q(x_t\vert x_{t-1})=\mathcal{N}(x_t;\sqrt{1-\beta_t}x_{t-1}, \beta_t\mathbf{I}) $ 을 정의할 수 있습니다. $\beta_1$은 데이터에 가까워지고, $\beta_T$는 노이즈에 가까워 집니다. 따라서, $\beta_1$은 작아야 하고 ($\sqrt{1-\beta_1}x_0 \approx x_0$) $\beta_T$는 커야합니다 ($\sqrt{1-\beta_T}x_{T-1} \approx 0$). 즉, forward process의 variance $\beta_t$에 대한 schedule을 $\beta_1 < \cdot\cdot\cdot < \beta_T$로 설정하고 conditional distribution가 정의된다면, foward process 를 

$$ q(x_{1:T} \vert x_0)=\prod_{t=1}^Tq(x_t \vert x_{t-1}),\  q(x_t \vert x_0) = \mathcal{N}(x_t;\sqrt{\bar{\alpha}_t}x_0, (1-\bar{\alpha}_t)I), \ \text{where} \ \alpha_t=1-\beta_t \ \text{and} \ \bar{\alpha}_t=\prod_{s=1}^t\alpha_s  $$ 로 어떤 데이터 $x_0$에 대한 $x_t$, 즉 perturbated data의 distribution을 정의할 수 있습니다.

 

$$\begin{align} x_t &= \sqrt{\alpha_t}x_{t-1}+\sqrt{1-\alpha_t}z_{t-1} \\ &=\sqrt{\alpha_t\alpha_{t-1}}x_{t-2} + \sqrt{\alpha_t}\sqrt{1-\alpha_{t-1}}z_{t-2}+\sqrt{1-\alpha_t}z_{t-1} \\ &= \sqrt{\alpha_t \alpha_{t-1}}x_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}} \bar{z}_{t-2}....(*) \\ &= \cdot\cdot\cdot \\&= \sqrt{\bar{\alpha_t}}x_0 + \sqrt{1 - \bar{\alpha_t}}z \end{align}$$

 

(*) $X \sim \mathcal{N}(\mu_X, \sigma_X^2)$와 $Y \sim \mathcal{N}(\mu_Y, \sigma_Y^2)$에서 $Z=X+Y$는 $Z \sim \mathcal{N}(\mu_X+\mu_Y, \sigma_X^2+\sigma_Y^2)$입니다.

 

Reverse Process

이제 반대로, Reverse process를 정의하는데, trainable model $\theta$에 대하여,

 

$$ p_\theta (x_{0:T}) = p(x_T) \prod_{t=1}^T p_\theta(x_{t-1} \vert x_t), \ \ p_\theta(x_{t-1}\vert x_t)=\mathcal{N}(x_{t-1};\mu_\theta(x_t, t), \Sigma_\theta(x_t,t)) $$

 

로 정의될 수 있습니다. 그리고 나서 $\theta$는 negative log likelihood의 variational lower bound를 optimized하기 위해서 학습됩니다. 

$$ \mathbb{E}[-\log p_\theta(x_0)] \leq \mathbb{E}_q[-\log {{p_\theta(x_{0:T})}\over{q(x_{1:T} \vert x_0)}}]=\mathbb{E}[-\log p(x_T)-\sum_{t\geq1} \log {{ {p_\theta(x_{t-1} \vert x_t)} }\over{q(x_t \vert x_{t-1})}} ]=L $$

 

지금의 식은 $q(x_t \vert x_{t-1})$에 대해 표현되어 있기 때문에 이 식을 다시 정리하면,

 

$$\begin{align}L &= \mathbb{E_q}[-\log p(x_T)-\sum_{t\geq1} \log {{ {p_\theta(x_{t-1} \vert x_t)} }\over{q(x_t \vert x_{t-1})}} ] \\&= \mathbb{E_q}[-\log p(x_T)-\sum_{t>1} \log {{ {p_\theta(x_{t-1} \vert x_t)} }\over{q(x_t \vert x_{t-1})}} -          \log{ \frac{p_{\theta}(x_0 \vert x_1)}{q({x_1 \vert x_0}) } }] \\&= \mathbb{E_q}[-\log p(x_T)-\sum_{t>1} \log {{ {p_\theta(x_{t-1} \vert x_t)} }\over{q(x_{t-1} \vert x_t, x_0)}}\cdot\frac{q(x_{t-1}\vert x_0)}{q(x_{t}\vert x_0)} - \log{ \frac{p_{\theta}(x_0 \vert x_1)}{q({x_1 \vert x_0}) } }] \\&= \mathbb{E_q}[-\log \frac{p(x_T)}{q(x_T\vert x_0)}-\sum_{t>1}\log \frac{p_\theta(x_{t-1}\vert x_t)}{q(x_{t-1}\vert x_t,x_0)}-\log p_\theta(x_0 \vert x_1)] \\&=  \mathbb{E}_q [ D_{KL}(q(x_T|x_0) || p(x_T)) + \sum_{t>1}D_{KL}(q(x_{t-1}|x_t, x_0) || p_\theta(x_{t-1}|x_t))-\log p_\theta (x_0|x_1) ] \end{align}$$

 

으로 $q(x_{t-1}|x_t, x_0)$에 관한 식으로 rewriting이 가능합니다. 여기서 $q(x_{t-1} \vert x_t, x_0)=\mathcal{N}(x_{t-1};\tilde{\mu}_t(x_t, x_0), \tilde{\beta}_t)$ 으로 정의합니다.

 

$q(x_{t-1} \vert x_t, x_0) = q(x_t \vert x_{t-1}) \frac{q(x_{t-1} \vert x_0)}{q(x_t \vert x_0)}$에서 $q(x_t \vert x_{t-1})$, $q(x_{t-1} \vert x_0)$, $q(x_t \vert x_0)$ 모두 정규분포이기 때문에, $q(x_{t-1} \vert x_t, x_0)$ 정규분포임을 알 수 있습니다. 따라서, 세 분포의 확률밀도함수로써 표현하자면,

$$\begin{align} q(x_t \vert x_{t-1})&=\frac{1}{\sqrt{2\pi\beta_t}}\exp{(-\frac{(x_t - \sqrt{1-\beta_t}x_{t-1})^2}{2\beta_t})} \\ q(x_t \vert x_0) &=\frac{1}{\sqrt{2\pi(1-\bar{\alpha}_t)}}\exp{(-\frac{(x_t - \sqrt{\bar{\alpha}_t}x_0)^2}{2(1-\bar{\alpha}_t)})} \\ q(x_{t-1} \vert x_0) &=\frac{1}{\sqrt{2\pi(1-\bar{\alpha}_{t-1})}}\exp{(-\frac{(x_{t-1} - \sqrt{\bar{\alpha}_{t-1}}x_0)^2}{2(1-\bar{\alpha}_{t-1})})} \end{align}$$

으로 표현이 가능합니다. 따라서,

 

$$\begin{align} \therefore q(x_{t-1} \vert x_t, x_0) &= \frac{1}{\sqrt{2\pi\beta_t (\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t})} } \exp{(-\frac{(x_t - \sqrt{1-\beta_t}x_{t-1})^2}{2\beta_t} - \frac{(x_{t-1} - \sqrt{\bar{\alpha}_{t-1}}x_0)^2}{2(1-\bar{\alpha}_{t-1})} + \frac{(x_t - \sqrt{\bar{\alpha}_t}x_0)^2}{2(1-\bar{\alpha}_t)} )} \\ &= \frac{1}{\sqrt{2\pi\beta_t(\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t})}}\exp{(-([\frac{1}{2(1-\bar{\alpha}_{t-1})}+\frac{1-\beta_t}{2\beta_t}]x_{t-1}^2 - [\frac{2\sqrt{1-\beta_t}}{2\beta_t}x_t+\frac{2\sqrt{\bar{\alpha}_{t-1}}}{2(1-\bar{\alpha}_{t-1})}x_0]x_{t-1} + C))} \\ &= \frac{1}{\sqrt{2\pi\beta_t (\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t})}} \exp{(- \frac{1}{2\beta_t (\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t})} [x_{t-1}^2 - (\frac{2\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\alpha_t}x_t + \frac{2\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}x_0)x_{t-1} + C ] )} \\ &\approx \frac{1}{\sqrt{2\pi\beta_t (\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t})}} \exp{(- \frac{1}{2\beta_t (\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t})}[x_{t-1} - (\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-1}}x_t )]^2 )}  \end{align}$$

 

이므로 $\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-1}}x_t $이고, $\tilde{\beta}_t=\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\beta_t$로써 정의됩니다.

 

따라서, 이전 연구[Ref. 1]에서는 diffusion process에서의 intermediate distributions의 확률 분포를 estimate하는 모델을 Diffusion Model이라고 정의했습니다. 따라서, diffusion model은 $q(x_{t-1} \vert x_t, x_0)$와 $\theta$를 통해 생성되는 $p_{\theta}(x_{t-1} \vert x_t)$를 matching 시키는 모델이고 각 $q(x_{t-1} \vert x_t, x_0)$이 Gussian distribution으로 정의했기 때문에, 계산의 용이성을 위해 Monte Carlo estimate 대신 Rao-Blackwellized fashion 으로 계산될 수 있다고 합니다. 간단하게, Gaussian distribution이기 때문에 VAE처럼 reparameterization trick으로 KL term을 학습하게 됩니다.

 

Diffusion models and Denoising autoencoder

그래서 본 논문이 이 Diffusion Model을 어떻게 발전시켰는지 알아보겠습니다. 그래서 Denoising Diffusion Probabilistic Model에서는 크게 2가지를 변경한다고 합니다.

 

1. Forward Process에서의 $\beta_t$

2. Reverse Process에서의 모델 아키텍쳐와 Gaussian distribution parameterization

 

를 선택함으로써 diffusion model과 denoising score matching 과의 새로운 connection을 만들었다고 합니다. 

 

 

Forward process

Forward process에서 $\beta_t$ 을 reparameterization을 통해 학습할 수 있지만, 상수로 고정시키고 간다고 합니다. 그렇게 된다면 $L_T = D_{KL}(q(x_T|x_0) || p(x_T))$에서 $q(x_T \vert x_0)$과 $p(x_T)$모두 trainable parameter가 없기 때문에 상수로써 무시할 수 있다고 합니다. 

 

Reverse process

Reverse process에서의 conditional distribution $p_\theta(x_{t-1}\vert x_t)=\mathcal{N}(x_{t-1};\mu_\theta(x_t, t), \Sigma_\theta(x_t,t))$에서 우선 $ \Sigma_\theta(x_t,t)=\sigma_t^2$로 설정합니다. 논문에서는 $\sigma_t^2$를 $\beta_t$ or ${{ 1-\bar{\alpha}_{t-1} }\over{1-\bar{\alpha}_{t}}}\beta_t$ 로 실험했었는데, 거의 비슷한 결과가 나왔다고 합니다. 그리고 두번째로, $\mu_{\theta}(x_t, t)$를 표현하기 위해서

 

$$ L_{t-1} = \mathbb{E}_q [ { {1}\over{2\sigma_t^2 }} || \tilde{\mu}_t(x_t,x_0)-\mu_\theta(x_t, t)||^2 ]+C $$

 

를 보게되면, $ x_t(x_0,\epsilon)=\sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_t} \epsilon $으로 reparameterizing할 수 있고, 이것을 $x_0$로 정리해서 대입하게 된다면

 

$$\begin{align} \tilde{\mu_t}(x_t(x_0, \epsilon), \frac{1}{\sqrt{\bar{\alpha}_t}}(x_t(x_0, \epsilon)-\sqrt{1-\bar{\alpha}_t}\epsilon) &= \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}(\frac{1}{\sqrt{\bar{\alpha}_t}}(x_t(x_0, \epsilon) - \sqrt{1-\bar{\alpha}_t}\epsilon)) + \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}x_t(x_0,\epsilon) \\ &= [\frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}\frac{1}{\sqrt{\bar{\alpha}_t}}+\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}]x_t(x_0, \epsilon)-[\frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}\frac{\sqrt{1-\bar{\alpha}_t}}{\sqrt{\bar{\alpha}_t}}]\epsilon \\ &= [ \frac{1}{\sqrt{\alpha_t}}(\frac{\beta_t}{1-\bar{\alpha}_t}+\frac{\alpha_t-\bar{\alpha}_t}{1-\bar{\alpha}_t}) ]x_t(x_0, \epsilon) - [\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\frac{1}{\sqrt{\alpha}_t}]\epsilon \\&= \frac{1}{\sqrt{\alpha_t}}(x_t(x_0, \epsilon) - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon) \end{align}$$

 

따라서, $\mu_{\theta}$는 $x_t$가 주어졌을 때, $\frac{1}{\sqrt{\alpha_t}}(x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon)$ 를 예측하도록 학습하게 됩니다. 결국, $x_t$가 model의 input이기 때문에 우리는 

 

$$ \mu_{\theta}(x_t, t)=\tilde{\mu}_t(x_t, \frac{1}{\sqrt{\bar{\alpha}_t}}(x_t - \sqrt{1-\bar{\alpha}_t} \epsilon_{\theta}(x_t))) = \frac{1}{\sqrt{\alpha_t}}(x_t - \frac{\beta}{\sqrt{1-\bar{\alpha}_t}}\epsilon_{\theta}(x_t, t)) $$

 

그리고 $\epsilon_{\theta}(x_t, t)$는 $x_t$에서 $\epsilon$을 추정하는 모델입니다. 따라서 우리는 $x_{t-1}$를 $p_{\theta}(x_{t-1}\vert x_t)$ 에서 sampling하는 것이 $x_{t-1} = \frac{1}{\alpha_t}(x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_{\theta}(x_t, t)) + \sigma_t\mathbf{z}$을 계산하는 것으로 가능합니다. 결국 우리가 학습시켜야 하는 모델은 $\epsilon_{\theta}$가 됩니다. 또한, 이것을 이용해 sampling하는 과정이 Langevin dynamic에서 data density의 gradient를 계산하는 것과 닮았다고 합니다. 또한, 

 

$$  \begin{align} L_{t-1}-C &= \mathbb{E}_{x_0,\epsilon} [ { {1}\over{2\sigma_t^2 }} || { {1}\over{\sqrt{\sigma_t}}} (x_t(x_0,\epsilon)-{{\beta_t}\over{\sqrt{1-\bar{\alpha}_t}}}\epsilon)-\mu_\theta(x_t(x_0,\epsilon),t)||^2 ] \\&= \mathbb{E}_{x_0,\epsilon}[ { {\beta_t^2}\over{ 2\sigma^2_t\sigma(1-\bar{\alpha}_t })  } || \epsilon - \epsilon_{\theta}( \sqrt{ \bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon,t) ||^2 ] \end{align} $$

 

으로 학습하는 것이 t에 따라서 noise scale을 다르게 하는 것도 denoising score matching과 닮았다고 말합니다. 마지막으로 위의 식이 Langevin-like reverse process 에서의 variational bound와 같으므로, denoising score matching과 닮은 objective를 optimizing하는 것은 Langevin dynamic과 닮은 sampling chain의 finite-time marginal을 fit하기 위해 variational inference를 사용하는 것과 같다고 말하고 있습니다. 

 

최종적으로 본 논문에서 제안한대로 objective을 종합해 Loss function은,

$$\mathcal{L}_{simple}(\theta)=\mathbb{E}_{t, x_0, \epsilon}[\Vert \epsilon - \epsilon_{\theta}(\sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon)\Vert^2]$$

으로 정리됩니다.

 

결국 위에 언급된 모든 내용을 summary하면, 우리는 reverse process의 mean fucntion approximator $\mu_{\theta}$가 $\tilde{\mu}_t$를 예측하도록 학습할 수 있고 parameterization을 수정하여 $\epsilon$을 predict하도록 학습하는 것이 가능하다고 합니다. 또한, $x_0$를 예측하도록 하는 것이 가능하지만, 이 경우에는 실험 초기에 sample quality가 안좋아진다는 것을 발견했다고 합니다.

다음으론 $\epsilon$을 추정하게끔하는 parameterization은

1. Langevin dynamic과 닮았고,

2. denoising score matching과 닮은 objective에 대한 diffusion model의 variational bound를 단순화하는 것

이 가능하다고 합니다. 그리고, 이것은 $p_{\theta}(x_{t-1} \vert x_t)$ 의 또 다른 표현이기 때문에, $\epsilon$을 예측하는 것과 $\tilde{\mu}_t$를 예측하는 것의 차이를 실험했다고 합니다.

 

결과적으로 보면, $\epsilon$을 예측하는 것이 더 성능이 좋다고 합니다. 저자들은 $\tilde{\mu}$를 학습할 경우 MSE대신 실제 variational bound를 학습할 때만 점수가 잘 나온다고 합니다. 위에 보시는 표와 같이 $\Vert \tilde{\mu}-\tilde{\mu}_{\theta} \Vert^2$ 으로 $\tilde{\mu}$를 추정할 경우에는 오히려 성능 자체가 나오지 않는다고 합니다. 또한, variational bound의 parameterized 의 diagonal $\Sigma$를 통합함으로써 diagonal $\Sigma$를 학습하는 것이 품질의 저하와 불안정한 학습으로 이어진다고 말하고 있습니다. 이 논문에서 제안한대로, 분산을 고정시키고 $\epsilon$을 추정하는 것이 훨씬 더 성능이 좋다고 강조하고 있습니다.

 

위에서 언급한 내용을 알고리즘으로 정리하면, 

 

이 두 알고리즘으로 설명이 가능합니다. t에 따라 noise level을 다르게 정의해 $x_t$에서 $\epsilon$을 추정하는 model $\epsilon_{\theta}$을 학습하고 (Algorithm 1), 우리가 정의한 $p_{\theta}(x_{t-1} \vert x_t)$를 사용해 reverse process에서의 sampling을 합니다. (Algorithm 2)

 

 

Experiments and Result

우선, $T=1000, \beta_1=10^{-4}, \beta_T=0.02$로 설정합니다. 이렇게 설정했을 때, $x_T$에서의 snr을 가능한 작게 유지하면서 reverse와 forward process가 거의 동일한 기능을 하도록 하는 것이 가능하다고 합니다. 또한, reverse process를 표현하기 위해 unmasked PixelCNN++과 비슷한 U-net backbone을 사용했고, time t에 따른 embedding을 사용했다고 합니다. 또한 데이터를 [0,255] 를 [-1, 1]로 전처리 했다고 합니다. 

 

위와 같이 Setting을 진행한다면, $D_{KL}(q(x_T|x_0)||\mathcal{N}(0,1))\approx10^{-5}$ 정도로 snr이 굉장히 비슷해진다고 합니다. 실제로 $\bar{\alpha}_T=0.9999$라는 값이 나오게 됩니다. 따라서, $L_T$는 거의 고려되지 않아도 될 정도의 scheduling이 필요하다는 뜻으로 해석됩니다.

 

그리고, Sampling 시에 가장 마지막 step에는 noise를 추가하지 않는다고 합니다. 또한 {0,1, ... 255}를 [-1, 1]로 변환한다면, 마지막 스텝에서 $p_{\theta}(x_0|x_1)$에서는 노이즈를 추가하지 않는다고 합니다.

이는 $\frac{1}{255}$가 $\bar{\sigma}_1$보다 월등하게 작기 때문에, lossless codelength of discrete data임을 보장한다고 합니다. 이는 마지막 step에서 noise를 추가하는 것이 크게 영향이 없다는 것을 의미합니다.

 

 

위에서도 언급한대로, Langevin dynamic과 Denoising score matching과 비슷한 형태의 알고리즘이기 때문에 생성되는 과정이 NCSN과 비슷하게 생성이 됩니다.

 

 

두번째로는, $p_{\theta}(x_0 \vert x_t)$ 에서 sampling을 하는 것입니다. t가 작을 수록 (데이터에 가까울 수록) 원본에 대한 정보가 많이 담겨있고, t가 클수록 (데이터에 멀 수록) 원본과 많이 다른 sampling이 가능합니다.

 

Quantitative Result

그 다음은 Quantitative Result입니다. Quantitative Result에서 FID score 3.17로 StyleGAN2+ADA를 이길 만큼의 좋은 모습을 보여주고 있습니다. 또, 논문에서는 train과 test의 codelength 차이가 0.03밖에 나지 않은 것으로 overfitting되어있지 않다라고 주장합니다. 또한, energy-based generative model과 score matching using annealed importance sampling보다는 수치가 좋게 잡히지만, 다른 모델들보다는 성능이 좋지 않다고 얘기하고 있습니다.

 

그럼에도 불구하고, 이미지 퀄리티가 훨씬 좋기 때문에 확산모델이 더 좋은 lossy compressors라고 주장합니다. $L_1 + L_2 + \cdot L_T$를 ratio, $L_0$를 distortion으로 정의했을 때, diffusion model은 1.78bit/dim의 ratio와 1.97bit/dim의 distortion으로 기록되고, 이 정도의 수치는 0-255 scale에서 0.95의 RMSE의 수치라고 설명하고 있습니다. NLL table을 보시면 Ours의 NLL은 3.75bits/dim(1.78bits/dim + 1.97bits/dim)의 결과가 나오는데 Sparse Transformer(2.80bits/dim), Gated PixelCNN(3.03bits/dim) 보다 결과가 안좋게 나오지만, DDPM의 NLL(lossless codelength) 절반 이상이 distortions수치이고 distortion은 감지할 수 없는 값이라고 언급하고 있습니다.

 

 

 

Interpolation

 

$x_0, x^{'}_0 \sim q(x_0)$ 를 sampling했을 때, q를 stochastic encoder로써 interpolation이 가능하다고 합니다. $x_t, x^{'}_t \sim q(x_t \vert x_0)$일 때, linear interpolation을 통해 $\bar{x}_t = (1-\lambda)x_t + \lambda x^{'}_t$ (논문에는 $x_0, x^{'}_0$ 으로 적혀있긴 한데 아마 오타이지 않을까 싶습니다.. 문맥상 $x_t, x^{'}_t$가 더 적합하다고 생각해서 이렇게 적었습니다.)

그리고 reverse process를 통해 $\bar{x}_0 \sim p(x_0 \vert \bar{x}_t)$ sampling을 할 수 있고, 또한 $\lambda$를 다르게 했을 때, 각각의 source에 가까워지게 reconstruction 되는 것을 확인 할 수 있습니다. 

 

 

 

Progressive lossy compression

DDPM의 ratio-distortion을 더 잘 보여주기 위해 progressive lossy code를 설명합니다. Minimal random code와 같은 procedure들의 접근을 가정한 Algorithm 3, 4는 q와 p의 KL Divergence를 이용해 q(x)에서의 sample x를 transmit할 수 있고, 이 경우 receiver는 p만 사용할 수 있다고 합니다. 이 경우에 Algorithm 3, 4는 $x_T, \cdot, x_1, x_0$을 순차적으로 transmit할 수 있다고 합니다. 

이때, 임의의 시간 t에서 receiver에서의 부분 정보인 $x_t$를 사용할 수 있으며 다음 식을 이용해 progressive하게 추정할 수 있다고 합니다.

$$ x_0 \approx \hat{x}_0=(x_t-\sqrt{1-\bar{\alpha}_t}\epsilon_{\theta}(x_t)) / \sqrt{\bar{\alpha}_t} $$

 

물론, $x_0 \sim p_{\theta}(x_0 \vert x_t)$에서 stochastic reconstruction이 가능하지만, 여기서는 고려하지 않는다고 합니다.

 

 

다음은 각 time t에서 distortion, rate를 계산한 plot입니다. distortion은 $\sqrt{\Vert x_0 - \hat{x}_0 \Vert^2/D }$ 로 $x_0$와 $x_0$를 estimate한 값과의 RMSE를 계산했고, rate는 수신된 누적 bit로 계산했다고 합니다. distortion은 distortion-rate plot에서 낮은 rate에서 빠르게 감소하고 이것은 bit의 대부분이 감지할 수 없는 왜곡에 존재한다고 말하고 있습니다.

 

[참고] Minimal Random Coding

 

 

Connection to autoregressive decoding

 

variational bound를 다시 써보면

$$\mathcal{L}=D_{KL}(q(x_T)\Vert p(x_T))+\mathbb{E}_q[\sum_{k\leq1} D_{KL}(q(x_{t-1}\vert x_t) \Vert p_{\theta}(x_{t-1} \vert x_t))] + H(x_0)$$

로 쓸 수 있는데, $q(x_t \vert x_0)$을 $x_0$의 t번째 coordinates로 정의해서 T를 데이터의 dimensionality으로 설정할 수 있다고 합니다. 예를 들어서, 원래 이미지의 dimension이 3 x 32 x 32 라고 한다면 T x 3 x 32 x 32로 dimensionality를 설정한다고 합니다. $t+1, t+2, \cdot, T$를 이용해서 $t$번째 coordinate를 예측하도록 모델을 학습시키는 것과 유사하다고 합니다. 결국 이 diffusion으로 학습한 모델은 autoregressive model을 훈련하는 것으로 볼 수 있습니다.

다시 말하자면, diffusion model은 data coordinate를 재정렬해서 표현될 수 없는 비트 순서를 갖는 autoregressive model의 일종이라고 볼 수 있습니다. 이전 연구에서 [Ref. 2] 이러한 reordering이 sample quality에 영향을 미치는 inductive bias를 가져올 수 있다는 것을 보여주었기 때문에, Gaussian 노이즈는 masking noise보다 더 자연스럽기 때문에 더 좋은 효과를 제공하는 것 같다고 추정합니다.

 

 

Discussion

 

본 논문에서 Discussion에서는 주로 DDPM과 NCSN과의 차이점을 설명하고 있습니다.

1. DDPM은 self-attention을 추가한 U-Net을 사용합니다. NCSN은 Dilated Convolution을 사용된 RefineNet을 사용합니다. 또한 NCSNv1(only normalization), NCSNv2(only output)에서 t에 대한 임베딩을 하는 것과는 다르게 모든 layer에서 t에 대한 임베딩을 합니다.
2. DDPM은 Forward process에 따라 데이터를 축소하여 Reverse process에서 데이터의 input을 스케일링된 input을 제공하지만, NCSN은 그렇지 못하다고 합니다. 
3. NCSN과는 달리 $D_{KL}(q(x_T \vert x_0) \Vert p(x_T))$에 대한 signal을 제거했기 때문에, prior와 posterior distribution간의 matching을 보장해준다고 합니다. 또한, NCSN과 달리 $\beta_t$가 매우 작고, 이는 forward process를 conditional Gaussian의 Markov chain에 의해 reversible하게 만든다고 합니다. 이 두 가지 이유때문에 sampling 중에 distribution shift을 예방할 수 있다고 합니다. 
4. DDPM은 forward process에서부터 엄밀하게 정의해서 나온 coefficient를 이용해 sampling을 진행합니다. 따라서, T step 이후에 실제 데이터 분포에 일치하도록 모델을 학습시킵니다. 반면에 NCSN은 sampling시에 sampler coefficient($\epsilon$ 의 Langevin dynamic or step size)을 이용해 post-hoc, 즉 결과를 보고 heuristics하게 정한다고 합니다. 이는 training시에 들어가지 않는 parameter이기 때문에, training procedure가 sampler의 최상의 결과를 내는 것을 직접적으로 보장하지 않는다고 합니다. 

 

 

Conclusion

 

Diffusion model을 이용해서 high-quality image를 보여주었으며, Markov chain, Denoising score matching과 annealed Langevin dynamic, autoregressive model 그리고 progressive lossy compression와의 연결성을 발견했습니다. 또한 Diffusion model은 이미지에 대한 좋은 inductive bias를 갖고 있기 때문에, 다른 유형의 generative model과 ML system과 다른 유형의 데이터 양식에서 그것들의 활용성이 기대된다고 합니다.

 

마지막으로 DDPM에 대한 튜토리얼 링크를 공유하면서 마무리하겠습니다. [Link]

다음에는 NCSNv2로 돌아오겠습니다.

 

 

Reference

[1] Deep Unsupervised Learning using Nonequilibrium Thermodynamics

Deep Unsupervised Learning using Nonequilibrium Thermodynamics

[2] Generating High Fidelity Images with Subscale Pixel Networks and Multidimensional Upscaling

 

 

이전글 : https://ivdevlog.tistory.com/4

Score-Based Generative Model[1] - Score Network

이전글 : https://ivdevlog.tistory.com/8

Score-Based Generative Model[2] - NCSN

다음글 : https://ivdevlog.tistory.com/15

Score-Based / Diffusion Model[4] - NCSNv2