[논문 리뷰] Understanding Diffusion Models: A Unified Perspective (2)

Understanding Diffusion Models: A Unified Perspective는 3가지 포스트로 나누어져 있습니다. 이전 포스팅을 참고해주세요.

Understanding Diffusion Models: A Unified Perspective (1)

 

Variational Diffusion Models

Variational Diffusion Models(VDM)

 

Variational Diffusion Models(VDM)은 이전 포스팅에서 다루었던 Hierarchical Variational Auto Encoder(HVAE)와 비슷합니다. 쉽게 말하자면 3가지의 다른점을 가진 HVAE라고 할수 있습니다:

 

• DM은 잠재 차원 (Latent Dimension) = 데이터 차원 (Data Dimension)입니다.

 

• DM은 이전 계층의 Latent를 중심으로 하는 Gaussian Noise입니다. 따라서 Encoder(Forward Process)에서는 따로 모델을 학습하지 않습니다.

 

• Latent Encoder의 Gaussian parameter는 최종 Time Step인 T에서 latent가 표준 가우시안 분포(완전한 Gaussian Noise)가 되도록 합니다.

 

말이 조금 어려우니 쉽게 풀어보겠습니다. 먼저 VAE, HVAE는 어떠한 data $x$에 대해서 잠재변수인 $z$로 Encoding 한뒤 다시 data와 같은 차원으로 Decoding을 하는 모델입니다. 반면 Diffusion Model은 data와 동일한 차원에서 분포를 $x_t$로 인코딩합니다. 즉 잠재차원과 데이터의 차원이 동일합니다.

 

