본문 바로가기

카테고리 없음

Variational AutoEncoder의 이해(feat.딥러닝의 확률 관점의 해석)

(Bayesian Theory과 Auto Encoder에 대한 기본적인 이해가 있다는 가정하에 작성되는 글입니다.)

 

머신러닝,딥러닝에서 Loss에 대한 확률적인 해석.

 

우리가 알고 있는 classic한 딥러닝 학습과정은 다음과 같습니다.

먼저 학습 데이터를 통해 training DB를 구성하고 $f_\theta(x)$로 표현되는 사용할 모델 구조를 정합니다. DNN,CNN,RNN 등을 예로 들 수 있습니다. 이때 모델의 파라미터 $\theta$는 사용하는 모델에 맞게 결정됩니다. 예를 들어 DNN을 사용한다면 각 층의 W와 b가 파라미터입니다.

Loss는 정답인 y와 출력인 $f_\theta(x)$의 다른 정도를 나타내며 Mean Square Error, Cross Entropy 와 같이 다양한 loss를 사용할 수 있습니다. 그 후 학습 과정을 통해 앞서 결정된 파라미터를 update하는데 이 과정은 $\theta^* = argmin_\theta L(f_\theta(x),y)$로 표현되고 일반적으로 Gradient Descent 방법을 통해 진행 됩니다. 학습이 완료되어 최적의 파라미터가 결정됐다면 새로운 입력에 대한 출력값을 얻어낼 수 있습니다. 이때 입력에 대한 출력은 항상 고정된 하나의 값을 얻습니다.

 

classic한 딥러닝의 경우 Loss를 정답과 출력이 다른 정도로 정의 했는데 이는 네트워크 출력값과 정답이 가깝기를 바랐기 때문입니다. 이제 Loss에 대해 확률적으로 생각하기 위해 "출력과 정답이 가깝기를 바란다"는 부분을 "네트워크의 출력이 given될 때 우리가 아는 정답이 나올 확률이 높기를 바란다"라는 해석으로 바꾸어 생각해보고자 합니다.  즉, likelihood $p(y|f_\theta(x))$값이 높기를 바라는 것으로 생각할 수 있습니다. Loss는 작을수록 좋아야 하기 때문에 이를 위해 $-\log(p(y|f_\theta(x)))$ 로 설정합니다.

해석을 위와 같이 바꿀경우 기존의 classic한 방법론과 달라지는 지점이 생깁니다. 사용할 모델 구조를 정하는 것에 추가적으로 도입한 likelihood가 어떤 확률 분포를 따르는지를 가정해주어야 합니다. 이 때 네트워크의 출력값은 가정된 확률분포의 파라미터 값으로 해석 됩니다. 예를 들어 likelihood를 정규분포로 가정했다면(편의상 분산은 무시) 해당 네트워크의 출력은 평균이 됩니다. 즉, $p(y|f_\theta(x))$는 평균값이 given되었을 때 y의 확률로 생각할 수 있습니다.  

따라서 위와 같은 상황에서 $p(y|f_{\theta_2}(x))$보다 $p(y|f_{\theta_1}(x))$ 값이 더 크기 때문에 $f_{\theta_1}(x)$을 더 좋은 모델로 생각할 수 있습니다.

 

위와 같이 Loss에 대해 확률적인 관점을 도입한 후 학습을 통해 최적의 모델 $f_{\theta^*}(x)$를 얻어냈다면 classic한 방법과 달리 최적 모델을 확률 분포 모델로 해석합니다. 즉, $y_{new} ~ p(y|f_{\theta^*}(x_{new}))$로 $y_{new}$는 새로운 입력 $x_{new}$가 모델을 통과하여 얻어진 확률 분포 $p(y|f_{\theta^*}(x_{new}))$에서 sampling하여 얻어지는 값으로 생각할 수 있게 됩니다.

 

