[CSE URP] Auto-Encoding Variational Bayes (ICLR 2014)

Original Paper https://arxiv.org/abs/1312.6114

Auto-Encoding Variational Bayes

---

GAN이라는 Generative Model를 톺아보면서 좀더 깊이 있게 공부하고자 쓰는 포스팅. 생성형 AI를 처음 접했던 만큼 논문도 천천히 읽어보면서 흡수해보려고 한다 마음처럼 잘 됐으면 좋겠다

---

Background 1: Unsupervised Learning, probability distribution

Generative model은 데이터 내부의 확률 분포를 찾는 일반적인 unsupervised learning 모델 중 하나이. 확률 분포와 unsupervised learning을 먼저 리뷰해보자. 

 

1. Supervised Learning (지도학습) 

Supervised learning의 대표적인 모델인 discriminative model은 일반적인 classification task에 해당. input에 해당하는 클래스를 맞추기 위해, 클래스가 🐱인지, 🐶인지를 구분하기 위해서 학습하게 되며 모델을 학습하고 loss를 구하기 위해 실제 label이 사용된다.

2. Unsupervised Learning

반면 unsupervised learning은 label 없이 data만으로 모델을 학습시킬 수 있는 학습방법으로, 생성 모델은 latent code를 가지고 image를 만들게 된다. 따라서 생성 모델은 class를 나누는 것이 아닌, training data를 학습(=분포를 학습)하는 것이 목적이 됨!

3. Probability Distribution

주사위의 각 눈이 나올 확률이 위와 같다고 가정하고 그래프를 그려보면... (당연히 모두 1/6이겠지만.. 예시를 위해 위와 같다고 해보자 ㅎㅎ) 그랬을 때 랜덤 변수 (=X, 1~6)의 값에 따라 나올 확률을 그린 것이 확률 분포이며, 여기서는 x가 여섯 개의 경우의 수를 가질 수 있는데, x가 이미지의 차원으로 커진다면 이미지의 확률 분포도 학습할 수 있음.

여기서 말하는 x는 이미지, 즉 n*n*3의 dimension을 가진 pixel들이 모여 매우 고차원의 x가 된다. 예시를 위해서 실제로는 매우 고차원의 x이지만, 1차원으로 생각해서 그래프(오른쪽)을 그려보면, 가장 높은 피크값이 데이터의 많은 부분을 차지하는 이미지를 의미하게 됨.

Background 2: Auto Encoder

Variational autoencoder에서 autoencoder이 무엇을 의미하는 걸까? Autoencoder(AE)의 구조를 먼저 살펴보고 시작하자. 

• 비지도학습(unsupervised learning)으로 입력 데이터의 표현을 효율적으로 학습할 수 있는 인공신경망 (feature extractor)
• Encoder - z (latent variable) - Decoder 형태
• output이 input의 형태로 나타나도록 하는 것이 목적이나, 그 과정에서 입력을 reconstruction해서 데이터를 효율적으로 representation 하는 방법 학습
• 차원을 축소(특징 추출)하기 위한 목적으로 만들어졌기 때문에 Encoder를 학습하기 위해 Decoder를 사용한 것으로 볼 수 있음

Variational Auto Encoder

• 대략적인 동작 과정
  1. Input 이미지 Encoder에 넣음
  2. 인코더를 거쳐 2개의 vector 평균과 표준편차를 output으로 배출
  3. 위의 평균과 표준편차를 통해 normal distribution 만듬
  4. distribution에서 데이터 sampling 사용해서 z 생성
• BUT!! Distribution에서 단순히 sampling 했을 경우에는 미분(=Backpropataion)이 불가능하다. 따라서 미분이 가능한 식으로 바꾸기 위해 추가로 trick을 사용했다. 그 Trick이 바로...
• 👉 Reparameterization Trick. Diffusion Model 공부하다 보면 나오는 바로 그 녀석이다 .....

 

Trick의 원리(?) 를 적어보자.

Before: Encoder의 출력인 평균과 표준편차로 바로 distribution 생성

After trick: 평균과 표준편차를 normal distribution에서 샘플링한 입실론과 연산을 수행하여 before에서 얻고자 하는 분포를 간접적으로 얻음

