Wasserstein Auto-encoders (vs VAE)

비전쪽을 잘 모르다보니 Wasserstein GAN이 나왔을 때도 어렵다는 얘기 정도만 듣고 따로 찾아보지는 않았는데 이번에 스토리 생성 관련 연구를 하면서 이 Wasserstein Distance를 활용한 Auto-encoder가 나와서 소개하고자 합니다. 

Backgrounds

1. Wasserstein Distance(편의상 W-dist)

W-dist는 두 확률 분포간의 거리를 나타내는 방식으로 Optimal Transport 문제에서 유래했다. 확률 분포를 흙더미라고 생각하고 $P(X)$를 $Q(X)$로 옮길 때 필요한 일의 양을 두 확률 분포의 거리라고 정의할 수 있다. 고전역학에서 일의 양은 힘 곱하기 거리로 표현되는데, 확률 분포의 무게는 합이 1로 제한되고, 거리는 단순히 euclidean distance로 표현할 수 있을 것이다. 이때 $P$의 $x$위치에서 $Q$의 $y$위치로 옮기는 확률 분포의 양을 $\gamma (x,y)$로 표현하고, 거리 함수(cost)를 $c(x,y)$로 표현해보자. 그러면 W-dist는 아래와 같이 정의할 수 있다.

$$\begin{gathered} W(P,Q)=\underset{\gamma \sim \mathrm{\Gamma }[P,Q]}{inf}{\mathbb{E}}_{X,Y\sim \gamma }[c(X,Y)], \\ \text{where }{\int }_x\gamma (x,y)dx=Q(Y),{\int }_y\gamma (x,y)dy=P(X),\gamma (x,y)\ge 0 \end{gathered}$$  

2. Variational Autoencoder

데이터셋의 Low dimension Latent Variable을 학습하기 위한 모델 중 하나인 VAE는 그 수학적 배경과 단순한 구현으로 인해 여전히 많이 사용되는 기법이다. 단순히 point estimate을 하는 auto encoder와 달리 VAE는 latent variable에 대한 gaussian prior 및 posterior를 가정하여 얻어진 tractable한 ELBO(Evidance Lower Bound)를 최대화 하는 방식으로 posterior 분포를 근사한다. 충분 통계량을 네트워크의 아웃풋으로 가정하는 amortized variational inference, 샘플링을 하면서도 gradient를 계산할 수 있도록 stochasticity를 고립시키는 reparametrization trick 등 이제는 매우 유명해진 기법들을 볼 수 있는 기념비적인 연구이기도 하다. 

VAE를 최적화하는 ELBO는 아래와 같이 계산된다.

$$L_{\text{VAE}}={\mathbb{E}}_{q_{\varphi }(z|x)}[\log p_{\theta }(x|z)]-D_{\text{KL}}(q_{\varphi }(z|x)|p(z))$$

첫번째 term은 보통 reconstruction loss로 계산하고(Mean Squared Error 등), 뒤의 KL divergence를 해석적으로 표현하기 위해 approximate posterior $q(z|x)$에 대해 normal 가정을 하게 되는데 이것이 latent variable의 표현력을 떨어뜨리는 한계가 될 수 있다. 단적으로 자연어를 모델링한다고 했을 때, '잘한다'라는 표현은 '칭찬'의 의미일 수도 있고 '비꼼'의 의미일 수도 있다. latent 공간에 이런 의미가 잘 모델링 되었다고 생각하면 분명 두 개 이상의 local maxima를 갖는 분포가 되어야 할 것이라고 생각할 수 있다. 하지만 VAE는 이것을 unimodal하게 가정해버리므로 suboptimal할 수밖에 없다. 이런 VAE의 근본적인 한계는 KL divergence를 해석적으로 계산하려고 하기 때문에 나온다. 그렇다면 다른 거리 metric을 가정하면 prior 및 posterior에 대한 자유도가 좀 더 높아질 수 있지 않을까?  

Waserstein Autoencoder

위에서 정의된 w-dist를 최소화하는 방식으로 Generator $G_X$를 만든다고 생각해보자. 우리의 태스크는 이제 $P_X$를 따르는 데이터셋 D를 가지고 있을 때, $G_X$와 $P_X$의 w-dist를 최소화하도록 G_X를 학습하는 일이다. 다행히도 $\gamma \sim \mathrm{\Gamma }[p,q]$를 찾는 대신, 사후 분포인 $Q(Z|X)$만 찾으면 된다는 것이 알려져 있다. (appendix 참조) 대신 $Q(Z)={\mathbb{E}}_{X\sim P_X}[Q(Z|X)]=P(Z)$라는 hard constraint가 있다. 즉 아래 식이 성립한다.

$$\underset{\gamma \in \mathrm{\Gamma }[p,q]}{inf}{\mathbb{E}}_{(x,y)\sim \gamma }[c(x,y)]=\underset{Q:Q_Z=P_Z}{inf}{\mathbb{E}}_{P_X}{\mathbb{E}}_{Q(Z|X)}[c(X,G(Z))]$$  

위의 hard constraint를 lagrange multiplier를 활용해서 soft constraint로 바꾸려면 양수 $\lambda$를 곱해서 아래와 같이 표현한다.

$$D_{\text{WAE}}(P_X,P_G)=\underset{Q(Z|X)\in \mathcal{Q}}{inf}{\mathbb{E}}_{P_X}{\mathbb{E}}_{Q(Z|X)}[c(X,G(Z))]+\lambda \cdot D_Z(Q_Z,P_Z)$$

