2017년 5월 6일 토요일

초짜 대학원생의 입장에서 이해하는 Auto-Encoding Variational Bayes (VAE) (3)


저번 글에 이어 오늘은 reparameterization trick에 대해 알아보고 VAE에 대한 소개를 마치도록 하겠습니다. 

The reparameterization trick: why?


우리가 풀고 싶어하는 문제가 이제 다음과 같은 lower bound를 maximize하도록 $\theta$와 $\phi$ parameter들을 조정해주는 것으로 바뀌었다는 것까지 지난 글에서 얘기를 했었습니다:
$$\cal{L}(\theta,\phi;x^{(i)})=-D_{KL}(q_\phi(z|x^{(i)})||p_\theta(z))+\mathbb{E}_{q_\phi(z|x^{(i)})}\left[\log p_\theta(x^{(i)}|z) \right].$$
$$(\theta^*,\phi^*)=\arg\max_{\theta,\phi} \cal{L}(\theta,\phi;x^{(i)}),$$
많은 경우, $D_{KL}(q_\phi(z|x^{(i)})||p_\theta(z))$는 적분 수식을 직접 풀어 계산할 수 있습니다. 왜 그런지에 대해서는 먼저 그렇다는 것만 알아두시고 (논문에서도 appendix로 빠져있습니다) 이 부분에 대해서는 뒤에 reparameterization trick 부분을 설명하고 난 이후에 부록처럼 설명을 달도록 하겠습니다.

위와 같이 KL divergence는 해결할 수 있다고 한다면, 이제 우리에게 남는 문제는 expected reconstruction error에 해당하는 $\mathbb{E}_{q_\phi(z|x^{(i)})}\left[\log p_\theta(x^{(i)}|z) \right]$를 어떻게 잘 풀어줄 것인가입니다.  저번 글 마지막에서 언급했듯이 이 부분이 문제가 되는 이유는 expectation을 풀기 위해 $z\sim q_\phi(z|x)$에 대한 sampling이 필요하기 때문입니다.

일단 가장 처음 쉽게 생각해볼 수 있는 문제로 $z$를 매우 많이 sampling해서 estimate를 구해야하고 따라서 시간이 많이 걸린다는 점이 떠오릅니다. 뭐 이 부분은 그래도 어짜피 우리가 stochastic gradient  descent를 사용하는 마당에 $x^{(i)}$ 하나당 $z^{(i)}$도 하나씩 뽑아서 $p_\theta(x^{(i)}|z^{(i)})$를 계산하여 마치 그 값이 $\mathbb{E}_{q_\phi(z|x^{(i)})}\left[\log p_\theta(x^{(i)}|z) \right]$에 대한 estimate라고 생각하면 별 문제가 없어보이기도 합니다.

한 번 그렇게 문제를 풀어보도록 하겠습니다. 원래 문제에서 $\mathbb{E}_{q_\phi(z|x^{(i)})}\left[\log p_\theta(x^{(i)}|z) \right]$ 부분을 위에서 구한 single value $x^{(i)}$와 $z^{(i)}$에 대한 esitimate으로 바꿔준 다음, 전체 데이터셋 $D$에 대해서 평균을 구하는 것으로 문제가 바뀝니다:
$$\mathbb{E}_D[\cal{L}(\theta,\phi;x^{(i)})]=\mathbb{E}_D\left[-D_{KL}(q_\phi(z|x^{(i)})||p_\theta(z))+\log p_\theta(x^{(i)}|z^{(i)})\right].$$
그럼 위 식에 대해 gradient를 구해서 하던대로 backpropagate하면 되겠군요. 심지어 sample을 많이 뽑아 gradient의 평균값을 구하게되면 당연하게도 원래 식과 같은 gradient로 수렴하게 되어있습니다. 어? 그러면 모든 문제가 해결된 것 아닌가요?

아쉽게도 이렇게 sampling을 하면서 생기는 문제가 있습니다. 본래 식에 있던 $\mathbb{E}_{q_\phi(z|x^{(i)})}\left[\log p_\theta(x^{(i)}|z) \right]$ estimate은 $q_\phi(\cdot)$에 dependent한 값인데 이를 대체한 $\log p_\theta(x^{(i)}|z^{(i)})$를 보시면 더이상 $q_\phi(\cdot)$와는 관계가 없어져 버립니다.

