Introduction

VAE에 대해 공부한 것을 정리했다. 가장 큰 도움이 된 자료는 Naver D2 세미나의 오토인코더의 모든 것이기에 이 글을 읽는 것보다 하루 정도 써서 위 강연을 듣는 것을 더 추천한다.

Basic Knowledge

PDF

확률밀도함수(Probability Density Function, PDF)는 continuous한 variable에 대해 확률을 표현하는 함수다.

from [1]

from [1]

따라서 당연히 전체에 대한 적분값은 1이 되고, 범위에 대한 적분 값은 그 범위에 걸릴 확률을 의미한다. 중요한 것은 특정한 x에 관한 값인데, 예를 들어 위 그래프의 경우 \( x=0\)일때 값은 \(0.4\)이지만 이는 당연히도 \(x=0\)일 확률이 \(0.4\)라는 의미는 아니다, 정확히 \(x=0.000\)일 확률은 \(0\)이기에.

따라서 \(p(x=0)\)은 상대적인 확률 값이라고 생각하면 좋다. 대충 \(p(x=0)\)이 높으면 높을 수록 다른 숫자에 비해서 \(x=0\)이 걸릴 확률이 ‘높다’라고만 생각해주자.

Function as Distribution

위: 전통적인 방식의 function, 아래: distribution으로 바꿔서 생각하기.

위: 전통적인 방식의 function, 아래: distribution으로 바꿔서 생각하기.

여태껏 나는 deep learning function을 fixed input에 대한 fixed output으로만 생각을 했지만, 조금만 달리 생각을 하면 distribution으로도 볼 수 있다. 예를 들어, 어떤 이미지 \(x\)에 대해서 \(\theta\)로 parameterized된 함수 \(p_\theta\)가 있다고 하자. 다시말하면 \(p_\theta\)라는 trained된 model이 있다고 생각하자. 이미지를 넣었을 때, \(p_\theta(x) = a\)이라고 하면 결과값이 \(a\)라고 볼 수 있지만, 조금만 돌려 생각해보아 결과를 \(N(\mu=a, \sigma=1)\)인 normal distribution이라고 생각하면 어떨까? (일단 \(\sigma\)는 무조건 1이라고 가정하면.) 즉 이렇게만 바꿔 생각해도 deep learning model이 fixed output이 아닌, distribution을 output하는 더 파워풀한 모델이 된 것이다!! 사실 그게 그거인데 말이다.

반대로 input 역시 distribution이라고 생각할 수 있다. 이건 약간 헷갈리는데, input인 \(z_i\)가 어느 distribution에서 샘플링 되어서 들어간고 생각하면 input 역시 distribution이라고 보면 된다. 꼭 알아야 되는 것은 이런 Deep-Learning as distribution의 경우 \(p_\theta(x|z=z_i)\)는 주어지지만 (= 계산이 가능하지만, 그냥 인풋에 때려 박으면 결과가 나오니까), 나머지는 확실하지 않다고 보자. 즉 input을 distribution으로 본다고 해도 진짜 분포가 통째로 들어가는거는 계산이 불가능하니까 그냥 이론적으로만 분포라고 생각하고 넘어가자.

이후 설명할 variational autoencoder 이해의 핵심은 (내가 생각하기에), model의 input, output이 fixed value가 아닌 distribution이라고 생각하는 것이라고 본다.

Generative Model

Notation Desc.
\(p_\theta\) generative model parameterized with \(\theta\)
\([x_1, …, x_n]\) \(n\)개의 이미지가 있는 트레이닝 데이터셋
\(z\) latent vector, generative model의 input

위에서 말했듯이 Generative Model이 distribution을 뱉는다고 생각하자. (생각하기 힘들면 model이 이미지를 뱉는데, 그게 normal distribution의 mean, std는 1이라고 생각을 해버리면 편하다.) 그럼 좋은 generative model은 무엇이냐하면 \(p_\theta(x=x_i)\)이 높아야겠다. (일단은 generative model이 비슷한 이미지가 아니라, 주어진 training dataset을 완벽하게 뱉어내는 model이라고 생각을 하자. Model이 완전히 overfitting되어 있지만 않으면 비슷한 사진도 많이 뱉어 내니까 이게 generative model의 목표여도 괜찮은 것 같다.)

물론 \(p_\theta\)는 input을 받아야되기 때문에 input \(z\)를 사용한 공식으로 바꿔 쓰면 \(p_\theta(x=x_i) = \int p_\theta(x=x_i|z)p(z)dz\)가 된다.

여기서 잠깐 \(p(z)\)는 뭐에요 인데, 그냥 \(N(\mu=0, \sigma=1)\)인 normal distribution이라고 생각하자. 이래도 충분하니까 이렇다고 하자.

아무튼 \([x_1, …, x_n]\)이 i.i.d.일 경우 \(p([x_1, …, x_n]) = \prod_i^n p(x_i) = \prod_i^n \int p_\theta(x=x_i|z)p(z)dz \) 이고 결과적으로 목표는

