Variational Diffusion Models
Variational Diffusion Model (VDM) 은 Markovian Hierarchical Variational Autoencoder 인데, 걍 3가지 제약사항이 있는거임
- latent dimension이 data dimension과 같다
- latent encoder 은 학습되지 않고, linear Gaussian model로 pre-defined됨.
- latent encoder 의 Gaussian parameters은 timestep마다 다르다. timestep T에서 latent는 표준정규분포임
- HVAE에서는 $q_\phi(z_t|x)$ 로 썼지만, 첫째 가정으로 인해 VDM에서는 $q_\phi(x_t|x_{t-1})$ 와 같이 쓴다.
- 두번째 가정에 의해, $q(x_t|x_{t-1)}$ 은 이전 단계의 output 을 mean으로 갖는 Gaussian distribution Gaussian distribution의 모수는 미리 정해진다. $\mu_t(x_t)= \sqrt{\alpha_t}x_{t-1}$ , $\Sigma_t(x_t)=(1-\alpha_t)I$. $\alpha_t$ 는 급격하게 변하지 않게 해서 latent variabl 의 분산이 비슷하게 유지되도록 정한다.
- → $p(x_{0:T})=p(x_T)\prod_{t=1}^Tp_\theta(x_{t-1}|x_t)$ , $p(x_T)=\mathcal{N}(x_T;0,I)$
이제 $q(x_t|x_{t-1})$ 은 더이상 $\phi$ 에 의해 정해지는 대신 미리 확정된다. 따라서 VDM에서는 되돌아오는 방향 $p_\theta(x_{t-1}|x_t)$ 에 관심을 둠
VDM에서 역시 ELBO를 최대화해서 모델을 최적화한다 :
- $\mathbb{E}{q(x_1|x_0)}[\log p\theta(x_0|x_1)]$ = reconstruction term = first step latent가 주어졌을 때 원본데이터의 확률. vanilla VAE에서도 볼 수 있는 항이다 (z 대신 $x_1$)
- $\mathbb{E}{q(x{T-1}|x_0)}[D_{KL}(q(x_T|x_{T-1}) || p(x_T))]$ = prior matching term = final latent $x_T$ 가 Gaussian prior 일 때 최소화됨. 최적화대상이 아니고 걍 고정값임. T가 크기만 하면 왠만해선 거의 0
- $\mathbb{E}{q(x{t-1}, x_{t+1}|x_0)}[D_{KL}(q(x_t|x_{t-1}) || p_\theta(x_t|x_{t+1}))]$ = consistency term = $x_{t+1}\to x_t$ 와 $x_{t-1} \to x_t$ 의 분포가 비슷할 때 최소화되는 값 → $x_t$ 의 분포를 일정하게 유지하게 만듦. denoising process도 Gaussian이 될 때 minimize된다.
1, 2, 3 중 3의 영향이 가장 dominant 함
모든 항이 expectation으로 계산되었으므로 Monte Carlo estimate 을 활용해서 계산 가능
근데, consistency term은 두개의 확률변수 $x_{t-1}, x_{t+1}$ 을 활용한 평균값으로 계산되기 때문에, Monte Carlo estimate의 분산은 하나의 확률변수를 썼을 때에 비해 높으며, T-1개의 consistency term을 더하는 식으로 계산되므로 T값이 커질수록 그 분산도 커진다.
따라서 이걸 하나의 확률변수만 쓰는 방식으로 계산하도록 바꿔보자. 이걸 위해서는 encoder transition을 $q(x_t|x_{t-1})=q(x_t|x_{t-1}, x_0)$ 로 바꿔 쓰면 된다. 어차피 Markov process이기 때문에 상관없음. 그러면 베이즈 공식에 의해 $q(x_t|x_{t-1},x_0) = \dfrac{q(x_{t-1}|x_t,x_0)q(x_t|x_0)}{q(x_{t-1}|x_0)}$이 되고, ELBO를 다시 전개하면 다음과 같다
이 공식은 각 term이 최대 1개의 확률변수의 기댓값에 의해 결정되므로, 더 낮은 분산을 가지는 ELBO에 대한 interpretation이 된다.
각 term을 해석해보면
- reconstruction term : vanilla VAE에서의 reconstruction term, Monte Carlo estimate으로 근사
- prior matching term : final noised input과 Gaussian prior가 얼마나 가까운 지
- denoising matching term : ground truth denoising transition $q(x_{t-1}|x_t,x_0)$ 과 learnable denoising transition $p_\theta(x_{t-1}|x_t)$ 가 얼마나 비슷한 지?
t=1이면 vanila VAE와 같음.
- p를 ground truth denoising step q에 최적화 시키는 게 목표임. 베이즈 공식으로 반대방향의 q값을 구함
여기서도 세번째 term이 최적화의 핵심인데, Gaussian transition 이라는 가정을 이용해서 최적화할 수 있다. 베이즈 공식에 의해
$q(x_{t-1}|x_t,x_0)=\dfrac{q(x_t|x_{t-1},x_0)q(x_{t-1}|x_0)}{q(x_t|x_0)}$ 이고
$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|x_{t-1})$를 활용해서 반대방향인 $q(x_{t-1}|x_t,x_0)$를 알 수 있고, reparameterization trick 에서 $x_t \sim q(x_t|x_{t-1})$ 을 $x_t=\sqrt{\alpha_t}x_{t-1}+\sqrt{1-\alpha_t}\epsilon$ 으로 바꿔서 처리 가능함. 그리고 이 값과 parameterized denoising step $p_\theta(x_{t-1}|x_t)$ 사이의 KL Divergence를 최소화함
$q(x_t|x_0)$ 는 reparameterization trick을 반복해서 수행해서 얻을 수 있음
noise variables $\{\epsilon_t^*,\epsilon_t\}_{t=0}^T \sim^{iid}\mathcal{N}(\epsilon;0,I)$ 을 활용해서 $x_t\sim q(x_t|x_0)$ 를 전개해보면 다음과 같음
Gaussian끼리의 합은 여전히 Gaussian이라는 공식을 활용하면 $\sqrt{\alpha_t-\alpha_t\alpha_{t-1}}\epsilon_{t-2}^* + \sqrt{1-\alpha_t}\epsilon_{t-1}^* = \sqrt{\sqrt{\alpha_t-\alpha_t\alpha_{t-1}}^2+\sqrt{1-\alpha_t}^2}\epsilon_{t-2}$
을 얻을 수 있음
이 유도식을 이용하면 $q(x_t|x_0)$의 Gaussian form을 안거니까 $x_0$에서 곧바로 $x_t$를 샘플링 할 수 있음
그리고 베이즈공식을 쓰면 $q(x_{t-1}|x_t,x_0)$ 의 Gaussian form도 도출 가능
이걸 통해서 알 수 있는 것 : Denoising step 또한 정규분포라는 사실 (평균은 $x_t, x_0$ 에 의해 정해지고, 분산은 $\alpha$ 에 의해 정해짐)
$\sigma_q^2(t)=\dfrac{(1-\alpha_t)(1-\bar\alpha_{t-1})}{1-\bar\alpha_t}$ 로 분산을 표현하고, denoising step은 ground truth $q(x_{t-1}|x_t,x_0)$ 를 따라가야 하므로 Gaussian으로 모델링하고, 분산도 $\sigma_q^2(t)I$ 로 설정하면 됨 (어차피 known임). 평균은 $x_0$를 모르므로 $x_t$ 에 대해 conditioned 됨
한편, 두개의 Gaussian distribution 사이의 KL Divergence를 계산하는 공식은 다음과 같음
$D_{KL}(\mathcal{N}(x;\mu_x,\Sigma_x) || \mathcal{N}(y;\mu_y,\Sigma_y)) = \frac{1}{2}[\log\dfrac{|\Sigma_y|}{\Sigma_x}-d+tr(\Sigma_y^{-1}\Sigma_x)+(\mu_y-\mu_x)^T\Sigma_y^{-1}(\mu_y-\mu_x)]$
우리가 하고싶은건 $q(x_{t-1}|x_t,x_0)$ 에다가 $p_\theta(x_{t-1}|x_t)$ 를 최적화시키는거임
한편, 위에서 $\mu_q$는 구했었음 → $\mu_q(x_t,x_0) = \dfrac{\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_\theta(x_t,t)$ 도 얘랑 비슷하게 형태를 정해보면
$\mu_\theta(x_t,t)=\dfrac{\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}$
여기서 $\hat{x}_\theta$ 는 $x_0$ 를 예측하는 용도로 neural network에서 학습하도록 $\theta$로 parameterize 한 애임
이렇게 형태를 갖춘 뒤 위의 최적화 식에 대입하면
→ 결 국 VDM을 최적화하는 건 original ground truth image x 를 예측하도록 신경망을 학습시키는 것과 같다.
→ 그리고, VDM의 ELBO인
의 denoising match term을 최소화하는 건 다음 기댓값
을 최소화하는 것과 같다. → 얘는 stochastic sampling을 통해 최적화할 수 있음