VAE에 여러 개의 random variable을 줄 수 있을까? 인공지능 및 기계학습 심화 2 : edwith

김동준 튜터 2023.03.20

안녕하세요, 조교입니다.

VAE는 보통 data variable $x$ 와 latent variable $z$ 하나씩 있는 형태로 구성됩니다.

inference network $q_\phi(z\vert x)$ 는 Gaussian distribution $N(\mu_\phi(x),\sigma_{\phi}(x)I)$ 의 mean과 covariance function을 NN의 output으로 설정하여 z distribution을 추정합니다.

generative network $p_\theta(x\vert z)$ 는 보통 Categorical distribution 혹은 dirac delta distribution으로 상정하여 NN의 output을 data x의 reconstruction으로 간주합니다.

그렇다면, latent variable이 $z_{1},z_{2}$ 두 개가 있는 형태의 VAE는 어떻게 만들 수 있을까요?

먼저 loss function의 측면에서 생각해 보도록 하겠습니다.

loss function은 이전 자료인 https://kooc.kaist.ac.kr/aiml-adv2/forum/129549를 참조하시면 결국

log-likelihood maximization은 KL optimization과 동일하다는 것을 설명드린 바가 있습니다.

즉, VAE는 $D_{KL}(p_{data}(x)\Vert p_{\theta}(x))\le D_{KL}(p_{data}(x)q_{\phi}(z\vert x)\Vert p_{\theta}(x\vert z)p(z))$ 에서 우변을 minimize하여 좌변이 그에 따라 최적화 되는 방식으로 학습됩니다.

여기서, latent variable이 두개가 된다면 VAE의 objective function은 $D_{KL}(p_{data}(x)\Vert p_{\theta}(x))\le D_{KL}(q_{\phi}(x,z_{1},z_{2})\Vert p_{\theta}(x,z_{1},z_{2}))$ 라고 할 수 있겠습니다.

이제 문제는, 우변을 어떻게 NN로 모델링 해야 latent variable이 1개인 vanilla VAE보다 더 유연한 모델링이 되겠느냐 입니다.

자, VAE에서 모델링 했던 방식을 그대로 여기에도 답습해 보겠습니다.

$q_\phi(x,z_{1},z_{2})=p_{data}(x)q_{\phi}(z_{1}\vert x)q_{\phi}(z_{2}\vert z_{1})$ 과 같이 NN를 짤 때 Markov assumption을 가정하여 $N N (x)$ 는 $z_{1}$ 에 대한 distribution을, $z_{2}$ 에 대한 distribution은 $z_{1}$ 만을 NN의 input으로 가지게 모델링 해보겠습니다.

그렇게 되면 좀 귀찮은 일이 벌어집니다. 일반적으로 latent variable을 여러개 사용하여 만든 Hierarchical VAE에서는 generative process는 $p_{\theta}(x,z_{1},z_{2})=p(z_{2})p_{\theta}(z_{1}\vert z_{2})p_{\theta}(x\vert z_{1})$ 을 가정하는데, 이렇게 되면 Joint KL divergence는 $D_{KL}(q_{\phi}(x,z_{1},z_{2})\Vert p_{\theta}(x,z_{1},z_{2}))=E_{p_{data}(x)}[-\log{p_{\theta}(x\vert z_{1})}+D_{KL}(q_{\phi}(z_{1}\vert z_{2},x)\Vert p_{\theta}(z_{1}\vert z_{2}))+D_{KL}(q_{\phi}(z_{2}\vert x)\Vert p(z_{2}))]$ 가 됩니다. 여기서 Bayes rule에 의해 $q_{\phi}(z_{1}\vert z_{2},x)=\frac{q_{\phi}(z_{1}\vert x)q_{\phi}(z_{2}\vert z_{1})}{q_{\phi}(z_{2}\vert x)}$ 이기 때문에 만약 $q_{\phi}(z_{1}\vert x), q_{\phi}(z_{2}\vert z_{1}), q_{\phi}(z_{2}\vert x)$ 을 알고 있으면 $q_{\phi}(z_{1}\vert z_{2},x)$ 를 계산할 수 있습니다.

