Variational Autoencoders (VAE)는 데이터를 latent 공간 상의 분포로 매핑하는 encoder와 latent 공간에서 데이터 공간 상의 분포로 매핑하는 decoder로 이루어진 확률 기반의 생성 모델로, 2013년에 발표된 이후로 많은 분야의 연구 대상이 되어 왔습니다. VAE는 직접적으로 다룰 수 없는 posterior probability $p(z|x)$를 encoder를 통해 근사시킴으로써 $q(z|x)$ 데이터의 likelihood에 대한 variational lower bound를 유도하여 생성 모델을 훈련합니다.
Importance Weighted Autoencoders (IWAE) 는 2016년 Burda에 의해 발표된 VAE의 개량 모델로 기존의 variational lower bound 보다 데이터 likelihood를 더 tight 하게 근사하는 lower bound를 제안했습니다. 이번 포스트에서는 VAE와 IWAE에 대해 간략히 알아보려 합니다.
Variational Autoencoders
VAE는 latent variable의 다룰 수 없는 posterior distribution을 근사하고자 제안되었습니다. 먼저 관측된 데이터 $x_i$는 다음 그림과 같이 관측되지 않은 latent variables $z$가 연관된 어떠한 generative process $p(x|\theta)$에 의해 샘플링되었다고 가정합니다.
이 가정에 의해 우리는 데이터 $x$에 대한 marginal likelihood를 다음과 같이 쓸 수 있고 데이터셋 $\mathcal{D}={x_i}_{i=1}^N$에 대해 $E_{p_{\mathcal{D}}}$ $[logp(x)]$를 최대화는 $\theta$를 찾는다면 데이터 분포 $p(x)$를 만들어내는 generative process를 알 수 있습니다.
하지만 marginal likelihood $p(x)$와 $z$의 posterior distribution $p(z|x)$를 보통 다룰 수 없는 경우가 많습니다. 이를 위해 VAE에서는 Gaussian 분포와 같은 다룰 수 있는 분포의 형식이면서 $p(z|x)$를 근사하는 $q(z|x)$를 사용합니다.
$q(z|x)$와 Jensen's inequality를 이용해 Equation 1에 대한 variational lower bound를 다음과 같이 유도할 수 있습니다.
결과적으로 marginal likelihood를 직접적으로 최대화 하는 것이 아니라 이에 대한 variational lower bound를 최대화 함으로써 $p(x)$를 모사하는 분포를 학습합니다.
Reparameterization trick
Variational lower bound의 최적화를 위해서는 이것의 편향되지 않고 분산이 적은 (unbiased and low-variance) 기울기가 필요합니다. 또한 Equation 2에서 $E_{q(z|x)}$로 되어 있는데 이 부분은 $q(z|x)$로부터 샘플링을 통해 계산하여야 합니다. 하지만 back-propagation 시 샘플링은 stochastic해서 미분이 불가능하므로 다른 방법이 필요합니다. 이를 해결하기 위해 VAE는 우리가 알고 있는 고정된 분포 $p(\epsilon)$으로부터 샘플링한 $\epsilon$을 통해 $z$에 대한 표현을 deterministic한 방법으로 우회합니다. 즉, $z$를 deterministic function $z=g(\epsilon, x, \phi)$로 다음과 같이 표현합니다.
이 방법으로 모델의 파라미터가 들어간 $q(z|x)$에 대한 샘플링을 $\epsilon$에 대한 샘플링으로 치환함으로써 기존의 back-propagation이 안되었던 문제를 해결했습니다. Reparameterization trick을 그림으로 보면 다음과 같습니다.
VAE 논문에서는 $q(z|x)$를 Gaussian 분포로 모델링하며 reparameterization trick을 이용해 다음 수식과 같이 표현할 수 있습니다. Gaussian 분포의 $\mu, \Sigma$는 feed-forward network인 encoder의 출력이 되고 $\epsilon \sim N(0,I)$로부터 샘플링합니다. 생성 과정을 나타내는 decoder $p(x|z)$ 또한 neural network로 나타내면, VAE는 최종적으로 encoder를 구성하는 $\phi$와 decoder를 구성하는 $\theta$를 variational lower bound의 최대화를 통해 구하는 작업으로 생각할 수 있습니다.
Interpretation
Autoencoder 입장에서 바라봤을 때, VAE는 $q(z)=E_{p_\mathcal{D}} q(z|x)$가 $z$의 사전 분포 $p(z)$에 비슷하도록 encoder를 학습시키면서 $p(x|z)$가 그럴듯한 샘플을 생성해내도록 decoder를 학습시킵니다. Equation 2의 수식을 다시 정리하면 다음과 같습니다.
첫 번째 항은 encoder/decoder를 통한 출력이 원래의 이미지로 원복되게끔 하여 data distribution $p_{\mathcal{D}}$의 구조를 유지하게끔 합니다. 위에 언급한 reparameterziation trick은 첫 번째 항에 적용됩니다. 두 번째 항은 approximate posterior $q(z|x)$가 $z$의 사전 분포 $p(z)$에 가깝도록 regularize하는 항으로 두 분포가 Gaussian 일경우 해석적으로 풀이가 가능합니다.
Importance Weighted Autoencoders
VAE는 posterior distribution을 추정하는 강력한 도구이지만 1) posterior distribution에 부여한 강력한 조건으로 인해 조건이 성립이 안되면 너무 간단한 latent representation을 얻게 되고 2) 데이터를 잘 설명하지 않을 것 같은 latent sample에 대해 지나치게 penalizing을 가합니다. (1번 내용에서의 가정은 VAE의 encoder가 데이터에 대해 non-linear regression으로 latent space의 분해된 파라미터 ($\mu, \Sigma$ 등)을 추정할 수 있다는 것입니다)
IWAE는 importance sampling을 통하여 VAE가 강제하는 posterior distribution에 대한 조건을 만족하지 않아도 lower bound를 잘 추정할 수 있도록 합니다. IWAE는 $q(z|x)$에서 $k$개의 샘플링을 통해 다음과 같은 새로운 lower bound를 제안합니다.
이때 $k$가 커진다면 $log p(x)$에 근접합니다. 따라서 $k$가 커질 수록 lower bound는 likelihood에 조금 더 tight 해지며, $k=1$일 경우 원래의 VAE로 돌아옵니다.
이때 IWAE가 VAE와 가장 다른 점은 $L_k (x)$의 기울기에 대한 Monte-Carlo estimator가 다음과 같이 importance weighted 평균이 된다는 점입니다. $w_i = w(x, g(\epsilon^i, x, \phi), \theta, \phi$=$frac{p(x, g(\epsilon^i, x, \phi)|\theta)}{q(g(\epsilon^i, x, \phi)|x, \phi)}$라 했을 때, $L_k (x)$의 기울기 추정은 다음과 같습니다.
위 수식에서 $\tilde{w}_i=\frac{w_i}{\sum_j w_j}$로 $w_i$의 normalized한 값입니다. VAE lower bound의 gradient와 비교해 보았을 때, IWAE에서는 gradient가 $\tilde{w}_i$만큼 가중 평균되어 있음을 볼 수 있습니다. ($\tilde{w}_i$를 $1/k$로 했을 때 본래 VAE로 돌아갑니다)
Interpretation
$log w_i=l_i$라 생각했을 때의 $L_k$의 gradient 추정은 다음과 같이 다시 표현할 수 있습니다.
따라서 $L_k$의 gradient는 가중 평균이 제일 높은, 데이터를 제일 설명을 잘 할 것 같은 샘플에 의해 주로 점유됩니다. 즉, $z_1, z_2, ..., z_k$의 샘플 중에서 데이터와 가장 연관이 높은 $z$에 대해서는 가중치가 높으므로 이들이 주로 gradient를 구성하게 됩니다. 따라서 위에서 언급했듯이 IWAE는 데이터에 연관이 적은 샘플 $z$에 대해 덜 penalize하게 되므로 multi-mode를 가진 true posterior distribution에 조금 더 근접하게 됩니다. 결론적으로 VAE에 비해 더 표현력 있는 데이터의 latent representation을 학습할 수 있게 되는 것이죠.
Experiment
VAE와 IWAE의 성능 비교는 다음과 같습니다. Lower bound가 얼마나 tight하지 확인하기 위한 지표로 negative log likelihood를 사용헀으며, encoder의 active unit의 수를 계산하여 주어진 모델의 capacity에 대해 얼마나 잘 latent space의 representation을 표현하는지 측정합니다. 밑의 결과를 보면 IWAE가 VAE에 비해 NLL이 더 낮으면서 active units이 더 많은 것을 확인할 수 있습니다.
두 번째 실험으로 VAE로 훈련을 한 후 IWAE로 훈련을 했을 때의 변화와 IWAE로 훈련을 한 후 VAE로 훈련을 했을 때의 지표 변화를 관측합니다. 예상대로 VAE로 훈련한 이후에 IWAE로 재 훈련을 수행했을 때, NLL이 낮아지고 active units이 늘어난 것을 볼 수 있습니다.
