VAE에 대한 잡설

정리 차원에서 남겨보는 VAE

https://hsaghir.github.io/data_science/denoising-vs-variational-autoencoder/

결론부터 말하자면 VAE는 모드 콜랩싱에 취약하다. ELBO의 

1) Likelihood term은 보통 reconstruction loss가 된다. CE를 쓰던 p-norm을 쓰던 그냥 오토인코더가 된다. 오토 인코더 학습만 해도 만만하지 않은데, 문제는 encoded z를 직접 쓰는게 아니라 여기에 z+std*e 를 이용해서 reconstruction을 한다는 것이다. 이게 일종에 데이터가 살고 있는 low-dimensional manifold로 보내주는 효과가 있긴 한데, reconstruction만 놓고 보면 쉽지가 않게 된다. 주어진 데이터를 열심히 인코딩해놨더니, 이걸 가지고 recon하는 것이 아니라 이걸 한번 흔든 값 (흔드는 정도는 현재 입력에 디펜데던트하고) 을 가지고 복원한 loss를 줄이려고 하니 일종의 ill-posed가 생긴다. 

여기에 ELBO의 두 번째 term인 prior fitting이 더해지면 문제가 더 심해지는데, 

2) 이 term은 인코딩된 z의 분포가 prior가 되게 하고, 보통 unit Gaussian에 피팅을 시켜버린다. 문제는 인코딩 된 z_mu를 그려보면 중앙에 몰리고, variance가 커져서 KLD를 맞추려고 하는 경우가 많다. 이러면 몇 가지 문제가 생기는데, 인코더의 역할이 없어져버리셔 p(z)에서 뽑은 z로 새로운 x를 만들어보면 완벽한 모드 콜랩싱이 일어나게 된다는 것이다. 

여기에 KLD의 동작 방식으로 생기는 문제가 하나 더 있다. ELBO이용한다는 것은 KL(q||p)에서 q를 학습하는 방식이다. 근데 요롷게 KLD에서 앞에 분포를 최적화하는 경우는 뒤에 있는 맞추고 싶은 분포에 특정 봉우리 하나만 맞추게 된다. 물론 q를 뒤에다 놓게되면 블러리한 애매한 분포를 학습하게 되어서, 그냥 이래저래 문제가 있다. 

결론을 적어보자면, VAE는 안될 수 밖에 없는 이유가 참 많다. 그래서 잘 안된다. 하하하. 

개인적으로 VAE에 대해서 참 좋아하는 이유 중 하는 이것이 infinite mixture model로 해석될 수 있다는 점이다. 즉 continuous한 z에 dependent한 어떤 믹스쳐 분포이기 때문에 MDN같이 믹스쳐의 수를 굳이 안 정해져도 된다는 장점이 있을 수 있지만, 실제로 써보면 별로다. 최근에 나온 월드 모델에서 VAE 대신에 MDN을 가지고 dynamic model을 구한 것은 이유가 있지 않을까, 생각하는데, 모 막상 물어보면 별거 없겠지. (CEM을 왜 썼냐 했더니, 그냥 썼다는 대답이..)

VAE를 학습시킬 때 한가지 잊으면 안되는 것은 우리가 나중에 샘플링을 한다는 것이다. 인코딩된 z들의 콜랩싱을 막는 한 가지 무식한 방법은 z_var의 최대치를 주는 것이다. 보통은 logvar를 모델링 하는데, 이걸 그냥 sigmnoid를 통과 시킨 후에 0.1 정도로를 곱해서 애초에 z_sample들이 멀리 못 벗어나게 하면, z_mu들이 좀 벌어지긴 하는데, 이 방식은 나중에 샘플링 할 때 다양한 z를 고려하지 못하므로 써서는 안되는 방법이다. 

몇 가지 논문들에 나온 휴리스틱으로는 

1) KLD 부분을 스케쥴링한다. 초반에는 reconstruction만 하고, 후반부에 KLD를 건다.

2) KLD의 minimum value를 정해서 일정 값 아래로 떨어지면 gradient가 reconstruction loss에만 dependent하게 만드는 것이다. 

물론 이 두 방법 모두 구현해봤는데, 사실 큰 도움이 안되더라. 

한가지 재밌는 것은 z_mu가 한 곳으로 모이는 것이 확실한 로컬 미니마라는 것이다. 여길 한번 벗어나면 곧잘 되는데, iteration이 적으면 여길 못 벗어나는 경우가 생긴다. 아놔 모르겠다.