$$ \theta^* = \argmax_\theta \bigg( \prod_i^n \int p_\theta(x=x_i|z)p(z)dz \bigg) = \argmax_\theta \bigg( \sum_i^n \log \int p_\theta(x=x_i|z)p(z)dz \bigg) $$

을 찾는게 된다.

Naïve Generative Model Finding doesn’t Work!

결국 원하는게

$$ \theta^* = \argmax_\theta \bigg( \sum_i^n \log \int p_\theta(x=x_i|z)p(z)dz \bigg) = \argmin_\theta \bigg( - \sum_i^n \log \int p_\theta(x=x_i|z)p(z)dz \bigg) $$ 이라면, 그냥 대충 CNN decoder network 때려박아서 대충 \([z_1, z_2, …, z_m] \sim N(0,1) \) 샘플링 한다음에,

$$ p_\theta(x_j) \approx \frac{1}{m}\sum_j p_\theta(x_i | z_j) $$

를 최대로 하는 \(\theta\)를 stochastic gradient ascent로 찾아버리면 안되나요? 라고 물어 볼 수 있다. 안타깝게도 이 방법은 안 된다. 간단히 설명하면 \(-\log p_\theta(x_i,z_j) \)는 결과적으로 MSE와 동일하기 때문에, 아래 그림과 같이, 의미론적으로 비슷한 이미지 보다는 Euclidean distance가 더 가까운 이미지를 더 선호하게 된다.

(a) 트레이닝셋의 \(x_i\). (b) 의미(perceptually)가 다른 generated image. (c) 대각선으로 1 pixel shift된 generated image. (c)의 경우가 더 의미론적으로 비슷하나, MSE는 (b)가 더 (a)와 가깝다. from [2]

(a) 트레이닝셋의 \(x_i\). (b) 의미(perceptually)가 다른 generated image. (c) 대각선으로 1 pixel shift된 generated image.
(c)의 경우가 더 의미론적으로 비슷하나, MSE는 (b)가 더 (a)와 가깝다.
from [2]

이 방법의 가장 큰 문제점은 무작정 샘플링을 할 경우 \(p_\theta(x=x_i|z_j)=0\)이 되는 \(z_j\)이 엄청 많은 데, 이들 역시 샘플링이 되기 때문에 그리 효율적이지 못하다는 점이다. 이제 드디어 VAE가 나온다. VAE는 위의 그냥 무작정 샘플링 \([z_1, z_2, …, z_m] \sim N(0,1) \) 보다 훨신 efficient 방법으로 샘플링을 하여 generative model을 만드는 데 도움을 준다.

Variational AutoEncoder

그러면 무작정 샘플링하지 말고, \(z_i \sim p_\theta(z|x_i) \), 즉 \(x_i\)가 나올 법한 \(z\) 중에서 \(z_i\)를 골라서 계산하면 좀 효율적이지 않을까? 그 외는 전부 \(p_\theta (x|z_i) = 0\)이니까 계산 하나 마나 잖아! 근데 문제는 \(p_\theta(z|x_i)\)는 계산이 불가능하다. 위에서 말했듯이, 우리는 \(p(x|z_i) \)만 효율적으로 계산 할 수 있다.

$$ p_\theta(z|x_i) = \frac{p_\theta(x_i, z)}{\int p_\theta(x,z)dz} = \frac{p_\theta(x_i | z)p(z)}{\int p_\theta(x,z)dz} $$

에서 분모인 \(\int p_\theta(x,z)dz\)는 계산이 불가능 하기 때문에, \(p_\theta(z|x_i)\)는 계산을 할 수가 없다.

따라서 \(p_\theta(z|x_i)\)를 approximate 해주는 새로운 function \(q_\phi(z|x_i)\)를 만들자.

VAE의 형태. VAE의 목표는 \(\theta\)를 찾는 것이지만, 이를 위해 \(p_\theta\)의 역함수인 \(q_\phi\)도 찾아야 한다. 이 점에서 Auto-Encoder과 비슷한 모양이 되기 때문에 VAE라는 이름이 붙었다.

VAE의 형태. VAE의 목표는 \(\theta\)를 찾는 것이지만, 이를 위해 \(p_\theta\)의 역함수인 \(q_\phi\)도 찾아야 한다. 이 점에서 Auto-Encoder과 비슷한 모양이 되기 때문에 VAE라는 이름이 붙었다.

approximation이 충분히 accurate하다면 \(q_\phi(z|x_i)\)에서 sampling 하면 된다. 이제 마지막으로 \(p_\theta\), \(q_\phi\) 둘 다 잘 fitting 시키는 loss function만 찾아내면 끝!

ELBO

Loss function은 다음과 같은 성질을 가지고 있어야 한다.

  1. \(q_\phi(z|x_i)\)이랑 \(p_\theta(z|x_i)\)이랑 비슷하다.
  2. \(p_\theta(x_i)\)가 높다.

