WGAN 리뷰(Wasserstein GAN)

오늘은 WGAN 논문을 읽어보았다. 여러가지 수학적 개념들이 등장해서 개인적으로 좀 반가운 논문이였다ㅎㅎ 수학적 배경이 탄탄한 논문이라서 WGAN이 학습이 잘 되는 이유에 대한 증명이 많이 담겨있다. 

 

Introduction

앞선 연구들은 real data distribution과 model distribution의 KL divergence를 최소화하는 방향으로 model distribution을 학습하고자 했다. 하지만 이러한 방법론은 상당한 문제점을 가지고 있는데, model manifold와 true distribution의 support는 잘 겹치지 않아서 이 둘 사이의 KL divergence를 계산하기 어렵다는 것이다.

 

이러한 문제를 극복하기 위해 simple한 distribution을 시작해 원하는 distribution으로 paramertic function을 이용해 구현하는 방법론이 각광받게 되었다. VAE와 GAN이 이러한 방법론인데, VAE 또한 likelihood를 근사해 ELBO를 최대화하고자 하는 방법론으로 앞선 연구와 같은 한계를 가지고 있다. GAN은 objective function을 정의하는 방법에 따라 유연한 형태를 취할 수 있다는 점이 장점이지만, 학습하기 너무 어렵다는 단점을 가지고 있다. 

 

이 논문에서는 model distribution와 data distribution 사이의 거리를 측정하는 distance에 대해 논할 예정이다. 이러한 distance는 probability distribution의 sequence의 수렴성과도 연관이 있는데, distance의 정의에 따라 수렴 가능성이 달라지기 때문에 이 부분을 눈여겨 보아야 한다.

 

이 논문에서는 weak topology를 지닌 distance가 sequence of distribution의 수렴을 더 쉽게 하기 때문에 유용하다고 말하는데, 그런 이유로 Earth-Mover distance를 고른게 아닌가 싶다. 

 

또한 저자는 loss function이 연속이기를 바라는데, 뒤에 나올 예제를 보면 loss function이 연속이지 않기 때문에 문제가 발생하기 때문이다.

 

Different Distances

Notation은 다음과 같이 정의하자.

 

우리가 앞으로 다룰 Total Variation distance(이하 TV), Kullback-Leibler divergence(이하 KL), Jensen-Shannon divergence(이하 JS), 그리고 WGAN에서 중요하게 다룰 Earth-Mover distance(이하 EM)의 정의는 아래와 같다.

The Total Variation distance

The Kullback-Leibler divergence

The Jensen-Shannon divergence

The Earth-Mover distance

예제

이제 아래의 예제를 생각해보자.

각각의 distance와 divergence의 정의에 의해 아래와 같이 계산된다.

 

EM distance를 제외한 모든 JS, KL, TV 모두 불연속하다는 것을 알 수 있다. 위와 같은 예제에서 EM distance를 제외한 나머지 distance와 divergence는 gradient descent 수행이 불가능하다!

 

뒤에서 알아볼 예정이지만, Wasserstein distance는 다른 distance나 divergence보다 weak하다! 

 

이 부분에 대한 증명은 차근차근 알아가보자.

증명

증명)

이것이 assumption 1

 

복잡한 증명의 결론으로 우리가 얻은 것은 EM distance를 최소화하는 것이 neural net상에서도 이론적으로 최적임을 증명할 수 있다는 사실이다.

 

또한 Theorem 2를 통해 EM이 JS, KL, TV보다 더 sensible한 cost function을 구성한다는 사실을 알 수 있으며, KL이 가장 strong하며, JS, TV, EM 순으로 약해진다는 사실을 알 수 있다.(사실 JS와 TV는 equivalent하다)

결-론 : EM distance가 암튼 짱임

Wasserstein GAN

앞에서 정의한 EM distance는 inf을 포함하고 있어 그 값을 찾기가 어렵다. 그래서 Kantorovich-Rubinstein duality 정리를 이용해 아래와 같이 EM distance를 재정의한다.

여기서 $f$는 1-Lipschitz function이며, Lipschitz constant가 1보다 작거나 같아야 한다.

 

위 식을 다시 변형해 Lipschitz constant가 K보다 작거나 같은 상황을 생각하고 $K \ times W(P_r, P_{\theta}$를 생각해도 된다. 이렇게 하면 $K$값을 변경함에 따라 제약 조건을 변경할 수 있기 때문에 좀 더 유연한 모델을 만들 수 있다. 

 

결국, K-Lipschitz function들을 파라미터 $w$로 나타내게 된다면, 아래와 같이 다시 쓸 수 있다.

위 식을 푸는 $w$가 존재한다는 것과 위 식의 backpropation이 가능하다는 점은 Theorem 3를 통해 수학적으로 증명이 가능하다.

 

여기서 $w$는 conpact space $W$의 원소이며, $w$가 compact space에 있도록 정의하기 위해 이 논문에서는 $w$의 값을 제한했다.(Weight clapping)

weight clapping은 상당히 까다로운데, 그 범위를 정하는 것은 상당히 어렵다

 

지금까지 설명한 과정을 알고리즘으로 나타내면 아래와 같다.

⭐이 논문에서는 GAN에서의 Discriminator를 Critic이라고 부른다. 

 

EM distance는 weak topology를 가지고 있으며 continous하고 differentiable하기 때문에(상당히 좋은 성질을 가지고 있다) critic을 optimal한 상태까지 학습시킬 수 있다는 장점이 있다. 

 

EM distance에서는 weight를 제한하기 때문에 gradient가 급격하게 바뀌는 feature를 어느정도 보정하였는데, 그렇기 때문에 critic이 지속적으로 gradient information을 generator에게 제공하고, generator가 잘 학습할 수 있도록 돕는다.

 

또한 WGAN을 사용하면 GAN에서 골칫거리였던 mode collapse 부분도 해결할 수 있다.

참고로 WGAN의 학습은 시간에 대해 unstable하기 때문에 mometum 기반의 optimizer(Adam 등)는 critic의 학습에 적절하지 않다. 그렇기 때문에 RMSProp등을 사용하는 것을 권장한다.

 

* 실험 결과 및 future work은 생략

 

읽고 도움이 된 블로그 글

1. https://ahjeong.tistory.com/7

2. https://www.slideshare.net/slideshow/wasserstein-gan-i/75554346

⭐3. https://jonathan-hui.medium.com/gan-wasserstein-gan-wgan-gp-6a1a2aa1b490