VAE ( Variational AutoEncoder )

이전 포스트에서 왜 Generative Models의 정의가 P(X,Y)를 구하는 게 아니라 왜 P(X)를 구한 후 샘플링하는 것이 목적이라고 하는 지 이해가 안 간다고 하였다. 여러 VAE를 설명하는 해외/국내 블로그들이 시작할 때 P(X)가 목표다 라고 하는데, 실제로는 Given Y에 대해 가장 그럴듯한 데이터 X를 생성하는 것이 목표이므로 P(X,Y)가 목표라고 생각한다. 곰곰이 생각해본 결과, 결과적으로 Generative Models 중 하나인 VAE가 결국 P(Y|X), P(X,Y), P(X|Y)를 모두 학습하기 때문에 P(X) / 데이터의 진정한 분포를 학습한다고도 할 수 있는 것 같다. GAN이 Implicit하게 P(X)를 구한다는 표현은 P(Y|X), P(X,Y), P(X|Y)를 따로 명시적으로 학습시키지 않고도 그럴듯한 데이터 X를 만들어내기 때문에 나왔을 것이다. 

그러면 VAE는 무엇일까? 여러 자료를 봤지만, Stanford cs231n Lecture13. Generative Models 가 가장 깔끔하고 논리적으로 소개한다.

VAE는 AutoEncoder의 확률모델적 변형으로 모델로부터 새로운 데이터를 샘플링할 수 있게 해준다. AutoEncoder란  Neural network that is trained to attempt to copy its input to its output인데, 실제로는 응용목적이 다르다. 더 깊게 들어갈 필요는 없고, 단순히 모델이 Encoder, Decoder구조를 따르기 때문에 "확률모델적 변형"이라고 부른다고 이해하면 된다.

VAE는 training data x(i), i=1~N, 들이 어떤 latent 표현 z로부터 생성된다는 것이다. PCA(주성분분석)과 비슷하게 어떤 도메인 데이터 X가 좀 더 단순한 변수 Z로 표현이 가능하다는 것과 같은 개념이다. 예를들어 x가 얼굴 이미지면, z는 얼굴의 크기, 모양, 방향 등의 특징들을 나타내는 vector일 것이다. VAE는 이런 latent vector Z의 분포를 알아내어 샘플링한 후 true conditional P(x|z(i))을 통해 그 Z에 대응되는 데이터 X를 샘플링하는 것이 목적이다. 

VAE의 매커니즘에 대한 개념적 요약

구체적인 VAE의 목표는 위 두 가지 분포 P(z)와 모든 z(i)에 대한 조건확률 P(x|z(i))을 알고 있는 모델 파라미터 θ를 학습하는 것이다. Prior인 P(z)는 단순하게 가우시안 분포로 가정하는데, 보통 합리적이라고 한다. 생성하고자 하는 데이터 X에 따라서 이것은 바뀔 수 있다. 예를 들어 X가 사람의 얼굴이고 z vector 중 하나가 나이를 나타낸다면 가우시안이 아니라 uniform distribution으로 가정할 수 있다. 아무튼 P(x|z)는 매우 복잡한데 이런 것은 Neural Network에게 맡긴다. 바로 Decoder다. 

그러면 이 모델을 어떻게 학습시킬까? 즉 어떻게 P(z)와 P(x|z) 분포들을 알고 있는 모델 파라미터 θ를 구할까? 하나의 방법은 모델이 training data x(i)를 생성할 확률 P(x(i))를 높이도록 학습하는 것이다. 수식적으로 P(x(i)) = integral of p(z)p(x|z)인데 아래 사진을 참고하면 된다. 그런데 문제는 무수히 많은 연속변수 z에 대해 integral을 계산할 수가 없다. 

P(x(i))를 최대화하는 모델 파라미터 θ를 학습하는 것은 결국 P(z), P(xz) 분포를 학습하는 것이다.