이때 VDM은 HVAE의 Encoder에 적용된 Markov Chain을 유지합니다. 쉽게 말해서, HVAE처럼 각 Step마다 $q_{\phi}(z_{t}|z_{t-1})$을 곱하여 계층적으로 latent를 변화시키는 Chain Rule을 적용한다는 말입니다. (다만 HVAE와는 다르게 $x->z$가 아닌 $x->x'$ 이므로 곱하는 값이 조금 달라지겠죠?)

 

수식으로 나타내자면, data $x$에 대한 잠재변수를 잠재변수 $x_t$로 나타낼 수 있습니다. t=0일때 $x_0$는 실제 data Sample을 의미하며, $t \in [1, T]$의 t는 각 계층의 time step을 의미합니다. ($x_0$는 원본 샘플 data, $x_1$은 x$_0$에 대한 잠재변수, $x_2$는 $x_1$에 대한 잠재변수, ..., $x_T$는 $x_{T-1}$에 대한 잠재변수가 되는 것입니다.) 즉 다음과 같은 수식으로 나타낼 수 있습니다:

 

$$q(x_{1:T}|x_0)= \prod^{T}_{t=1}q(x_t|x_{t-1}) \tag{30}$$

 

이때 latent에 대한 Gaussian Noise를 추가하는 것으로 Encoding을 하기 때문에 HVAE 등과 달리 각 time step t에서 따로 모델을 학습하지 않습니다. 평균과 표준편차에 대한 하이퍼 파라미터나 Linear Gaussian Model로 고정하여 사용합니다.

 

Forward Process에서 사용하는 Gaussian Encoder는 하이퍼 파라미터로 평균: $\mu_{t}(x_{t}) = \sqrt{\alpha_{t}}x_{t-1}$, 분산:$\sum_{t}(x_{t}) = (1-\alpha_{t})I$를 사용합니다. 프로세스를 진행하는 동안 잠재변수의 분산은 보존됩니다(Variance-Preserving). 

$\alpha_{t}$는 Time Step t에 따라 잠재적으로 학습가능한 계수라는 것을 유의해야 합니다. (근데 그럼 분산은 어떻게 보존하지..?)

VDM의 Forward Process에 대한 Gaussian Noise 생성 수식은 다음과 같습니다:

 

$$q(x_{t}|x_{t-1}) = \mathcal{N}(x_{t};\sqrt{\alpha_{t}}x_{t-1}, (1-\alpha_{t})I) \tag{31}$$

 

식을 해석해보자면, 먼저 $x_{t}$는 $x_{t-1}$을 중심으로 하는 Markov Chain입니다. 이러한 Process는 Gaussian Distribution을 따르기 때문에 (이전 출력값을 기준으로 하는 Gaussian Noise) $\mathcal{N}(x_{t};\mu, \sum)$으로 나타낼 수 있습니다. 이때 평균은 이전 출력값의 $\sqrt{\alpha_{t}}$만큼 곱한 값을 사용하며, 공분산은 단위행렬 I에 $(1-\alpha_{t})$만큼 곱한 값을 사용합니다.

 

3번째 가정을 통해 마지막 latent인 $p(x_{t})$는 표준 Gaussian Distribution을 따른 다는 것을 알 수 있습니다. HVAE에 대한 수식을 VDM으로 변경하면 다음과 같습니다:

 

$$p(x_{0:T}) = p(x_{T})\prod^{T}_{t=1}p_{\theta}(x_{t-1}|x_{t}) \tag{32}$$

 

HVAE와 다른점은 딱 한가지입니다. HVAE는 입력 data $x$를 latent variable인$z$로 축소한 뒤 값을 계산하는 반면 VDM은 원본의 차원을 그대로 가져갑니다.

이때 $p(x_T)$는 다음과 같이 표준 Gaussian Distribution입니다.

 

$$p(x_T)=\mathcal{N}(x_T;0, I) \tag{33}$$

 

결국 Forward Process는 Time Step t에 따라 점진적으로 진행되는 Gaussian Distribution이기 때문에 더이상 $q$는 $\phi$로 파라미터화 되지 않습니다. 따라서 VDM은 $p_{\theta}(x_{t-1}|x_{t})$만 학습합니다.

 

이제 HVAE처럼 ELBO를 유도하여 VDM을 최대화 해 봅시다.

 

$$\begin{align}
\log p({x}) &=\log \int p\left({x}_{0: T}\right)dx_{1: T} \tag{34}
\\
&=\log\int\frac{p\left({x}_{0: T}\right) q\left({x}_{1: T} \mid {x}_0\right)}{q\left({x}_{1: T} \mid {x}_0\right)} d {x}_{1: T} \tag{35}
\\
& =\log \mathbb{E}_{q\left({x}_{1: T} \mid {x}_0\right)}\left[\frac{p\left({x}_{0: T}\right)}{q\left({x}_{1: T} \mid {x}_0\right)}\right] \tag{36}
\\
& \geq \mathbb{E}_{q\left({x}_{1: T} \mid {x}_0\right)}\left[\log \frac{p\left({x}_{0: T}\right)}{q\left({x}_{1: T} \mid {x}_0\right)}\right] \tag{37}
\\
& =\mathbb{E}_{q\left({x}_{1: T} \mid {x}_0\right)}\left[\log \frac{p\left({x}_T\right) \prod_{t=1}^T p_{{\theta}}\left({x}_{t-1} \mid {x}_t\right)}{\prod_{t=1}^T q\left({x}_t \mid {x}_{t-1}\right)}\right] \tag{38}
\\
& =\mathbb{E}_{q\left({x}_{1: T} \mid {x}_0\right)}\left[\log \frac{p\left({x}_T\right) p_{{\theta}}\left({x}_0 \mid {x}_1\right) \prod_{t=2}^T p_{{\theta}}\left({x}_{t-1} \mid {x}_t\right)}{q\left({x}_T \mid {x}_{T-1}\right) \prod_{t=1}^{T-1} q\left({x}_t \mid {x}_{t-1} \right)}\right] \tag{39}
\\
& =\mathbb{E}_{q\left({x}_{1: T} \mid {x}_0\right)}\left[\log \frac{p\left({x}_T\right) p_{{\theta}}\left({x}_0 \mid {x}_1\right) \prod_{t=1}^{T-1} p_{{\theta}}\left({x}_t \mid {x}_{t+1}\right)}{q\left({x}_T \mid {x}_{T-1}\right) \prod_{t=1}^{T-1} q\left({x}_t \mid {x}_{t-1}\right)}\right] \tag{40}
\\
& =\mathbb{E}_{q\left({x}_{1: T} \mid {x}_0\right)}\left[\log \frac{p\left({x}_T\right) p_{{\theta}}\left({x}_0 \mid {x}_1\right)}{q\left({x}_T \mid {x}_{T-1}\right)}\right]+\mathbb{E}_{q\left({x}_{1: T} \mid {x}_0\right)}\left[\log \prod_{t=1}^{T-1} \frac{p_{{\theta}}\left({x}_t \mid {x}_{t+1}\right)}{q\left({x}_t \mid {x}_{t-1}\right)}\right] \tag{41}
\\
& =\mathbb{E}_{q\left({x}_{1: T} \mid {x}_0\right)}\left[\log p_{{\theta}}\left({x}_0 \mid {x}_1\right)\right]+\mathbb{E}_{q\left({x}_{1: T} \mid {x}_0\right)}\left[\log \frac{p\left({x}_T\right)}{q\left({x}_T \mid {x}_{T-1}\right)}\right]+\mathbb{E}_{q\left({x}_{1: T} \mid {x}_0\right)}\left[\sum_{t=1}^{T-1} \log \frac{p_{{\theta}}\left({x}_t \mid {x}_{t+1}\right)}{q\left({x}_t \mid {x}_{t-1}\right)}\right] \tag{42}
\\
& =\mathbb{E}_{q\left({x}_{1: T} \mid {x}_0\right)}\left[\log p_{{\theta}}\left({x}_0 \mid {x}_1\right)\right]+\mathbb{E}_{q\left({x}_{1: T} \mid {x}_0\right)}\left[\log \frac{p\left({x}_T\right)}{q\left({x}_T \mid {x}_{T-1}\right)}\right]+\sum_{t=1}^{T-1} \mathbb{E}_{q\left({x}_{1: T} \mid {x}_0\right)}\left[\log \frac{p_{{\theta}}\left({x}_t \mid {x}_{t+1}\right)}{q\left({x}_t \mid {x}_{t-1}\right)}\right] \tag{43}
\\
& =\mathbb{E}_{q\left({x}_1 \mid {x}_0\right)}\left[\log p_{{\theta}}\left({x}_0 \mid {x}_1\right)\right]+\mathbb{E}_{q\left({x}_{T-1}, {x}_T \mid {x}_0\right)}\left[\log \frac{p\left({x}_T\right)}{q\left({x}_T \mid {x}_{T-1}\right)}\right]+\sum_{t=1}^{T-1} \mathbb{E}_{q\left({x}_{t-1}, {x}_t, {x}_{t+1} \mid {x}_0\right)}\left[\log \frac{p_{{\theta}}\left({x}_t \mid {x}_{t+1}\right)}{q\left({x}_t \mid {x}_{t-1}\right)}\right] \tag{44}
\\
& =\underbrace{\mathbb{E}_{q\left({x}_1 \mid {x}_0\right)}\left[\log p_\theta\left({x}_0 \mid {x}_1\right)\right]}_{\text {reconstruction term }}-\underbrace{\mathbb{E}_{q\left({x}_{T-1} \mid {x}_0\right)}\left[D_{\mathrm{KL}}\left(q\left({x}_T \mid {x}_{T-1}\right) \| p\left({x}_T\right)\right)\right]}_{\text {prior matching term }} \tag{45}-\sum_{t=1}^{T-1} \underbrace{\mathbb{E}_{q\left({x}_{t-1}, {x}_{t+1} \mid {x}_0\right)}\left[D_{\mathrm{KL}}\left(q\left({x}_t \mid {x}_{t-1}\right) \| p_\theta\left({x}_t \mid {x}_{t+1}\right)\right)\right]}_{\text {consistency term }} \\
\end{align}$$

 

식이 좀 많지만 그렇게 어려운 수식은 아닙니다. (저번 포스트의 자료를 참고하면 충분히 풀수 있는 수식입니다.)

Eq 34는 Eq 1과 같은 수식입니다. ($p(x)=\int{p(x,z)}dz$)  즉 $x_0$에 대한 나머지 t들의 확률에 대한 합과 같다는 의미입니다.

Eq 35~37은 ELBO의 수식유도와 같으며 Eq38은 Eq 30, Eq 32를 수식에 대입하였습니다.

Eq 39는 수식의 분자는 $\prod$의 t=1번째 값, 분모는 $\prod$의 t=T번째 값을 추출하여 곱으로 나타내었으며, Eq 40은 분모와 분자를 $\prod^{T-1}_{t=1}$로 맞춰주기 위해 분자를 $\prod^T_{t=2} = \prod^{T-1}_{t=1}$로 바꿉니다.

Eq 41은 log의 특성에 따라 $\prod$를 바깥으로 빼고 Eq 42도 마찬가지로 곱을 더하기로 바꾸어 나열하였습니다. 이때 $\prod$는 곱집합이기 때문에 log 바깥으로 나가면서 $\sum$이 됩니다.

Eq 43은 단순히 $\sum$을 Expectation 바깥으로 보냈습니다.

Eq 44는 Expectation에 대한 확률밀도함수를 각 식에 맞게 축소합니다. (근데 이렇게 해도 되나?)

마지막으로 Eq45는 2, 3번항을 2개의 KL Divergence로 바꾸는데 이 수식을 한번 살펴보겠습니다.

 

$$\begin{aligned}
\mathbb{E}_{q(x_{1:T}|x_0)} \bigg[ \log{{p(x_T)}\over{q(x_T|x_{T-1})}} \bigg] &= \int{q(x_{1:T}|x_0) \bigg(\log{{p(x_T)}\over{q(x_T|x_{T-1})}}\bigg)} dx_{1:T}
\\
&=\int{\prod_{t=1}^{T}q(x_{t}|x_{t-1}) \bigg(\log{{p(x_T)}\over{q(x_T|x_{T-1})}}\bigg)} dx_{1:T}
\\
&=\int{\prod_{t=1}^{T-1}q(x_{t}|x_{t-1}) \bigg(q(x_T|x_{T-1}) \log{{p(x_T)}\over{q(x_T|x_{T-1})}}\bigg)} dx_{1:T}
\\
&=-\int{\prod_{t=1}^{T-1}q(x_{t}|x_{t-1}) \bigg(q(x_T|x_{T-1}) \log{{q(x_T|x_{T-1})}\over{p(x_T)}}\bigg)} dx_{1:T}
\\
&=-\int{\int{\prod_{t=1}^{T-1}q(x_{t}|x_{t-1}) \bigg(q(x_T|x_{T-1}) \log{{q(x_T|x_{T-1})}\over{p(x_T)}}\bigg)}} dx_{T} \ dx_{1:T-1}
\\
&=-\int{\bigg( \prod_{t=1}^{T-1}q(x_{t}|x_{t-1})\int{\bigg(q(x_T|x_{T-1}) \log{{q(x_T|x_{T-1})}\over{p(x_T)}}\bigg)}dx_{T}\bigg)} dx_{1:T-1}
\\
&=-\int{\bigg( \prod_{t=1}^{T-1}q(x_{t}|x_{t-1})D_{KL}(q(x_T|x_{T-1})||p(x_T)) \bigg)}dx_{1:T-1}
\\
&=-\int{\bigg( q(x_{1:T-1}|x_0)D_{KL}(q(x_T|x_{T-1})||p(x_T)) \bigg)}dx_{1:T-1}
\\
&=-\mathbb{E}_{q(x_{1:T-1}|x_0)}[ D_{KL}(q(x_T|x_{T-1})||p(x_T))]
\end{aligned}$$

 

결합 확률 함수의 기댓값

 

결합 확률함수 $f(x,y)$일 때, 확률 변수 $g(x,y)$의 기댓값은 다음과 같습니다:

 

• 이산 확률 변수:

$$\mathbb{E}_{g(x,y)}[f(x,y)]=\sum_x \sum_y g(x,y)f(x,y)$$

 

• 연속 확률 변수:

$$\mathbb{E}_{g(x,y)}[f(x,y)]=\int \int g(x,y)f(x,y)dxdy$$

 

---

 

Eq 43에서 다음과 같은 식을 도출하고 Eq 44처럼 항을 바꿔주면

$$-\mathbb{E}_{q(x_{1:T-1}|x_0)}[ D_{KL}(q(x_T|x_{T-1})||p(x_T))] = -\mathbb{E}_{q(x_{T-1}|x_0)}[ D_{KL}(q(x_T|x_{T-1})||p(x_T))]$$

 

위의 식들은 2번째 항에 대한 식으로 3번째 항 또한 같은 방식으로 식을 정리할 수 있습니다. 3번째항에 대한 식은 생략하겠습니다.

 

이때 Eq 45의 ELBO 식은 다음의 3가지의 Component로 해석할 수 있습니다:

 

• $reconstruction \ term$은 모델의 1계층의 reverse process에서 $x_0$  sample의 log 확률을 예측합니다. (VAE에서 z->x의 재구성에 대한 확률 분포를 예측하는 것과 같은 개념)

 

• $prior \ matching \ term$은 T 단계의 마지막 잠재변수가 Gaussian Prior과 일치할때 최소화합니다. 이때 $q(x_T|x_{T-1})$은 식에서 보는것과 같이 최적화할 파라미터가 딱히 없기 때문에 (VAE는 Encoder를 통해 잠재변수에 대한 파라미터 $\phi$가 있었지만, VDM의 $q(x_T|x_{T-1})$은 Gaussian Noise를 추가하는 개념이기 때문에 따로 파라미터가 없습니다.) 사실상 값이 0에 수렴합니다.

 

• $consistency \ term$은 forward process와 reverse process가 양방향에서 이루어질때 서로 일관성 있게 분포를 생성하는지를 비교합니다. 즉 $x_{t-1} -> x_{t}$로 Gaussian Noise를 추가한 것과 $x_{t+1} -> x_{t}$로 noise를 제거한 것이 최대한 같도록 학습합니다. 모든 time step에 대해 계산해야되기 때문에 Cost가 많이 듭니다.

 

consistency term은 분홍색 forward process와 연두색 reverse process의 분포가 같을 때 최소화됩니다.

위의 식에 대한 Monte Carlo 추정치를 이용하여 ELBO를 근사할 수 있지만, 한가지 문제점이 있습니다. 3번째 Component에 의해 diffusion process는 매 time step마다 $x_{t-1}$과 $x_{t+1}$의 기대값을 이용해 $x_{t}$를 근사하기 때문에 2가지 랜덤 변수를 가진 추정치는 1가지 랜덤변수를 가진 값보다 더 분산이 높을 수 밖에 없습니다(high variance). (분산이 높다는 것은 곧 과대적합을 의미합니다.) 즉 T의 값이 커질수록 점점 높은 분산을 가지게 됩니다.

 

그렇다면 이번에는 1개의 랜덤변수만을 사용하여 ELBO를 유도해봅시다. Markov Chain에 의해 $q(x_{t}|x_{t-1}) = q(x_{t}|x_{t-1}, x_0)$를 사용해서 식을 구성해봅시다. 이 Markov Chain은 Bayes Rule에 의해 다음과 같이 정의할 수 있습니다.

 

$$q(x_t|x_{t-1}, x_0) = {{q(x_{t-1}|x_t, x_0)q(x_t|x_0)}\over{q(x_{t-1}|x_0)}} \tag{46}$$

 

 

Bayesian Rules (Bayes' Theorem)

 

Bayesian Rules은 통계에서 가장 중요한 법칙중 하나입니다. Bayesian Rules, Bayes Rules, Bayesian Theorem, Bayes Theorem 등 여러가지로 불립니다. 각 사건에 대한 확률 P(A), P(B)에 대한 조건부 확률 P(A|B)는 다음과 같습니다:

 

$$P(A|B) = {{P(B|A)P(A)} \over {P(B)}}$$

 

위를 해석하면, 사건 A의 확률이 사건 B에 의해 갱신된 확률을 의미합니다. 

수식은 결합 확률 $P(A,B) = P(B,A)$이기 때문에 이를 통해 도출할수 있는 수식입니다. Bayesian Rules을 좀더 확장하면 다음과 같은 수식도 얻을 수 있습니다.

 

$$P(A|B,C) = {{P(C|A,B)P(A|B)}\over{P(C|B)}}$$

 

사건A의 확률이 사건 B와 C에의해 갱신된 확률입니다. ($P(A|B,C) = P(A|(B \cap C))$)

 

6.6 베이즈 정리 — 데이터 사이언스 스쿨

---

 

위 Eq 46을 적용해 Eq37의 ELBO를 다시 유도해봅시다.

 

$$
\require{cancel}\begin{align}
\log p({x}) &\geq \mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_{0: T})}{q({x}_{1: T} \mid {x}_0)}\bigg] \tag{47}
\\
&=\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_T) \prod_{t=1}^T p_{{\theta}}({x}_{t-1} \mid {x}_t)}{\prod_{t=1}^T q({x}_t \mid {x}_{t-1})}\bigg] \tag{48}
\\
&=\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_T) p_{{\theta}}({x}_0 \mid {x}_1) \prod_{t=2}^T p_{{\theta}}({x}_{t-1} \mid {x}_t)}{q({x}_1 \mid {x}_0) \prod_{t=2}^T q({x}_t \mid {x}_{t-1})}\bigg] \tag{49}
\\
&=\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_T) p_{{\theta}}({x}_0 \mid {x}_1) \prod_{t=2}^T p_{{\theta}}({x}_{t-1} \mid {x}_t)}{q({x}_1 \mid {x}_0) \prod_{t=2}^T q({x}_t \mid {x}_{t-1}, {x}_0)}\bigg] \tag{50}
\\
&=\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p_{{\theta}}({x}_T) p_{{\theta}}({x}_0 \mid {x}_1)}{q({x}_1 \mid {x}_0)}+\log \prod_{t=2}^T \frac{p_{{\theta}}({x}_{t-1} \mid {x}_t)}{q({x}_t \mid {x}_{t-1}, {x}_0)}\bigg] \tag{51}
\\
&=\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_T) p_{{\theta}}({x}_0 \mid {x}_1)}{q({x}_1 \mid {x}_0)}+\log \prod_{t=2}^T \frac{p_{{\theta}}({x}_{t-1} \mid {x}_t)}{\frac{q({x}_{t-1} \mid {x}_t, {x}_0) q({x}_t \mid {x}_0)}{q({x}_{t-1} \mid {x}_0)}}\bigg] \tag{52}
\\
& =\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_T) p_{{\theta}}({x}_0 \mid {x}_1)}{q({x}_1 \mid {x}_0)}+\log \prod_{t=2}^T \frac{p_{{\theta}}({x}_{t-1} \mid {x}_t)}{\frac{q({x}_{t-1} \mid {x}_t, {x}_0) \cancel{q({x}_t \mid {x}_0)}}{\cancel{q({x}_{t-1} \mid {x}_0)}}}\bigg] \tag{53}
\\
& =\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_T) p_{{\theta}}({x}_0 \mid {x}_1)}{\cancel{q({x}_1 \mid {x}_0)}}+\log \frac{\cancel{q({x}_1 \mid {x}_0)}}{q({x}_T \mid {x}_0)}+\log \prod_{t=2}^T \frac{p_{{\theta}}({x}_{t-1} \mid {x}_t)}{q({x}_{t-1} \mid {x}_t, {x}_0)}\bigg] \tag{54}
\\
& =\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_T) p_{{\theta}}({x}_0 \mid {x}_1)}{q({x}_T \mid {x}_0)}+\sum_{t=2}^T \log \frac{p_{{\theta}}({x}_{t-1} \mid {x}_t)}{q({x}_{t-1} \mid {x}_t, {x}_0)}\bigg] \tag{55}
\\
& =\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}[\log p_{{\theta}}({x}_0 \mid {x}_1)]+\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_T)}{q({x}_T \mid {x}_0)}\bigg]+\sum_{t=2}^T \mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p_{{\theta}}({x}_{t-1} \mid {x}_t)}{q({x}_{t-1} \mid {x}_t, {x}_0)}\bigg] \tag{56}
\\
& =\mathbb{E}_{q({x}_1 \mid {x}_0)}[\log p_{{\theta}}({x}_0 \mid {x}_1)]+\mathbb{E}_{q({x}_T \mid {x}_0)}\bigg[\log \frac{p({x}_T)}{q({x}_T \mid {x}_0)}\bigg]+\sum_{t=2}^T \mathbb{E}_{q({x}_t, {x}_{t-1} \mid {x}_0)}\bigg[\log \frac{p_{{\theta}}({x}_{t-1} \mid {x}_t)}{q({x}_{t-1} \mid {x}_t, {x}_0)}\bigg] \tag{57}
\\
& =\underbrace{\mathbb{E}_{q({x}_1 \mid {x}_0)}[\log p_{{\theta}}({x}_0 \mid {x}_1)]}_{\text {reconstruction term }}-\underbrace{D_{\mathrm{KL}}(q({x}_T \mid {x}_0) \| p({x}_T))}_{\text {prior matching term }}-\sum_{t=2}^T \underbrace{\mathbb{E}_{q({x}_t \mid {x}_0)}\bigg[D_{\mathrm{KL}}(q({x}_{t-1} \mid {x}_t, {x}_0) \| p_{{\theta}}({x}_{t-1} \mid {x}_t))\bigg]}_{\text {denoising matching term }} \tag{58}
\end{align}
$$

 

