[DL] 딥러닝 모델: VAE (Variational Autoencoder) 이해하기

---

생성형 모델(generative model)은 단순히 데이터를 분류하거나 예측하는 것이 아니라 새로운 데이터를 생성하는 모델입니다. VAE (Variational autoencoder)는 2013년 제안된 대표적인 생성형 모델입니다. VAE는 학습 데이터의 분포를 잘 표현하는 잠재 공간(latent space)를 학습하고자 하고, 학습된 latent space를 통해 새로운 샘플을 생성할 수 있습니다.

 

Generative Model이란? 

생성형 모델(generative model)은 학습 데이터의 분포를 모델링하여, 학습 데이터와 유사한 샘플을 모델로부터 생성하는 것을 목표로 합니다. 따라서 아래의 그림처럼 학습 데이터의 분포 $p_\text{data}(\mathbf{x})$와 유사한 생성형 모델 $p_\theta(\mathbf{x})$를 학습하고자 합니다. 

 

Figure 1. 학습 데이터의 분포와 생성형 모델

이를 위해 생성형 모델은 maximum likelihood에 기반하여, 학습 데이터 $\mathbf{x}$에 대한 likelihood $p_\theta(\mathbf{x})=p(\mathbf{x}|\theta)$를 최대화하는 방식을 주로 사용합니다. Maximum likelihood를 사용하는 이유는 모델이 학습 데이터를 설명할 확률이 커질수록, 그 모델이 그럴듯한 새로운 샘플을 생성할 수 있을 것이라는 직관에서부터 비롯합니다.

 

VAE (Variational Autoencoder)의 구조

1) Latent Variable Model

VAE는 학습 데이터의 분포를 직접 모델링하지 않고, latent variable $\mathbf{z}$를 이용한 latent variable model을 사용합니다. Latent variable $\mathbf{z}$은 관측할 수 없는 내재된 변수로서 데이터 $\mathbf{x}$와 연관된 상태, 특징 등을 의미하는 압축된 정보일 수 있습니다.

Figure 2. Latent variable model

Latent variable model에서 새로운 데이터는 아래 두 단계에 거쳐 생성될 수 있습니다.

 

• (Step 1) Latent variable $\mathbf{z}$가 prior distribution $p(\mathbf{z})$로부터 생성됩니다.

$$ \mathbf{z}\sim p(\mathbf{z}) $$

• (Step 2) Data $\mathbf{x}$가 conditional distribution $p(\mathbf{x}|\mathbf{z})$로부터 생성됩니다.

$$ \mathbf{x}\sim p(\mathbf{x}|\mathbf{z}) $$

 

여기서 VAE는 prior distribution $p(\mathbf{z})$을 단순한 Gaussian distribution $\mathcal{N}(0,I)$로 설정합니다. 또한 조건부 분포 $p_\theta(\mathbf{x}|\mathbf{z})$는 신경망을 통해 모델링되며 이는 VAE의 decoder에 해당하는 부분입니다.

 

Figure 3. VAE의 generative process

Decoder에 해당하는 모델을 학습하기 위해서, 우리가 최대화하고자 하는 데이터 $\mathbf{x}$에 대한 marginal likelihood는 아래와 같이 표현될 수 있습니다.

 

$$ p_\theta(\mathbf{x})=\int p_\theta(\mathbf{x}|\mathbf{z})p(\mathbf{z})d\mathbf{z} $$

 

하지만 이를 직접 구하는 것은 intractable하다는 문제가 발생합니다. Latent variable $\mathbf{z}$가 고차원이고, VAE는 신경망으로 모델링 된 복잡한 likelihood $p_\theta(\mathbf{x}|\mathbf{z})$을 가지기 때문에 closed-form으로 적분을 계산할 수 없기 때문입니다.

 

직접적인 계산이 어렵다면 Monte Carlo sampling을 기반한 근사를 생각해볼 수 있습니다. 예를 들어, $\mathbf{z}$를 prior distribution $p(\mathbf{z})$에서 많이 샘플링하여 $p_\theta(\mathbf{x})\approx\frac{1}{n}\sum_{i}p_\theta(\mathbf{x}|\mathbf{z}_i)$로 근사할 수 있습니다. 하지만 이 방법에도 문제가 존재합니다. Latent space에 존재하는 대부분의 $\mathbf{z}$는 데이터 $\mathbf{x}$를 생성하는데 기여하지 않습니다. 따라서 $p_\theta(\mathbf{x})$를 정확하게 근사하기 위해서는 매우 많은 수의 샘플링이 필요합니다.

 