→ z와 평균, 표준편차를 식으로 나타낼 수 있게 되어 미분이 가능해짐! 

 

여기서 AE와 VAE의 차이점으로 돌아가보자. 

→ 먼저 AE와 VAE의 동작 과정을 통해 latent vector z를 찾아내는 방법에서 차이점이 있음을 확인할 수 있다. AE는 단순히 차원 축소를 했다면, VAE는 데이터의 분포를 찾은 후 샘플링을 통해 z를 생성했음을 알 수 있다.

→ 위에서 언급된 것처럼, sampling하는 과정은 미분, 즉 backpropagation이 불가능하다는 문제가 있기 때문에 Reparameterization Trick을 사용하는 것

Loss Function

두 가지 function의 합(Reconstruction Error + Regularization)으로 나타냄

• Reconstruction Error

• input과 output의 차이를 최소화 하는 것
• normal distribution의 경우 MSE, bernoulli distribution의 경우 Cross entropy로 나타남 → VAE는 기본적으로 bernoulli를 따르는 Decoder 사용
• Regularization 

VAE는 normal distribution을 만들어 sampling하기 때문에 z는 normal distribution을 따른다는 가정이 들어가있음

그 과정을 최적화에 넣기 위해, encoder를 통과해서 나오는 z 값의 확률분포와 정규분포와의 거리가 최소가 되도록 KL divergence 최소화. KL Divergence는 두 확률 분포 차이를 계산하기 위해 사용하는 함수로, 아래와 같이 작성된다.

 

그렇다면 Loss Function을 구하기 위해서 Objective Function이 무엇일지 생각해보자. 

아래 그림과 같이 내가 갖고 있는 데이터(z)로 x에 대한 확률 분포를 정의한다면, 데이터 x가 나올 확률이 가장 커지는 distribution을 만들고자 할 것이다.

따라서, training data의 likelihood(가능도)를 최대화, x가 나올 확률 p(x)를 최대화 시키는 것이 objective function ! 티스토리 수식을 잘 모르겠어서 직접 써봤다 글씨는 흐린눈해주자...

주인장 손목이 아파서...밑으로 갈수록....못보겠는 글씨체...

최대화 시켜야 할 Data Likelihood 식을 적어본 것이다. 그러나 디코더만으로 z에 대한 x를 뽑아내는 것에는 수식적으로 계산하기 힘든, intractable한 연산이 너무 많다(위 그림의 빨간 부분).

해결방법

⇒ p(x|z) 구하기 위해서 추가적인 인코더 네트워크 q(z|x) 정의해주자! 여기서 이 q(z|x)는 실제 p(z|x)를 뉴럴 네트워크로 최대한 근사하는 분포이다.

⇒즉, 원래 목적은 디코더만으로 z에서 x를 만드는 것이지만, 디코더만으로는 학습할 수 없어서 인코더의 도움을 받아야 한다.

이를 가지고 최대화시켜야 할 likelihood p(x)를 다시금 정리해보자 :)

• p(x)의 최대를 구하기 위해 log(p(x))의 최대를 구해도 됨→ z가 encoder의 q(z|x) distribution을 따를 때 log(p(x))의 기대값을 구함 → 베이즈 정리 + 수식 계산 → KL divergence 형태로 나타냄
• 정리하면, 아래의 식 maximize
•  

각 항의 의미는 아래와 같다. 

• 첫번째 항 : decoder를 통해 복원될 수 있도록 만들어주면 됨
• 두번째 항 : encoder를 통과한 확률 분포가 z의 분포와 비슷하게 만들어주면 됨
• 세번째 항 : x → z 방향의 조건부 확률 값은 intractable 하기 때문에 정확한 KL divergence 값 구할 수 없다. 단지 KL divergence 값 자체가 0 이상이라는 것만 알 수 있음
• 즉, 빨간 박스 부분(ELBO)을 maximize 하면 됨
• θ는 decoder의 파라미터, Φ는 encoder의 파라미터

최종 Loss Function을 도출하면, 결과적으로 아래 이미지와 같아진다.

• Reconstruction Error : x가 Encoder(Φ라는 network) 지나서 z가 나오고, 해당 z를 Decoder(θ라는 network)에 넣었을 때 x가 나올 확률 maximize.
  • 즉, input과 output의 차이를 최소화 하는 것