앞서 나온 수식들과 형태나 방법은 유사합니다. (Eq 47~52는 앞서나온 수식의 대입과 곱집합을 분리하는 process입니다.)

Eq 53부터 조금 달라지는데 Eq 53~55까지의 유도는 다음과 같습니다.

 

$$
\begin{aligned}
\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log& \frac{p({x}_T) p_{{\theta}}({x}_0 \mid {x}_1)}{q({x}_1 \mid {x}_0)}+\log \prod_{t=2}^T \frac{p_{{\theta}}({x}_{t-1} \mid {x}_t)}{\frac{q({x}_{t-1} \mid {x}_t, {x}_0) q({x}_t \mid {x}_0)}{q({x}_{t-1} \mid {x}_0)}}\bigg]
\\
&=\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_T) p_{{\theta}}({x}_0 \mid {x}_1)}{q({x}_1 \mid {x}_0)}+ \log \prod^T_{t=2}{\frac{p_{{\theta}}({x}_{t-1} \mid {x}_t)q({x}_{t-1} \mid {x}_0)}{q({x}_{t-1} \mid {x}_t, {x}_0) q({x}_t \mid {x}_0)}} \bigg]
\\
&=\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_T) p_{{\theta}}({x}_0 \mid {x}_1)}{q({x}_1 \mid {x}_0)} +\log{\prod^T_{t=2}\frac{q({x}_{t-1}\mid{x}_0)}{q({x}_t \mid {x}_0)}}+ \log \prod^T_{t=2}{\frac{p_{{\theta}}({x}_{t-1} \mid {x}_t)}{q({x}_{t-1} \mid {x}_t, {x}_0)}} \bigg]
\\
&=\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_T) p_{{\theta}}({x}_0 \mid {x}_1)}{q({x}_1 \mid {x}_0)} +
\log{\frac{q({x}_{1}\mid{x}_0)}{q({x}_2 \mid {x}_0)}\frac{q({x}_{2}\mid{x}_0)}{q({x}_3 \mid {x}_0)}...\frac{q({x}_{T-1}\mid{x}_0)}{q({x}_T \mid {x}_0)}} +
\log \prod^T_{t=2}{\frac{p_{{\theta}}({x}_{t-1} \mid {x}_t)}{q({x}_{t-1} \mid {x}_t, {x}_0)}} \bigg]
\\
& =\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_T) p_{{\theta}}({x}_0 \mid {x}_1)}{q({x}_1 \mid {x}_0)}+\log \frac{q({x}_1 \mid {x}_0)}{q({x}_T \mid {x}_0)}+\log \prod_{t=2}^T \frac{p_{{\theta}}({x}_{t-1} \mid {x}_t)}{q({x}_{t-1} \mid {x}_t, {x}_0)}\bigg]
\\
& =\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_T) p_{{\theta}}({x}_0 \mid {x}_1)q({x}_1 \mid {x}_0)}{q({x}_1 \mid {x}_0)q({x}_T \mid {x}_0)}+\log \prod_{t=2}^T \frac{p_{{\theta}}({x}_{t-1} \mid {x}_t)}{q({x}_{t-1} \mid {x}_t, {x}_0)}\bigg]
\\
& =\mathbb{E}_{q({x}_{1: T} \mid {x}_0)}\bigg[\log \frac{p({x}_T) p_{{\theta}}({x}_0 \mid {x}_1)}{q({x}_T \mid {x}_0)}+\sum_{t=2}^T \log \frac{p_{{\theta}}({x}_{t-1} \mid {x}_t)}{q({x}_{t-1} \mid {x}_t, {x}_0)}\bigg]
\end{aligned}
$$

이후 Eq 56~58까지는 이전 ELBO 유도와 같은 방식으로 진행됨으로 생략하겠습니다.

 

