Understanding Diffusion Models (Part 2: VAE)

Part2는 VAE에 관한 내용이다. Part1에서 생성모델의 개념과 ELBO에 대해서 살펴보았다. 여기서 본 ELBO를 가지고 VAE, Variational Autoencoders를 수행할 수 있다. 여기서는 이 VAE에 대해서 자세히 다룬다. 반드시 Part1을 이해하고 이 포스팅을 참고하는 것이 도움이 될 것이다.

2024.11.20 - [논문리뷰/Diffusion models] - Understanding Diffusion Models: A Unified Perspective (Part 1: 생성 모델의 개념과 ELBO)

Understanding Diffusion Models: A Unified Perspective (Part 1: 생성 모델의 개념과 ELBO)

 

2.2. Variational Autoencoders

Variational Autoencoder(VAE)는 기본적으로 직접 ELBO를 최대화한다. 이러한 접근은 posterior 중 최적점을 찾아내는 것이므로 variational이라고 표현한다. 아래 그림은 autoencoder인데, 입력 값을 인코더를 통해 z로 변환한 후, 이를 다시 디코더로 복원하여 원래 입력 값과 유사하게 만드는 과정을 보여준다. VAE는 이 autoencoder와 유사하게 입력 데이터를 latent variable로 변환한 후 다시 복원하는 구조를 가진다.

https://www.mdpi.com/2076-3417/13/22/12413

 

우리는 ELBO를 분해해서 이것을 이해할 수 있다. 첫 번째 식은 part1 마지막에 정리된 ELBO의 식이다.

$\mathbb{E}_{q_{\phi}(z|x)}\Bigg[\text{log}\cfrac{p(x,z)}{q_\phi(z|x)}\Bigg] = \mathbb{E}_{q_{\phi}(z|x)}\Bigg[\text{log}\cfrac{p_\theta(x|z)p(z)}{q_\phi(z|x)}\Bigg]$                                                (17)

$\quad\quad\quad\quad\quad\quad\quad\quad\quad =\mathbb{E}_{q_{\phi}(z|x)} [\text{log}\,p_{\theta}(x|z)] + \mathbb{E}_{q_{\phi}(z|x)}\Bigg[\text{log}\cfrac{p(z)}{q_{\phi}(z|x)}\Bigg] $                (18)

$\quad\quad\quad\quad\quad\quad\quad\quad\quad =\underbrace{\mathbb{E}_{q_\phi(z|x)}[\text{log}\,p_\theta(x|z)]}_\text{reconstruction term} - \underbrace{D_{KL}(q_\phi(z|x)||p(z))}_\text{prior matching term} $                   (19)

 

(17) 의 좌변은 loglikelihood $p(x)$의 ELBO이다. 여기서 결합 확률 $p(x,z)$는 $p(x|z)\cdot p(z)$로 나타낼 수 있다. (→ 조건부 확률의 정의에 의해서) 따라서 이 부분을 모델의 파라미터에 의존하는 $p_\theta(x|z)p(z)$로 바꿔 쓸 수 있다.

(18) 곱을 합으로 분해할 수 있는 로그함수의 성질을 이용하여 분리한다.

(19) 좌변의 $\text{log}p_{\theta}(x|z)$는 어떤 잠재 변수가 주어졌을 때 $x$가 나타날 확률이다. 그런데 여기에 들어가는 잠재 변수, 즉 $z$가 $x$에 의해서 생성되는 값이고, 이 값에 의해 다시 $x$가 나타날 확률의 기대값을 구하는 부분이므로 $z$의 전체 공간($\mathbb{E}_{q\phi(z|x)}$)에서 $x$가 나타날 확률은 원래 input이 복원되는 확률이다. 따라서 이를 reconstruction term이라고 한다. 우변은 part1에서  정의한 KLD 공식에 의하면 $KL(p||q) = \sum_i\,p_i\,\text{log}(\cfrac{p_i}{q_i})$ 이지만, 여기서 마이너스(-)를 곱하게 되면 다음과 같이 변경된다(로그의 성질에 의해). $KL(p||q) = -\sum_i\,q_i\,\text{log}(\cfrac{p_i}{q_i})$. 따라서 KLD term으로 표현할 수 있고, 이 부분을 잘 살펴보면 $x$에서 잠재 변수가 나오는 분포를 $p(z)$, 즉 encoder를 prior($p(z)$)에 유사하도록 맞춰주는 것이다.

 

결과적으로 이 부분에서 $q_\phi(z|x)$는 encoder로 이해할 수 있다. 즉 $q_\phi(z|x)$는 encoder $p(z|x)$를 근사하며 우리는 이것을 학습해야 한다. 동시에 $p_\theta(x|z)$를 학습해야 하는데, 이것은 잠재 변수에서 관측치로 변환해주는 decoder가 된다. (19) 의 식에 따라서 VAE에서 ELBO를 최대화 하는 것은, reconstruction term을 최대로 하고 prior matching term을 최소로 하는 것을 의미한다.

 

VAE의 특징은 (19)번 식에서 볼 수 있듯이 $\phi, \theta$를 모두 최적화 한다는 것이다. 여기서 $\phi$는 encoder의 파라미터이고, $\theta$는 decoder의 파라미터이다. 즉, encoder와 decoder를 함께 최적화한다. Encoder는 일반적으로 diagonal 공분산을 가진 다변량 가우시안 분포를 모델링하며, prior는 종종 표준 가우시안 분포로 선택된다.

$q_\phi(z|x) = \mathcal{N}(z;\mu_\phi(x), \sigma^2_\phi(x)\mathrm{I})$                      (20)

