[논문리뷰] Auto-Encoding Variational Bayes(VAE) 모든 수식 알아보기

0. 들어가며

제목은 모든 수식을 알아본다고 호기롭게 썼으나 제가 이해한 만큼만 이 포스팅에 담길 예정입니다. 학습이 어떤식으로 이루어 지는지에 초점을 맞추기 보다는 왜 loss가 이런 식으로 구성되었고, Autoencoder를 사용한 이유 등 남들은 조금 궁금해 하지 않았던 내용도 다룰 예정입니다.

 

Autoencoder는 아주 예전부터 등장했던 아이디어 입니다. 보통 차원 축소나 denoising등을 위하여 많이 쓰였으나 Kingma는 이 구조를 Likelihood들 approximation하는데 사용합니다.

 

위 그림에서 보면 Variational Autoencoder는 Explicit density안의 approximate density에 분류되어 있습니다. 즉, 모든 Generative model의 목표인 Input dataset의 분포를 우리가 알고있는 어떠한 분포로 Explicit하게 approximate하는 방법론 이라는 뜻입니다. VAE에서는 간단하게 설명해서 Input dataset에 어떤 latent distribution Z가 있고, 이 Z를 우리가 아는 gaussian distribution 등으로 근사시킵니다.

 

지금 VAE 자체는 많이 사용되지 않지만, 현재 많은 연구가 이루어지고 있는 Diffusion Model이 결국 VAE를 여러 층 쌓은 모델이기 때문에 반드시 알아야 하는 모델입니다.

 

1. DeepLearning의 Maximum Likelihood 관점 해석

일반적으로 딥러닝의 classification, regression model을 생각하면, Input이 입력으로 들어와 blackbox model 안에서 가중치 업데이트가 일어나 고정된 하나의 값을 output으로 주는 모델이라고 생각합니다. 하지만, 단순히 이렇게 해석하는 것이 아니라 output 값을 어떤 분포의 Likelihood를 Maximize하게 하는 파라미터로 볼 수 있습니다. Maximum Likelihood는 다음 포스팅에서 보실 수 있습니다.

 

https://rla020.tistory.com/36

Maximum Likelihood Estimation(MLE) 이해해보기

 

예를 들어 어떤 아파트의 집값을 예측하는 모델을 만들고자 한다고 가정하겠습니다. 이 때 Input X로는 (평 수, 주변 편의시설 수, 위치, 거주자의 소득 수준...) 등이 들어 갈 수 있고, ground truth Y로는 실제 집값이 들어갈 것입니다. 그러면 Supervise방법으로 간단히 모델을 training 시킬 수 있습니다. MSE Loss를 통해 모델을 학습시킨다고 하면 실제 y와 모델의 output의 Loss $(y-f(x))^2$이 계산되어 가중치를 업데이트 합니다.

하지만 여기에서 y가 어떤 deterministic한 값이 아니라 특정 파라미터를 갖는 분포라고 '가정'을 해 볼 수 있습니다. 간단히 말해서 특정 Input x가 주어졌을 때 집값 y가 $\mu,\sigma$를 갖는 가우시안 분포라고 가정을 하는 것 입니다. 이렇게 가정하고 어떤 x에 대한 y의 maximum log likelihood(= minimize negative log likelihood)는 다음과 같습니다.

신기하게도 y를 gaussian이라고 가정한다면 MSE, Bernoulli로 가정한다면 Cross-entropy가 나오게 됩니다. 직관적으로 해석해보면 다음과 같습니다.

• 가중치가 $\theta_2$일 때의 model $f_{\theta_2}(x)$에서는 그림과 같이 y의 likelihood가 굉장히 낮게 나온다. 
• Gaussian에서 Likelihood는 평균($\mu$)일 때 maximize 된다.
• negative log likelihood를 이용하여 likelihood를 minimize시키려면 model의 output $f_{\theta}(x) = \mu = y$가 되어야 한다.
• 즉, model의 output은 y가 되어야 하고 MSE를 이용하여 가중치를 업데이트 한다.

결국 model의 output과 y를 같게 만들어 주는것이 모델의 목표가 되는데, 이것을 deterministic하게 보면 모델이 strict한 output $\hat{y}$를 예측하는 것으로 볼 수 있고, Maximum Likelihood관점에서 보면 어떤 distribution의 parameter를 예측하는 것으로도 볼 수 있는 것 입니다. 이렇게 어떤 분포를 모델링하면 그 분포에서 Sampling 또한 진행 할 수 있게 됩니다.

 

