[논문 정리] Auto-Encoding Variational Bayes

Variational Auto Encoder (VAE) 에 대한 수학적 정리를 해보려 한다. 이해가 쉽지 않은 논문이라 논문을 읽고 정리하기 보다는, Clova 이활석 님 발표와 ratsgo 님, 마지막으로 GAN tutorial, VAE tutorial 자료를 많이 참고하였고 이미지 또한 이에 따른 것이다. 또한, 이미지는 잘못 이해하였거나 틀린 부분이 있는 경우에 계속적으로 공부하면서 수정할 예정이다.

Generative Model

VAE 모델은 우선 Generative Model 이다. 따라서, Generative Model이 무엇인지 간략하게 짚고 넘어가려 한다. 아래 두 개의 이미지가 나름 직관적이라 이를 통해 이야기를 해보려 한다.

generative model architecture

Generative Model 은 Latent Variable z 를 input 으로 한다. z 는 p(z) 분포에서 sampling 되는 random variable 이며, 이 때 p(z) 는 sampling하기 쉬운 normal distribution 같은 분포가 주로 쓰인다. random variable z가 들어가면, deterministic function인 Generator (ex. MLP) 가 원래 데이터 X의 분포를 학습하게끔 업데이트된다. 

generative architecture

결과적으로 우리는 z가 샘플링되는 p(z) 분포의 파라미터를 학습해야 하고, z를 원 데이터 X로 맵핑해주는 generator 를 학습해야 한다. 이를 식으로 표현하면 아래와 같다. 

 

marginalize

좌변의 앞부분은 sampling 된 z가 generator를 통해 나온 값을 x로 만드는 것을 의미하며, 뒷 부분은 z의 분포 p(z)를 의미한다. 이 값들을 maximize 하는 것은, 결국 데이터 분포를 더 잘 학습하는게 된다.

괜히 복잡해졌는데, z sampling을 통한 training process 를 진행하면서, (z marginalize out) 데이터 x 에 대한 분포를 학습하는 것으로 이해하면 좋을 것 같다.

 

Variational Auto-Encoder

marginalize

위의 식을 최대화하면, generative model이 학습되지만, VAE는 이렇게 간단하게 끝나지 않는다. 그 이유는 p(z)를 normal distribution으로 하게 되면, 학습이 잘 되지 않기 때문이라고 하는데, 학습하려는 데이터 차원이 커질 수록, 모든 데이터에 대해 하나의 분포를 가지고 있는 distribution of latent variable은 부족해보이므로, X마다 reconstruct에 적합한 파라미터를 주게 된다. 이에 따라, 기존의 P(z) 분포에 X라는 조건을 추가해서, P(z|x)가 된다. 

 

근데 여기서 또 간단하지 않은게, P(z|x)가 데이터 갯수만큼 있게 되면, 수도 없는 z에 대한 분포를 가정하고 학습해야 한다. 예로 들어, 파라미터가 백만 개가 있는 gaussian mixture 인 셈이다. 이에 따라, 수많은 P(z|x_{i}) 를 근사하는 분포 q(z) 를 구하게 되는데, 이 때 변분 추론이라 불리는 variational inference 가 등장한다.

 

Variational Inference

Variational Inference (VI)의 컨셉은 굉장히 간단한데, 근사하려고 하는 분포와 우리가 가지고 있는 분포의 KL divergence 를 최소화하는 것이다. 첫 번째 줄은 1) 이라고 표기하여 아래 식을 하나씩 살펴보겠다.

 

variational inference

1) 식은 KL diveregence 에 대한 식의 표기이다.

 

2) 식은 bayes rule 로 분모의 p(z|x) 가 p(x|z)로 바뀌면서, 나머지 term(evidence, normalize term) 들이 채워진다. encoder 부분을 decoder 로 바꿔주기 위함으로 보인다.

 

3) 식은 2)식을 3가지로 풀어준 것으로 2) 식의 뒷 부분이 log 로 씌워졌기 때문에, 덧셈으로 분배된다. 

 