$p(z) = \mathcal{N}(z; 0, \mathrm{I})$                                              (21)

 

KLD term은 analyrically 하게 계산 가능하고, reconstruction term은 몬테카를로 시뮬레이션으로 근사할 수 있다.

$\underset{\phi,\theta}{\text{arg max}}\;\mathbb{E}_{q\phi(x|z)}[\text{log}\,p_\theta(x|z)-D_{KL}(q_\phi(z|x)||p(z))\approx\underset{\phi,\theta}{\text{arg max}}\displaystyle\sum_{l=1}^L\text{log}\,p_\theta(x|z^{(l)})-D_{KL}(q_\phi(z|x)||p(z))$ (22)

여기서 $\{z^{(l)}\}{l=1}^L$은 데이터셋의 모든 관찰값 $x$에 대해 $q\phi(x|z)$에서 샘플링된다. 그런데 여기서 loss가 미분 불가능한 확률적 샘플링으로 이루어지므로 backpropagation이 어렵다. 하지만 $q_\phi(x|z)$가 다변량 가우시안을 모델링하므로 reparameterization trick을 통하여 해결할 수 있다. Reparameterization trick은 예를 들어 정규 분포를 가지는 샘플이라면 다음과 같이 표현할 수 있다.

$x = \mu +\sigma\epsilon\quad\text{with}\;\epsilon\sim\mathcal{N}(\epsilon;\;0,\mathrm{I})$

다시 말해, 임의의 가우시안 분포에서 샘플링하는 것을 표준 가우시안 분포에서 샘플링 하고, 표준 편차로 스케일링 한 후 평균으로 이동시키는 것이다.

결국 VAE에서는 $z$가 인풋 $x$와 노이즈 $\epsilon$의 deterministic한 함수로 표현할 수 있다.

$z = \mu_\phi(x)+\sigma_\phi(x)\odot\epsilon$

이 함수의 그라디언트 계산을 통해 $\mu_\phi, \sigma_\phi$를 최적화 할 수 있다.

 

2.3. Hierarchical Variational Autoencoders

Hierarchical Variational Autoencoder (HVAE)는 latent variables 위에 여러 계층으로 확장한 VAE의 일반적인 형태이다. 그래서 여기서는 latent variable 자체가 더 높은 수준의, 더 추상적인 latent variable에 의해서 생성되었음으로 해석한다. 일반적으로 T개의 계층을 가진 HVAE는 각각의 latent가 이전 latent를 조건부로 갖을 수 있지만, 여기서는 Markovian HVAE (MHVAE)에만 초점을 맞춘다. 따라서 이 markovian chain에 의해서 latent $z_t$만이 오로지 latent $z_{t+1}$의 condition이 된다.

 

수학적으로는 다음과 같이 표현할 수 있다. (markovian chain의 특성에 의해 모든 확률의 곱이 되므로 아래 식을 전개해보면 확인할 수 있음)

$p(x, z_{1:T}) = p(z_T)p_\theta(x|z_1)\displaystyle\prod_{t=2}^Tp_\theta(z_{t-1}|z_t)$             (23)

$q_\phi(z_{1:T}|x) = q_\phi(z_1|x)\displaystyle\prod_{t=2}^Tq_\phi(z_t|z_{t-1})$                      (24)

 

이 개념을 ELBO에 쉽게 적용해 볼 수 있다. Part1에서 식 (5) ~ (8)에서 했던 것과 동일하다.

$\text{log}\,p(x) = \text{log}\displaystyle\int p(x, z_{1:T})dz_{1:T}$                                        (25)

$\quad\quad\quad\;\;=\text{log}\displaystyle\int\cfrac{p(x,z_{1:T})q_\phi(z_{1:T}|x)}{q_\phi(z_{1:T}|x)}dz_{1:T}$                     (26)

$\quad\quad\quad\;\;=\text{log}\,\mathbb{E}{q\phi(z_{1:T}|x)}\Bigg[\cfrac{p(x, z_{1:T})}{q_\phi(z_{1:T}|x)}\Bigg]$                          (27)

$\quad\quad\quad\;\;\geq\mathbb{E}{q\phi(z_{1:T}|x)}\Bigg[\text{log}\cfrac{p(x, z_{1:T})}{q_\phi(z_{1:T}|x)}\Bigg]$                           (28)

 

(25) (1)의 식과 동일하게 전체 $z$공간에 대해 적분하면 (25)와 같이 쓸 수 있다. ($z$ → $z_{1:T}$)

(26) 분자와 분모에 $q_\phi(z_{1:T}|x)$를 곱한다.

(27) 기대값의 정의에 따라 적분을 기대값 형태로 바꿔준다.

(28) Jensen’s Inequality에 의해서 성립한다. (log 함수는 오목 함수 이기 때문에, $f(E[X])\geq E[f(X)]$

 

그런 다음 결합 분포 식 (23)과 posterior (24)를 식 (28)에 넣어서 아래 결과를 얻을 수 있다.

$\mathbb{E}_{q\phi(z_{1:T}|x)}\Bigg[\text{log}\cfrac{p(x, z_{1:T})}{q_\phi(z_{1:T}|x)}\Bigg] = \mathbb{E}_{q\phi(z_{1:T}|x)}\Bigg[\text{log}\cfrac{p(z_T)p_\theta(x|z_1)\prod_{t=2}^Tp_\theta(z_{t-1}|z_t)}{q_\phi(z_{1}|x)\prod_{t=2}^Tq_\phi(z_t|z_{t-1})}\Bigg]$            (29)

Variational Diffusion Model들을 연구할 때, 이런 objective를 해석 가능한 component로 더 분리할 수 있다.