[논문리뷰] VAE(Auto-Encoding Variational Bayes)

Auto-Encoding Variational Bayes

너무나도 유명한 논문.. 공부할 겸 다시 읽었다. 
읽을 때마다 새로운 느낌이 드는 이유는 뭘까... 신비한 머신러닝의 세계~~ 
VAE 는 직관적으로 내용을 전부 이해하기에는 익숙하지 않은 사람에게는 수식이 너무 많아서,, 여러 번 읽어야 할 것 같은 논문이다. 

1. Problem Suggestion 

How can we perform efficient approximate inference and learning with directed probabilistic models whose continuous latent variables and/or parameters have intractable posterior distributions?

VAE 의 목표는,, 아주 다루기 어려운 posterior distribution 을 가진 연속 잠재변수를 가진 모델로부터 효율적으로 inference 및 학습을 시키는 것이다. 

모델을 도식화 해 보면 다음과 같다. 검은색 실선은 일반적인 generative model, pθ(z)pθ(x|z)  을 나타내는 것이다. 하지만 이 경우 실제 dataset 의 posterior 분포 pθ(z|x) 를 알아내기란,, 불가능하다 ^^ ;; pθ(x)를 모르기 때문에. 그래서 이것을 근사하여(variational approximation) qφ(z|x)(검정색 점선) 을 두어 대리로 posterior 를 모델링할 수 있도록..! 하는 것이다. 여기서 φ와  θ는 모두 학습되는 parameter 이다! 

2. 1 Problem scenario

이건 논문에서 나온 문제 scenario 이다. 자 다음과 같은 iid dataset이 있다고 해보자. (continuous 일 수도, discrete일 수도 있음) 

이 dataset 이 생성되는 과정은 다음과 같다. 
1. 잠재 변수 z가 어떤 Prior distribution pθ*로부터 생성된다. 
2. dataset xi 가 어떤 conditional distribution pθ*(x|z)로부터 생성된다.
이 때 이 prior와 likelihood는 parametric한 distribution인 pθ(z) 와 pθ(x|z)로 나타낼수 있다고 생각하고, 이 pdf들은 θ나 z에 관해서 미분 가능하다고 생각하자. 

하지만 불행하게도...  우리는 실제 parameter θ*를 모르고, latent variable z 또한 모르니.. 모델링은 매우 어려워진다. 

아주 중요한 점은, marginal probability p(x) 를 쉽게 이용하기 위한 가정은 하지 않을 것이라는 것이다. 그러니까, 우리의 모델은 다음과 같은 경우에도 동작한다는 것. 

1. marginal likelihood를 구할 수 없는 경우 

• pθ(x) = pθ(z)pθ(x|z) dz 를 구할 수 없어서 true posterior density 또한 구할 수 없다. 
• 이 경우는 nonlinear hidden layer를 사용하는 neural net 도 해당하는 경우인데 아주 빈번하게 발견된다. 

2. 데이터셋이 큰 경우

• 이건 monte carlo approximation을 염두에 두고 쓴 구절 같은데,, monte carlo의 경우는 dataset이 너무커지면 각 datapoint로부터 sampling 을 여러 번 해야하기 때문에 비효율적이다. 하지만 이 사람들이 말한방법을 쓰면 큰 데이터셋 으로도 모델링이 가능하다는 것. 

자 위와 같은 문제를 풀기 위해서 recognition model qφ(z|x)를 도입한다. 이 친구는 구할 수 없는 true posterior pθ(z|x)에 대한 approximation이다. mean field 변분 추론과는 다르게, 여기서의 φ는 factorial 하거나 parameter가 closed form으로 계산될 필요가 없다. 대신에 recognition model parameter φ는 generative model parameter θ와 함께 Jointly 하게 학습되는 것이다. 

어떤 면에서는 recognition model qφ(z|x)를 인코더 라고 볼수도 있을 것이다. 어떤 datapoint x로부터 그것의 분포를 모델링 하도록(가우시안..이 많이 활용될 것이당)! 비슷한 면으로 pθ(x|z)를 디코더 라고 볼 수 있을 것이다. 

요런 느낌이다. 

2.2 Variational Bound 

우리가 원하는 건 결국 x에 관한 marginal likelihood 값이다. 이것은 다음과 같이 나타낼 수 있는데, 

이 부분은 오른쪽의 첫번째 term인 KL divergence 를 푸는 것을 먼저 해보면 쉽다. 

이렇게.. 한번 해보세여 

 

저렇게 해 놓으면 DKL(qφ(z|x (i) )||pθ(z|x (i) )) 이 Term 은 항상 Nonnegative 이니까 식은 다음과 같이 변하는데, 

이렇게 된다. 여기서의 우변을 variational lower bound라고 부른다. 근데 이 부분을 또 이렇게 바꿀수 있다 (위의 증명 참고) 

