Generative model 기초 2. VAE 정리
VAE(Variational Auto Encoder)는 GAN과 같이 Generative model(생성 모델)의 한 종류이다. 생성모델의 근본이라고 볼 수 있는 모델이다. VAE의 구조와 목적, VAE의 Loss Function 유도 과정까지 정리해보자.
Variational Auto Encoder
VAE는 input image를 Encoder를 통해 feature를 추출하여 Latent Vector z로 만들고, 이 Latent Vector z를 다시 Decoder를 통해 output image를 생성하는 구조이다. 이 때 AE(Auto Encoder)는 input과 output을 같게 만들도록 학습시키는데, VAE에서는 Latent Vector z를 decoding할 때 input image와는 비슷하지만 다른 이미지를 generation한다. 이 과정에서 우리가 원하는 이미지를 생성할 수 있다. VAE의 목표는 Generative model과 같다. 즉, Input data와 같은 분포를 가지는 분포에서 sample을 추출하고 여기서 새로운 이미지를 생성해내는 것이 목표이다.
Encoder를 통해 생성된 feature, 즉 latent vector $z$는 Gaussian distribution을 따른다고 가정한다. Encoder output으로 $\mu$와 $\sigma$를 생성하는데, 이 값을 이용하여 normal distribution을 생성한 뒤 샘플링을 통해 확률 분포에서 값을 추출한다. 이 값을 Decoder에 넣어 Input Image와는 유사하지만 다른 이미지를 생성하는 것이다.
학습된 VAE에서 Encoder는 차원 축소(Manifold learning)의 역할을 하고, Decoder(Generative model learning)는 생성 모델의 역할을 한다.
AE vs VAE
그렇다면 VAE에 대해 정리하기 전에, AE와 VAE의 차이점은 무엇일까?
AE(Auto Encoder)는 input image x를 동일하게 x로 복원하여, Latent vector z를 잘 임베딩하는 것이 목적이다. 즉, Encoder의 차원 축소(manifold learning)가 학습의 목표이다. Unsupervised 방법을 Supervised 방법으로 바꿔(Decoder의 output image를 이용하여) 원본 이미지를 잘 복원할 수 있도록 학습시키는 방법이다. Encoder의 feature extractor로서의 역할이 중요하다고 볼 수 있다.
반면 VAE(Variational Auto Encoder)는, input image x를 가지고 이와 유사한 새로운 output을 generation하는 것이 목적이다. 즉, Decoder의 Generation modeling이 학습의 목표이다. Encoder가 생성한 Latent variable을 이용해 sampling을 진행하고, Decoder를 이용해 새로운 이미지를 생성하는 Generation model이다.
VAE Architecture
이미지 출처:https://wikidocs.net/152474
VAE의 구조는 다음과 같다. 크게 세 부분으로 나누어 볼 수 있다.
- Encoder
- input $x_i$ → $q_{\phi}(x)$ → output ${\mu}_i, {\sigma}_i$
- Gaussian Encoder
- input $x$를 받아 데이터의 latent variable을 생성한다. 생성된 variable의 분포는 정규분포를 따른다.
- Sampling(Reparameterization Trick)
- Encoder에서 생성한 ${\mu}_i, {\sigma}_i$를 사용해 Gaussian distribution을 생성하고, 만들어진 분포로부터 샘플링을 통해 $z$를 추출하여 Decoder에 넣어주면 입력 $x$와 유사하지만 다른 이미지를 생성할 수 있다.
- 하지만 sampling과정은 미분불가능하므로 back propagation을 수행할 수 없다. 따라서 Reparameterization이라는 트릭을 통해 샘플링을 수행한다. Reparameterization Trick은 뒤에서 다시 논의한다.
- Decoder
- input $z_i$ → $g_{\theta}(z)$ → output $p_i$
- Bernoulli Decoder
- Sampling된 latent vector를 입력으로 받아 이미지를 생성한다. 생성된 이미지는 bernoulli 분포를 따른다.
Loss Function은 위와 같다. 유도 과정에서 더 자세히 설명하겠지만, Loss는 두 부분으로 나누어 볼 수 있다. $x_i$와 생성된 이미지에 대한 Reconstrunction Error, Encoder에서 생성한 Latent variable이 우리가 가정한 분포를 따라야한다는 Regularization으로 구성되어있다.
VAE Loss Function
KL Divergence
KL Divergence에 대해 자세히 다루지는 않겠지만, 간단하게 정리해보자.
KL Divergence는 두 확률 분포의 차이를 계산하는데 사용되는 함수이다. 두 확률분포 P와 Q를 비교하고 싶을 때 KL Divergence를 사용할 수 있다.
그림을 보면 알 수 있듯이, 두 확률분포가 가까우면 그 값이 점점 작아지고(같아지면 0), 두 확률분포가 멀어지면 KLD 값이 커진다.
수식은 아래와 같이 나타낼 수 있다.
그럼 VAE Loss function의 유도 과정을 살펴보자. 먼저 원래 VAE의 목적이었던, decoder에서 시작해보자.
구하고자 하는 것은 $x$에 대한 확률 분포이다. Decoder의 output이 우리가 원하는 $x$가 나올 확률$(p_{\theta}(x))$이 가장 높아하므로, $x$의 likelihood를 최대화하는 방식으로 학습을 진행해야한다.
그래서 위와 같이 식을 전개해 보았을 때, 첫번째 term $p_{\theta}(z)$는 우리가 gaussian distribution이라고 가정해놓은 부분이고, 두번째 term은 Decoder에 해당하는 부분이라 Neural network로 구현할 수 있다. 하지만 모든 $x$에 대해 적분해야하는 것은 불가능하다.
따라서 Decoder만을 사용해서 학습을 진행하는 것은 불가능하다고 생각이 된다. 그래서 여기에 Encoder를 추가해서 문제를 해결해 보고자 한다.
여전히 $p_{\theta}$를 maximize하는 것이 목표이다. $p_{\theta}$ 에 log를 씌워주고 Expectation 형식으로 식을 바꿔주었다. 첫 줄을 살펴보면, $z$ ~ $q_{\phi}(z|x^{(i)})$, 즉 z의 분포가 x가 Encoder를 거쳐 나오는 확률 분포라는 조건이 추가되었다.
첫 번째 줄의 식에 Bayes’ rule을 적용하여 두번째 줄과 같이 변경해주었다. 그 뒤에 세번째 줄에서 분모 분자에 $q_{\phi}$를 곱해주었다. 그리고 전개해주면 네번째 줄과 같은 식으로 분리할 수 있다.
네번째 줄의 식은 KL Divergence식으로 표현이 가능하고, 결국 마지막 줄의 식과 같이 유도된다.
그럼 마지막에 나온 식을 살펴보자. 우리의 목적은 $p_{\theta}(x^{(i)})$ 를 maximize 하는건데, 그러기 위해서는 마지막 줄의 식을 maximize하면 된다는 뜻이다. 이 식을 차례로 살펴보면, 3개의 term으로 쪼개 볼 수 있다.
① Decoder Network. Latent vector $z$에서 $x$가 나올 확률을 최대화해주는 부분.
② Encoder를 통과하는 결과 $(q_{\phi}(z|x^{(i)}))$ 가 내가 정의하는 distribution $(p_{\theta}(z))$ 와 비슷하도록 만들어주는(KL Divergence) 부분.
③ 문제는 세번째 term이다. $q_{\phi}$ 부분은 Encoder, $p_{\theta}$ 부분은 Decoder부분으로 볼 수 있는데, Decoder 에서 $x$ → $z$ 로 오는 조건부 확률은 알 수 없다. 따라서 이 부분은 계산이 불가능하다. 다만, KL Divergence이기 때문에 0보다 크다는 것을 알 수 있다.
결국 세번째 term은 우리가 계산할 수 없는 부분이다. 그렇다면 첫번째, 두번째 term을 Lower bound로 하여, $\log p_{\theta}$를 maximize하고 싶다면 최소한 이 부분을 maximization하는 것으로 문제로 축소한다. 세번째 term은 최소한 0보다 크다는 점을 알고 있으니 이부분은 고려하지 않기로 한다.
따라서 파란색 박스 부분을 Variational lower bound 즉, ELBO 라고 부르고, 이 식을 maximize해주는 Encoder parameter $\phi$ 와 Decoder parameter $\theta$ 를 찾는 문제로 바뀌었다.
위 식을 다시 정리하면 위와 같다. Minimization problem으로 바꾸기 위해 (-)를 붙여줬다. 이제 위 함수를 두 부분으로 나누어 볼 수 있다. 첫번째 term은 Reconstruction Error, 두번째 term은 Regularization 으로 나누어진다. 두 부분에 대해 정리해보면 아래와 같이 정리할 수 있다.
- Reconstruction Error:
- 현재 샘플링용 함수에 대한 negative log likelihood
- $x_i$에 대한 복원 오차(AutoEncoder 관련)
- Regularization:
- 현재 샘플링용 함수에 대한 추가 조건
- 샘플링의 용이성/생성 데이터에 대한 통제성을 위한 조건을 prior에 부여하고 이와 유사해야한다는 조건을 부여
VAE Optimization
그럼 이렇게 정리된 수식을 실제 구현을 위해 Optimization 해보자.
Regularization
먼저 Regularization term을 살펴보자.
Regularization 부분은 Encoder를 통해 나온 latent vector z가 우리가 가정해 놓은 distribution과 일치하도록 만들어져야한다는 조건이다. 그렇다면 여기서 우리는 몇가지 가정을 추가해야한다.
Assumtion 1. Encoder output $q_{\phi}(z|x)$ 는 Multivariate gaussian distribution을 따르고 diagonal covariance를 갖는다.
Assumtion 2. prior $p(z)$는 Multivariate normal distribution이다.
이렇게 가정했을 때, 이 둘 사이를 갖게 만들어준다 ⇒ KL Divergence를 이용한다
KLD가 multivariate normal distribution일 때, 위의 식에서 맨아랫줄과 같이 전개될 수 있다. (tr은 대각행렬의 합을 의미한다) 이 때, p(z)가 N(0,I)를 따르므로 평균과 분산에 각각 0과 1을 대입하여 풀어주면 위와 같이 유도가 된다.
결과적으로 Encoder를 통과한 $\mu$와 $\sigma$를 가지고 최적화 할 수 있는 식을 유도해냈다.
Reconstruction Error
다음으로는 Reconstruction Error를 최적화해보자.
Reconstruction Error는 Expactation식이기 때문에 적분식으로 풀어쓸 수 있다. 하지만 해당 적분식에서 모든 z에 대해 적분한다는 것은 쉽지않다.
그래서 대신, Monte-carlo technique을 적용하기로 했다.
여기서 Monte-carlo technique이란? 어떤 분포를 가정하고, 이 분포에서 데이터를 샘플링 할 때 반복적으로 여러번 많은 수의 샘플을 추출하여 평균을 내면, 이 값이 전체 데이터 분포의 true 기댓값과 거의 동일해진다는 알고리즘을 일컫는 말이다. 따라서 우리는, 확률분포에서 $z$를 계속 추출한 뒤 각각의 $z$에 대하여 log값을 계산하고 평균을 내는 문제로 바꾸어서 풀고자한다.
하지만 딥러닝에서 Monte-carlo technique를 적용하기는 역시 어렵다. 반복적으로 샘플링하여 계산을 하기엔 계산량이 지나치게 많아진다. 따라서 한번 더 문제를 축소하여 L=1로 적용하기로 한다. 즉, 샘플링을 한 번만 진행한다. 랜덤하게 하나를 샘플링하여 이것을 대표값으로 쓰는 것이다.
샘플링을 사용하면 미분이 불가능하므로, back propagation이 불가능하게 된다. 이 문제를 해결하기 위해 하나의 trick을 추가 적용하게 되는데, 바로 reparameterization trick이다. 이 부분은 reconstruction error를 마저 이해한 뒤 뒤에서 살펴보기로 하자.
우여곡절 끝에 여기까지 왔다. 정리하자면, 원래는 Expectation에 대해 적분을 수행해야하는데, 적분이 불가능해서 Monte-carlo technique를 적용하여 샘플링 평균을 내고자 했으나, 이 방법마저도 계산이 불가능하여 샘플링을 한번만 해서 값을 구하자는 흐름으로 식이 전개되었다.
그럼 이제 축소한 식에서, 확률 분포 $p_{\theta}$를 구하는 문제로 넘어가보자.
여기서 하나의 가정이 추가된다. $p_{\theta}$는 어떤 확률 분포를 가지는가에 대한 가정이다.
Assumption 3. Decoder의 likelihood $p_{\theta}$의 확률 분포 → Multivariate Bernoulli or Gaussian distribution
여기서는 $p_{\theta}$가 Bernoulli 라고 가정한다.
$z^i$를 입력으로 받아 $p_i$로 복원할 때, $D$ 만큼 확률을 계산하게 되는데, 확률은 곱으로 표현할 수 있고 log가 붙기 때문에 이를 다시 합으로 바꿔준다(첫번째 줄 식). 여기서 $p_{\theta}$식을 bernoulli 식으로 바꿔주고, 이를 전개하면 위와 같은 결과를 얻을 수 있다. 즉, Cross Entropy 식으로 유도하였다.
추가로, $p_{\theta}$를 Gaussian 으로 가정해보자.
만약 Gaussian으로 가정한다면, Decoder가 $\mu$와 $\sigma$를 output으로 생성하게 된다. 그렇게 되면 위와 같이 식을 전개할 수 있고, 여기에 identity covariance를 가정하여 $\sigma$에 1을 대입하면 최종적으로 맨 아래 식과 같이 MSE로 유도할 수 있다.
이렇게 Reconstruction Error, Regularization에 대해 전부 계산 가능한 식으로 유도를 마쳤다. 이제 마지막으로 Reparameterization 과정을 살펴보자.
Reparameterization
Encoder output으로 $\mu$와 $\sigma$를 만들고, 이 값을 이용해 normal distribution을 만든 뒤 샘플링하는 방식은 미분이 불가능하다(back propagation 불가능). 따라서 이 식을 미분가능한 식으로 만들어줘야 한다.
Reparameterazation 방식은 다음과 같다. 평균이 0이고 표준편차가 1인 표준 정규 분포에서 샘플링을 수행한 뒤($\epsilon$), Encoder output $\sigma$ 에 곱해주고 $\mu$ 에 더해준다. 이 값을 새로운 샘플링 값 $z$로 사용한다. 해당 방법을 사용하면 $z$에 대한 식을 표현할 수 있고, 미분가능하게 되었다.
Summary
최종적으로 유도한 과정을 포함하여 VAE에 대해 정리해보면 아래와 같다.
최종적으로 VAE는 Gaussian Encoder + Bernoulli Decoder로 이루어진 구조이고, Loss function은 Reconstruction Error, Regularization 두개의 term으로 구성되어있다.
VAE Loss Function의 유도과정은 [2]Smart Design Lab 강의를 인용했다. 깔끔 명료한 강의라 자세히 들어보고 싶은 분은 꼭 한번 들어보길 추천한다.
Reference
[1] 한땀한땀 딥러닝 컴퓨터 비전 백과사전: https://wikidocs.net/152474
[2] Smart Design Lab @KAIST: https://www.youtube.com/watch?v=GbCAwVVKaHY&t=546s
[3] Auto-Encoding Variational Bayes(paper): https://arxiv.org/abs/1312.6114