Data likelihood인 P(x)를 이용하는 방법 대신에 P(z|x)를 이용하는 방법도 있다. 그러나 P(z|x) = P(x|z)P(z) / P(x)인데 P(x)때문에 결국 이것도 intractable하다. 그래서 대안으로 P(z|x)를 근사하는 Q(z|x)를 학습하는 Neural Network를 만들 수 있다. 바로 Encoder다. 이 Encoder Network가 학습하는 Q(z|x)는 Data likelihood를 maximize하는 과정 중 하나인 Lower Bound를 유도할 수 있게 하고, 이 Lower Bound는 tractable해서 문제가 풀린다. 이제부터는 수식이다. 그전에 실제 모델인 Encoder, Decoder가 각각 뭐하는 지 아래의 그림을 통해 알아보자. Encoder는 z|x(i)에 대한 Mean과 Covariance를 output하고, 그것을 가지고 우리는 z를 샘플링한다. Encoder가 샘플링하진 않으며 뒤에 다시 나온다. 샘플링 된 z를 가지고 Decoder는 x|z에 대한 Mean과 Covariance를 output하고 샘플링한다. 즉 새 데이터 X'를 output한다. 

Decoder가 output하는 것은 결국 Input data X에 대응되는 이미지 X'이다.

수식은 캡쳐사진으로 대체하겠다. 보기만 해도 이해가 되고, 수식 부분은 다른 블로그에도 잘 설명이 되어있다. 

맨 밑의 변에서 마지막 항은 intractable은 하지만 >=0임으로 자동으로 앞의 두항이 logp(x(i))의 lower bound가 된다

결국 VAE에서 Data Likelihodd P(x)를 maximize하도록 학습하는 것은 lower bound를 학습하는 것으로 대체할 수 있다. 정리하자면 아래의 캡처와 같다. LowerBound의 첫항은 Reconstruction Error로 이것을 최소화되도록, 두번째항은 q(z|x(i))와 true prior P(z), 가우시안 분포,의 KL Divergence로 이 차이를 최소화하도록 VAE를 학습하면 된다. 두번째항을  KL-Divergence와 Gaussian 분포함수의 정의에 대입하면 DKL[qλ(z|x)||p(z)]=DKL[N(μ,σ)||N(0,1)] = 0.5(μ2+σ2−log(σ2)−1) 가 된다. (유도과정) 결국 우리에게 익숙하도록 바꾸자면 Neural Network VAE의 최종 Loss는 다음과 같다. 이 블로그에서 가져왔다 . 내 생각엔 첫째 항도 ||X - f(z)||^2로 바꾸주는 게 좋겠다. X는 Input Data, f(z)는 Decoder의 Output이다.

그런데 첫째항 log(p(x|z) 을 최대화(=Reconstruction Error 최소화)를 하는 데 문제가 있다. VAE는 Encoder의 Output, z|x(i)에 대한 Mean과 Covariance를 가지고 z를 Random Sampling한다. 만약 Encoder가 z를 그냥 output하면 그것은 분포를 학습하지 않기 때문에 그냥 AutoEncoder, 즉 학습한 데이터만 기억하는 Neural Network일 것이다. 그런데 Random Sampling은 Back Propagation이 가능하지 않다. 따라서 Reconstruction Error를 최소화하기 위한 Back Propagation이 Encoder에 반영되지 않아서 Encoder가 데이터 X를 Generate하기 적합한 latent code Z에 대한 분포를 제대로 학습할 수 없다. 이때 Back Propagation이 가능하도록 하는 Trick이 Reparameterization Trick이다. 노이즈 엡실론 ϵ를 normal distribution에서 샘플링해 z의 분산과 곱해 평균에 더하는 것이다.

오른쪽에서는 이제 Encoder로 Back Prop이 가능하다

이 분블로그에서 뭔가 이론적 배경을 제시하시는데 궁금한 분은 봐도 될 것 같다.

Back Prop이 된다!

 결론적으로 VAE는 P(z|x)를 근사하는 Q(z|x) 분포를 Encoder를 통해 학습하고, Decoder를 통해 latent code z에 대해 가장 가능성 높은 x를 output하기 위한 P(x|z) 분포를 학습한다. true prior인 P(z)를 normal distribution N(0,1)로 가정하고 KL Divergence항을 통해서 학습한다. 따라서 Test time에는 Encoder는 떼어버리고 N(0,1)에서 노이즈를 샘플링해 Decoder에 Input하면 새로운 데이터 X'를 생성할 수 있다. 

새로운 데이터(이미지) 생성

 

VAE의 장점 중 하나는 q(z|x)를 inference하기 때문에, 즉 주어진 이미지에 대해서 가능한 latenet code의 분포를 예측하기 때문에 feature representation task들에 적용될 수 있다. 

 

INDEX) kl divergence 정의