DDPM 논문 리뷰 및 수식 설명
Diffusion Model이란?
디퓨전 모델이란 VAE, GAN과 같은 생성형 모델 중 하나로 데이터(ex 이미지)에 가우시안 노이즈를 추가하여 완전한 가우시안 노이즈 이미지를 만드는 과정(forward process)과 그러한 가우시안 노이즈 이미지에서 역으로 노이즈를 제거하여 온전한 이미지를 만드는 과정(reverse process)를 통해 데이터를 생성한다.
Diffusion이라고 표현하는 이유는 노이즈를 추가하는 과정이 물체가 확산하는 과정과 유사하기 때문이다. 예를 들어 잉크가 물에 퍼지는 상황을 생각해보자. 잉크 분자는 점점 물에 퍼지면서 무질서한 상태가 될 것이다. 좀 더 자세히 들여다본다면 아주 짧은 시간 t에서 잉크 분자의 움직임은 현재 위치를 평균으로 가지는 3차원 가우시안 확률분포에서 결정될 수 있다. 이미지에 노이즈를 추가하는 과정 또한 마찬가지로 현재 픽셀 값을 평균으로하는 가우시안 분포에서 샘플링을 하는 것으로 볼 수 있다.
Foward Process & Reverse Process
Foward Process
DDPM에서 언급하는 Foward Process란 이미지에 가우시안 노이즈를 추가하는 과정으로 마르코프 성질을 따른다고 가정한다.
$$q(\mathbf{x}_{1:T}|\mathbf{x}_0) := \prod_{t=1}^T q(\mathbf{x}_t|\mathbf{x}_{t-1})$$
위의 식은 원본 이미지 $\mathbf{x_0}$에서 시작하여 점진적으로 노이즈를 추가해가는 전체 Forward Process를 의미하는 결합 확률 분포이다.
미래의 상태를 결정하는 요소가 오로지 현재의 상태인 마르코프 성질을 활용해 통해 위의 식을 아래와 같은 과정을 통해 유도할 수 있다.
$$q(\mathbf{x}_t|\mathbf{x}_{t-1}) := \mathcal{N}(\mathbf{x}_t; \sqrt{1-\beta_t}\mathbf{x}_{t-1}, \beta_t\mathbf{I})$$
Foward Process에 사용되는 확률 분포를 가우시안 분포로 가정했기 때문에 $q(\mathbf{x}_t|\mathbf{x}_{t-1})$를 $\mathcal{N}(\mathbf{x}_t; \sqrt{1-\beta_t}\mathbf{x}_{t-1}, \beta_t\mathbf{I})$로 정의할 수 있다.
여기서 $\beta_t$는 노이즈의 크기 파라미터이다.(0< $\beta_1$< $\beta_2$<…< $\beta_T$<1)
$\sqrt{1-\beta_t}\mathbf{x}_{t-1}$은 이전 상태(t-1)의 이미지의 픽셀값들을 미세하게 줄인 것을 의미하며, 이 값을 가우시안 분포의 평균으로 사용한다. 공분산의 경우에는 $\beta_t$를 곱한 값을 사용한다.
($\sqrt{1-\beta_t}$로 스케일링하는 이유는 더보기를 참고)
더보기
💡 $\sqrt{1-\beta_t}$로 스케일링하는 이유 → variance가 발산하는 것을 막기 위함 reparameterization trick를 사용해 샘플링하는 것을 표현하면 아래와 같이 표현할 수 있다.
$$ \mathbf{x_t} \sim \mathcal{N}(\sqrt{1-\beta_t}\mathbf{x_{t-1}}, \beta_t\mathbf{I})\\ \mathbf{x_t}=\sqrt{1-\beta_t}\mathbf{x_{t-1}}+\sqrt{\beta_t}\odot{\boldsymbol{\epsilon}},\; where\;\boldsymbol{\epsilon}\sim \mathcal{N}(0,\mathbf{I}) $$
여기서 분산의 기본 성질을 이용하여 $Var(x_t)$를 아래와 같이 계산할 수 있다.
$\epsilon$는 분산이 1인 표준 정규분포에서 샘플링했으므로 확률변수 $\epsilon$의 분산은 1로 고정된다.
아래는 분산이 수렴하는 과정을 직관적으로 알 수 있게 표현한 식이다.
하지만 스텝의 수 t가 매우 클 경우, 매 스탭마다 확률분포를 계산하고 샘플링하는 과정의 비용이 많이 든다. 이러한 문제를 해결하기 위해 reparameterization trick을 이용하여 $\mathbf{x}_0$가 주어졌을 때 $\mathbf{x}_t$의 확률분포 $q(\mathbf{x}_t|\mathbf{x}_{t-1})$를 아래와 같은 과정을 통해 구할 수 있다.
Reverse Process
Reverse Process는 Foward Process의 역과정으로 노이즈를 걷어내는 역할을 한다.(denoising 과정) 즉, $\mathbf{x_t}$상태의 이미지에서 $\mathbf{x_{t-1}}$이미지로 가는 과정이고, 확률 분포로 나타내면 $q(\mathbf{x}_{t-1}|\mathbf{x_t})$로 표현할 수 있다. 하지만 $q(\mathbf{x}_{t-1}|\mathbf{x_t})$는 직접 계산할 수 없는 확률 분포이기 때문에 딥러닝 모델 $p_\theta(\mathbf{x_{t-1}}|\mathbf{x_{t}})$을 학습시켜 해당 확률분포로 근사할 것이다.
PS. $q(\mathbf{x}_{t-1}|\mathbf{x_t})$는 가우시안 분포로 가정함
$$p_\theta(\mathbf{x}_{0:T}) := p(\mathbf{x}_T) \prod_{t=1}^T p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)$$
위의 수식은 Reverse Process를 표현한 식으로, $\mathbf x_0,\mathbf x_1,\mathbf x_2,...,\mathbf x_T$의 결합확률 분포를 마르코프 체인 성질를 이용해 정의할 수 있음을 보여준다. 아래와 같이 마르코프체인 성질을 이용해 결합확률분포를 유도할 수 있다.
$$p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t) := \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t))$$
우리는 Reverse Process 또한 가우시안 분포라고 가정을 했으므로, $p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)$는 위와 같이 신경망으로 예측한 평균과 공분산을 가지는 가우시안 분포로 정의 할 수 있다.
$\boldsymbol{\mu}_\theta(\mathbf{x}_t, t)$은 신경망이 예측하는 조건부 분포의 평균으로 $\mathbf{x}_t$와 $t$를 입력으로 받아 계산된다.
$\boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t)$은 신경망이 예측하는 조건부 분포의 공분산으로 $\mathbf{x}_t$와 $t$를 입력으로 받아 계산된다. 하지만 DDPM에서는 공분산 값을 $\beta_t$로 고정한다.
목적함수
디퓨전 모델의 손실함수는 이미지($x_0$) 확률분포의 negative-log-liklihood를 최소화하는 것을 목표로 한다. 변분 추론을 통해 ELBO를 도출하여 최소화하게 되면 negative-log-liklihood를 최소화하게 된다. 아래는 전체 과정을 수식으로 나타낸 것이다.
$$ \begin{align*} &\mathbb{E}{x{1:T} \sim q(x_{1:T}|x_0)}[-\log p_\theta(x_0)] \\ &= \mathbb{E}{x{1:T} \sim q(x_{1:T}|x_0)}\left[-\log \frac{p_\theta(x_0, x_1, x_2, \ldots, x_T)}{p_\theta(x_1, x_2, x_3, \ldots, x_T|x_0)}\right] \\ &= \mathbb{E}{x{1:T} \sim q(x_{1:T}|x_0)}\left[-\log \frac{p_\theta(x_0, x_1, x_2, \ldots, x_T)}{p_\theta(x_1, x_2, x_3, \ldots, x_T|x_0)} \cdot \frac{q(x_{1:T}|x_0)}{q(x_{1:T}|x_0)}\right] \\ &\leq \mathbb{E}{x{1:T} \sim q(x_{1:T}|x_0)}\left[-\log \frac{p_\theta(x_0, x_1, x_2, \ldots, x_T)}{q(x_{1:T}|x_0)}\right] \\ &= \mathbb{E}{x{1:T} \sim q(x_{1:T}|x_0)}\left[-\log \frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)}\right] \\ &= \mathbb{E}{x{1:T} \sim q(x_{1:T}|x_0)}\left[-\log \frac{p_\theta(x_T)\prod_{t=1}^T p_\theta(x_{t-1}|x_t)}{\prod_{t=1}^T q(x_t|x_{t-1})}\right] \\ &= \mathbb{E}{x{1:T} \sim q(x_{1:T}|x_0)}\left[-\log p_\theta(x_T) - \sum_{t=1}^T \log \frac{p_\theta(x_{t-1}|x_t)}{q(x_t|x_{t-1})}\right] \\ &= \mathbb{E}{x{1:T} \sim q(x_{1:T}|x_0)}\left[-\log p_\theta(x_T) - \sum_{t=2}^T \log \frac{p_\theta(x_{t-1}|x_t)}{q(x_t|x_{t-1})} - \log \frac{p_\theta(x_0|x_1)}{q(x_1|x_0)}\right] \\ &= \mathbb{E}{x{1:T} \sim q(x_{1:T}|x_0)}\left[-\log p_\theta(x_T) - \sum_{t=2}^T \log \frac{p_\theta(x_{t-1}|x_t)}{q(x_{t-1}|x_t, x_0)} - \log \frac{p_\theta(x_0|x_1)}{q(x_1|x_0)}\right] \\ &= \mathbb{E}{x{1:T} \sim q(x_{1:T}|x_0)}\left[-\log p_\theta(x_T) - \sum_{t=2}^T \log \frac{p_\theta(x_{t-1}|x_t)}{q(x_{t-1}|x_t, x_0)} - \sum_{t=2}^T \log \frac{q(x_{t-1}|x_0)}{q(x_t|x_0)} - \log \frac{p_\theta(x_0|x_1)}{q(x_1|x_0)}\right] \\ &= \mathbb{E}{x{1:T} \sim q(x_{1:T}|x_0)}\left[-\log p_\theta(x_T) - \sum_{t=2}^T \log \frac{p_\theta(x_{t-1}|x_t)}{q(x_{t-1}|x_t, x_0)} - \log \frac{q(x_1|x_0)}{q(x_T|x_0)} - \log \frac{p_\theta(x_0|x_1)}{q(x_1|x_0)}\right] \\ &= \mathbb{E}{x{1:T} \sim q(x_{1:T}|x_0)}\left[-\log p_\theta(x_T) - \sum_{t=2}^T \log \frac{p_\theta(x_{t-1}|x_t)}{q(x_{t-1}|x_t, x_0)} - \log \frac{p_\theta(x_0|x_1)}{q(x_T|x_0)}\right] \\ &= \mathbb{E}{x{1:T} \sim q(x_{1:T}|x_0)}\left[-\log \frac{p_\theta(x_T)}{q(x_T|x_0)} - \sum_{t=2}^T \log \frac{p_\theta(x_{t-1}|x_t)}{q(x_{t-1}|x_t, x_0)} - \log p_\theta(x_0|x_1)\right] \end{align*} $$
지금부터는 이 전개 과정을 차근차근 풀어볼 것이다.
ELBO 도출
- 베이즈 룰을 이용하여 기존 $p_\theta(x_0)$를 재표현
- $q(x_{1:T}|x_0)$를 분자와 분모에 곱한다. 이 과정에서 $-D_{KL}(p_\theta(x_{1:T}|x_0)||q_\theta(x_{1:T}|x_0))$를 도출하게 되는데 이 값은 0보다 항상 같거나 작다.
- 위에서 도출한 $-D_{KL}$ term를 제거하면 항상 원래식보다 항상 크거나 작게 된다. 즉, ELBO가 도출된것이다. 하지만 negative-log-liklihood이기 때문에 여기서는 하한이 아니라 상한으로 생각하면 된다.
- $\mathbf x_1,\mathbf x_2,\mathbf x_3,...,\mathbf x_T$의 결합확률 분포를 마르코프 체인 성질를 이용해 새롭게 정의한다.
- 이후 로그 성질을 이용해 정리하면 두개의 term으로 나뉜다.
- 빨간색으로 칠해진 term을 보면 $p_\theta$는 $x_t$가 주어졌을 때 $x_{t-1}$의 확률 분포 즉, Reverse Process를 의미한다. 하지만 $q$는 $x_{t-1}$가 주어졌을 때 $x_t$의 확률 분포 즉, Foward Process를 의미한다. 이로인해 신경망이 오히려 노이즈해지는 과정으로 배우는 것처럼 보인다. 그러나 앞으로 남아있는 수식전개 과정을 통해 그렇지 않다는 것을 알 수 있을 것이다.
- t=1일 때 term을 따로 분리
- $q(x_t|x_{t-1})$를 마르코프 체인 성질을 이용하여 변형
- 로그 성질을 이용하여 별개의 term으로 분리
- 파란색으로 밑줄 친 term을 log 성질을 이용하여 정리
- 오른쪽 기준 첫번째와 두번째 텀을 로그성질을 이용하여 하나로 합침
- log성질을 이용하여 마지막 텀의 분모를 맨 처음 term의 분모로 바꿔줌
수식 전개 과정을 통해 ELBO를 도출하게되고, 마지막에 ELBO는 세가지 term으로 나뉘게 된다. 각각의 텀들은 아래와 같은 의미를 가진다.
$L_T$: 이 term은 최종 노이즈가 추가된 입력의 분포가 표준 가우시안 사전 분포에 얼마나 가까운지를 나타낸다. $x_{T}$는 가우시안 분포에서 샘플링된다고 가정했으므로 0에 가깝다. VAE에서 Regularization 항과 비슷하다고 생각할 수 있다.
$L_{t-1}$: 이 term은 Denoising 과정을 나타낸다. $p_\theta$는 $x_t$가 주어졌을 때 $x_{t-1}$에 대한 확률분포를 예측해야하는데 이걸 확률분포 $q(x_{t-1}|x_t,x_0)$와 닮도록 학습한다는 의미를 담고 있다.
$L_{0}$: 이 term은 VAE에서 Reconstruction 항과 비슷하다고 생각할 수 있다.(원본이미지($x_0$)에 대한 likelihood이기 때문). t =1일 때 단 한 순간의 likelihood값으로 매우 작은 값을 가진다.
실제 훈련할 때는 $L_T$, $L_0$는 매우 작은 값이라 무시되고, $L_{t-1}$를 간소화하여 손실함수를 구성하게 된다.
손실함수
$$L_{\text{simple}}(\theta) := \mathbb{E}_{t,x_0,\epsilon} \left[ \left\| \epsilon - \epsilon_\theta \left( \sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t} \epsilon, t \right) \right\|^2 \right]$$
DDPM에서는 목적 함수에서 나왔던 $L_{t-1}$ term을 간소화하여 손실함수를 위와 같이 도출한다.
이제부터는 어떻게 간소화했는지 그 과정을 설명하고자 한다.
KL-Divergence
$$ \begin{equation*} \begin{aligned} p(x) &\sim \mathcal{N}(\mu_1, \sigma_1^2) \\ q(x) &\sim \mathcal{N}(\mu_2, \sigma_2^2) \\ KL(p\|q) &= \log\left(\frac{\sigma_2}{\sigma_1}\right) + \frac{\sigma_1^2 + (\mu_1 - \mu_2)^2}{2\sigma_2^2} - \frac{1}{2} \end{aligned} \end{equation*} $$
$L_{t-1}$는 $q(x_{t-1}|x_t,x_0)$와 $p(x_{t-1}|x_t)$의 KL 다이버전스 값을 의미한다. 두 가우시안 분포의 KL 다이버전스 값은 위와 같은 공식을 통해 구할 수 있다.
$$L_{t-1} = E_q \left[ \frac{1}{2\sigma_t^2} \|\tilde{\mu}_t(x_t, x_0) - \mu_\theta(x_t, t)\|^2 \right] + C$$
공식을 사용하여 KL 다이버전스 값을 구하면 위와 같은 식이 나오게 되는데 상수항의 경우에는 전부 합쳐서 C로 표현한다. 그 이유는 분산은 신경망을 통해 나오는 출력값이 아니라 사전에 정해진 값($\beta_t$를 의미)을 사용하기 때문이다. 즉, 학습 하는데 필요하지 않은 값들은 C로 두어 간소화한다.
$\tilde \mu$ 와 $\mu$ 값 구하기
KL 다이버전스 값을 구할 때 각각의 확률분포 평균을 $\tilde \mu(x_t, x_0)$, $\mu(x_t,t)$로 표현했다.
여기서는 각각의 평균을 구하는 식을 구함과 동시에 어떻게 $x_t,x_0,t$로 그 값을 구할 수 있는지 알아볼것이다.
$\tilde \mu(x_t, x_0)$값 구하기
$\tilde \mu(x_t, x_0)$은 $q(x_{t-1}|x_t,x_0)$의 평균을 의미한다. 평균을 구하기 위해서는 식을 전개하여 확률 밀도 함수로 표현해야한다. 아래는 그 전개과정이다.
전개를 하면 $x_{t-1}$에 대한 2차식으로 정리할 수 있다.
정규분포를 전개하여 얻은 식의 $x^2$과 $x$계수와 $q(x_{t-1}|x_t,x_0)$을 전개했을 때 나오는 식의 $x_{t-1}^2$, $x_{t-1}$의 계수를 비교하여 분산과 평균을 구할 수 있다.(아래 식 참조)
여기서 구한 $\tilde \mu(x_t,x_0)$은 $x_t, x_0$에 대한 식으로 이루어져 있다. 하지만 이것을 오직 $x_t$에 대해서만 의존하도록 변형할 수 있다. 아래는 그 과정이다.
$q(x_{t-1}|x_t, x_0)$의 평균과 분산은 최종적으로 아래와 같이 나오게 된다.
$$ \tilde{\mu}_t = \frac{1}{\sqrt{\alpha_t}} \left(x_t - \frac{(1 - \alpha_t)}{\sqrt{1 - \bar{\alpha}_t}} \epsilon\right) \qquad\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t $$
$\mu_\theta(x_t, t)$값 구하기
$p_\theta(x_{t-1}|{x_t})$ 또한 가우시안 분포이기 때문에, $\tilde \mu_t$, $\tilde\beta_t$에 사용된 식을 그대로 활용할 수 있다.
다만 Reverse Process를 진행할 때, $\epsilon$를 제외한 나머지값들은 주어지기 때문에 신경망을 이용하여 노이즈를 예측한다. $\epsilon_\theta(x_t,t)$
$t$를 입력으로 같이 넣는 이유는 신경망을 시간에 따라 학습시키기 위함이다.(시간에 따라 노이즈 주입정도가 달라짐)
그러면 아래와 같이 $p_\theta(x_{t-1}|x_t)$의 평균을 구할 수 있다.
$$\mu_{\theta}(\mathbf{x}_t, t) = \frac{1}{\sqrt{\alpha_t}} \left(\mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_{\theta}(\mathbf{x}_t, t)\right)$$
최종 손실함수
위에서 구한 $\tilde \mu$와 $\mu$를 이용하여 DDPM에서 제안한 최종 손실함수를 아래와 같이 유도할 수 있다.
여기서 앞에 붙어 있는 weight term을 제거하면 DDPM에서 제시한 최종 손실함수가 도출된다.
$$ L_{\text{simple}}(\theta) := \mathbb{E}_{t,x_0,\epsilon} \left[ \left\| \epsilon - \epsilon_\theta \left( \sqrt{\alpha_t} x_0 + \sqrt{1-\alpha_t} \epsilon, t \right) \right\|^2 \right] $$
즉, 시간 t에서 이미지에 주입된 노이즈를 모델이 예측한 값과 시간 t에서 실제로 주입된 노이즈와의 차이를 의미한다고 생각할 수 있다.
참고자료
https://process-mining.tistory.com/182
https://lilianweng.github.io/posts/2021-07-11-diffusion-models/