위의 식에서 우리는 $Q(Z|X)$를 $\varphi$로 매개화하고 그것을 학습할 것이므로 infimum을 제거한 부분을 loss로 사용할 수 있다. (이 부분도 엄밀한 equality는 아니고 근사값으로 볼 수 있다.) 이 식에서 cost 함수 부분을 L1이나 L2 같은 일반적인 거리로 둔다면 우리가 알고 있는 L1, L2 reconstruction term이 나오는 것을 쉽게 확인할 수 있다. 

VAE와 달라지는 부분은 결국 regularization term이라는 얘기다. VAE는 모든 Z들의 posterior가 P(Z)와 같아지게 하지만 WAE는 posterior의 mixture( $={\mathbb{E}}_X[Q(Z|X)]$ )가 P(Z)와 같아지게 만든다. 따라서 아래 그림처럼 latent code의 표현력을 더 잘 보존할 수 있게 된다.

regularization term으로 prior와 posterior mixture의 거리($D$)를 최소화해야 하는데, 

여기에서 이 거리라는 개념를 MMD(Maximum Mean Discrepancy)로 정의할 수도 있고, GAN discriminator로 정의할 수도 있다. 그 두 가지 모두 성능은 유사하며 알고리즘은 아래와 같다.  

Divergence Modeling

1. WAE-GAN

WAE GAN은 divergence 텀을 prior-posterior discriminator의 negative log likelihood로 봤다고 생각하면 된다. 먼저 prior 분포의 z와 posterior 분포의 z를 구분하는 binary classifier를 만들고 그 classifier에서 prior로 분류할 확률을 divergence로 사용하는 것이다. 아주 직관적인 방식이다. 대신 이 방식은 classifier학습을 위해 AE 학습 외에 별도 backprop을 해줘야 한다. 

2. WAE-MMD

MMD(mean maximum discrepancy)는 두 분포 사이의 거리를 나타내는 지표 중 하나로 아래와 같은 정의를 갖는다. 

$$\text{MMD}(P,Q)=\parallel {\mathbb{E}}_{X\sim P}[\varphi (X)]-{\mathbb{E}}_{Y\sim Q}[\varphi (Y)]{\parallel }_{\mathcal{H}}$$

여기에서 $\varphi (\cdot )$는 피쳐 스페이스로 보내는 맵핑 함수라고 생각하면 된다. 간단하게 생각해서 $\varphi (x)=x$ 인 경우에는 MMD는 단순히 두 분포의 평균의 차이가 된다. 만약 $\varphi (x)=(x,x^2)$이라면 MMD는 $\sqrt{(\mathbb{E}[X]-\mathbb{E}[Y]{)}^2+(\mathbb{E}[X^2]-\mathbb{E}[Y^2]{)}^2}$이 된다. 이 경우 두 분포의 평균과 분산을 모두 반영한 거리라고 볼 수 있다. 이런 관점에서 임의의 $n$차 모먼트까지 고려할 수 있다면 더 표현력이 좋아질 것이다. 피쳐 스페이스에서 무한차원을 갖는 맵핑함수를 implicit하게 만들어낼 수 있는 방법이 커널 트릭이다. 

위의 MMD의 제곱 값을 구하면 다음과 같다. 

저 커널함수를 RBF $k(x,y)=\exp (-\frac{\parallel x-y{\parallel }^2}{2{\sigma }^2})$  등으로 정의하면 feature space $\mathcal{H}$에서의 implicit dimension은 무한대가 된다. (그걸 보는 방법은 단순히 $\parallel x-y\parallel$를 중심으로 tayler expansion을 해보면 된다.)

다만 gaussian kernel은 tail이 너무 얕아서 heavy tail을 가진 inverse multiquadratic kernel을 사용했다고 한다. 위 알고리즘처럼 샘플을 $n$개 뽑았다고 가정하면 kernel trick을 사용한 결과는 첫번째, 두번째 텀에서는 자기 자신과의 유사도를 제외한 $n(n-1)$개, joint 분포(마지막 텀)에서는 $n^2$개가 나와서 위와 같은 평균식이 나오게 된다. 

이 경우 prior나 posterior에 대해서 특정한 분포적 가정을 두지 않기 때문에 샘플링만 가능하다면 더 자유롭게 모델링할 수 있다는 장점이 있다. VAE에서 지겹게 쓰던 isotropic gaussian 가정을 버려도 된다는 것이다! 그래서 위에서 언급했던 Multimodal Prior를 표현하기 위해 Mixture of Gaussian을 사용한 논문들이 나온 것들도 확인할 수 있다. 물론 posterior를 분포로 가져가려면 amortized variational inference를 써야 하기 때문에 충분 통계량을 모델링하기 위해 결국 어떤 분포를 가정 할 수밖에 없긴 하다. 그래도 posterior가 point estimate인 경우에도 학습이 가능하다는 것은 장점이라고 할 수 있다. 

3. Summary

formulation은 다르게 시작했지만 Variational Autoencoder와 동일하게 reconstruction term + regularization term으로 구성이 동일하게 되는 것을 볼 수 있고, regularization term을 어떻게 정의하느냐에 따라 두 가지 형태가 나올 수 있다. MMD 같은 경우는 다양한 함수적 가정을 가진 prior를 정의할 수 있어 VAE 보다 자유도가 높다고 할 수 있다.

4. Experimental Result

GAN 계열의 모델과 비교하면 FID는 한참 밀리지만(최근 CelebA SOTA FID는 4..) 같은 latent variable modeling 기법 안에서 생각했을 때는 VAE보다 훨씬 좋은 표현력을 보여주는 것을 확인 할 수 있다.

References

• Wasserstein Auto-Encoders
• MMD