이 문제를 해결하기 위해 실제 학습 데이터 $\mathbf{x}$와 관련된 영역에서만 latent variable $\mathbf{z}$를 효과적으로 샘플링 하는 것이 필요합니다. 다시 말해 posterior distribution $p(\mathbf{z}|\mathbf{x})$를 구하고 여기서 추출된 의미 있는 (즉, 데이터 $\mathbf{x}$를 생성하는데 기여하는) latent variable $\mathbf{z}$를 통해 $p_\theta(\mathbf{x})$를 구하는 것이 필요합니다.

 

728x90

 

2) Variational Inference

하지만 posterior distribution $p(\mathbf{z}|\mathbf{x})$를 구하는 것 역시 intractable하다는 문제를 가지고 있습니다. Posterior는 Bayes theorem에 따라 $p(\mathbf{z}|\mathbf{x})=\frac{p(\mathbf{x}|\mathbf{z})p(\mathbf{z})}{p(\mathbf{x})}$로 정의되는데, 앞서 언급했던 것과 같이 분모에 있는 marginal likelihood $p_\theta(\mathbf{x})$의 계산이 불가능하므로 posterior 또한 계산이 불가능합니다.

 

여기서 VAE의 핵심적인 아이디어가 나옵니다. 바로 true posterior $p(\mathbf{z}|\mathbf{x})$를 tractable한 분포인 approximate posterior $q_\phi(\mathbf{z}|\mathbf{x})$로 근사하는 변분 추론(variational inference)를 통해 최적화하는 것입니다.

 

이때 VAE에서 approximate posterior $q_\phi(\mathbf{z}|\mathbf{x})$는 신경망을 통해 모델링되고, 이는 VAE의 encoder에 해당하는 부분입니다. 즉, VAE에서 encoder는 decoder의 학습을 위해 도입되었습니다.

 

Figure 4. VAE의 training process

 

3) ELBO (Evidence Lower Bound)

학습 데이터 $\mathcal{D}=\{\mathbf{x}^{(1)}, \mathbf{x}^{(2)},\dots, \mathbf{x}^{(N)}\}$가 주어졌을 때, marginal likelihood는 $p_\theta(\mathcal{D})=\prod_{x\in\mathcal{D}}p_\theta(\mathbf{x})$입니다. 이를 개별 데이터 $\mathbf{x}$에 대한 합으로 표현하기 위해서 log-likelihood를 계산하면 $\log p_\theta(\mathcal{D})=\sum_{\mathbf{x}\in\mathcal{D}}\log p_\theta(\mathbf{x})$입니다. 따라서 이후 수식은 $p_\theta(\mathbf{x})$ 대신 $\log p_\theta(\mathbf{x})$를 최대화하는 것을 목표로 합니다. $\log p_\theta(\mathbf{x})$는 아래와 같이 전개됩니다.

 

$$ \begin{align} \log p_\theta(\mathbf{x})&=\mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})}\left[\log p_\theta(\mathbf{x})\right] \\ &=\mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})}\left[\log\left[\frac{p_\theta(\mathbf{x},\mathbf{z})}{p(\mathbf{z}|\mathbf{x})}\right]\right] \\ &=\mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})}\left[\log\left[\frac{p_\theta(\mathbf{x},\mathbf{z})}{ q_\phi(\mathbf{z}|\mathbf{x})}\frac{q_\phi(\mathbf{z}|\mathbf{x})}{p(\mathbf{z}|\mathbf{x})}\right]\right] \\ &=\underbrace{\mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})}\left[\log\left[\frac{p_\theta(\mathbf{x},\mathbf{z})}{q_\phi(\mathbf{z}|\mathbf{x})}\right]\right]}_{\mathcal{L}(\theta,\phi)~(\text{ELBO})}+ \underbrace{\mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})}\left[\log\left[\frac{q_\phi(\mathbf{z}|\mathbf{x})}{p(\mathbf{z}|\mathbf{x})}\right]\right]}_{D_{KL}(q_\phi(\mathbf{z}|\mathbf{x})||p(\mathbf{z}|\mathbf{x}))} \end{align} $$

 

$\mathcal{L}(\theta,\phi)$는 ELBO (evidence lower bound)로 정의되어 있는 부분이고, $D_{KL}(q_\phi(\mathbf{z}|\mathbf{x})||p(\mathbf{z}|\mathbf{x}))$는 approximate posterior $q_\phi(\mathbf{z}|\mathbf{x})$와 true posterior $p(\mathbf{z}|\mathbf{x})$의 분포 차이를 나타내는 KL divergence입니다.

 