이를 통해 Loss에 확률적인 관점을 도입하면 고정 입력에 대해 고정 출력을 얻는 classic한 방법과 달리 고정 입력에 대해 sampling이 가능해진다는 장점을 가지고 있습니다. 이러한 관점은 이어서 설명드릴 Variational AutoEncoder에서 중요하게 사용되는 관점입니다.

 

Variational AutoEncoder

VAE는 AutoEncoder와 마찬가지로 앞단과 뒷단으로 구성된 대칭적 구조를 가지지만 해당 구조를 가지게 된 배경이 AE와 전혀 다릅니다. 

먼저 뒷단에 해당하는 generation model(generator)이 이미지를 생성하는 딥러닝 모델이라고 가정하겠습니다. generator는 input 이미지 x와 동일한 이미지를 생성하는 것을 목표로 합니다. generation model에 대한 전체적인 notation은 아래 이미지와 같습니다.

 

 

generator $g_\theta (\cdot)$은 latent variable z를 입력 받아 이미지를 생성합니다. z는 prior $p(z)$로부터 sampling 됩니다. 앞서 확률론적 해석에서 설명드렸듯이 모델의 출력 $g_\theta(z)$는 likelihood $p(x|g_\theta(z))$ (이하 $p_\theta(x|z)$)에 대해 가정된 확률분포의 파라미터값입니다. 따라서 생성되는 이미지 x는  $g_\theta(z)$을 파라미터 값으로 가지는 가정된 확률분포 $p(x|g_\theta(z))$에서 sampling되었다고 해석 가능합니다.

 

이때 latent variable z를 생성하는 prior $p(z)$를 정규 분포와 같은 다루기 쉬운 확률 분포를 사용함으로써 이미지 생성 과정을 조절하고자 합니다. 한가지 의문점은 '복잡한 이미지를 생성하는 manifold가 단순한 분포를 통해 표현되는가'인데 이는 sample z가 generator $g_\theta(\cdot)$를 통과하는 것으로 해결됩니다. 예를들어 generator가 여러개의 layer로 구성된 Multi-layered neural network라면 초반 몇개의 layer가 단순한 분포에서 sampling된 z를 적절한 manifold로 mappping하는 역할을 하고 그 이후의 layer들에서 가정된 확률분포의 파라미터 값을 얻어내는 과정을 진행한다고 생각할 수 있습니다. 적절한 manifold로 mapping하는 것은 VAE 논문에 있는 아래의 그림으로 예를 들 수 있습니다.

 

이러한 generator를 학습하고자 하는데 일반적으로 네트워크를 학습할 때와 동일하게 Loss인 $-\log(p_\theta(x|z))$를 최소화 즉, Maximum Likelihood Estimation(이하 MLE)를 진행하게 되면 문제가 발생합니다. 예를 들어 아래와 같은 이미지가 존재한다고 생각해보겠습니다.

 

(a)는 실제 sample이고 (b)는 일부분이 변형된 이미지 (c)는 단순히 실제 sample을 한 픽셀만큼 오른쪽으로 옮긴 이미지 입니다. 

generator의 likelihood $p_\theta(x|z)$를 Gaussian 분포로 가정을 한다면 해당 loss는 MSE로 해석이 가능합니다. 이 경우, 실제 (c)가 (b)보다 의미적으로 실제 sample과 유사함에도 MSE는 (b)가 더 작기 때문에 generator가 의미적으로 다른 (b)와 더 유사한 이미지를 생성하게끔 학습이 되는 문제가 발생합니다.

 

따라서 latent variable z를 prior에서 sampling하여 사용하는 대신에 이상적인 sampling함수를 도입 하고 해당 함수에서 z를 sampling하여 사용하고자 합니다. 이상적인 sampling함수는 input 이미지 x와 비슷한 생성물을 잘 만들어내도록 하는 z를 sampling 해주는 함수를 의미합니다. x를 evidence로 받아 generator를 거쳤을 때 적어도 x는 잘 생성되게 하는 z를 sampling하도록 해주기 때문에 $p(z|x)$로 표현됩니다.  우리는 이러한  함수(확률분포)가 무엇인지 모르므로 이를 추정해야하는데 이때 Variational Inference 이하 VI라는 개념이 사용됩니다.

 