우리는 이 lower bound를 varaiational parameter 인 φ와 generative parameter인  θ에 대해서 모두 최적화 하고 싶은 것. 하지만 이 lower bound로 gradient 를 흘려보내 주는 것은 약간 문제가 있는데, 이런 경우에 일반적으로 사용하는 monte carlo gradient estimator는 분산이 커서 우리의 목적에는 실용적이지 않다고 한다. 

2.3 The SGVB estimator and AEVB algorithm

최적화 하고 싶은 식은 2.2에서 찾았는데, 이제 이 식을 어떻게 optimize 할 건지 estimator 를 찾아야 한다. 
2.4에서와 같은 mild한 조건이 있다면, posterior를 다음과 같이 reparameterize 할 수 있다. z ∼ qφ(z|x)

이제 function f(x) 에 대해서 다음과 같이 Monte Carlo estimation을 할 수 있다. 

이 estimator를 variational lower bound에 적용하여, 우리의 일반적인, Stochastic Gradient Variational Bayes estimator를 만들어 낼 수 있는 것이다! (variational lower bound에 관한 식 중 첫번째 식 참고.) 

variational lower bound 두 번째 식에서, kl divergence는  φ를 regularization 하는 term이라고 생각할 수 있다. Approximate posterior가 prior pθ(z) 를 따르게 한다는 점에서..! 그래서 estimator 를 다음과 같이 구성하는 것도 가능하다. 

N datapoint 를 가진 dataset X 가 주어진다면, 전체 dataset 에 대한 Lower bound 또한 구할 수 있다. 

흔히 사용되는 방법인 graident based optimization이 가능하단 이야기. 다음 알고리즘을 참고하도록 하자. 

2.4 The reparametrization trick 

위에서도 나오지만 다시 설명한다면, z 자체를 sampling 한다면 random 연산은 미분이 불가능하니까 gradient backward 가 불가능하다. 
그렇기 때문에 reparametrization trick 을 사용하는데 이게 뭐냐면  다음과 같은 deterministic matpping 이 주어졌을 때,  $z=g_{\phi}(\epsilon, x)$, 우리는 $q_{\phi}(z|x)\Pi_{i}dz_{i} = p(\epsilon)\Pi_{i}d\epsilon_{i}$ 임을 알 수 있기 때문에

$ \int \! {q_{\phi}(z|x)f(z)} \, dx = \int \! {p_{\epsilon}f(z)} \, d\epsilon = \int \! p(\epsilon)f(z) \, d\epsilon = \int \! {p(\epsilon)f(g_{\phi}(\epsilon, x))} \, d\epsilon$

 임을 증명할 수 있다. 그래서 이제 미분 가능한 estimator 로 만들어 낼 수 있는 것.  

저자들이 첨부해 놓은 예시를 생각해 본다면, univariate gaussian case를 생각해보자. $z \sim p(z|x) = N(\mu, \sigma^2)$. 이 경우에는 다음과 같이 reparametrization을 할 수 있을 것이다. $ z = \mu + \sigma\epsilon$. $\epsilon$은 normal distribution을 따르는 noise variable이다. 따라서 다음과 같이 된다.

그렇다면, 어떤 경우에 다음과 같은 reparameterization trick 이 유효할까? 논문에서는 다음과 같은 세 가지 경우를 예시로 들고 있다.

1. Tractable한 inverse CDF. $\epsilon ~ U(0, I)$ 이고, $g_{\phi}(\epsilon, x)$ 가 $q_{\phi}(z|x)$ 의 inverse CDF인 경우. 
2. 위에서의 gaussian 예제처럼, 다음과 같은 형식으로 표현이 가능한 distribution g(.) = location + scale · $\epsilon$
3. Composition, random variable을 보조 variable들의 다른 transformation처럼 생각하는 것 (Log-Normal 이나 gamma 등) 

3. Example : Variational Auto-Encoder 

자 이제 위에서 나온 모든 trick들을 적용한 VAE에 대해서 살펴볼 차례다.. 
맨 위에서 봤었던 그림처럼 encoder 와 decoder 를 가지고, AEVB algorithm (2.3 설명) 을 통해 각 parameter를 최적화 하는 것. 

그림으로 표현하자면 다음과 같다. prior를 isotropic multivariate gaussian으로 가정했을 때의 경우이다. 이 경우에 true posterior는 또한 interactable 하기 때문에 posterior를 multivariate gaussian으로 approximation을 한다. 

휴, loss function을 식으로 표현하면 다음과 같다. 이 경우에 KL divergence는 closed form으로 계산이 가능하다. 

decoding term $logp_{\theta}(x^{(i)}|z^{(i, l)}) $ 은 베르누이나 가우시안 MLP 로 가정한다.. 우리가 모델링하려는 데이터의 타입에 따라 결정해 주면 될 것이다. 

5. Experiment 

4. 는 related work 라 생략 ㅎ 

MNIST 와 Frey face dataset 을 가지고 generative model을 학습시켰다고한다. 

Continuous Latent vector에 따라서 sample이 자연스럽게 변하는 모습을 확인할 수 있다!