$$\underbrace{\mathbb{E}_{q({x}_1 \mid {x}_0)}[\log p_{{\theta}}({x}_0 \mid {x}_1)]}_{\text {reconstruction term }}-\underbrace{D_{\mathrm{KL}}(q({x}_T \mid {x}_0) \| p({x}_T))}_{\text {prior matching term }}-\sum_{t=2}^T \underbrace{\mathbb{E}_{q({x}_t \mid {x}_0)}\bigg[D_{\mathrm{KL}}(q({x}_{t-1} \mid {x}_t, {x}_0) \| p_{{\theta}}({x}_{t-1} \mid {x}_t))\bigg]}_{\text {denoising matching term }}$$

 

이 수식은 2가지 랜덤변수를 1가지로 줄임으로써 분산이 줄어든 ELBO를 도출 할 수 있습니다. 각각의 Component는 다음과 같은 의미를 지닙니다:

 

• $reconstruction \ term$: 1번째 계층에대한 log 확률을 예측하는, 이전의 reconstruction term과 동일합니다. 이수식은 Monte Carlo 근사를 통해 정리할 수 있습니다.

 

• $prior \ matching \ term$: time step T에서 최종적으로 만들어지는 Gaussian Noise($q(x_T|x_0)$)와 실제 $p(x_T)$의 차이를 최소화 합니다. 다만 이때 두 값은 거의 유사하며 q는 파라미터가 없이 단순히 정규분포를 계속 쌓아올리는 구조이기 때문에 거의 0에 수렴하는 값입니다. (무시해도 되는 상수취급되는 값입니다.)

 

• $denoising \ matching \ term$: Ground-Truth denoising step인 $q(x_{t-1}|x_t,x_0)$와 근사분포 $p_\theta()x_{t-1}|x_t)$의 차이를 최소화 합니다. 즉 기존처럼 t-2 step에서 t-1 step의 gaussian 분포를 얻는것이 아니라, t step과 $x_0$를 이용해 t-1 step의 gaussian 분포를 얻어 reverse process를 통해 근사한 값과 비교를 하는 것입니다.

 

3번에서 어떻게 $q(x_{t-1}|x_t,x_0)$를 구할수 있을지 궁금하실수 있습니다. 실질적으로 Gaussian Noise가 있는 이미지를, 없는 이미지로 바꿔주는 수식이기 때문입니다. 다만 원본 이미지와 Gaussian Noise가 있는 이미지를 조건으로 주기 때문에 Bayes rule에 의해 수식의 유도가 가능합니다. 

 

$$q(x_t|x_{t-1}, x_0) = {{q(x_{t-1}|x_t, x_0)q(x_t|x_0)}\over{q(x_{t-1}|x_0)}}$$

즉
$$q(x_{t-1}|x_t, x_0)=\frac{q(x_t|x_{t-1}, x_0)q(x_{t-1}|x_0)}{q(x_t|x_0)}$$

 

 

즉 forward process에 대한 값 만으로도 수식의 유도가 가능합니다.

 

유도한 두 ELBO 식은 Markov Chain 이외에는 기존 ELBO와 크게 바뀐 부분이 없습니다. (많이 바뀐거 같은데..?) 그렇기 때문에 이 수식은 임의의 HVAE와 같을 뿐 아니라 T=1일 때 VAE의 ELBO와 완전히 동일합니다.

 

저자는 reconstruction term의 cost가 매우 높다고 하는데, 제생각에는 denoising matching term을 잘못 쓴것 같습니다. 다만 Encoder까지 따로 학습해야하는 HVAE와 다르게 VDM은 Decoder part(reverse process)만 학습하면 되기 때문에 훨씬 안정적이라고 합니다. 

 

이전에 다룬것 처럼 계층을 지나면서 Gaussian Noise를 추가할 때 분포는 $q(x_t|x_{t-1}, x_0) = q(x_{t}|x_{t-1}) = \mathcal{N}(x_{t};\sqrt{\alpha_{t}}x_{t-1}, (1-\alpha_{t})I)$라는 것을 알고 있습니다. 따라서 $q(x_{t-1}|x_t, x_0)=\frac{q(x_t|x_{t-1}, x_0)q(x_{t-1}|x_0)}{q(x_t|x_0)}$에서 우리는 $q(x_t|x_0)$와 $q(x_{t-1}|x_0)$의 분포를 알면 수식을 구할수 있습니다.

Sample $x_t~q(x_t|x_{t-1})$는 reparameterization trick에 의해 다음과 같이 나타낼수 있습니다:

 

$$x_t = \sqrt{\alpha_t}x_{t-1} + \sqrt{1-\alpha_t}\epsilon  \ \rm{with} \ \epsilon \sim  \mathcal{N}(\epsilon;0,I) \tag{59}$$

 

그리고 $x_{t-1}~q(x_{t-1}|x_{t-2})$에 대해서도 적용해보면 다음과 같습니다:

 

$$x_{t-1} = \sqrt{\alpha_{t-1}}x_{t-2} + \sqrt{1-\alpha_{t-1}}\epsilon \ \rm{with} \ \epsilon \sim \mathcal{N}(\epsilon;0,I) \tag{60}$$

 

이때 $q(x_t|x_0)$는 reparameterization trick을 통해 재귀적으로 정의할 수 있습니다. 2T개의 Random 노이즈 변수 $\{ \epsilon^*_t, \epsilon_t \}^T_{t=0} \overset{iid}{~} \mathcal{N}(\epsilon;0,I)$에 대한 $x|t ~ q(x_t|x_0)$ sample을 재귀적으로 나타내 봅시다.

 

$$
\begin{align}
\boldsymbol{x}_t & =\sqrt{\alpha_t} \boldsymbol{x}_{t-1}+\sqrt{1-\alpha_t} \boldsymbol{\epsilon}_{t-1}^* \tag{61}
\\
& =\sqrt{\alpha_t}\left(\sqrt{\alpha_{t-1}} \boldsymbol{x}_{t-2}+\sqrt{1-\alpha_{t-1}} \epsilon_{t-2}^*\right)+\sqrt{1-\alpha_t} \boldsymbol{\epsilon}_{t-1}^* \tag{62}
\\
& =\sqrt{\alpha_t \alpha_{t-1}} \boldsymbol{x}_{t-2}+\sqrt{\alpha_t-\alpha_t \alpha_{t-1}} \boldsymbol{\epsilon}_{t-2}^*+\sqrt{1-\alpha_t} \boldsymbol{\epsilon}_{t-1}^* \tag{63}
\\
& =\sqrt{\alpha_t \alpha_{t-1}} \boldsymbol{x}_{t-2}+\sqrt{{\sqrt{\alpha_t-\alpha_t \alpha_{t-1}}}^2+\sqrt{1-\alpha_t} \boldsymbol{\epsilon}_{t-2}} \tag{64}
\\
& =\sqrt{\alpha_t \alpha_{t-1}} \boldsymbol{x}_{t-2}+\sqrt{\alpha_t-\alpha_t \alpha_{t-1}+1-\alpha_t} \boldsymbol{\epsilon}_{t-2} \tag{65}
\\
& =\sqrt{\alpha_t \alpha_{t-1}} \boldsymbol{x}_{t-2}+\sqrt{1-\alpha_t \alpha_{t-1}} \boldsymbol{\epsilon}_{t-2} \tag{66}
\\
& =\ldots \tag{67}
\\
& =\sqrt{\prod_{i=1}^t \alpha_i \boldsymbol{x}_0+\sqrt{1-\prod_{i=1}^t \alpha_i \boldsymbol{\epsilon}_0}} \tag{68}
\\
& =\sqrt{\bar{\alpha}_t} \boldsymbol{x}_0+\sqrt{1-\bar{\alpha}_t} \boldsymbol{\epsilon}_0 \tag{69}
\\
& \sim \mathcal{N}\left(\boldsymbol{x}_t ; \sqrt{\bar{\alpha}_t} \boldsymbol{x}_0,\left(1-\bar{\alpha}_t\right) \mathbf{I}\right) \tag{70}
\end{align}
$$

 

Eq 62는 재귀적으로 봤을 때, $x_{t-1} = \sqrt{\alpha_{t-1}}x_{t-2} + \sqrt{1-\alpha_{t-1}}\epsilon^*_{t-2}$이기 때문에 이를 대입합니다. 여기서 Eq 64가 이해가 잘 안될 수 있는데, $\sqrt{\alpha_t-\alpha_t \alpha_{t-1}} \boldsymbol{\epsilon}_{t-2}^*$는 $\mathcal{N}(0, (\alpha_t - \alpha_t\alpha_{t-1})I)$, $\sqrt{1-\alpha_t} \boldsymbol{\epsilon}_{t-1}^*$는 $\mathcal{N}(0, (1-\alpha_t)I)$로 나타낼 수 있습니다. 즉 정규분포의 가법성에 따라 $\mathcal{N}\left(\mathbf{0},\left(1-\alpha_t+\alpha_t-\alpha_t \alpha_{t-1}\right) \mathbf{I}\right)=\mathcal{N}\left(\mathbf{0},\left(1-\alpha_t \alpha_{t-1}\right) \mathbf{I}\right)$이므로 Eq 67이 되게 됩니다.  이때 $\epsilon$은 모두 $\mathcal{N}(\epsilon;0, I)$를 따르는 정규분포이므로 모두 같은 값이라 통일 한 것 같습니다. 그리고 $\bar{\alpha}_t$는 모든 timestep t에 대한 $\alpha_t$의 곱집합입니다.

 

 

정규분포의 가법성

독립된 두 정규분포 $X \sim \mathcal{N}(\mu_1, \sigma^2_1)$ $Y \sim \mathcal{N}(\mu_2, \sigma^2_2)$에 대하여 다음과 같이 새로운 정규분포를 만들 수 있습니다.

 