한편 marginal likelihood $p_\theta(\mathbf{x})$는 evidence라고 부르기도 합니다. 따라서 evidence의 최소 하한인 ELBO (evidence lower bound)를 정의할 수 있으며, ELBO를 최대화하면 $\log p_\theta(\mathbf{x})$의 최적화와 동일한 효과를 얻을 수 있습니다. ELBO를 기준으로 수식을 정리하면 아래와 같습니다. 

 

$$ \mathcal{L}(\theta,\phi)= \underbrace{\log p_\theta(\mathbf{x})}_{\color{blue}{\text{objective}}}- \underbrace{D_{KL}(q_\phi(\mathbf{z}|\mathbf{x})||p(\mathbf{z}|\mathbf{x}))}_{\color{red}{\text{penalty}}} $$

 

이는 최대화하고자 하는 marginal log-likelihood가 objective로 주어지고, 최소화하고자 하는 KL divergence가 penalty로 더해진 형태입니다. 따라서 ELBO를 최대화하는 것은 marginal likelihood $p_\theta(\mathbf{x})$를 최대화하면서, approximate posterior $q_\phi(\mathbf{z}|\mathbf{x})$는 true posterior $p(\mathbf{z}|\mathbf{x})$와 유사하게 근사시키는 의미로 볼 수 있습니다.

 

그렇다면 ELBO를 어떻게 최대화할까요? 일단 위 수식은 intractable한 성분들로 구성되어 있습니다. Marginal likelihood는 우리가 구하고자 하는 모르는 부분이고, KL divergence 역시 true posterior를 모르므로 구할 수 없습니다. 대신 ELBO를 다음과 같이 정의에 따라 전개해보면 tractable한 성분으로 분해할 수 있습니다.

 

$$ \begin{align} \mathcal{L}(\theta,\phi)&= \mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})}\left[\log\left[\frac{p_\theta(\mathbf{x},\mathbf{z})}{q_\phi(\mathbf{z}|\mathbf{x})}\right]\right] \\ &=\mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})}\left[\log\left[\frac{p_\theta(\mathbf{x}|\mathbf{z})p(\mathbf{z})}{q_\phi(\mathbf{z}|\mathbf{x})}\right]\right]\\ &=\mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})}\left[ \log p_\theta(\mathbf{x}|\mathbf{z})-\log\left[\frac{ q_\phi(\mathbf{z}|\mathbf{x})}{p(\mathbf{z})}\right]\right] \\\\ &=\underbrace{\mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})}\left[\log p_\theta(\mathbf{x}|\mathbf{z})\right]}_{\text{Reconstruction term}}- \underbrace {D_{KL}(q_\phi(\mathbf{z}|\mathbf{x})||p(\mathbf{z}))}_{\text{Regularization term}}\end{align} $$

 

첫번째 reconstruction term은 encoder를 통해 추론된 $\mathbf{z}$를 decoder에 입력했을 때, 데이터 $\mathbf{x}$가 생성될 log-likelihood의 기대값으로 계산됩니다. Decoder 파라미터 $\theta$를 최적화하여 이 항이 최대화되면, decoder는 $\mathbf{z}$로부터 실제 데이터와 유사한 샘플을 복원(생성)합니다.

 

두번째 regularization term은 encoder로부터 추론된 $\mathbf{z}$의 분포와 prior distribution와의 KL divergence로 계산됩니다. VAE에서 prior distribution은 Gaussian으로 설정되고, $q_\phi(\mathbf{z}|\mathbf{x})$ 역시 Gaussian으로 가정하기 때문에 KL divergence를 계산할 수 있습니다. Encoder 파라미터 $\phi$를 최적화하여 이 항이 작아지면, encoder가 학습한 분포가 prior에서 크게 벗어나지 않게 제약을 줄 수 있습니다.

 

4) Reparameterization Trick

ELBO의 reconstruction term를 최적화하려면 \(\mathbf{z}\sim q_\phi(\mathbf{z}|\mathbf{x})\)의 샘플링이 필요합니다. 그러나 샘플링 연산은 미분이 불가능하기 때문에, encoder 파라미터 $\phi$에 gradient를 전파할 수 없는 문제가 발생합니다. 이를 해결하기 위해 VAE는 reparameterization trick을 사용합니다. Encoder는 데이터 $\mathbf{x}$로부터 평균 $\mu_\phi(\mathbf{x})$과 표준편차 $\sigma_\phi(\mathbf{x})$를 출력합니다. 이를 수식으로 표현하면 다음과 같습니다.

 

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

 

