Wasserstein GAN

구현은 상당히 쉽지만, GAN을 구현하기 위해 알아야 하는 배경지식과 STARGAN을 구현하면서 느낀 GAN의 수렴이 굉장히 불안정한 문제를 풀기 위해 가장 먼저 적합한 솔루션을 제시하고 수학적으로 증명한 논문. 구현과 모델보다는 개념위주로 보면 좋겠다.

Introduction

생성이란 무엇인가에 대하여 저자들은 아래와 같이 답한다. 생성은 확률분포를 학습하여 확률분포에 있는 변수를 출력하는 것이고, 이는 또한 우리가 가진 데이터에서 가능도를 최대화 하는 것과 같다고 말한다.
즉, $$\max_{\theta \in R_d} \frac{1}{m}\sum_{i=1}^{m}logP_{\theta}(x^{(i)})$$
이 논문에서는 continuity의 의미를 $\theta_t$ 가 $\theta$로 수렴하면, $P_{\theta_t}$ 또한 $P_{\theta}$로 수렴한다고 정의한다. 또한 이 continuity는 우리가 확률간의 거리를 어떻게 정의하느냐에 따라 달라진다고 한다. 정확한 수렴을 위해서는 continuity가 성립해야하고 이를 성립시키기 위해 distance함수를 잘 골라야 한다는 말이다. 따라서 이 논문에서는 이 distance를 잘 구하기 위해 여러 metric을 비교해보고 왜 EM distance (Wasserstein L1)을 선택하였는지 설명한다.

Differnet Distances

4개의 distance 를 소개하기 전에 notation을 잘 익혀둘 필요가 있다.

• $\mathcal{X}$ : compact metric set. 1. bound가 되어있고, 2. 이 bound 경계를 포함한다.
• $\Sigma$ : borel subset. 측정 가능한 집합으로 1. $P_r$, $P_g$와 같은 확률 분포가 확률값이 계산된다. 2. 연속함수의 기댓값을 계산하기 위한 최소조건

1. Total Variation (TV)
   $$\delta(\mathbb{P}_r, \mathbb{P}_g) = \sup_{A \in \Sigma}|\mathbb{P}_r(A) - \mathbb{P}_g(A) |$$.
   A는 엄청 큰 범위 $\Sigma$안에 있는 범위라고 생각하면 된다. 즉, 우리가 전체를 모르기 때문에 그것의 일정 집합 A에서 논다고 생각. 따라서 이 식은 "큰 범위에서 어던 범위 A를 잡았을 때 가장 큰 확률값의 차이" 라고 할 수 있다. 또한 두 확률 분포가 겹치지 않는다면 이 metric은 항상 1이다. 아래 그림 참조
   