2. Manifold 가정

200x200 RGB 채널에서 random하게 값을 뽑아서 plot하면 어떻게 될까요?

10만,100만 번을 sampling해도 대부분 그림과 같은 noise가 나오게 됩니다. 하지만 이 분포 안에는 어떤 인물의 이미지, 스케치, 글자, 숫자 등등이 포함되어 있습니다. 하지만 random sampling을 했을 때 그런 이미지가 나오지 않는 이유를 생각해보면 " 어떤 특정 공간에 데이터가 밀집해 있다"라고 생각할 수 있고 이것을 Manifold 가정이라고 합니다.

위의 그림을 swiss roll이라고 하는데, 저렇게 3차원 공간 상에서 데이터는 2차원의 manifold에 몰려 있다고 생각 할 수 있습니다. 그리고 manifold를 잘 찾았다면 manifold내의 좌표에 따라 이미지는 유의미하게 변하게 됩니다.

Manifold는 수학 쪽에서 더 심도있게 다루지만 저희는 "어떤 고차원의 공간이 존재할 때 데이터는 저차원의 manifold에 존재 한다는 것을 manifold 가정이라고 한다" 라고만 알고 넘어가겠습니다.

 

3.Variational AutoEncoder

Variational AutoEncoder는 어떤 Input dataset의 분포 $p(x)$를 Variational approximation을 통해 구하고자 합니다. Variational approximation으로 manifold상의 z를 Gaussian 분포를 통해 근사하는데, 이때 모델(Encoder)은 gaussian의 파라미터를 예측합니다. 그리고 근사한 Gaussian을 통해 z를 sampling하여 모델(Decoder)은 x를 다시 생성 해내게 됩니다.

 

3.1 Why variational approximation?

그런데 굳이 variational approximation이라는 방법을 써야하나? 라는 의문이 들 수 있습니다. 그냥 z를 gaussian noise 같은 noise input으로 받아서 Decoder를 통과시켜서 x를 복원하는 모델을 써도 될 것 같습니다.(MSE 사용)

이 때 $p(x|g_\theta(z))$를 gaussian이라고 가정해서 파라미터를 찾는다고 생각하면,

위의 수식을 통해 $p(x)$를 구할 수 있고,

몬테카를로 방법으로 적분을 근사 할 수 있습니다. 

 

하지만 당연히 이런 간단한 방법에는 문제가 있습니다.

위 그림을 보면, (a)를 기준으로 (b)는 앞쪽이 잘려있는 형태이고, (c)는 한 pixel씩 오른쪽으로 shift 돼있는 이미지 입니다. 단순히 MSE를 사용한다면, (b)가 (c)보다 MSE가 작게 나올 것 입니다. 다른 부분이 앞쪽 밖에 없기 때문이죠. 그렇기 때문에 '의미상' 가까운 (c)가 아닌 (b) 쪽으로 학습이 되기 때문에 올바른 학습이 진행되지 않게됩니다.

 

3.2 Variational approximation

그래서 VAE에서는 z를 gaussian같은 분포가 아니라 x와 유사한 sample이 나올 수 있는 분포 $P(z|x)$에서 sampling하는 것을 목표로 합니다. 하지만 $P(z|x)$의 실제 분포는 알 수가 없습니다. $p(x), p(z)$등의 분포를 모르기 때문이죠. 따라서 variational approximation을 사용합니다. 해당 방법은 어떤 분포를 우리가 아는 분포로 최대한 가깝게 만드는 방법입니다.

위와 같이 우리가 원하는 실제 분포 $P(Z|X)$가 있다면 이 분포를 우리가 아는 gaussian같은 분포 $Q(z)$로 근사할 수  있고 이를 variational approximation이라고 합니다. 이렇게 되면 sampling도 쉽게 할 수 있게 되죠. 그러면 어떻게 variational approximation을 진행하는지 알아보겠습니다.

 

1. log likelihood $log(p(x))$에 $\int q_\phi(z|x)dz$를 추가합니다. 이 때 $\int q_\phi(z|x)dz$는 z에 대한 적분이기 때문에 1입니다.
2. $p(x) = \frac{p(x,z)}{p(z|x)}$공식을 이용해서 풀어줍니다.
3. 해당 식을 쭉 풀어 보면 위와 같이 $ELBO(\phi)$ term과 $KL$ term이 나오게 됩니다.