$X+Y \sim \mathcal{N}(\mu_1+\mu_2, \sigma^2_1+\sigma^2_2)$

$X-Y \sim \mathcal{N}(\mu_1-\mu_2, \sigma^2_1-\sigma^2_2)$

---

$q(x_t|x_0)$의 가우시안형태를 유도했으니, $q(x_{t-1}|x_0)$를 통해 $q(x_{t-1}|x_t, x_0)$를 유도할 수 있습니다.



$$
\begin{align}
q(x_{t-1} \mid x_t, x_0) &= \frac{q(x_t \mid x_{t-1}, x_0) q(x_{t-1} \mid x_0)}{q(x_t \mid x_0)} \tag{71}
\\
& =\frac{\mathcal{N}(x_t ; \sqrt{\alpha_t} x_{t-1},(1-\alpha_t) \mathbf{I}) \mathcal{N}(x_{t-1} ; \sqrt{\bar{\alpha}_{t-1}} x_0,(1-\bar{\alpha}_{t-1}) \mathbf{I})}{\mathcal{N}(x_t ; \sqrt{\bar{\alpha}_t} x_0,(1-\bar{\alpha}_t) \mathbf{I})} \tag{72}
\\
& \propto \exp \bigg\{-\bigg[\frac{(x_t-\sqrt{\alpha_t} x_{t-1})^2}{2(1-\alpha_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)}\bigg]\bigg\} \tag{73}
\\
& =\exp \bigg\{-\frac{1}{2}\bigg[\frac{(x_t-\sqrt{\alpha_t} x_{t-1})^2}{1-\alpha_t}+\frac{(x_{t-1}-\sqrt{\bar{\alpha}_{t-1}} x_0)^2}{1-\bar{\alpha}_{t-1}}-\frac{(x_t-\sqrt{\bar{\alpha}_t} x_0)^2}{1-\bar{\alpha}_t}\bigg]\bigg\} \tag{74}
\\
& =\exp \bigg\{-\frac{1}{2}\bigg[\frac{(-2 \sqrt{\alpha_t} x_t x_{t-1}+\alpha_t x_{t-1}^2)}{1-\alpha_t}+\frac{(x_{t-1}^2-2 \sqrt{\bar{\alpha}_{t-1}} x_{t-1} x_0)}{1-\bar{\alpha}_{t-1}}+C(x_t, x_0)\bigg]\bigg\} \tag{75}
\\
& \propto \exp \bigg\{-\frac{1}{2}\bigg[-\frac{2 \sqrt{\alpha_t} x_t x_{t-1}}{1-\alpha_t}+\frac{\alpha_t x_{t-1}^2}{1-\alpha_t}+\frac{x_{t-1}^2}{1-\bar{\alpha}_{t-1}}-\frac{2 \sqrt{\bar{\alpha}_{t-1}} x_{t-1} x_0}{1-\bar{\alpha}_{t-1}}\bigg]\bigg\} \tag{76}
\\
& =\exp \bigg\{-\frac{1}{2}\bigg[(\frac{\alpha_t}{1-\alpha_t}+\frac{1}{1-\bar{\alpha}_{t-1}}) x_{t-1}^2-2(\frac{\sqrt{\alpha_t} x_t}{1-\alpha_t}+\frac{\sqrt{\bar{\alpha}_{t-1}} x_0}{1-\bar{\alpha}_{t-1}}) x_{t-1}\bigg]\bigg\} \tag{77}
\\
& =\exp \bigg\{-\frac{1}{2}\bigg[\frac{\alpha_t(1-\bar{\alpha}_{t-1})+1-\alpha_t}{(1-\alpha_t)(1-\bar{\alpha}_{t-1})} x_{t-1}^2-2(\frac{\sqrt{\alpha_t} x_t}{1-\alpha_t}+\frac{\sqrt{\bar{\alpha}_{t-1}} x_0}{1-\bar{\alpha}_{t-1}}) x_{t-1}\bigg]\bigg\} \tag{78}
\\
& =\exp \bigg\{-\frac{1}{2}\bigg[\frac{\alpha_t-\bar{\alpha}_t+1-\alpha_t}{(1-\alpha_t)(1-\bar{\alpha}_{t-1})} x_{t-1}^2-2(\frac{\sqrt{\alpha_t} x_t}{1-\alpha_t}+\frac{\sqrt{\bar{\alpha}_{t-1}} x_0}{1-\bar{\alpha}_{t-1}}) x_{t-1}\bigg]\bigg\} \tag{79}
\\
& =\exp \bigg\{-\frac{1}{2}\bigg[\frac{1-\bar{\alpha}_t}{(1-\alpha_t)(1-\bar{\alpha}_{t-1})} x_{t-1}^2-2(\frac{\sqrt{\alpha_t} x_t}{1-\alpha_t}+\frac{\sqrt{\bar{\alpha}_{t-1}} x_0}{1-\bar{\alpha}_{t-1}}) x_{t-1}\bigg]\bigg\} \tag{80}
\\
& =\exp \bigg\{-\frac{1}{2}(\frac{1-\bar{\alpha}_t}{(1-\alpha_t)(1-\bar{\alpha}_{t-1})})\bigg[x_{t-1}^2-2 \frac{(\frac{\sqrt{\alpha_t} x_t}{1-\alpha_t}+\frac{\sqrt{\bar{\alpha}_{t-1}} x_0}{1-\bar{\alpha}_{t-1}})}{\frac{1-\bar{\alpha}_t}{(1-\alpha_t)(1-\bar{\alpha}_{t-1})}} x_{t-1}\bigg]\bigg\} \tag{81}
\\
& =\exp \bigg\{-\frac{1}{2}(\frac{1-\bar{\alpha}_t}{(1-\alpha_t)(1-\bar{\alpha}_{t-1})})\bigg[x_{t-1}^2-2 \frac{(\frac{\sqrt{\alpha_t} x_t}{1-\alpha_t}+\frac{\sqrt{\bar{\alpha}_{t-1}} x_0}{1-\bar{\alpha}_{t-1}})(1-\alpha_t)(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t} x_{t-1}\bigg]\bigg\} \tag{82}
\\
& =\exp \bigg\{-\frac{1}{2}(\frac{1}{\frac{(1-\alpha_t)(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}})\bigg[x_{t-1}^2-2 \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) x_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) x_0}{1-\bar{\alpha}_t} x_{t-1}\bigg]\bigg\} \tag{83}
\\
& \propto \mathcal{N}(x_{t-1} ; \underbrace{\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) x_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) x_0}{1-\bar{\alpha}_t}}_{\mu_q(x_t, x_0)}, \underbrace{\frac{(1-\alpha_t)(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t} \mathbf{I})}_{\mathbf{\Sigma}_q(t)} \tag{84}
\end{align}
$$



앞서 나왔던 것처럼 정규 분포를 평균과 분산에 대한 식(그래프)으로 나타내면 다음과 같습니다.

$$\mathcal{N}(x;\mu,\sigma^2) = \frac{1}{\sqrt{2\pi\sigma^2}}\exp{(-\frac{(x-\mu)^2}{2\sigma^2})}$$



즉 분포는 $ \exp{(-\frac{(x-\mu)^2}{2\sigma^2})}$에 비례하므로 Eq 73과 같이 exp에 대한 식으로 바꿀 수 있습니다. 또한 $x_{t-1}$에 대한 식으로 나타낼 경우 $x_t, x_0$로 이루어진 식을 상수 취급 $(C(x_t, x_0))$할 경우 Eq 76과 같은 비례 식을 얻을 수 있습니다.



마지막으로 Eq 83은 위의 정규분포 식의 $\exp{(-\frac{(x-\mu)^2}{2\sigma^2})}$에 대입했을 때, 



$$
\begin{aligned}
\mathcal{N}\bigg(x_{t-1} ; & \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) x_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) x_0}{1-\bar{\alpha}_t}, \frac{(1-\alpha_t)(1-\bar{\alpha}_{t-1}))}{1-\bar{\alpha_{t}}}I \bigg) 
\\
&\propto \exp{\bigg( -\frac{1}{2}(\frac{1}{\frac{(1-\alpha_t)(1-\bar{\alpha}_{t-1}))}{1-\bar{\alpha_{t}}}}) \bigg[ x_{t-1} - \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) x_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) x_0}{1-\bar{\alpha}_t} \bigg]^2 \bigg)}
\\
&= \exp{\bigg( -\frac{1}{2}(\frac{1}{\frac{(1-\alpha_t)(1-\bar{\alpha}_{t-1}))}{1-\bar{\alpha_{t}}}}) \bigg[ x^2_{t-1} -2 \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) x_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) x_0}{1-\bar{\alpha}_t}x_{t-1} + \bigg(\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) x_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) x_0}{1-\bar{\alpha}_t}\bigg)^2 \bigg] \bigg)}
\\
&\propto \exp{\bigg( -\frac{1}{2}(\frac{1}{\frac{(1-\alpha_t)(1-\bar{\alpha}_{t-1}))}{1-\bar{\alpha_{t}}}}) \bigg[ x^2_{t-1} -2 \frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) x_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) x_0}{1-\bar{\alpha}_t}x_{t-1} \bigg] \bigg)}
\end{aligned}
$$



위의 식을 통해 Eq 84를 얻을 수 있습니다. 식의 $\mu_q(x_t, x_0)$는 $x_t$와 $x_0$를 변수로 갖는 함수이며, $\Sigma_q(t)$는 $\alpha$를 계수로 가지는 함수 입니다. 이때 계수 $\alpha$는 매 timestep t에 따라 값이 고정되고, 하이퍼 파라미터로  값이 설정되거나, 모델링 네트워크의 출력을 통해 특정할 수 있습니다. 또한 Eq84의 $\Sigma_q(t) = \sigma^2_q(t)I$라 할때 $ \sigma^2_q(t)$의 값은 다음과 같습니다.



