Variational Auto Encoder(VAE)

Auto Encoder(AE)는 결정론적인 모델이라 생성 모델로 쓰기에는 결과가 좋지 않았음. 더 나은 데이터 생성을 결과를 얻기 위해 AE를 확률적 버전으로 구성한 것이 VAE라고 볼 수 있다.

AE를 확률론적 버전으로 구성할 때 encoder를 단순히 $p(\bold{z})$로 구성하면 생성 결과가 좋지 못하기 때문에 $p(\bold{z}|\bold{x})$와 같이 조건부 형태로 구성 한다. 이때 $p(\bold{z}|\bold{x})$의 posterior를 계산하기 위한 marginal likelihood $p(\bold{x})$를 직접 계산하는 것이 어렵기 때문에 Variational Inference을 사용하여 계산하기 쉬운 (일반적으로 가우시안인) $q(\bold{z}|\bold{x})$를 정의하고 KL divergence로 이 분포를 $p(\bold{z}|\bold{x})$와 가깝게 만드는 방법을 사용함. (물론 가우시안인 아닌 $p(\bold{z}|\bold{x})$를 가우시안인 $q(\bold{z}|\bold{x})$로 근사하는 것은 한계가 있다. 이것이 Variational 방법의 한계 중 하나이다.)

최종적으로 VAE의 encoder는 $q_{\boldsymbol{\phi}}(\bold{z}|\bold{x})$의 형식을 갖고 decoder는 $p_{\boldsymbol{\theta}}(\bold{x}|\bold{z})$의 형식을 갖게 됨.

Cap 2024-02-14 14-16-18-959.jpg

VAE의 손실 함수는 여러가지 형태로 기술할 수 있지만 간단히 정리하면 다음과 같다.

$$ \mathcal{L} = -\mathbb{E}_{q(\bold{z}|\bold{x})}[\log p(\bold{x}|\bold{z})] + \text{KL}(q(\bold{z}|\bold{x})\|p(\bold{z})) $$

여기서 $-\mathbb{E}_{q(\bold{z}|\bold{x})}[\log p(\bold{x}|\bold{z})]$는 reconstruction error로 encoder가 생성한 잠재 변수 $\bold{z}$를 기반으로 원본 $\bold{x}$를 얼마나 잘 재구성하는지를 측정하는 것으로 디코더의 성능을 평가하는 부분으로 볼 수 있다. 일반적으로 cross entropy를 사용함.

$\text{KL}(q(\bold{z}|\bold{x})\|p(\bold{z}))$는 $q(\bold{z}|\bold{x})$가 근사하려는 $p(\bold{z}|\bold{x})$에 얼마나 가까운지를 측정하는 것으로 인코더의 성능을 평가하는 부분으로 볼 수 있다. 일반적으로 $p(\bold{z}) \sim \mathcal{N}(\bold{0}, \bold{I})$를 사용. $q(\bold{z}|\bold{x})$가 근사하려던 원래 분포가 $p(\bold{z}|\bold{x})$였는데 KL divergence에 $p(\bold{z})$가 쓰이는 이유는 VAE의 손실 함수 유도과정 때문이다. 아래 참고 자료의 <시각적 이해를 위한 머신러닝> 참조.

참고로 VAE의 encoder가 출력하는 잠재 변수 $\bold{z}$의 평균 $\mu$과 분산 $\sigma^2$을 이용하여 KL divergence를 다음과 같이 계산할 수 있다. 아래 식에서 $J$는 잠재 변수의 차원이다.

$$ \text{KL}(q(\bold{z}|\bold{x})\|p(\bold{z})) = -{1\over2} \sum_{j=1}^J (1 + \log(\sigma_j^2) - \mu_j^2 - \sigma_j^2) $$

$\beta$-VAE

VAE의 단점은 생성된 이미지가 blur 이미지가 되는 경향이 있다는 것이다. 이것을 보완하기 위해 여러 방법을 사용할 수 있지만, 간단한 방법은 KL 항에 페널티를 감소시켜서 모델을 결정론적 autoencoder에 더 가깝게 만드는 것이다.

$$ \mathcal{L}\beta(\boldsymbol{\theta},\boldsymbol{\phi}|\bold{x}) = \underbrace{-\mathbb{E}{q_{\boldsymbol{\phi}}(\bold{z}|\bold{x})}[\log p_{\boldsymbol{\theta}}(\bold{x}|\bold{z})]}{\mathcal{L}E} + \beta\underbrace{D\text{KL}(q{\boldsymbol{\phi}}(\bold{z}|\bold{x})\|p_{\boldsymbol{\theta}}(\bold{z}))}_{\mathcal{L}_R} $$

여기서 $\mathcal{L}_E$는 reconstruction 에러(negative log likelihood)이고 $\mathcal{L}_R$은 KL regularizer이다. 이것을 $\beta$-VAE 목적이라고 한다. $\beta=1$을 설정하면 표준 VAEs에서 사용되는 목적을 복구한다. $\beta=0$을 설정하면 표준 autoencoder에서 사용되는 목적을 복구한다.

$\beta$를 0에서 무한대로 변화시키면 rate distortion curve의 다양한 지점에 도달하게 된다. 이러한 점들은 reconstruction error(distortion)과 입력에 관한 잠재에 저장된 정보의 양(코드에 해당하는 rate) 사이의 다양한 tradeoff를 만든다. $\beta < 1$을 사용하면 각 입력에 대해 더 많은 bit를 저장하므로 reconstruct 이미지는 덜 blur가 되고, $\beta>1$을 사용하면 더 압축된 표현을 얻는다.

InfoVAE

VAE를 학습하는데 디코더가 강력하면 잠재 코드가 무시되는 경향이 있고 데이터 공간과 잠재 공간에서 KL 항 사이의 불일치 때문에 빈곤한 posterior 근사를 하는 경향이 있다. 다음 형식의 일반화된 목적을 사용하여 어떤 정도로 이 문제를 수정할 수 있다.

$$ Ł(\boldsymbol{\theta},\boldsymbol{\phi}|\bold{x}) = -\lambda D_\text{KL}(q_{\boldsymbol{\phi}}(\bold{z})\|p_{\boldsymbol{\theta}}(\bold{z})) - \mathbb{E}{q{\boldsymbol{\phi}}(\bold{z})}[D_\text{KL}(q_{\boldsymbol{\phi}}(\bold{x}|\bold{z})\|p_{\boldsymbol{\theta}}(\bold{x}|\bold{z}))] + \alpha \mathbb{I}_q(\bold{x};\bold{z}) $$