2. Kullback-Leibler (KL) divergence
   $$KL(\mathbb{P}_r || \mathbb{P}_g) = \int \log (\frac{P_r(x)}{P_g(x)}) P_r(x) d(\mu (x)$$
   생략

3. Jensen-SHannon (JS) divergence
   $$JS(\mathbb{P}_r, \mathbb{P}_g) = KL(\mathbb{P}_r || \mathbb{P}_m) + KL(\mathbb{P}_g || \mathbb{P}_m)$$
   생략

4. Earth-Mover (EM) distance or Wasserstein-1
   $$W(\mathbb{P}r, \mathbb{P}g) = \inf{\gamma \in \Pi (\mathbb{P}_r, \mathbb{P}_g)} \mathbb{E}{(x,y)~\gamma}[||x-y||]$$
   $\Pi$는 두 확률분포의 결합 분포이고, 아래 그림을 참조하여
   
   나올 수 있는 $w$ 값들에 대하여 매핑된 점들의 거리 기댓값의 최솟값을 구하는 것이다. 이 때의 변수는 확률분포 $\mathbb{P}_r, \mathbb{P}_g$이다. 아래 예제에서 이해하지.

Example 1

$Z ~ U[0,1]$인 uniform distribution이고, $\mathbb{P}0$는 $(0, Z) in \mathbb{R}^2$를 따른다고 하자. 또한 $\mathbb{P}{\theta}$는 $(\theta, Z)$ 를 따른다고 하자.

풀이는 위 사진 참고

논문의 Theorem들은 EM이 가장 strong한 특성을 갖고 있기 때문에 수렴에 관해서도 EM이 성립하지 않으면 나머지도 성립하지 않는 것을 증명한 것이다.

Wasserstein GAN

Kantorovich-Rubinstein duality를 사용해서
$$W(\mathbb{P}r, \mathbb{P}g) = \sup{||f||L \leq 1} \mathbb{E}{x\mathbb{P}r} [f(x)] - \mathbb{E}{x\mathbb{P}_g}[f(x)]$$
가 성립하고, 이는 $\mathbb{P}_r$이 real dataset, $\mathbb{P}_g$가 생성한 dataset인것을 생각하면 이 거리 함수가 최대화가 되려면 real에 대해서는 최대화, fake에 대해서는 최소화가 되어야 한다. 따라서 discriminator의 loss function은 아래와 같다.
$$\max{w \in W} \mathbb{E}{x~\mathbb{P}_r} [f_w(W)] - \mathbb{E}-{z ~ p(z)}[f_w(g{\theta}(z))]$$.
위의 식을 보면 $f_w$는 discriminator고, $g_{\theta}$는 generator임을 알 수 있다. 즉, discriminator에서 나온 아웃풋을 그냥 mean취한 것으로 구현 하면 된다.

Theorem 3에서는 generator에 대하여 loss function을 어떻게 하는지에 대해 논한다. Theorem 3의 결론부터 말하면 $$\max_{||f||_L \leq 1} \mathbb{E}_{x \sim \mathbb{P}_r}[f(x)] - \mathbb{E}_{x \sim \mathbb{P}_{\theta}}[f(x)]$$
를 만족하는 해 $f$가 존재하고, 이를 $\theta$에 대하여 미분하면
$$\nabla_{\theta}W(\mathbb{P}_r, \mathbb{P}_g) = -\mathbb{E}_{z \sim p(z)}[\nabla_{\theta}f(g_{\theta}(z))]$$
가 된다. 또한 두개의 식은 well-defined되어 있다.
식을 잘 보면 $\theta$에 대하여 미분 했으므로 generator를 업데이트 한다고 할 수 있다.

알고리즘 부분은 먼저 D를 $n_{critic}$번 학습하고, G를 1번 학습 하는 것 외에는 업데이트는 위의 식으로 한다.

Improved stability

Star GAN을 구현 하면서 WGAN이 쓰여 이렇게 포스팅을 하지만 STARGAN에 쓰일 만큼 영향을 많이 끼친 이유가 있다. 여기서 장점 3가지를 드는데

1. 논문의 그림을 보면 기존의 GAN과는 달리 loss 가 줄어들수록 성능이 좋아지는 것을 볼 수 있다. 즉, 기존에는 loss가 왔다갔다 하면서 언제 수렴했는지도 몰랐는데 WGAN은 성능 지표를 알 수 있는 것에 굉장한 장점이 있다.
2. WGAN은 수렴성이 강해서 모델에 민감하게 반응하지 않는다. 하지만 내 경험상 STARGAN을 구현 할 때 generator의 upsampling 부분에서 layer를 하나 추가되어 수렴이 안되는 것을 볼 수 있었다. STARGAN은 WGAN-GP를 썼는데도 이렇게 민감한 것을 볼 수 있었다.
3. 저자들은 WGAN을 실험하면서 어떤 mode collapse (discriminator를 속이기 위해 쉬운 것만 생성하는 것.) 을 본 적이 없다고 한다. 그런데 이후 GAN논문에서 mode collapse 문제가 계속 발생하는 것을 보면 이 문제를 완전히 해결 한 것 같지는 않다.