Variational Inference   

 

우리가 모르는 true posterior을 추정하고자 할때 쓰이는 방법입니다. 다루기 쉬운 확률 분포를 하나 정하고 그 확률 분포의 parameter를 잘 조절해가며 최대한 true posterior에 가깝게 만들고 이렇게 approximated 분포를 기존의 true posterior 대신에 사용하는 방식을 사용합니다. 즉, 앞서 이상적인 sampling함수 $p(z|x)$에 대해 VI를 적용하여 얻은 approximated 분포 $q_\phi(z|x)$를 대신하여 사용할 수 있습니다.

VI를 적용하기 위해 $p(x)$와 $p(z|x)$ 그리고 $q_\phi(z|x)$와의 관계를 알아보겠습니다.

 

다음과 같이 $\log{p(x)}$가 Evidence Lower BOund와 $p(z|x)$와 $q_\phi(z|x)$의 KL Divergence로 이루어져 있음을 확인 가능합니다.

이 때 $\log{p(x)}$는 모르는 상수값이 때문에 ELBO값과 KL값이 상보적인 관계를 이룸을 알 수 있습니다. 

$p(z|x)$를 잘 approximate하는 $q_\phi(z|x)$를 만들기 위해서는 두 분포간의 KL값이 작아야 하는데 이와 같은 상보적인 관계가 있음을 알기 때문에 KL값을 낮추는 것은 ELBO값을 높이는 것과 같음을 알 수 있습니다. 즉, VI의 optimization 문제는 max $ELBO(\phi)$로 표현 됩니다.

이렇게 이상적인 sampling함수에 VI를 적용하는 과정은 VAE의 앞단 부분을 이루게 됩니다.

 

VAE의 내용을 정리하면 generator로 input x를 generation하고자 하는 부분과 이러한 generation이 잘 되도록 하는 generator의 입력을 잘 만들어주도록 하는 이상적인 sampling함수를 approximate하는 부분으로 나누어져 있고 각각 모델의 뒷단과 앞단을 구성한다고 볼 수 있습니다.

 

 

VAE가 optimize해야하는 문제는 총 2가지 입니다. 

첫번째는 이상적인 sampling함수를 approximate하는 부분입니다. 앞서 VI부분에서 설명드렸듯이 $ELBO(\phi)$를 maximize함으로 얻어낼 수 있습니다. 해당 optimization문제는 $\phi$에 대한 optimization문제입니다.

두번째는 추정한 sampling함수에서 sampling된 z를 통해 x를 생성하는 generator를 학습하는 부분입니다. 확률론적인 관점에서 네트워크(generator)를 바라보고 있기 때문에 likelihood $p_\theta(x|z)$를 maximize함으로 얻어낼 수 있습니다. 해당 optimization문제는 $\theta$에 대한 optimization문제입니다. 여기서 주목할 점은 $p_\theta(x|z)$가 $ELBO(\phi)$식에 포함되어 있다는 점입니다. 앞서 나온 $ELBO(\phi)$식을 전개하면 다음과 같습니다.

 

 

 

이와 같이 $E_{q_\phi(z|x)}[\log{p(x|z)}]$로 표현되는 likelihood(좀 더 자세하게는 likelihood 평균)가 식에 포함됨을 알 수 있습니다.따라서 두 optimization문제 모두 EBLO를 최적화의 목적식으로 하되, 최적화 대상만 각각 $\phi$와 $\theta$로 다르므로 최종 VAE의 Loss는 다음과 같이 표현 됩니다.

 

이때 $L_i(\phi,\theta,x_i)$에서 앞 term은 $x_i$를 evidence로 줬을때 지금까지의 appoximated 샘플링 함수에서 z를 샘플링하고 그 z를 입력으로 하는 generator를 통해 $x_i$가 생성될 확률로, 의미적으로는 $x_i$를 줬을 때 다시 $x_i$가 나오느냐의 reconstruction error로 생각할 수 있고 뒷 term은  $q_\phi(z|x_i)$와 $p(z)$의 KL Divergence로 같은 reconstruction error를 갖는 $q_\phi(z|x_i)$가 여러 개라면 기왕이면 prior와 같은 모양이 되라는 뜻으로 regularization으로 생각할 수 있습니다.

 