정리해보면 $log(p(x))$는 ELBO와 KL Divergence를 합친 값 입니다. 우리가 원하는 상황은 우리의 모델 $q_\phi(z|x)$가 실제 $p(z|x)$와 같아지는 상황이기 때문에 KL term을 최소화 하면 됩니다. 하지만 안타깝게도 우리는 KL Divergence term을 구할 수 없습니다. 두 분포 사이의 유사도가 KL Divergence이지만 $p(z|x)$라는 실제 분포는 알지 못하기 때문에 variational approximation을 이용하게 된건데, 이용하는 과정에 또 $p(z|x)$가 나온 상황이죠. 하지만 다시 생각해보면 log likelihood $log(p(x))$가 어떤 상수로 고정되어 있다고 하면,KL Divergence를 최소화 하는 것은 반대로 ELBO를 최대화 하는 것과 같게됩니다.

따라서 우리의 목표는 ELBO를 최대화 하는 것으로 변환되었습니다.

그리고 ELBO를 $p(x,z) = p(x|z)p(z)$를 이용해서 풀면 맨 밑과 같은 식이 도출됩니다. Loss라는 것은 결국 최소화 해야하는 것이기 때문에 minus를 붙여 살펴보면

다음과 같이 Reconstructioin term과 Regularization term이 나오게 됩니다.

Reconstruction term은 encoder $q_\phi(z|x)$로 부터 나온 z를 Decoder에 통과시켜 예측한 gaussian의 파라미터가 되고,

Regularization term은 encoder가 다루기 쉬운 확률 분포 $p(z)$(gaussian 같은)가 되게 합니다. variational approximation을 위해서 모델이 gaussian이여야 한다고 이미 앞에서 한차례 말을 했습니다.

 

이 때 Reconstruction term은 output이 gaussian이라고 가정했기 때문에, MSE Loss를 사용하여 모델을 학습시키면 됩니다. 그리고 Regularization term 같은 경우 두 분포가 모두 gaussian이기 때문에 다음과 같이 쉽게 계산이 가능합니다.

따라서 전체적인 도식을 보면 다음과 같습니다.

1. Input x가 입력으로 들어가서 $p(z|x)$에 대한 variational approximation을 이용하기 위해 gaussian의 $\mu,\sigma$를 예측한다. 
2. approximation된 분포 $q_\phi(z|x)$로 부터 z를 sampling한다.
3. sampling된 z를 이용해서 gaussian이라고 가정한 x의 분포에 대한 $\mu,\sigma$를 예측한다.
4. 분포를 알기 때문에 분포에서 sampling을 진행할 수 있게 된다.
5. 최종 loss는 Reconstruction Error + Regularization Error가 된다.

베르누이 분포에 대해서는 다음과 같습니다.

 

3.3 Reparameterization Trick

마지막으로 우리는 z를 Encoder가 예측한 $\mu,\sigma$를 가진 gaussian에서 sampling한다고 했습니다. 하지만 여기서 문제가 발생하는데 $N(\mu,\sigma)$에서 바로 sampling을 하게 된다면 backpropagation을 할 수 없게 됩니다. 그래서 나온 방식이 reparameterization trick입니다. 이 방식은 $N(\mu,\sigma)$에서 바로 sampling하는 것이 아니라 먼저 normal gaussian에서 $\epsilon$을 sampling한 다음 $\sigma$를 곱해주고 $\mu$를 더해줍니다. 이런 방식으로 normal gaussian의 분포를 이동시킬 수 있기 때문에 $N(\mu,\sigma)$에서 바로 sampling한 것과 같은 효과를 주고 back propagation도 가능하게 됩니다.

 

4. 마치며

VAE는 log likelihood를 이용하여 분포를 approximation하는 모델입니다. GAN은 특정 data class만 잘 생성하는 mode collapse 문제가 있지만 VAE는 모든 data class를 잘 생성한다고 알려져 있습니다. 하지만 각 픽셀별 MSE나 cross entrophy를 사용하기 때문에 이미지가 blurry하게 나온다는 문제 또한 가지고 있습니다. 

 

 

출처: 이활석 - 오토인코더의 모든 것