4) 식은 덧셈 연산으로 풀어진 3) 식을 정리한 것으로, 4-1) 식은 KL div 로 정리하였고, 4-2) 식은 q(z) 가 pdf 이므로 적분하여 log p (x) 만 남은 모습이며, 마지막 4-3) 식은 4-2) 식과 유사하지만, log 안에 있는 값이 적분해주는 값 z에 dependent 하여, 기댓값의 형태가 된다.

 

 

결과적으로 우리에게 남은 식은 위와 같다. 첫 번째와 두 번째 식은 위치만 조금 바꿨을 뿐, 정확히 같은 것이다. 이때 KL divergence 자체적 특성인 0 보다 크거나 같다를 이용해, 아래 식의 첫 번째 부분을 제거하고 등식을 부등식으로 바꿔줄 수 있다. 아래 식의 첫 번째 부분을 없애는 이유는, 우리가, p(z|x)를 모르기 때문에 (알기 어려워서 VI를 시작했다.) 이에 대한 KL divergence 를 계산하려고 하기 보다는, 이를 없애고, 아래의 부등식에 대해 우변을 최대화하는 문제가 바꾸기 위함이다. 아래 식의 우변을 최대화하는 것은 좌변을 최대화하는 것과 같다. 이를 Evidence Lower Bound (ELBO)라고 한다.

 

위의 ELBO 식을 다시 나눠서 해석해보자. 우선, 위의 VI 식을 전개하면서 q(z)로 계속해서 표기하였는데, VAE에서 샘플링되는 z는 X마다 다르기 때문에, 우리에게는 q(z|x) 가 필요하다. 이렇게 바뀌는 것이 VI 전개식에 영향을 주지 않으므로, 바꿔준다.

아래의 식은 reconstruction error 에 대한 것이다. 다만 VAE의 경우 deterministic function인 generator 가 있기 때문에 아래와 같이 바뀐다.

아래 식은 VI를 통해 근사된 q(z) 분포에 대한 regularization term 이다. 위의 VI 식을 전개하면서 q(z)로 계속해서 표기하였는데, VAE에서 샘플링되는 z는 X마다 다르기 때문에, 우리에게는 q(z|x) 가 필요하다. 이렇게 바뀌는 것이 VI 전개식에 영향을 주지 않으므로, 바꿔준다.

위에서 언급했던 것과 같이, generative model 에서 z를 생성하는 분포는 샘플링하기 쉬운 것을 사용하기 때문에, gaussian distribution p(z)를 가정하여, q(z)가 이와 유사하게 형성되도록 한다. 해당 penality term으로 VAE를 통해 생성된 manifold z의 분포는 auto-encoder에 비해 안정적인 모습을 보인다고 한다. 또한, q(z)를 multi-variate normal distribution, p(z) 를 stand normal 가정하게되면, 위의 식이 간단하게 계산된다고 한다.

Reparametrization Trick

VI 를 통한 ELBO 식으로 loss function 세팅이 끝나고, 학습을 하려고 하면, 문제가 생긴다. 이유는, q(z|x) 에 따른 z 가 stochastic 하기 때문에, gradient 를 계산할 수 없어, 분포에 대한 파라미터가 학습이 될 수 없기 때문이다. 이에 따라, 약간의 트릭이 들어간다. 이해에 도움이 되는 Clova 이활석 님 발표자료를 보면, trick 을 적용함으로써, stochastic 한 부분을 epsilon 분포가 발생하게 함으로써, q(z|x) 의 파라미터가 업데이트가 가능하게 해준다. 해당 트릭에는 기존 분포에 사용되는 파라미터를 어떻게 "조합 및 변환" 할 것인지를 정하는 함수와 noise 를 넣어 "stochastic term"을 넣어주는 epsilon 부분이 필요하다. 어떤 함수와 noise 분포가 적합한지에 대해서는 추후에 보강토록 하겠다.