Variational Auto-Encoder (Kingma, 2013 논문 리뷰)

VAE(Variational Auto-Encoder)는 생성형 AI의 역사에 중요한 획을 그은 모델로, 입력 데이터의 확률 분포를 학습하여 새로운 데이터를 생성하는 모델이다. 이 모델은 2013년 D.P.Kingma와 Max Welling의 논문(1)으로 처음 공개되었으며, 2024년 3월 현재 33,556회 인용되었다. 본 포스팅에서는 이 논문을 리뷰하면서 VAE의 구조를 알아보고자 한다.
 
(1) Kingma, D. P. & Welling, M. (2014). Auto-Encoding Variational Bayes. 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, .

---

VAE의 구조

들어가기에 앞서, VAE의 구조에 대해 살펴보자.

VAE는 위와 같은 구조를 가진다. 각각의 요소를 살펴보자.

• Input: vector x를 입력받는다.
• Encoder: Input vector x를 encoding하여 latent vector z로 변환한다. encoder의 parameter는 phi로 표기하며, input vector에 대한 latent vector의 분포를 q(z|x)로 나타낸다. Encoder의 목적은 이 분포를 approximation 하는 것이다.
• Latent Space: Encoding되어 latent vector 상태로 존재하는 입력 정보이다. 예를 들어, 사과 사진이 input으로 들어온다면, 기존 이미지보다 단순화된 정보가 latent vector로 저장된다. 각 unit은 사과의 색상, 크기, 모양 등의 속성을 의미할 수 있다. 이때, latent space에는 z의 "분포"가 저장되기 때문에, 여기에 Gaussian noise를 추가할 수 있다.
• Decoder: Decoder는 latent vector를 다시 input vector와 비슷한 형태로 복원한다. decoder의 parameter는 theta로 표기하며, 복원을 위해 p(x|z)를 이용한다. 따라서 decoder의 목적은 이 분포를 approximation하는 것이다.
• Output: input과 같은 형식으로 새로 생성된 데이터이다.

AE(Autoencoder)와 VAE(Variational Autoencoder)는 바로 이 latent space에서 noise의 주입 여부에 따라 바뀐다. AE는 input으로 들어온 x가 encoder, latent space, decoder를 거쳐 output으로 나갈 때 x와 최대한 비슷해지는 것이 목적이다. 반면, VAE는 latent space에 noise를 추가하여, decoder를 거쳤을 때 입력 데이터에 없었던, 새로운 출력을 만드는 것이 목표이다. 이것이 VAE가 생성형 AI (Generative AI)인 이유이다.

Introduction

Approximate inference & Learning

 
Approximate inference란, 복잡한 모델에서 정확한 확률을 계산하는 것이 불가능할 때 확률의 결과를 추정하기 위해 사용하는 기법이다. 특히, 이 논문에서 문제 상황으로 설정한 것은, 다루기 어려운 Posterior를 가진 *directed probabilistic model을 이용한 추정을 효율적으로 하는 방법이다. 기존 방법으로는 이러한 추정을 위해 **mean-field approach를 사용했었는데, 이는 posterior 자체를 근사해야 하고, 일반적인 경우에는 이것이 어렵다.
 
*Directed probabilistic model: direct graphical model이라고도 한다. 이는 아래 그림처럼 노드를 연결하는 edge(화살표)의 방향성이 한쪽으로만 움직이는 그래프를 말한다. 화살표의 방향은 보통 인과 관계를 나타내는데, 이러한 형태로 복잡한 확률 분포를 간단한 확률 분포의 조합으로 나타낼 수 있다.

출처: Carnegie Mellon Univ.

 
** Mean-field approach: 자유도를 평균화하여 원본에 근사하는 방법.
 
본 논문에서는 posterior를 계산하기 어려워 기존 방법으로는 접근이 어려운 문제에 대한 해법을 제시한다. 이러한 새로운 해법을 AEVB(Auto-Encoding Variational Bayes)라고 하는데, 여기서 핵심적으로 들어가는 알고리즘이 SGVB estimator이다. SVGB estimator를 이용하여, 과도한 연산 없이 아주 효율적으로 posterior를 근사할 수 있으며, 이를 이용해 parameter를 학습시킬 수 있다.
 
Method

위 그림에서 점선은 encoder를, 실선은 decoder를 나타낸다. 즉, 점선은 q_phi (z|x)에 따라 "계산하기 어려운 posterior "에 해당하는 p_theta (z|x)를 만드는 과정이고, 실선은 p_theta(z)p_theta(x|z)를 이용하여 x를 generation하는 과정을 나타낸다. 여기서 parameter인 phi와 theta는 공동으로 학습된다. 
 

상황 설정

 
N개의 i.i.d.(independent & identically distributed) 샘플로 구성된 아래와 같은 데이터셋을 생각하자. 이때 X는 이산변수일 수도, 연속변수일 수도 있다.

이 데이터셋의 원소들은, 랜덤한 연속변수 z에 의해 generation 된다고 가정하자. 즉, 이 generation 과정은 아래의 두 스텝을 따른다.
 
1. x_i로부터 z가 생성된다. 이때 다음의 prior 분포가 이용된다.

2. 생성된 z로부터 다시 x_i가 생성된다. 이때 다음의 likelihood가 이용된다.

 
이 두 가지 과정에 사용된 분포의 확률밀도함수는, z와 theta에 따라 변한다. 이러한 상황 설정은 parameter (theta)와 latent variable(z)를 직접적으로 알기는 어렵게 하지만, 두 가지 장점이 있다. 첫째, P(x), P(z|x) 를 직접 다루기 힘든 경우에도 사용할 수 있다. 둘째, 큰 데이터셋에 대해서도 적용할 수 있다.
 