이게 어떤 문제인지 가만히 생각을 해보겠습니다. 우리가 원하는 것은 모델 $q_\phi(\cdot)$에 대한 parameter $\phi$를 "잘" update하는 것입니다: $$(\theta^*,\phi^*)=\arg\max_{\theta,\phi} \cal{L}(\theta,\phi;x^{(i)}).$$
여기서 "잘"이라 함은  $q_\phi(\cdot)$가 $p_\theta(x|z)$로부터도 피드백을 받아서 좋은 encoder가 되기를 바라는 것인데 위와 같이 관계가 없어져 버린 이상 미분을 하면 더이상 제대로 학습을 할 수가 없게 되어버립니다.

이 것을 매우 잘 설명해주는 그림이 있습니다(분명 전 시간에 그림을 많이 넣겠다고 했는데...이제야 첫 그림이라니...orz...):

약간의 notation 차이가 있습니다만 알아보는데 그리 어렵지는 않을 것입니다. 여기서 $\mu$와 $\Sigma$는 모델 $Q$를  Gaussian으로 정했을 때 자연스래 따라오는 parameter $\phi$입니다. 왼쪽에 있는 모델이 현재 sampling을 했을 때의 VAE 모델이 되겠습니다. 빨간색 박스로 $z$가 normal distribution으로부터 sampling되고 있다는 것을 보여주고 있죠.

이 문제를 (좀 다른 관점으로) 저번 글의 설명을 가져와서 재언하자면 "sampling"은 미분이 가능한 연산이 아니기 때문에 생기는 문제라 볼 수 있습니다. Gradient descent로 문제를 푼다는 것은 모델이 parameter에 대해 미분이 가능하고 이는 모델이 어떤 의미에서는 deterministic하다는 것을 의미합니다. 즉, fixed parameter들에 대해 stochasticity는 input에만 있고 같은 input에 대해서는 항상 같은 output이 나와야 하는데 이 "sampling"이란 녀석은 모델 자체에 stochasticity를 넣어버리기 때문에 문제가 됩니다.

이제 문제를 알았으니 해결할 방법을 알아봐야겠죠. 원래 문제를 잘 아는 것부터 시작이라 학교에서 배웠습니다!! (글 두 편을 문제 설명하는데만 날린 자의 변명...)

이 문제를 교묘하게 잘 피하고자 사용된 방법이 바로 reparameterization trick입니다. 여기서 "trick"이라고 했듯이 약간 technical tip과 같은 느낌의 단순한 식 변환으로 위에 설명한 문제를 모두 깔끔히 해결합니다.

The reparameterization trick


사실 이 녀석도 이름이 다 말해주고 있습니다. 지금 문제가 되고 있는 변수 $z\sim q_\phi(z|x)$를 잘 변환해서 미분이 가능한 녀석으로 만들어보자: 즉, "reparameterization"을 해서 뭔가 해결을 보겠다는 것이죠. 이 때 사용되는 변환 함수를$g_\phi (\epsilon,x)$라는 미분 가능한 함수와 보조 noise variable $\epsilon$로 나타내겠습니다:
$$\tilde{z}=g_\phi (\epsilon,x)~~with~~\epsilon\sim p(\epsilon).$$
물론 어떤 $q$에 대해 $g$를 사용하여 reparameterize 할 수 있는가에는 몇 가지 조건들이 더 필요합니다만 이 부분은 큰 틀을 보기 위해 잠시 넘어가겠습니다.

어찌 되었든 $g_\phi (\epsilon,x)$를 사용하면 어떻게 문제가 해결된다는 것인지 살펴봐야겠죠. 임의의 함수 $f(z)$의 $q_\phi(z|x)$에 대한 Monte Carlo expectation estimate은 다음과 같습니다:
$$\mathbb{E}_{q_\phi(z|x^{(i)})}\left[f(z)\right]=\mathbb{E}_{p(\epsilon)}\left[f(g_\phi (\epsilon,x^{(i)}))\right]=\frac{1}{L}\sum_{l=1}^L f(g_\phi (\epsilon^{(l)},x^{(i)})),~~where~~\epsilon^{(l)}\sim p(\epsilon).$$
이걸 앞에서 single value estimate을 넣었듯이 원 식에 갈아끼면 다음과 같이 변합니다:
$$\tilde{\cal{L}}(\theta,\phi;x^{(i)})=-D_{KL}(q_\phi(z|x^{(i)})||p_\theta(z))+\frac{1}{L}\sum_{l=1}^L (\log p_\theta(x^{(i)}|z^{(i,l)})), \\where~~z^{(i,l)}=g_\phi (\epsilon^{(i,l)},x^{(i)})~~and~~\epsilon^{(l)}\sim p(\epsilon).$$
오~~~뭔가 바뀐거 같은데 이게 수식으로만 봐선 잘 와닿지도 않고 그래서 뭐? 이런 느낌이죠.