Loss의 목적식인 ELBO가 계산이 된다면 쉽게 코딩으로 구현이 가능합니다.

 

ELBO 계산 과정

(1)Regularization term

먼저 prior 와 approximation sampling 함수에 대한 가정을 진행합니다. 

$p(z) \sim N(0,I)$ 그리고 $q_\phi(z|x_i) \sim N(\mu_i , \sigma^2_iI)$ 로 둘다 Gaussian 분포를 따른다고 가정합니다. 다른 분포로도 가정이 가능하지만 Gaussian 분포들 사이의 KL divergence를 closed-form으로 구하는 방법이 이미 존재하여 위와같이 가정합니다. 해당 가정에 따른 KL Divergence값은 아래와 같습니다.

 

 

(2)Reconstruction Error term

 

reconstruction error인 $E_{q_\phi(z|x_i)}[\log{p_\theta(x_i|z)}]$ 값은 Monte-carlo sampling을 통해 실제 적분값을 구하지 않고 근사하여 사용합니다. 이 때 $q_\phi(z|x_i)$로부터 총 L개의 샘플을 뽑아 평균값을 구합니다.

reconstruction error는 $p_\theta(x_i|z)$에 대해 어떤 확률분포를 가정하는지에 따라 다르게 해석됩니다. 

 

(2)-1 베르누이 분포 가정

베르누이 분포를 가정했을 경우 reconstruction error를 전개해보면 cross entropy와 같음을 확인 가능합니다. 확률론적인 관점에서 네트워크의 출력은 가정한 확률분포의 parameter 임으로 베르누이 분포의 parameter인 $p_i$가 출력됨을 볼 수 있습니다.

 

(2)-2 가우시안 분포 가정

 가우시안 분포를 가정했을 경우 reconstruction error를 전개해보면 MSE와 같음을 확인 가능합니다. 역시 확률론적인 관점에서 네트워크의 출력은 가우시안 분포의 parameter인 평균과 분산임을 볼 수 있습니다.

 

이렇게 Regularization과 Reconstruction Error의 합으로 ELBO를 계산할 수 있습니다.

loss인 -ELBO가 계산이 되었으므로 학습을 진행할 수 있는데 Backpropagation에서 한가지 문제가 발생합니다. z가 chain rule을 통해 backprop되어야 하는데 z는 random variable이기 때문에 이러한 과정이 불가능합니다. 따라서 Reparameterization Trick을 통해 Backpropagation을 가능하게 해줍니다. 

 

 

새로운 random node인 error항을 도입하여 원래의 분포를 유지하면서 z를 deterministic하게 바꾸어주기 때문에 $\phi$에 대한 Backprop가 문제 없이 진행이 가능합니다.

 

최종적인 VAE의 예시 그림을 보여주며 마무리 짓겠습니다.

 

예시는 MNIST 이미지를 생성하는 VAE로, MLP encoder와 MLP decoder로 구성되며 이에 대한 확률 분포 가정은 각각 가우시안 분포와 베르누이 분포입니다. 원활한 학습을 위해 reparameterization trick도 반영 됐음을 볼 수 있습니다. 이러한 VAE를 학습한 결과 적절하게 MNIST 데이터를 잘 생성함을 확인할 수 있습니다. 만약 이러한 VAE에 대해 추가적인 성능 향상을 원한다면 latent variable z의 차원인 J를 조절하거나 encoder와 decoder의 모델구조를 MLP가 아닌 다른 구조로 변경해 볼 수 있겠습니다.

 

 

 

reference:https://www.slideshare.net/NaverEngineering/ss-96581209

이활석님의 오토인코더 강의를 참고하여 작성하였습니다.