이 상황에서 추정해야 할 변수는 아래 세 가지이며, 각 변수를 추정할 방법을 표로 만들었다.

parameters (theta)	ML or MAP estimation
latent variable (z)	posterior inference
variable (x)	marginal inference

 

---

ELBO

 
VAE 모델에서 optimization해야 하는 값은 무엇인가? 이는 ELBO(Evidence of Lower Bound = Variational Lower Bound) 라는 식이다. ELBO는 다음과 같이 정의된다.

식 (1)

식이 복잡해 보이는데, 무슨 뜻인지 하나하나 뜯어 보자.

먼저, 위의 두 식은 VAE의 구조에서도 간략히 설명하였다. q는 encoder에 의해 정의되는 분포이다. Input vector x가 들어왔을 때 latent vector z의 분포를 나타내며, phi는 encoder의 parameter이다.  p는 decoder에 의해 정의되는 분포이다. latent vector z가 들어왔을 때 x의 분포를 나타내며, theta는 decoder의 parameter이다. 

 
다음으로, 위의 KL divergence (KL divergence의 개념: https://cascade.tistory.com/46) 는 아래 두 분포 간의 차이를 측정한다.
(1) encoder에 의해 정의된 latent vector의 분포 (q(z|x))
(2) latent vector의 사전 분포 (p(z))

 
마지막으로, 이 항은 재구성 손실(Reconstruction loss)를 나타내는 항이다. 이는 latent vector를 통해 생성된 데이터가 원본 데이터와 얼마나 유사한지를 측정한다. Log likelihood의 꼴로 표현된 이 식이 최대가 될수록 생성된 데이터가 원본 데이터와 유사한 것이다.
 
종합하면, ELBO는 아래와 같이 정리할 수 있다.

ELBO의 최적화는 (1)재구성 손실(Reconstruction loss)의 최소화 + (2) KL divergence의 최소화 이다

---

 

Variational Bound

 
각 x의 marginal likelihood를 계산하기 위해, 아래 식을 사용하자.

여기에 식 (1) 을 적용하면,

 
KL divergence의 non-negativity를 적용하면 아래와 같다. 여기서 log p_theta (x, z)는 결합 로그 확률 분포로, 기댓값을 x와 z를 모두 고려하여 나타내는 것이다.

 
이제, optimization을 통해 ELBO를 parameter로 미분해야 한다. 여기서 phi로 직접 미분하는 것은 어렵다. 일반적으로 이런 문제에서는 ***naive Monte Carlo gradient estimator라는 방법을 사용해 아래와 같이 "근사적으로" 미분을 수행할 수 있다.

이때 각 z값은

와 같이 추출하였다. 하지만, 이 방법은 variance가 너무 크고, 위에서 열거한 세 가지 목적을 달성하는 데 실용적이지 않아 새로운 방법을 찾아야 한다.
 
*** naive Monte Carlo gradient estimator
 
naive Monte Carlo 방법은 "분포를 알기 위해 수행하는 노가다"로 볼 수 있다. 예를 들어, 원주율을 계산하기 위해 아래와 같이 단위원을 하나 그려 놓고 무작위로 점을 찍어 안으로 들어오는 비율을 보는 것이다. ELBO의 최적화를 위해 실용적인 알고리즘은 아니다. 이후에 나올 reparametrization trick은 이 노가다를 어떻게 하면 "효율적으로" 할지를 위해 고안되었다.

출처: geeksforgeeks

 

SGVB estimator & AEVB algorithm

 
이제, "실용적인" ELBO와 그 미분값을 찾기 위하여 새로운 방법을 적용해 보자. 여기서 사용되는 것이 reparametrization trick이다. 기존 방법에서는 아래의 근사 분포를 직접 이용했었다.

Reparametrization trick은 이것을 직접 이용하는 것이 아니라, 위 분포에 포함되는 z_tilde라는 새로운 랜덤한 변수를 잡아 아래와 같은 새로운 분포를 만드는 것이다. 이때 epsilon은 noise distribution을 따른다.

 
이제, 위에서 시도했던 Monte Carlo Estimation을 f(z)에 대해 직접 하는 것이 아니라 g_phi 분포로 돌려서 수행할 수 있다. 

 
이를 ELBO에 적용하여, SGVB(Stochastic Gradient Variaitional Bayes) estimator를 아래와 같이 만들자.

 
위 식에서 우변이 ELBO이고, 좌변이 이를 근사한 SGVB estimator이다.

 
이제 위와 같이 monte carlo estimation을 적용하여, ELBO가 계산 가능한 형태로 바뀌었다. 이는 이제 SGD(Stochastic Gradient Descent)나 Adagrad와 같은 알고리즘을 이용하여 최적화할 수 있다.

SGD는 전체 N개의 데이터셋(batch)으로 loss function을 계산하는 것이 아니라, M개짜리 minibatch로 쪼개어 loss function을 계산하는 방법이다. 이는 batch gradient descent보다 수렴하는 데 필요한 step 수는 더 많을 수 있지만, 계산이 빠르고 같은 시간에 더 많은 step을 갈 수 있다. 또한, local이 아닌 global minimum에 수렴할 가능성이 더 높아진다.
 
이러한 최적화 알고리즘을 적용하기 위해, ELBO를 아래와 같이 정리할 수 있다.

실험

저자들은 이 생성모델을 검증하기 위하여 MNIST (전에 소개했던 CNN https://cascade.tistory.com/40 에서도 사용했던 손글씨 데이터셋이다) 와 Frey Face 얼굴 데이터셋을 사용하였다.

Frey Face 데이터셋

 
기존의 모델인 wake-sleep 알고리즘과 비교한 결과, 모든 경우에서 VAE가 우월한 결과를 보였다. 또한, latent variable의 개수가 많다고 해서 overfitting이 나타나지 않았다.