$$\sigma^2_q(t) = \frac{(1-\alpha_t)(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}} \tag{85}$$



Denoising Transition Step인 $p_\theta(x_{t-1}|x_t)$가 gt denoising step인 $q(x_{t-1}|x_t,x_0)$를 따르도록 학습합니다.

$\mathcal{N}(x_{t-1} ; \underbrace{\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) x_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) x_0}{1-\bar{\alpha}_t}}_{\mu_q(x_t, x_0)}, \underbrace{\frac{(1-\alpha_t)(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t} \mathbf{I})}_{\mathbf{\Sigma}_q(t)}$에서 $\alpha$는 각 time step마다 고정되므로, $\Sigma_q(t)$는 쉽게 알수 있으며, $p_\theta(x_{t-1}|x_t)$는 $x_0$가 조건으로 들어가 있지 않기 때문에 분산을 $x_t$에 대한 함수 $\mu_\theta(x_t,t)$로 나타낼 수 있습니다.



$p_\theta(x_{t-1}|x_t) = \mathcal{N}(x_{t-1};\mu_\theta(x_t, t), \Sigma_q(t))$



이제 $ p_\theta(x_{t-1}|x_t) $와 $ q(x_{t-1}|x_t, x_0)$의 차이를 최소화 하기 위해 KL Divergence를 계산해봅시다. 두 정규분포(Gaussian Distributions)의 KL Divergence는 다음 수식을 통해 구할 수 있습니다.



$$D_{\mathrm{KL}}(\mathcal{N}({x} ; {\mu}_x, {\Sigma}_x) \| \mathcal{N}({y} ; {\mu}_y, {\Sigma}_y))=\frac{1}{2}[\log \frac{|{\Sigma}_y|}{|{\Sigma}_x|}-d+\operatorname{tr}({\Sigma}_y^{-1} {\Sigma}_x)+({\mu}_y-{\mu}_x)^T {\Sigma}_y^{-1}({\mu}_y-{\mu}_x)] \tag{86}$$



위의 공식을 p와 q에 적용합니다. 이때 KL Divergence를 최소화 해야 하므로(p와 q의 분포가 일치하도록)이를 고려하여 수시을 정리해 봅시다.



$$
\begin{align}
& \underset{{\theta}}{\arg \min } D_{\mathrm{KL}}\left(q\left({x}_{t-1} \mid {x}_t, {x}_0\right) \| p_{{\theta}}\left({x}_{t-1} \mid {x}_t\right)\right)
\\
= & \underset{{\theta}}{\arg \min } D_{\mathrm{KL}}\left(\mathcal{N}\left({x}_{t-1} ; {\mu}_q, {\Sigma}_q(t)\right) \| \mathcal{N}\left({x}_{t-1} ; {\mu}_{{\theta}}, {\Sigma}_q(t)\right)\right) \tag{87}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2}\left[\log \frac{\left|{\Sigma}_q(t)\right|}{\left|{\Sigma}_q(t)\right|}-d+\operatorname{tr}\left({\Sigma}_q(t)^{-1} {\Sigma}_q(t)\right)+\left({\mu}_{{\theta}}-{\mu}_q\right)^T {\Sigma}_q(t)^{-1}\left({\mu}_{{\theta}}-{\mu}_q\right)\right] \tag{88}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2}\left[\log 1-d+d+\left({\mu}_{{\theta}}-{\mu}_q\right)^T {\Sigma}_q(t)^{-1}\left({\mu}_{{\theta}}-{\mu}_q\right)\right] \tag{89}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2}\left[\left({\mu}_{{\theta}}-{\mu}_q\right)^T {\Sigma}_q(t)^{-1}\left({\mu}_{{\theta}}-{\mu}_q\right)\right] \tag{90}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2}\left[\left({\mu}_{{\theta}}-{\mu}_q\right)^T\left(\sigma_q^2(t) \mathbf{I}\right)^{-1}\left({\mu}_{{\theta}}-{\mu}_q\right)\right] \tag{91}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2 \sigma_q^2(t)}\left[\left\|{\mu}_{{\theta}}-{\mu}_q\right\|_2^2\right] \tag{92}
\end{align}
$$

 

$\mu_q$는 줄여서 쓴것이므로 다음과 같이 풀어서 쓸 수 있습니다.

$$
\begin{align}
{\mu}_q\left({x}_t, {x}_0\right)=\frac{\sqrt{\alpha_t}\left(1-\bar{\alpha}_{t-1}\right) {x}_t+\sqrt{\bar{\alpha}_{t-1}}\left(1-\alpha_t\right) {x}_0}{1-\bar{\alpha}_t} \tag{93}
\end{align}
$$

 

이것과 마찬가지로 $\mu_\theta$를 비슷하게 풀어서 쓸수 있는데, $p_\theta(x_{t-1}|x_t)$는 $x_0$에 대한 식으로 나타낼 수 없으므로 임의의 $\hat{x}_\theta$를 통해 나타내 봅시다.

 

$$
\begin{align}
{\mu}_{{\theta}}({x}_t, t)=\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) {x}_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) \hat{{x}}_{{\theta}}({x}_t, t)}{1-\bar{\alpha}_t} \tag{94}
\end{align}
$$



$\hat{x}_\theta(x_t,t)$는 neural network에 의해 paramterized 되어 index t에 대한 $x_t$가 $x_0$를 예측할 수 있도록 하는 역할입니다. 이제 각 값을 대입하여 식을 다시 고쳐봅시다.



$$
\begin{align}
& \underset{{\theta}}{\arg \min } D_{\mathrm{KL}}(q({x}_{t-1} \mid {x}_t, {x}_0) \| p_{{\theta}}({x}_{t-1} \mid {x}_t))
\\
= & \underset{{\theta}}{\arg \min } D_{\mathrm{KL}}(\mathcal{N}({x}_{t-1} ; {\mu}_q, {\Sigma}_q(t)) \| \mathcal{N}({x}_{t-1} ; {\mu}_{{\theta}}, {\Sigma}_q(t))) \tag{95}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2 \sigma_q^2(t)}[\|\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) {x}_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) \hat{{x}}_{{\theta}}({x}_t, t)}{1-\bar{\alpha}_t}-\frac{.\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) {x}_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) {x}_0 \|^2]}{1-\bar{\alpha}_t}\|_2] \tag{96}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2 \sigma_q^2(t)}[\|\frac{\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) \hat{{x}}_{{\theta}}({x}_t, t)}{1-\bar{\alpha}_t}-\frac{\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) {x}_0}{1-\bar{\alpha}_t}\|_2^2] \tag{97}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2 \sigma_q^2(t)}[\|\frac{\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t)}{1-\bar{\alpha}_t}(\hat{{x}}_{{\theta}}({x}_t, t)-{x}_0)\|_2^2] \tag{98}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2 \sigma_q^2(t)} \frac{\bar{\alpha}_{t-1}(1-\alpha_t)^2}{(1-\bar{\alpha}_t)^2}[\|\hat{{x}}_{{\theta}}({x}_t, t)-{x}_0\|_2^2] \tag{99}
\end{align}
$$



Eq 99를 보게되면, 결국 diffusion model을 최적화 하는 것은 t번째 noise image를 이용해 예측한 Image $\hat{x}_\theta(x_t, t)$가 $x_0$와 유사하도록 학습하는 것을 알 수 있습니다. 즉 모든 time step에 따른 noise image가 reverse process를 거치며 도출되는 ELBO의 총합을 최소화 하도록 합니다.



$$ \underset{\theta}{\arg \min}{\mathbb{E}_{t \sim U \{2, T \}}[\mathbb{E}_{q(x_t|x_0)}[D_{KL}(q(x_{t-1}|x_t,x_0) || p_\theta(x_{t-1}|x_t))]]} \tag{100}$$

 

Learning Diffusion Noise Parameter

이번에는 VDM의 noise parameter를 학습하는 방법을 살표보겠습니다. 만약 매개변수 $\eta$를 갖는 신경망 $\hat{\alpha}_\eta(t)$를 사용하여  $\alpha_t$를 모델링 한다면 $\bar{\alpha}_t$를 계산하기 위해 각 timestep에 대한 $\alpha_t$에 대한 추론을 여러번 해야되므로 효율적이지 못합니다. Eq 99에 Eq 85를 대입하여 식을 $\alpha_t$로 고쳐봅시다.