그렇지만, 짜잔! 우리는 $x\rightarrow z_{1}\rightarrow z_{2}$ 가 되는 방향으로 NN를 주었기 때문에 $q_{\phi}(z_{2}\vert x)=\int q_{\phi}(z_{2},z_{1}\vert x)dz_{1}=\int q_{\phi}(z_{1}\vert x)q_{\phi}(z_{2}\vert z_{1})dz_{1}$ 을 통해서만 $q_{\phi}(z_{2}\vert x)$ 를 학습할 수 있게 되는 것입니다. 그래서 이러한 방식으로 NN를 주게 된다면 $q_{\phi}(z_{2}\vert x)$ 를 exact하게 계산이 불가능합니다.

이는 결국 Joint KL divergence를 근사적으로만 계산이 가능하게 되는데, 이는 큰 loss variance에 시달립니다.

또 다른 문제는 이러한 방식으로 NN를 주는 방식은 보통은 posterior collapse 문제를 그대로 가지고 있는 것으로 알려져 있습니다. 이 문제는 decoder가 너무 좋아서 $x\rightarrow z_{1}$ NN가 $x$ 의 어떤 information이라도 담지 않도록 학습하여, 결국 임의의 normal 분포에서 뽑은 latent를 decoder에 태워도 data를 충분히 잘 생성한다는 것을 의미합니다.

Data만 잘 생성하는 것이 아니라 $x$ 의 information을 많이 함유하는 latent vector $z$ 를 생성하는 것 또한 VAE의 목표이기 때문에 posterior collapse는 큰 문제가 되고, 이러한 문제를 극복하기 위하여 많은 기법들이 발전되어 왔습니다.

NN가 $x\rightarrow z_{1}\rightarrow z_{2}$ Markov 방식으로 짜여져 있다면 posterior collapse가 더 심화되어 $z_{1}\rightarrow z_{2}$ 는 사실상 아무것도 학습이 되지 않습니다. 그래서 이 문제를 해결하기 위해 우리는 NN 구조를 바꿔주게 됩니다.

학습 loss function으로부터 힌트를 얻어서 NN 구조를 바꿔보겠습니다.

학습 loss에서 우리는 두 가지 계산을 할줄 알아야 했습니다. 첫번재는 $q_{\phi}(z_{1}\vert z_{2},x)$ 이고 두번째는 $q_{\phi}(z_{2}\vert x)$ 입니다. 이 두 가지 계산을 Bayes rule같은 거추장스럽고 번거로운 방법 말고 NN의 feed-forward 계산으로 한꺼번에 끝내버린다면 얼마나 좋을까요? 즉, NN는 input으로 $x$ 를 받으면 $z_{2}$ distribution에 대한 mean과 variance를 output으로 내뱉어 $q_{\phi}(z_{2}\vert x)=N(\mu_{\phi}(x),\sigma_{\phi}(x)I)$ 를 구성하고, $x$ 와 $z_{2}$ 를 input으로 받으면 $z_{1}$ distribution에 대한 mean과 variance를 output으로 내뱉어 $q_{\phi}(z_{1}\vert z_{2},x)=N(\mu_{\phi}(x,z_{2}),\sigma_{\phi}(x,z_{2})I)$ 를 구성하게 되는 것입니다.

이렇게 되면 장점이 그냥 바로 Joint KL divergence를 계산할 수 있다는 점입니다. 또 posterior collapse 또한 잘 해결하는 것으로 알려져 있어서 매우 좋은 inference 및 generation을 동시에 할 수 있게 됩니다.

위와 같은 이유로 Hierarchical VAE에서는 전부 위와 같은 NN 구조를 차용하여 inference distribution을 정의한 후, 학습하고 있습니다.

오늘은 latent variable의 개수를 늘리는 방법에 대하여 알아보았습니다.

정리하자면, 보통은 generative process는 direct extension하는 반면, inference process는 조금의 treatment가 더 필요하다는 것이었습니다.

다음 시간에 또 뵙겠습니다.

감사합니다

인공지능 및 기계학습 심화 2

comment