대학원생의 풀이


위에 수식을 잘 보시면 더이상 $z$가 아니라 $\epsilon$에  stochasticity가 더해지는 것을 보실 수 있습니다. 즉, 이 부분이 중요합니다: $z^{(i,l)}=g_\phi (\epsilon^{(i,l)},x^{(i)}),~~where~~\epsilon^{(l)}\sim p(\epsilon)$. 따라서 이전에는 $z\sim q_\phi(z|x)$ model로부터 직접 sampling되었기 때문에 모델 자체에 stochasticity가 들어가는 문제가 있었지만 이를 이제는 우회할 수 있게 됩니다.

좀 더 쉬운 이해를 위해 구체적인 예시를 보여드리겠습니다. 예를 들어 우리가 고른 모델이 Gaussian이라서 $Q(Z|X)=\cal{N}(\mu(X),\Sigma(X))$라 나타낼 수 있다면 이전에는  이 모델 distribution으로부터 바로 sampling한 stochastic variable이었던 $z$가 이제는 $\epsilon\sim \cal{N}(\mathbb{0},I)$에서 sampling한 $\epsilon$에 대해 $z=\mu(X)+\Sigma^{1/2}*\epsilon$이라고 나타낼 수 있게 됩니다.

이...뭐....눈 가리고 아웅?.느낌일 수 있습니다. 하지만! 둘의 차이는 확실하죠. 앞에 $z$는 parameter $(\mu,\Sigma)$에 대해 미분이 가능한 형태가 아닌 반면, reparameterization trick을 거친 $z$는 $(\mu,\Sigma)$에 대해 미분이 가능합니다. 그래서 첫 그림 오른쪽과 같은 모델을 보셔도 되지만 편의를 위해 또 다른 그림을 가져오면 다음과 같습니다.


이제는 이해가 되시려나요. 오른쪽을 보시면 $z$의 구성요소에서 $\epsilon$이 들어가면서  $\phi$에 대해 backpropagation이 가능해지는 것을 보실 수 있습니다. 이를 위해서 맨 처음에 얘기했듯 $g$라는 함수가 미분이 가능한 연속 함수여야 한다는 제약조건과 함께 Gaussian의 예시에서 보여드린 것과 같이 이렇게 reprameterization을 했을 때 동치가 이루어지는 함수 모델들이어야만 한다는 조건이 붙습니다.

그럼 어떤 $q_\phi (z|x)$에 대해 어떤 $g_\phi (\cdot)$ 변환 함수와 $\epsilon\sim p(\epsilon)$를 골라야 적절한 선택일까요? 이에 대해 P.Kingma느님이 정리를 다음과 같이 해주셨습니다:


이 세가지 옵션이 모두 다 안 통할때는 "good approximations to the inverse CDF exist requiring computations with time complexity comparable to the PDF (see e.g. [Dev86] for some methods)."라고 친절히 말씀해주셨으나 불민한 초짜 대학원생은 되는 경우만 사용하기로 했습니다....(사실 레퍼까지 읽어보지 않았네요..ㅋㅋ 후에 시간과 잉여력이 남으면 더 파보도록 하겠습니다.)

*덧: inverse CDF라는 것을 보면서 음 이건 또 뭐야..하실 수 있는데 단순하게 설명하자면  난수 생성기라고 생각하시면 됩니다. 즉, 확률 변수 $X$의 누적 분포 함수(CDF) $F(x)$의 역함수 $F^{−1}$을 알 수 있다면 기본 난수 생성기를 이용하여 확률 변수 $X$에 대한 난수 생성기를 만들 수 있는데 기본 난수 생성기를 확률 변수 $U$로 정의하면 이 확률 변수는 $Unif(0,1)$을 따르며, 아래와 같이 정리할 수 있습니다:
$$X=F^{−1}(U).$$
이 때, 확률 변수 $X$의 난수는 자기 자신의 누적 분포 함수의 역함수에 기본 난수를 대입하여 얻을 수 있습니다.

같은 맥락으로 이를 이용해서 $q_\phi (z|x)$의 inverse CDF를 $g_\phi (\epsilon,x)$로 두고 값을 뽑아내면 된다는 얘기입니다.

결과


요즘 GAN이 워낙 판을 주름잡고 있지만 VAE의 결과도 상당히 괜찮습니다. 당시의 상황을 생각해보면 sensational한 결과였을 것으로 보입니다. ImageNet과 같은 큰 모델에도 어느 정도 사용이 가능한 generative model은 거의 유일무이 하지 않았을까 싶네요.