$$
\begin{align}
\frac{1}{2 \sigma_q^2(t)} \frac{\bar{\alpha}_{t-1}(1-\alpha_t)^2}{(1-\bar{\alpha}_t)^2}[\|\hat{{x}}_{{\theta}}({x}_t, t)-{x}_0\|_2^2] & =\frac{1}{2 \frac{(1-\alpha_t)(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}} \frac{\bar{\alpha}_{t-1}(1-\alpha_t)^2}{(1-\bar{\alpha}_t)^2}[\|\hat{{x}}_{{\theta}}({x}_t, t)-{x}_0\|_2^2] \tag{101}
\\
& =\frac{1}{2} \frac{1-\bar{\alpha}_t}{(1-\alpha_t)(1-\bar{\alpha}_{t-1})} \frac{\bar{\alpha}_{t-1}(1-\alpha_t)^2}{(1-\bar{\alpha}_t)^2}[\|\hat{{x}}_{{\theta}}({x}_t, t)-{x}_0\|_2^2] \tag{102}
\\
& =\frac{1}{2} \frac{\bar{\alpha}_{t-1}(1-\alpha_t)}{(1-\bar{\alpha}_{t-1})(1-\bar{\alpha}_t)}[\|\hat{{x}}_{{\theta}}({x}_t, t)-{x}_0\|_2^2] \tag{103}
\\
& =\frac{1}{2} \frac{\bar{\alpha}_{t-1}-\bar{\alpha}_t}{(1-\bar{\alpha}_{t-1})(1-\bar{\alpha}_t)}[\|\hat{{x}}_{{\theta}}({x}_t, t)-{x}_0\|_2^2] \tag{104}
\\
& =\frac{1}{2} \frac{\bar{\alpha}_{t-1}-\bar{\alpha}_{t-1} \bar{\alpha}_t+\bar{\alpha}_{t-1} \bar{\alpha}_t-\bar{\alpha}_t}{(1-\bar{\alpha}_{t-1})(1-\bar{\alpha}_t)}[\|\hat{{x}}_{{\theta}}({x}_t, t)-{x}_0\|_2^2] \tag{105}
\\
& =\frac{1}{2} \frac{\bar{\alpha}_{t-1}(1-\bar{\alpha}_t)-\bar{\alpha}_t(1-\bar{\alpha}_{t-1})}{(1-\bar{\alpha}_{t-1})(1-\bar{\alpha}_t)}[\|\hat{{x}}_{{\theta}}({x}_t, t)-{x}_0\|_2^2] \tag{106}
\\
& =\frac{1}{2}(\frac{\bar{\alpha}_{t-1}(1-\bar{\alpha}_t)}{(1-\bar{\alpha}_{t-1})(1-\bar{\alpha}_t)}-\frac{\bar{\alpha}_t(1-\bar{\alpha}_{t-1})}{(1-\bar{\alpha}_{t-1})(1-\bar{\alpha}_t)})[\|\hat{{x}}_{{\theta}}({x}_t, t)-{x}_0\|_2^2] \tag{107}
\\
& =\frac{1}{2}(\frac{\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t-1}}-\frac{\bar{\alpha}_t}{1-\bar{\alpha}_t})[\|\hat{{x}}_{{\theta}}({x}_t, t)-{x}_0\|_2^2] \tag{108}
\end{align}
$$

 

여기서 $q(x_t|x_0) = \mathcal{N}(x_t;\sqrt{\bar{\alpha}}x_0, (1-\bar{\alpha}I)$입니다. Signal-tonoise Ratio(SNR)의 분포에 대한 정의에 의하면 $SNR=\frac{\mu^2}{\sigma^2}$이므로 $q(x_t|x_0)$의 분포를 SNR로 나타내면 다음과 같습니다.

$$
\operatorname{SNR}(t) = \frac{\bar{\alpha}_t}{1-\bar{\alpha}_t } \tag{109}
$$

이를 Equation 108에 적용해봅시다.

$$
\frac{1}{2}(\frac{\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t-1}}-\frac{\bar{\alpha}_{t}}{1-\bar{\alpha}_{t}})[\|\hat{{x}}_{{\theta}}({x}_t, t)-{x}_0\|_2^2]=\frac{1}{2}(\operatorname{SNR}(t-1)-\operatorname{SNR}(t))[\|\hat{{x}}_{{\theta}}({x}_t, t)-{x}_0\|_2^2] \tag{110}
$$

 

SNR은 signal과 noise 사이의 비율을 나타냅니다. 값이 커질수록 signal이 많아지고 값이 작아질수록 noise가 많아집니다. 즉 t의 값이 커질수록 noise가 증가하여야하므로 SNR은 단조 감소함수여야 합니다. 이를 위해 다음과 같이 함수를 만들어 봅시다.

$$
\operatorname{SNR} (t)= \exp{(-\omega_{\eta}(t))} \tag{111}
$$

$\omega_{\eta}(t)$는 파라미터 $\eta$를 사용해 점진적으로 증가하는 신경망입니다. 이때 $\exp(-x)$에 대한 그래프를 보면 다음과 같습니다.

exp(-x)의 그래프

결국 t의 값이 커질수록 SNR의 값이 지속적으로 하락하는 그래프가 그려지게 되는 것을 볼 수 있습니다. Equation 109를 Equation 111에 대입하면 다음과 같습니다.

 

$$
\begin{align}
\frac{\bar{\alpha}_t}{1-\bar{\alpha}_t} = \exp{(-\omega_\eta(t))} \tag{112}
\\
\therefore \bar{\alpha}_t = \operatorname{sigmoid}(-\omega_\eta(t)) \tag{113}
\\
\therefore 1-\bar{\alpha}_t = \operatorname{sigmoid}(\omega_\eta(t)) \tag{114}
\end{align}
$$

 

Three Equivalent Interpretations

 

앞서 공부한 내용을 정리하자면, Variational Diffusion Model은 각 timestep t에서 주어진 noise image $x_t$를 통해 $x_0$를 예측하며 훈련합니다. 이 때 $x_0$는 두가지의 방식으로 파라미터화 할 수 있습니다.

 

1. Reparameterization Trick

$x_0$를 $x_t$에 대한 표준 분포로 나타내면 2가지의 매개변수를 가지게 됩니다. Eq 69를 정리하면 다음과 같습니다.

 

$$
\begin{align}
x_t = \sqrt{\bar{\alpha}}x_0 + \sqrt{1-\bar{\alpha}}\epsilon_0 \tag{69}
\\
x_0 =  \frac{x_t-\sqrt{1-\bar{\alpha}}\epsilon_0}{\sqrt{\bar{\alpha}}} \tag{115}
\end{align}
$$

 

이를 $\mu_q(x_t, x_0)$에 대입해 봅시다.

 

$$
\begin{align}
{\mu}_q({x}_t, {x}_0) & =\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) {x}_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) {x}_0}{1-\bar{\alpha}_t} \tag{116}
\\
& =\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) {x}_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) \frac{{x}_t-\sqrt{1-\bar{\alpha}_t} {\epsilon}_0}{\sqrt{\bar{\alpha}_t}}}{1-\bar{\alpha}_t} \tag{117}
\\
& =\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) {x}_t+(1-\alpha_t) \frac{{x}_t-\sqrt{1-\bar{\alpha}_t} {\epsilon}_0}{\sqrt{\alpha_t}}}{1-\bar{\alpha}_t} \tag{118}
\\
& =\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) {x}_t}{1-\bar{\alpha}_t}+\frac{(1-\alpha_t) {x}_t}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}}-\frac{(1-\alpha_t) \sqrt{1-\bar{\alpha}_t} {\epsilon}_0}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}} \tag{119}
\\
& =(\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}+\frac{1-\alpha_t}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}}) {x}_t-\frac{(1-\alpha_t) \sqrt{1-\bar{\alpha}_t}}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}} {\epsilon}_0 \tag{120}
\\
& =(\frac{\alpha_t(1-\bar{\alpha}_{t-1})}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}}+\frac{1-\alpha_t}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}}) {x}_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}} {\epsilon}_0 \tag{121}
\\
& =\frac{\alpha_t-\bar{\alpha}_t+1-\alpha_t}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}} {x}_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}} {\epsilon}_0 \tag{122}
\\
& =\frac{1-\bar{\alpha}_t}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}} {x}_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}} {\epsilon}_0 \tag{123}
\\
& =\frac{1}{\sqrt{\alpha_t}} {x}_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}} {\epsilon}_0 \tag{124}
\end{align}
$$

 

$\mu_q(x_t,x_0)$는 실제 $q(x_{t-1}|x_t,x_0)$에 대한 분포의 평균을 나타내므로 우리는 이에 근사시킬 $p_\theta(x_{t-1}|x_t)$의 평균 $\mu_\theta(x_t, t)$에 대한 식으로 나타내봅시다.

 

$$
\mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}}x_t - \frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}\sqrt{\alpha_t}}\hat{\epsilon}_\theta(x_t, t) \tag{115}
$$

 

앞서 언급했던것 처럼 $\bar{\alpha}_t$는 모든 timestep t에 대한 $\alpha_t$의 곱집합입니다. 앞서나온 KL Divergence 수식에 항을 대입하면 다음과 같이 정리 할 수 있습니다.

 

$$
\begin{align}
& \underset{{\theta}}{\arg \min } D_{\mathrm{KL}}(q({x}_{t-1} \mid {x}_t, {x}_0) \| p_{{\theta}}({x}_{t-1} \mid {x}_t))
\\
= & \underset{{\theta}}{\arg \min } D_{\mathrm{KL}}(\mathcal{N}({x}_{t-1} ; {\mu}_q, {\Sigma}_q(t)) \| \mathcal{N}({x}_{t-1} ; {\mu}_{{\theta}}, {\Sigma}_q(t))) \tag{126}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2 \sigma_q^2(t)}[\|\frac{1}{\sqrt{\alpha_t}} {x}_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}} \hat{{\epsilon}}_{{\theta}}({x}_t, t)-\frac{1}{\sqrt{\alpha_t}} {x}_t+\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}} {\epsilon}_0\|_2^2] \tag{127}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2 \sigma_q^2(t)}[\|\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}} {\epsilon}_0-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}} \hat{{\epsilon}}_{{\theta}}({x}_t, t)\|_2^2] \tag{128}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2 \sigma_q^2(t)}[\|\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}}({\epsilon}_0-\hat{{\epsilon}}_{{\theta}}({x}_t, t))\|_2^2] \tag{129}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2 \sigma_q^2(t)} \frac{(1-\alpha_t)^2}{(1-\bar{\alpha}_t) \alpha_t}[\|{\epsilon}_0-\hat{{\epsilon}}_{{\theta}}({x}_t, t)\|_2^2] \tag{130}
\end{align}
$$

 

$\hat{\epsilon}_\theta(x_t, t)$는 $x_0$에서 $x_t$를 결정하는 noise $\epsilon \sim \mathcal{N}(\epsilon;0,I)$를 예측하는 방법을 학습하는 신경망입니다. 즉 $x_0$를 예측하도록 학습하는 것은 noise를 예측하는 것과 동일하다는 것을 증명 할 수 있습니다.

 

