VAE Tutorial 1
in VAE on Deep-Learning
- VAE Tutorial 목차
VAE Tutorial 1: CS231n Lecture 13
cs231n 강의 내용과 Kingma의 논문을 통해 Variational Auto-Encoder를 정리해봅니다. 그리고 그 이후에 sequential data에 VAE를 적용한 사례인 SentenceVAE와 MusicVAE를 다룹니다. 다음과 같이 이 VAE 튜토리얼은 진행합니다.
- VAE Tutorial 목차
이 글에서는 1. CS231n 강의 내용을 다룰 것입니다. CS231n lecture 13에서는 여러가지 generative model을 다루는데 그 중에 VAE가 들어있습니다.
1. Supervised Learning vs Unsupervised Learning
13강 이전까지는 supervised learning 문제에 대한 딥러닝을 다뤘습니다. Supervised Learning은 데이터와 정답이 있는 경우 모델의 output이 정답과 같아지도록 loss function을 정의하고 학습합니다. Supervised Learning을 통해 Classification, Regression, Detection, Segmentation, image-captioning 등 다양한 task를 학습할 수 있습니다.
하지만 이 경우 항상 label이 있어야합니다. 그래서 데이터에 대한 cost가 큽니다. 딥러닝의 경우 모델의 성능이 데이터의 양과 질에 크게 의존하기 때문에 이에 대한 부담이 있습니다. 그렇다면 label이 없이 학습하는 모델은 없을까요? 그게 바로 Unsupervised Learning 입니다. Unsupervised Learning은 input이 되는 x 데이터만으로 학습할 수 있습니다. 주 목적은 x 데이터의 underlying hidden structure을 학습하는 것입니다. 개인적으로는 data의 meaningful representation을 학습하는 것이 중요하다고 생각해서 Unsupervised Learning을 관심있게 보고 있습니다.
Unsupervised Learning으로는 Clustering, dimensionality reduction, feature learning, density estimation 등을 할 수 있습니다. 지도학습과는 달리 아직 미개척 분야이며 풀어야할 문제가 많고 어쩌면 그 안에 보물을 담고 있을지도 모릅니다. Unsupervised Learning 중에 가장 유명한 것이 K-means와 Auto-Encoder 입니다. 하지만 Auto-Encoder를 그냥 사용하는 경우는 없고 보통 VAE(Variational Auto-Encoder)를 많이 사용합니다. VAE의 경우 Auto-Encoder와는 달리 데이터의 의미있는 representation을 학습할 수 있습니다. 하지만 Auto-Encoder를 처음 접하고 나서 VAE로 넘어갈 때 상당히 큰 벽이 있는 것 같습니다.
VAE를 제대로 이해하려면 코드부터 보는 것이 아니라 네 가지를 알아야합니다.
- VAE는 Generative model이라는 것
- latent variable이라는 것이 있으며 이것을 바탕으로 데이터를 생성한다는 것(Decoder)
- 문제를 더 쉽게 만들기 위해 latent variable 이라는 것을 Encoder를 통해 추출한다는 것
- VAE의 학습과정은 MLE라는 것
2. Generative Model
VAE는 일종의 Generative model이라고 봐야합니다. Generative model이란 training data가 주어졌을 때 이 data가 sampling된 분포와 같은 분포에서 새로운 sample을 생성하는 model입니다. 즉 \(p_{model}(x)\)가 최대한 \(p_{data}(x)\)에 가깝게 만드는 것이 목표입니다. 이것을 어떻게 할 수 있을까요? 결국 얼마나 기존 모델과 가까운 것인가에 대한 지표를 만들어야하고 그 차이를 최소화하도록 gradient를 계산해서 업데이트할 것입니다.
다음 그림은 Generative Model에 대해 Ian Goodfellow가 정리한 도표입니다. Generative Model은 크게 Explicit Density와 Implicit Density 두 가지로 나눌 수 있습니다. Explicit Density 모델을 data를 샘플링한 모델의 구조를 명확히 정의를 합니다. 그 모델로부터 data를 sampling 하는 것입니다. 하지만 Implicit Density에서는 모델에 대한 explicit하게 정의하지 않습니다. 예를 들어, GAN의 경우 noise로부터 바로 data로의 transformation을 학습합니다. VAE는 data의 model의 density model을 explicit하게 정의해서 직접적으로 학습하는 경우라고 볼 수 있습니다.
다시 말하자면 다음과 같습니다. “density estimation”은 x라는 데이터만 관찰할 수 있을 때, 관찰할 수 없는 x가 샘플된 확률밀도함수(probability density function)을 estimate하는 것입니다. 예를 들어 다음 그림의 첫번째 그림처럼 데이터 포인트가 찍혀있을 경우에 대해서 생각해봅니다. Density estimation이 하는 일은 이 데이터 포인트가 샘플링 됐을 것 같은 확률밀도함수를 찾아내는 것입니다. 데이터 포인트가 많은 곳은 확률이 높고 데이터 포인트가 없는 곳은 확률이 낮을 것입니다. 이러한 확률분포(파란색 선)가 있다면 현재 데이터 포인트와 같지는 않지만 비슷한 데이터를 생성해낼 수 있을 것입니다.
즉, 확률밀도함수 즉 최대한 \(p_{data}(\mathcal{x})\)에 가까운 \(p_{model}(x)\)를 찾아냈을 때, \(p_{model}(x)\)을 통해 새로운 데이터를 “generate” 할 수 있습니다. 이렇게 새로운 data를 generate를 할 수 있으므로 Generative Model이라고 부르는 것입니다. 이미지의 경우 각 pixel은 0에서 255의 값을 가질 수 있고 rgb 3차원이므로 나타낼 수 있는 이미지의 양은 어마어마하게 많습니다. 하지만 우리가 “realistic”이라고 느끼는 이미지들은 그 중에 정말 일부입니다. 또한 데이터셋에 있는 이미지들은 그보다 더 일부일 것입니다. Generative Model은 그 엄청나게 큰 space에서 realistic한 이미지를 샘플링할 수 있는 확률밀도함수입니다.
Generative Model 중에 가장 선명하고 진짜같은 이미지를 생성하는 것은 GAN(Generative Adversarial Network)입니다. GAN은 여러 변형체들이 많은데 그 중에 하나의 학습된 모델을 가지고 생성한 데이터 예시입니다. 하지만 GAN의 경우 random noise로부터 데이터를 생성해내므로 데이터의 의미있는 representation을 학습할 수 없습니다. 이제 VAE에 대해서 살펴보겠습니다.
3. Variational Auto-Encoder
앞으로 MNIST 데이터 셋을 생성하는 Generative Model에 대해서 이야기하도록 하겠습니다. 일단 probability density function을 정의합니다. \(p_{\theta}(x)\)라고 정의합니다. dataset은 \(X = [x^{(i)}]_{i=1}^N\)이라고 할 수 있습니다. 이 때, dataset의 datapoint(\(x^{(i)}\))는 i.i.d. 하다고 가정합니다. 이 density function은 \(\theta\)를 parameter로 가지고 있습니다. 이 정의의 의미는 \(\theta\)라는 parameter가 정해졌을 때 \(x\)라는 데이터가 나올 확률입니다. 이 확률을 최대화하는 것이 Generative Model 혹은 density estimation의 목표입니다. 이 식은 \(z\)라는 latent variable을 사용해서 다음과 같이 쓸 수 있습니다. 앞으로 우리는 이 식을 미분해서 그 미분값에 따라 stochastic gradient ascent를 할 것입니다. \[p_{\theta}(x)=\int p_{\theta}(z)p_{\theta}(x|z)dz - (1)\]
이 식은 다음 그림으로 보면 더 이해가 잘 갑니다. VAE에서는 Auto-Encoder와는 달리 latent variable을 정의합니다. 왜 갑자기 “latent variable”이라는게 등장했을까요? 우리가 생성하고 싶은 데이터들은 상당히 차원이 높고 예를 들어 datapoint의 pixel 사이에 복잡한 관계가 있습니다. 따라서 pixel 사이의 관계를 확률모델로 모델링하는 것이 아니라 데이터를 표현하는 \(z\) (mnist의 경우 어떤 숫자인지)가 있으면 그 \(z\)로부터 데이터를 생성하는 graphical model을 생각해보는 것입니다.
(1) 식의 우변에서 \(p_{\theta}(z)\)는 latent variable \(z\)를 sampling 할 수 있는 확률밀도함수입니다. 그리고 \(p_{\theta}(x\vert z)\)는 \(z\)가 주어졌을 때 \(x\)를 생성해내는 확률밀도함수입니다. 중요한 것은 이 \(p_{\theta}(x\vert z)\)가 \(\theta\)에 대해 미분가능해야한다는 것입니다. 보통 output을 Gaussian 분포나 Bernoulli 분포로 설정합니다. (1) 식은 intractable 한데 그 이유는 가능한 모든 z에 대해 integral을 취하기 때문입니다. 따라서 이 식에 대해서 직접적으로 미분할 수 없습니다. 하지만 이 MLE 문제를 풀기 위해서 우리는 미분을 해야합니다.
문제는 두 가지입니다. (1) 어떻게 \(z\)라는 놈을 구성할 것인가? (2) intergral을 어떻게 해결할 것인가? 입니다. VAE는 이 두 가지 문제를 해결합니다. 일단 CS231n 강의의 그림으로 다시 보자면 다음과 같습니다. \(\theta^*\)는 dataset이 sampling된 true distribution의 parameter입니다(이 분포는 parameterized 되었다고 가정하는 것입니다). simple한 Gaussian 분포로부터 \(z^{(i)}\)를 sampling 합니다. decode network \(p_{\theta^*}(x\vert z)\)는 \(z^{(i)}\)로부터 데이터를 생성합니다. 우리가 하고 싶은 것은 \(\theta^*\)를 estimate 하는 것입니다.
\(z\)는 사실 mnist에서 어떤 숫자인지만 나타내야하는 것이 아니라 각 숫자마다도 다른 복잡한 특징들을 나타내야 합니다. 이것을 단순한 normal distribution으로 만들고 decoder network의 layer들이 알아서 data를 생성해내는데 필요한 정보를 추출하도록 하는 것입니다. 즉, 여러 층의 layer 들이 있다면 앞의 층들은 normal distribution을 latent value로 변환해주는 일을 하는 것이고 뒤의 층들은 이 latent value를 가지고 realistic한 digit를 sampling 할 수 있는 확률밀도함수를 만들어내는 것입니다.
두 번째 문제인 integral의 경우 integral을 다 계산하지 않고 Monte-Carlo estimation을 통해 estimate 할 것입니다. 여기서 Bayesian이 등장합니다. 대부분의 \(z\)에 대해서는 \(p_{\theta}(x \vert z)\)는 거의 0의 값을 가질 것입니다. 따라서 sampling이 상당히 많이 필요합니다. 데이터셋이 클 경우에 이것은 너무 cost가 큽니다. 좀 더 efficient하게 이 sampling 과정을 진행하려면 data에 dependent하게 \(z\)를 sampling 할 필요가 있습니다. \[p_{\theta}(x)=\int p_{\theta}(z)p_{\theta}(x|z)dz \approx \frac{1}{N}\sum_{i=1}^N p_{\theta}(x|z^{(i)})- (2)\]
따라서 \(p_{\theta}(z \vert x)\)를 생각해보는 것입니다. \(p_{\theta}(z \vert x)\)가 하는 역할은 \(x\)가 주어졌을 때 이 \(x\)를 생성해낼 것 같은 \(z\)에 대한 확률분포를 만드는 것입니다. 이것은 Bayes’ rule에 의해 다음과 같이 쓸 수 있습니다. 하지만 (1) 식에서 보듯이 \(p_{\theta}(x)\)를 계산하는 것은 intractable 하므로 이 posterior 또한 intractable posterior가 됩니다. 따라서 이 poterior를 approximate 하는 새로운 함수를 정의합니다.
\(\phi\)라는 새로운 parameter로 표현되는 \(q_{\phi}(z \vert x)\)는 일종의 encoder라고 볼 수 있습니다. 원래의 posterior를 approximate 했기 때문에 error가 존재합니다. 따라서 원래의 objective function에 대한 lower bound를 정의할 것입니다. 그 전에 VAE의 네트워크 구조를 살펴보겠습니다. “encoder”는 \(q_{\phi}(z \vert x)\)이며 \(x\)를 input으로 받아서 \(z\) space 상에서 확률분포를 만듭니다. 이 확률분포는 gaussian이라고 가정해서 만듭니다. 이 data dependent한 gaussian 분포로부터 \(z\)를 sampling 합니다. sampling된 \(z\)를 가지고 “decoder” \(p_{\theta}(x|z)\)는 \(x\)의 space 상의 gaussian distribution 혹은 Bernoulli distribution을 output으로 내놓습니다. 그러면 \(x\)를 이 분포로부터 sampling 할 수 있습니다. 이러한 구조를 가지기 때문에 “Auto-Encoder”가 되는 것이며 학습이 되고 나면 latent variable \(z\)라는 data의 의미있는 representation을 얻을 수 있습니다.
4. ELBO(Evidence Lower Bound)
이제 VAE를 어떻게 학습시키는지를 살펴보기 위해 objective function을 변형시켜보겠습니다. log likelihood는 다음과 같습니다. 이 값을 최대화시키는 것이 목표입니다. 이 식 자체는 intractable 하기 때문에 변형이 필요합니다. \[log p_{\theta}(x^{(i)})\]
이 log likelihood를 \(q_{\phi}(z \vert x)\)로부터 sampling한 latent variable \(z\)에 대한 expectation 식으로 바꿀 수 있습니다. \(p_{\theta}(x^{(i)}\)가 \(z\)에 dependent하지 않기 때문입니다. \[log p_{\theta}(x^{(i)}) = \mathbb{E}_{z\sim q_{\phi}(z \vert x^{(i)})}[log p_{\theta}(x^{(i)})] - (3)\]
여기서 Bayes’ Rule을 적용해보겠습니다. \[p_{\theta}(z \vert x^{(i)}) = \frac{p_{\theta}(x^{(i)} \vert z)p_{\theta}(z)}{p_{\theta}(x^{(i)})}\] \[p_{\theta}(x^{(i)}) = \frac{p_{\theta}(x^{(i)} \vert z)p_{\theta}(z)}{p_{\theta}(z \vert x^{(i)})} -(4)\]
(3)식에 (4)를 넣어서 다시 쓰면 다음과 같습니다. \[log p_{\theta}(x^{(i)}) = \mathbb{E}_{z\sim q_{\phi}(z \vert x^{(i)})}[log \frac{p_{\theta}(x^{(i)} \vert z)p_{\theta}(z)}{p_{\theta}(z \vert x^{(i)})}] - (5)\]
그 다음에 expectation 안의 항에 같은 값을 곱하고 나눕니다. \[log p_{\theta}(x^{(i)}) = \mathbb{E}_{z\sim q_{\phi}(z \vert x^{(i)})}[log \frac{p_{\theta}(x^{(i)} \vert z)p_{\theta}(z)}{p_{\theta}(z \vert x^{(i)})}\frac{q_{\phi}(z \vert x^{(i)})}{q_{\phi}(z \vert x^{(i)})}] - (6)\]
이 때, \(p_{\theta}(z)\)와 \(q_{\phi}(z \vert x^{(i)})\)를 하나로 묶고 \(p_{\theta}(z \vert x^{(i)})\)와 \(q_{\phi}(z \vert x^{(i)})\)를 하나로 묶어서 별도의 expectation으로 내보냅니다. \[log p_{\theta}(x^{(i)}) = \mathbb{E}_{z}[log p_{\theta}(x^{(i)} \vert z)] - \mathbb{E}_{z}[log\frac{q_{\phi}(z \vert x^{(i)})}{p_{\theta}(z)}] + \mathbb{E}_{z}[log\frac{q_{\phi}(z \vert x^{(i)})}{p_{\theta}(z \vert x^{(i)})}] - (7)\]
(7)식에서 우변을 살펴보겠습니다. 우변의 두번째 항과 세번째 항은 잘 보면 KL-Divergence의 형태인 것을 알 수 있습니다. 따라서 KL의 형태로 바꿔쓰면 다음과 같습니다. \[log p_{\theta}(x^{(i)}) = \mathbb{E}_{z}[log p_{\theta}(x^{(i)} \vert z)] - D_{KL}(q_{\phi}(z \vert x^{(i)}) \Vert p_{\theta}(z)) + D_{KL}(q_{\phi}(z \vert x^{(i)}) \Vert p_{\theta}(z \vert x^{(i)}))\]
이 식의 의미를 살펴보면 다음과 같습니다. 우변의 첫번째 항은 reconstruction을 의미합니다. \(q_{\phi}(z \vert x^{(i)})\)로부터 sampling한 \(z\)가 있으며 그 \(z\)를 가지고 \(p_{\theta}(x^{(i)} \vert z)\)가 \(x^{(i)}\)를 생성할 log likelihood입니다. 우변의 두번째 항은 prior인 \(z\)와 근사된 posterior인 \(q_{\phi}(z \vert x^{(i)})\) 사이의 KL-divergence 입니다. 즉, 근사된 posterior의 분포가 얼마나 normal distribution과 가깝운지에 대한 척도입니다(prior를 normal distribution으로 가정했을 때). 마지막 항은 원래의 posterior와 근사된 posterior의 차이로서 approximation error라고 볼 수 있습니다. 하지만 앞에서 살펴봤듯이 \(p_{\theta}(z \vert x^{(i)})\)는 intractable 해서 세번째 항을 계산하기 어렵습니다. 하지만 KL의 성질대로 이 세번째 항은 무조건 0보다 크거나 같습니다. 
따라서 첫번째 항과 두번째 항을 하나로 묶어주면 원래의 objective function에 대한 tractable한 lower bound를 정의할 수 있습니다. MLE 문제를 풀기 위해 objective function을 미분해서 gradient ascent 할 것입니다. Lower bound가 정의된다면 이 lower bound를 최대화하는 문제로 바꿀 수 있고 결국 lower bound의 gradient를 구하게 될 것입니다. lower bound의 두 항은 모두 미분가능하기 때문에(어떻게 미분가능한건지는 뒤에서 살펴보겠습니다) 이제 우리는 최적화를 할 수 있습니다. 
lower bound를 다시 정의하자면 다음과 같습니다. \[\mathcal{L}(x^{(i)}, \theta, \phi) = \mathbb{E}_{z}[log p_{\theta}(x^{(i)} \vert z)] - D_{KL}(q_{\phi}(z \vert x^{(i)}) \Vert p_{\theta}(z))\]
이 lower bound 식은 evidence의 log 값인 \(p_{\theta}(x^{(i)})\)의 lower bound이기 때문에 Evidence Lower Bound, ELBO라고 부릅니다. \[log(p_{\theta}(x^{(i)})) \geq \mathcal{L}(x^{(i)}, \theta, \phi)\]
따라서 원래 \(p_{\theta}(x)\)를 최대화하는 문제는 다음과 같이 바뀝니다. \[\theta^*, \phi^* = argmax_{\theta, \phi}\sum_{i=1}^N \mathcal{L}(x^{(i)}, \theta, \phi)\]
여기까지의 식을 한 번에 보자면 다음과 같습니다.
이 ELBO를 구하는 과정은 다음 그림을 통해 이해해볼 수 있습니다. x를 encoder의 input으로 집어넣으면 encoder는 latent space 상에서의 mean과 variance를 내보냅니다(이 때, mean과 variance는 latent vector의 dimension마다 하나씩입니다). 그러면 이 mean과 variance가 posterior를 나타내게 되고 prior와의 KL을 구할 수 있습니다. 그 이후에 \(z\)로부터 decoder는 data의 space 상의 mean과 variance를 내보냅니다(만약 decoder의 output을 gaussian이라고 가정했다면. Bernoulli 분포라고 가정했다면 다른 형태). 그러면 ELBO의 첫번째 항 값을 구할 수 있고 ELBO가 구해집니다. 구한 값에 Backprop을 해서 업데이트하면 VAE의 학습과정이 완성됩니다.
하지만 코드로 짜기에는 ELBO의 두 항을 구하는 과정이 명확하지 않습니다. 이 과정은 VAE의 원논문인 “Auto-Encoding Variational Bayes” 논문에 나와있습니다. 따라서 다음 포스트에서는 논문의 내용을 살펴보도록 하겠습니다.