위 결과는 2D latent space에 대해 학습시킨 후 data manifold를 시각화한 것입니다. 상당히 재미있죠. 그럴듯하게 있음직한 위치에 데이터들이 서로 모여있는 것을 보실 수 있고 상당히 부드럽게 이미지들이 변화하는 것을 확인할 수 있습니다. 


위는 MNIST 데이터에 대해 각각 다른 latent space dimension으로 학습시키고 random하게 값을 sample해본 결과입니다. 역시 dimension이 작은 model에서도 잘 나오는 것을 볼 수 있습니다.

내 멋대로 정리


전통적인 Variational Bayes 방식에서는 같은 문제를 풀 때, 모델인 posterior $q(z|x)$를 세워 intractability 문제를 해결하였으나 lower bound를 좀 더 maximize해서 gap을 더 줄이기  위해서는 likelihood $p(x|z)$를 사람이 좀더 똑똑한 방식으로 열심히 디자인해서 새로 갈아껴보고 하는 방식으로 연구들이 진행되어 왔습니다. 따라서 적절한 모델을 디자인하기 위해 domain knowledge가 매우 중요했었죠.

VAE는 DNN이 이런 optimization 문제를 일반적으로 잘 푼다는 것이 알려져있기 때문에,
  1. variational inference의 아이디어를 이용해서 기존 문제를 optimization으로 바꿔주고, 
  2. 이렇게 바뀐 문제를 위와 같은 방식 대신 $q(z|x)$와 $p(x|z)$가 각각 encoder와 decoder인 Auto-Encoder로 모델링하여 
  3. 이를 gradient descent를 사용하여 둘 다 한 방에 update하였다는 점에서 
당시로서는 매우 획기적인 발상의 전환이었을 것으로 보입니다.

물론 VAE의 결과들이 전에 언급한대로 약간 이미지들이 다 흐리멍텅하게 보이는 단점이 있습니다. 이 부분에 대해서 앞서 제가 설명했듯이 MLE를 사용한 것으로 설명하는 쪽도 있지만 꼭 이렇다할 명확한 이유를 제대로 밝힌 연구는 없는 것 같습니다.

마찬가지로 왜 GAN이 보다 Sharp한 이미지를 생성할 수 있는지에 대한 논의도 여전히 진행 중인 것 같구요. 그럼에도 불구하고 VAE의 이론적 배경과 결과들을 볼 때 매우 아름다운 느낌을 받는 것은 저 하나만이 아닌 것 같습니다. Generative model들이 그렇지만 VAE가 매우 매력적인 방법임에는 틀림 없는 것 같군요.

* 조금씩 남아있는 자잘한 정리 및 증명들은 후에 부록처럼 글로 정리해보겠습니다. (연휴 같은 것 없었던 대학원생은...orz...잉여력이 부족하군요..ㅎㅎ)

다음 읽을거리



참고문헌:



댓글 3개:

  1. 좋은 글 잘 읽고있습니다. 감사합니다!!
    첫번째 그림에서 하나 이해가 잘 안가는게 있어서 질문드립니다.
    KL[N(u(X),Sigma(X))||N(0,I)] 이부분... 이 전체 데이터 셋 D에 대해 평균을 구하는(저희가 reparameterization trick 을 쓰고자 하는 식의) KL(q(z|x)||p(z)) 부분인것 같은데 p(z)를 임의로 N(0,I)로 두신건지 궁금합니다.
    음 다시말해 원래 p(z)는 저희가 모르는 분포이기 때문에 q distribution으로 estimate 하는거로 알고있는데 이 그림에서는 실제 p(z)가 N(0,I)를 따른다고 임의로 설정해두신건지 (실제로 저희가 데이터를 보고있을때 p(z)는 모르니까요) 궁금합니다.

    답글삭제
    답글
    1. 네 맞습니다. 다만 이유가 있는데요. p(z)가 나타내는 latent space를 Gaussian으로 정의를 하면 이를 바탕으로 적절한 함수 변환을 통해 어떤 함수든지 표현할 수 있기 때문에 Gaussian distribution으로 나타낸다는 설명이 있습니다.참고문헌으로 제가 올려둔 것 중에 Tutorial on Variational Autoencoders 를 읽으시면 더 이해가 잘 되실겁니다.

      삭제
    2. 아하 정말 감사합니다ㅠㅠ 도움이 많이 되네요!

      삭제