2. Tweedie's Formula

수학적으로 Tweedie's Formula는 다음과 같이 정의됩니다.

Gaussian variable  $z \sim \mathcal{N}(z;\mu_z,\Sigma_z)$에 대해서,
$$\mathbb{E}[\mu_z|z] = \underbrace{z}_{MLE}+\underbrace{\Sigma_z \nabla_z \log{p(z)}}_{Bayes\ correction}$$

자세한 내용은 아래를 참고하세요.
https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3325056/

Tweedie’s Formula and Selection Bias

아래 pdf 파일에도 설명이 잘 되어 있습니다.

https://efron.ckirby.su.domains/talks/2010TweediesFormula.pdf

 

Eq 70에서 $q(x_t|x_0)$는 다음과 같이 나타낼 수 있습니다.

$$q(x_t|x_0) = \mathcal{N}(x_t;\sqrt{\bar{\alpha}}x_0, (1-\bar{\alpha})I)$$

이같은 정규분포에 대하여 Tweedie's Formula를 통해 다음이 성립합니다.

$$\mathbb{E}[\mu_{x_t}|x_t] = x_t + (1-\bar{\alpha}_t)\nabla_{x_t}\log{p(x_t)} \tag{131}$$

 

이때 $ \nabla_{x_t}$는 $x_t$에 대한 미분을 의미합니다. $\nabla_{x_t}\log(p(x_t))$는 score function입니다. (관련된 내용은 Score based model 파트에서 다뤼보도록 하겠습니다.)

 

 

Score Function

 

Score Function이라는 개념이 생소할 수 있는데, Score Function은 우리가 인공신경망을 처음 배울때 나오는 개념입니다. log likelihood function의 1차 미분 값을 Score Function이라 부릅니다. 미분의 기준은 파라미터($\theta$)입니다. MLE(Maximum Liklihood Estimation)은 이 Score Function이 0이 되는 값을 찾는 것입니다.

 

https://en.wikipedia.org/wiki/Informant_(statistics)

Informant (statistics) - Wikipedia

---

 

또한 Tweedie's Formula에 의하면 관측된 $x_i$에대한 최적의 평균 ($\mu_{x_t}=\sqrt{\bar{\alpha}_t}x_0$)은 다음과 같이 정의 할 수 있습니다. ($\nabla_{x_t} = \nabla$로 간소화했습니다.)

 

$$
\begin{align}
\sqrt{\bar{\alpha}_t} x_0 &= x_t + (1-\bar{\alpha}_t) \nabla\log{p(x_t)} \tag{132}
\\
\therefore x_0 &= \frac{x_t + (1-\bar{\alpha}_t) \nabla\log{p(x_t)}}{\sqrt{\bar{\alpha}}} \tag{133}
\end{align}
$$

 

다시한번 $\mu_q(x_t, x_0)$에 적용을 해봅시다.

 

$$
\begin{align}
{\mu}_q({x}_t, {x}_0) & =\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) {x}_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) {x}_0}{1-\bar{\alpha}_t} \tag{134}
\\
& =\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) {x}_t+\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t) \frac{{x}_t+(1-\bar{\alpha}_t) \nabla \log p({x}_t)}{\sqrt{\bar{\alpha}_t}}}{1-\bar{\alpha}_t} \tag{135}
\\
& =\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) {x}_t+(1-\alpha_t) \frac{{x}_t+(1-\bar{\alpha}_t) \nabla \log p({x}_t)}{\sqrt{\alpha_t}}}{1-\bar{\alpha}_t} \tag{136}
\\
& =\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1}) {x}_t}{1-\bar{\alpha}_t}+\frac{(1-\alpha_t) {x}_t}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}}+\frac{(1-\alpha_t)(1-\bar{\alpha}_t) \nabla \log p({x}_t)}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}} \tag{137}
\\
& =(\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_t}+\frac{1-\alpha_t}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}}) {x}_t+\frac{1-\alpha_t}{\sqrt{\alpha_t}} \nabla \log p({x}_t) \tag{138}
\\
& =(\frac{\alpha_t(1-\bar{\alpha}_{t-1})}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}}+\frac{1-\alpha_t}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}}) {x}_t+\frac{1-\alpha_t}{\sqrt{\alpha_t}} \nabla \log p({x}_t) \tag{139}
\\
& =\frac{\alpha_t-\bar{\alpha}_t+1-\alpha_t}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}} {x}_t+\frac{1-\alpha_t}{\sqrt{\alpha_t}} \nabla \log p({x}_t) \tag{140}
\\
& =\frac{1-\bar{\alpha}_t}{(1-\bar{\alpha}_t) \sqrt{\alpha_t}} {x}_t+\frac{1-\alpha_t}{\sqrt{\alpha_t}} \nabla \log p({x}_t) \tag{141}
\\
& =\frac{1}{\sqrt{\alpha_t}} {x}_t+\frac{1-\alpha_t}{\sqrt{\alpha_t}} \nabla \log p({x}_t) \tag{142}
\end{align}
$$

 

$\bar{\alpha}_t$는 모든 timestep t에 대한 $\alpha_t$의 곱집합인 점을 유의하면 단순한 수식입니다. 그렇다면 $\mu_q(x_t, t)$는 다음과 같이 정의할 수 있습니다.

$$\mu_{\theta}(x_t, t)=\frac{1}{\sqrt{\alpha_t}}x_t+\frac{1-\alpha_t}{\sqrt{\alpha_t}}s_{\theta}(x_t, t) \tag{143}$$

 

$ s_{\theta}(x_t, t) $는 time step $t$에 대한 $x_t$의 Score function(기울기) $\nabla_x \log{p(x_t)}$를 예측하는 신경망을 뜻합니다. 위 수식을 대입해 식을 다시 정리해봅시다.

 

$$
\begin{align}
& \underset{{\theta}}{\arg \min } D_{\mathrm{KL}}(q({x}_{t-1} \mid {x}_t, {x}_0) \| p_{{\theta}}({x}_{t-1} \mid {x}_t))
\\
= & \underset{{\theta}}{\arg \min } D_{\mathrm{KL}}(\mathcal{N}({x}_{t-1} ; {\mu}_q, {\Sigma}_q(t)) \| \mathcal{N}({x}_{t-1} ; {\mu}_{{\theta}}, {\Sigma}_q(t))) \tag{144}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2 \sigma_q^2(t)}[\|\frac{1}{\sqrt{\alpha_t}} {x}_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}} \hat{{\epsilon}}_{{\theta}}({x}_t, t)-\frac{1}{\sqrt{\alpha_t}} {x}_t+\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}} {\epsilon}_0\|_2^2] \tag{145}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2 \sigma_q^2(t)}[\|\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}} {\epsilon}_0-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}} \hat{{\epsilon}}_{{\theta}}({x}_t, t)\|_2^2] \tag{146}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2 \sigma_q^2(t)}[\|\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \sqrt{\alpha_t}}({\epsilon}_0-\hat{{\epsilon}}_{{\theta}}({x}_t, t))\|_2^2] \tag{147}
\\
= & \underset{{\theta}}{\arg \min } \frac{1}{2 \sigma_q^2(t)} \frac{(1-\alpha_t)^2}{(1-\bar{\alpha}_t) \alpha_t}[\|{\epsilon}_0-\hat{{\epsilon}}_{{\theta}}({x}_t, t)\|_2^2] \tag{148}
\end{align}
$$

 

잘 보시면 Eq 130과 비슷한 모양을 하고 있는것을 볼수 있습니다. 다만 반대로 매핑해야 합니다. 우리의 목표인 $\epsilon_0$는 $\nabla \log{p(x_t}$와 매핑되고, 이를 추정하는 신경망인 $\hat{\epsilon}_\theta$는 $ s_{\theta}(x_t, t)$ 와 매핑됩니다. (음수 양수 차이는 어짜피 제곱을 하기 때문에 의미가 없어집니다.)

 

그렇다면 Tweedie's Formula (Eq 133)로 유도된 수식을 Reparameterization trick(Eq 115)에 대입해볼 수 있습니다.

 

$$
\begin{align}
{x}_0=\frac{{x}_t+(1-\bar{\alpha}_t) \nabla \log p({x}_t)}{\sqrt{\bar{\alpha}_t}} & =\frac{{x}_t-\sqrt{1-\bar{\alpha}_t} {\epsilon}_0}{\sqrt{\bar{\alpha}_t}} \tag{149}
\\
\therefore(1-\bar{\alpha}_t) \nabla \log p({x}_t) & =-\sqrt{1-\bar{\alpha}_t} {\epsilon}_0 \tag{150}
\\
\nabla \log p({x}_t) & =-\frac{1}{\sqrt{1-\bar{\alpha}_t}} {\epsilon}_0 \tag{151}
\end{align}
$$

 

식을 해석해보면, reparameterization trick의 경우는 이미지를 noise 변수인 $\epsilon$을 통해 점차 손상시키고 이를 복원하는 방향을 모델링합니다. score function의 개념으로 봤을때는 log 확률을 Maximize 하는 방향으로 모델링합니다. 여기서 결국 손상된 이미지를 denoising 하는 방법은 위의 Eq 151에 의해 score function과 같이 log 확률을 maximize하는 것과 동일하다(일정한 계수의 곱만큼 다르긴 하지만)는 것을 증명할 수 있습니다.

 

즉 VDM은 3가지의 동등한 목표(three equivalent objective)를 가지고 있습니다.

 

• 원본 이미지 $x_0$를 예측하도록 학습.
• source noise인 $\epsilon_0$를 예측하도록 학습.
• 이미지에 대한 임의의 noise 단계의 score $\nabla \log p({x}_t)$를 예측하도록 학습.