이때 $\mathbf{z}$를 위 분포에서 샘플링하지 않고, 아래와 같이 모델링하는 것이 reparameterization trick입니다.

 

$$ \mathbf{z} = \mu_\phi(\mathbf{x}) + \sigma_\phi(\mathbf{x}) \odot \varepsilon, \quad \varepsilon \sim \mathcal{N}(0,I) $$

 

여기서 $\varepsilon$은 모델과 무관하게 Gaussian에서 추출된 랜덤 값으로, $\sigma_\phi(\mathbf{x})$와 element-wise product 연산이 수행됩니다. 또한 $\mu_\phi(\mathbf{x})$와 $\sigma_\phi(\mathbf{x})$는 encoder의 출력이기 때문에 미분 가능한 연산입니다. 따라서 $\mathbf{z}$를 재파라미터화하면 미분이 가능해져 gradient가 encoder $\phi$까지 전파될 수 있으며, encoder-decoder 전체를 학습할 수 있게 됩니다.

 

Figure 5. Reparameterization trick

 

5) Loss Function

신경망 학습을 위해 ELBO에 -1를 곱하여 loss function으로 정의할 수 있으며, 이는 reconstruction loss와 regularization loss로 나눠 계산할 수 있습니다. 

 

■  Reconstruction loss

Decoder가 Bernoulli output을 가지는 경우 $p_\theta(\mathbf{x}|\mathbf{z})$는 multivariate Bernoulli distribution를 따릅니다. 데이터 $\mathbf{x}$의 차원을 $D$라고 할 때, reconstruction loss은 아래와 같이 binary cross entropy의 합으로 계산됩니다.

 

$$ -\log p_\theta(\mathbf{x}|\mathbf{z})=-\sum_{i=1}^D x_i\log y_i+(1-x_i)\cdot\log(1-y_i) $$

 

■ Regularization loss

Prior distribution $p(\mathbf{z})=\mathcal{N}(0,I)$와 approximate posterior $q_\phi(\mathbf{z}|\mathbf{x})$ 모두 Gaussian을 따른다고 가정하면, latent variable $\mathbf{z}$의 차원을 $J$라고 할 때 regularization loss는 아래와 같이 closed-form으로 계산됩니다. 자세한 증명은 논문 [1]을 참고바랍니다.

 

$$ D_{KL}(q_\phi(\mathbf{z}|\mathbf{x})||p(\mathbf{z}))=-\frac{1}{2}\sum_{j=1}^J\left(1+\log\sigma_j^2-\mu_j^2-\sigma_j^2\right) $$

 

Autoencoder와의 차이점 

Autoencoder와 VAE는 encoder-decoder 구조를 가진다는 공통점이 있습니다. 하지만 두 모델의 목적과 지향하는 바는 조금 다릅니다. Autoencoder는 입력 데이터를 encoder를 통해 latent vector로 압축했다가 decoder를 통해 다시 복원합니다. 압축된 latent vector는 deterministic한 고정값이므로, 새로운 데이터를 생성하기 위한 확률적 구조를 가지지 않습니다. Autoencoder는 차원 축소, 데이터 압축, 특징 추출 등의 압축 및 복원에 활용됩니다. VAE (Variational Autoencoder)는 확률적인 모델입니다. Encoder를 통해 latent space의 확률 구조를 학습하고, 샘플링된 latent vector를 통해 decoder는 원본 데이터와 유사한 새로운 데이터를 생성합니다. VAE는 encoder-decoder를 통한 압축-복원보다는, decoder를 통한 새로운 데이터를 생성하는데 더 강점을 가지고 있습니다.

 

Figure 6. Autoencoder와 VAE 비교

 

References

[1]  D. P. Kingma et al., "Auto-encoding Variational Bayes," arXiv, 2013.

[2] D. P. Kingma et al., "An Introduction to Variational Autoencoders," arXiv, 2019.

[3] C. Doersch et al., "Tutorial on Variational Autoencoders," arXiv, 2016.

[4] https://deepgenerativemodels.github.io/notes/vae

---

오타나 잘못된 부분 있으면 댓글 부탁드립니다. 도움이 되셨다면 공감 눌러 주시면 감사하겠습니다 :)