그럼 일단, 1번을 해결하기 위해 KL-발산[3]을 이용하자. KL-발산은 대충 말하면 두 distribution이 얼마나 비슷한지 측정하는 함수로 생각하면 된다. 즉

$$ D_{KL}(q_\phi(z|x_i)||p_\theta(z|x_i)) $$ 을 줄이는 게 목표가 된다. 근데 왜 \(D_{KL}(p||q)\)가 아닌 \(D_{KL}(q||p)\)를 하는 지에 대해 생각해 볼 수 있는 데, 그건 [4]에서 참고해 보자.

따라서 위 식을 전개해 보면 [5]:

$$ \begin{aligned} & D_\text{KL}( q_\phi(\mathbf{z}\vert\mathbf{x}) | p_\theta(\mathbf{z}\vert\mathbf{x}) ) & \\ &=\int q_\phi(\mathbf{z} \vert \mathbf{x}) \log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{z} \vert \mathbf{x})} d\mathbf{z} & \\ &=\int q_\phi(\mathbf{z} \vert \mathbf{x}) \log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})p_\theta(\mathbf{x})}{p_\theta(\mathbf{z}, \mathbf{x})} d\mathbf{z} & \scriptstyle{\text{; Because }p(z \vert x) = p(z, x) / p(x)} \\ &=\int q_\phi(\mathbf{z} \vert \mathbf{x}) \big( \log p_\theta(\mathbf{x}) + \log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{z}, \mathbf{x})} \big) d\mathbf{z} & \\ &=\log p_\theta(\mathbf{x}) + \int q_\phi(\mathbf{z} \vert \mathbf{x})\log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{z}, \mathbf{x})} d\mathbf{z} & \scriptstyle{\text{; Because }\int q(z \vert x) dz = 1}\\ &=\log p_\theta(\mathbf{x}) + \int q_\phi(\mathbf{z} \vert \mathbf{x})\log\frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{x}\vert\mathbf{z})p_\theta(\mathbf{z})} d\mathbf{z} & \scriptstyle{\text{; Because }p(z, x) = p(x \vert z) p(z)} \\ &=\log p_\theta(\mathbf{x}) + \mathbb{E}_{\mathbf{z}\sim q_\phi(\mathbf{z} \vert \mathbf{x})}[\log \frac{q_\phi(\mathbf{z} \vert \mathbf{x})}{p_\theta(\mathbf{z})} - \log p_\theta(\mathbf{x} \vert \mathbf{z})] & \\ &=\log p_\theta(\mathbf{x}) + D_\text{KL}(q_\phi(\mathbf{z}\vert\mathbf{x}) | p_\theta(\mathbf{z})) - \mathbb{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x})}\log p_\theta(\mathbf{x}\vert\mathbf{z}) & \end{aligned} $$ 라고 나온다. 마지막 줄을 다음과 같이 변형시키면

$$ \log p_\theta(\mathbf{x}) - D_\text{KL}( q_\phi(\mathbf{z}\vert\mathbf{x}) | p_\theta(\mathbf{z}\vert\mathbf{x}) ) = \mathbb{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x})}\log p_\theta(\mathbf{x}\vert\mathbf{z}) - D_\text{KL}(q_\phi(\mathbf{z}\vert\mathbf{x}) | p_\theta(\mathbf{z})) $$

좌변이 우리가 maximize하고 싶어하는 값이 나오며 (1번과 2번 성립), 우리는 이 것을 ELBO라고 부른다. loss function은 위의 negative가 될 것이고,

$$ \phi^*, \theta^* = \argmin_{\phi, \theta}\bigg( - \mathbb{E}_{\mathbf{z}\sim q_\phi(\mathbf{z}\vert\mathbf{x_i})}\log p_\theta(\mathbf{x_i}\vert\mathbf{z}) + D_\text{KL}(q_\phi(\mathbf{z}\vert\mathbf{x_i}) | p_\theta(\mathbf{z})) \bigg) $$

여기서 좌항은 Reconstruction Error로 볼 수 있고, 우항은 Regularization Term으로 볼 수 있다. 특히 우항 덕분에 \(z\)가 조금 더 \(N(0,1)\)에 가까워지기 때문에 VAE를 generative model로 사용할 수 있게 된다.

Overall

\(p_\theta, q_\phi \) 둘 다 gaussian distribution을 output한다고 가정했을 경우

\(p_\theta, q_\phi \) 둘 다 gaussian distribution을 output한다고 가정했을 경우

Code

https://github.com/hwalsuklee/tensorflow-mnist-VAE 참조.

References

[1] http://ko.wikipedia.org/wiki/확률_밀도_함수

[2] https://arxiv.org/pdf/1606.05908.pdf

[3] https://ko.wikipedia.org/wiki/쿨백-라이블러_발산

[4] [https://blog.evjang.com/2016/08/variational-bayes.html](https://blog.evjang.com/2016/08/variational-bayes.html

[5] [https://lilianweng.github.io/posts/2018-08-12-vae/](https://lilianweng.github.io/posts/2018-08-12-vae