• Regularization : x가 Encoder(Φ라는 network)를 지났을 때 z의 분포가 가정한 z의 분포와 일치하도록 만듬

---

Optimization

• Regularization 계산
•  

Assumption

1. 가우시안 encoder를 통과해서 나오는 분포가 정규분포 + diagonal한 covariance를 가진다고 가정 할 경우 아래와 같이 표현 가능

2. decoder를 통과해서 나오는 p의 분포는 평균이 0이고, 모든 표준편차가 1로 표현되는 표준 정규 분포를 따른다고 가정→ N(0, I) 로 표현 가능

해당 내용을 regularization을 나타내는 KL divergence 식에 대입하면,

Reconstruction Error 계산

위와 같은 적분 방법으로 기댓값을 계산할 수 있지만, 계산이 어렵다는 단점이 있음 :(

그리하여 등장한 방법을 소개한다. Monte Carlo - Reparameterization trick을 활용한 계산이다. 

1. Monte-Carlo (몬테카를로) technique 사용

특정 분포에서 샘플링을 무한대로 하고 평균을 내면, 전체의 기대값과 거의 동일해진다!

즉, z를 계속 뽑아서 각 z를 샘플링한 값을 평균 → 계산량이 많아지기 때문에 딥러닝에서 사용하기 어려움 → L을 1로 가정 (trick : 샘플링을 한번만 해서 대표값으로 사용)

 

Reparameterization Trick 계산

평균이 0이고, 표준편차가 1인 표준정규분포에서 입실론을 샘플링해서 위의 식에 대입 → z에 대한 식이 나오고, 미분 가능해진다 !정리하면, 아래와 같은 식 계산하면 됨. 

p가 확률 분포이기 때문에 베르누이를 따른다고 가정하고 계산 ( 가우시안을 따른다고 가정하고 계산할 수도 있음)

이때, 위 식의 Bernoulli처럼 베르누이 식에 의해 아래와 같은 cross entropy 식이 도출된다. 

Back to paper

VAE의 결과 사진을 통해서 더 큰 차원의 latent vector으로 축소를 했을 때 더 선명해진 것을 볼 수 있다. 하지만 너무 큰 차원의 latent vector를 사용하면 차원을 축소한 의미가 없어질 것.

Summary

마지막으로 코드와 요약정리를... 해보자. 요약이 요약이 아니겠지만 (!)

• Encoder이후 z를 sampling하는데, 그냥 sampling하면 back propagation이 불가능해서 reparameterization trick을 사용했다. 표준 정규분포에서 샘플링을 한 후 그 값에 평균, 표준편차 연산을 해서 인코더의 최종 출력 z를 만든다.
• real data X가 가우시안 인코더로 들어가면 평균, 표준편차를 output으로 내보낸다. (가우시안 인코더: output은 normal을 따른다고 가정하기 때문)
• Decoderloss function에서 베르누이 디코더에 의한 reconstruction error는 cross entropy, 그리고 Regularization은 KL divergence 식을 쓴다. KL divergence는 가우시안 인코더에 해당하는 q가 normal dist 따르도록 만든다.
• 그 후 베르누이를 따르는 디코더를 통과시킨다.
• 알고리즘 ( Code link https://colab.research.google.com/github/smartgeometry-ucl/dl4g/blob/master/variational_autoencoder.ipynb#scrollTo=t0rdg3FU7KRw )

1. M 개의 데이터(사진)을 가져온다. (M: batch size)
2. noise distribution으로부터 랜덤 샘플 입실론을 뽑아낸다. (입실론은 reparameterization trick에 사용되는 상수)
3. Neural network의 파라미터들(세타, 파이)과 데이터, 입실론으로 Loss를 구하고, Loss를 세타와 파이에 대해 미분하여 gradient를 계산
4. Gradient를 사용하여 Neural network의 파라미터들을 업데이트
5. repeat!

---

Diffusion Model 상에 나오는 VAE, KL 분포, Reconstruction Error, Monte Carlo 기법 등을 총괄적으로 다루는 다른 논문을 읽어보았는데... 이제 웅얼웅얼은 가능한 갓난애기가 된 기분입니다 허허