[논문리뷰]Wasserstein GAN

Introduction

머신러닝에 있어서 확률분포를 학습하는 것은 어떤 의미인가? 고전적인 답은 pdf를 학습하는 것.

pdf가 주어지고, 우리의 real data와 likelihood를 최대화 하는 것이다. 또한 KL-divergence를 통해 두 pdf 간의 거리를 최소화 하는 것이다. \(\begin{aligned} \max_{\theta \in R^d} \frac {1}{m} \sum_{i=1}^{m} P_\theta (x^{(i)})\\ \\ KL(\mathbb{P}_r||\mathbb{P}_\theta) = \int \log(\frac {P_r} {P_\theta})\ P_r(x) \mu(dx) \\ \\ \mathbb{P}_\theta : parameterized\ density\\ \mathbb{P}_r : real\ data\ distribution \end{aligned}\)

• 하지만 실제 데이터 pdf를 구하기는 굉장히 어렵다.
• 존재하지 않는 실제 확률밀도함수를 측정하기 보다, 새로운 확률 변수 Z와 그에 해당하는 고정된 pdf인 $p(z)$를 정의하고 특성함수 $g_\theta$를 통해 Z를 $\chi$에 매핑한 후, 특정한 pdf인 $P\theta$으로 생성해낸다.

• $\theta$가 바뀜에 따라, pdf를 실제 데이터 분포에 가깝게끔 새로 만든 pdf를 변경할 수 있다.

• 또한 두 pdf 사이의 상관계수도 고려한다.

The contributions of this paper are:

• 우리는 “Earth Mover(EM) distance”가 다른 distance나 divergence와 어떻게 다르게 동작하는지 설명할 것이다.
• EM distance 방법을 사용한 GAN을 W-GAN이라고 명명한다.
• WGAN을 훈련시키는 데에 있어 discriminator와 generator의 balance를 걱정하지 않아도 된다.
• WGAN의 가장 실용적인 이익은 discriminator를 최적으로 훈련시킴으로써 EM distance를 연속적으로 측정하는 것이다.

Different Distances

아래 내용은 확률 거리 척도를 구하는 4가지 공식에 대한 내용에 대해서 나옵니다.

수식이 워낙 많고 복잡해서 Wasserstein GAN 수학 이해하기 1 - 임성빈을 참고하였습니다.

Total Variation은 두 확률측도(확률분포)의 측정값이 벌어질 수 있는 값 중 가장 큰 값을 말합니다.

수식을 보면 ‘sup’ 이라는 것이 있는데, Supremum의 준말이며 the least upper bound 라고 부릅니다. 해석하면 상한에서 (upper bound)에서 가장 작은 값(minimum)입니다.

아래 그림을 보면, 이해 하기 쉽습니다. A라는 특정 구간 내에서 값을 대입 했을 때, 두 pdf의 차이가 가장 큰 부분을 뜻합니다.

쿨백-라이블러 발산은 두 확률분포의 차이를 계산하는 데에 사용하는 함수이며, 두 개의 확률분포를 계산하는데 있어 가장 많이 사용하는 방법입니다.

KL-divergence를 사용하는데 문제가 있는데, 그것이 바로 Symmetry하지 않는 다는 것입니다. 이를 해결하기 위해 나온 개념이 Jensen-Shannnon Divergence입니다.

여기서 $\Pi(\mathbb{P},\ \mathbb{Q})는 두 확률 분포 P, Q의 joint pdf들을 모든 집합이고 $\gamma$는 그들 중 하나입니다. 또한 inf라는 것이 있는데, 위에서 설명한 sup의 반대되는 개념입니다. infimum의 준말이며, 하한선 중 가장 큰 값을 나타냅니다.

즉, 모든 joint pdf들 중에서 X와 Y의 차이에 대한 기댓값(평균)을 가장 작게 추정한 값을 의미합니다.

초록색 선이 뜻하는 것이 두 joint pdf의 값의 거리에 대한 차이를 나타낸 것입니다. 쉽게 말해서 초록색 선들의 분포의 기댓값을 가장 작게 추정한 것입니다.

Example of distance

Example 1은 각 거리 함수의 예를 나타내고 EM distance의 타당성을 증명하는 단계입니다. 보시다시피 각 거리 함수의 추정값을 볼 수 있는데, EM distance는 일정한 상수값을 가지는 반면, 나머지 거리 함수의 추정값은 특정한 조건이 있거나 무한대로 수렴하는 것을 볼 수 있습니다.

위 4개의 Example들을 하나씩 증명해 봅시다.

EM distance

KL-divergence

JS divergence

Total Variation

그래서 뭐가 다르냐?

왼쪽 그래프는 EM distance를 나타낸 그래프이고, 오른쪽 그림은 JS divergence를 나타낸 것입니다. EM plot은 연속적이고 어디서나 gradient가 존재하기 때문에 Loss function으로 사용하는데 제약이 없지만, JS plot은 불연속적이고 0인 지점에서 미분이 불가능하기 때문에 gradient를 구하지 못합니다.

논문에서 ‘weaker’이라는 말이 나오는데, 분포수렴을 알아야합니다.

분포수렴이란 확률분포 수렴 종류 중 하나로서 제일 약한 수렴을 말합니다. 확률 분포의 개별적인 특징보다 전체적인 모양을 중시하는 수렴입니다.

즉, EM(Wasserstein distance)는 TV, KL, JS보다 약한(weak) metric으로 수렴을 판정하는데 soft한 성질을 가집니다.

다음은 EM-distance가 연속성이 있는 Loss function인지 아닌지 증명하고 어떤 조건을 만족해야 하는지 논문에서 밝힌 것입니다.

Theorem 1.의 Corollary(필연적인 결과)

$\theta$를 파라미터로 가지는 $g_{\theta}$를 feedforward neural network(순방향 신경망)가 된다고 하고, $p(z)$에서 추출한 임의의 z의 거리의 Expectation 유한하다면,

assumption 1은 만족하고 EM-distance는 어디서나 연속이며 거의 모든 곳에서 미분 가능하다는 것을 보여줍니다.

그러므로 기존에 사용하던 JS-divergence보다는 섬세한 cost function이라고 합니다.

확실한 사실은 TV, JS, KL은 low dimensional manifolds(특징의 차원이 낮은 manifolds)에서 cost function으로 사용하기에는 부적절합니다. 다음 섹션에서 EM distance가 어떤 방법으로 optimize되는지 알아봅니다.

Wasserstein GAN

식을 다시 한번 살펴보면, 결국 우리는 $P_r$과 $P_\theta$의 상한에서 가장 작은 값을 골라 내야 합니다. 하지만 $P_r$은 우리가 추정하고자 하는 값이기에 이 값을 구하는 것은 intractable(아주 다루기 힘든)합니다..

그래서 * Kantorovich-Rubinstein duality*를 사용하여 계산을 단순화시킬 수 있습니다.

여기서 sup에 적용되는 것은 $f(x)$이며, $f(x)$는 1-Lipschitz function입니다.

1-Lipschitz function

립시츠 연속 함수는 두 점 사이의 거리를 일정 비 이상으로 증가시키지 않는 함수이다.

결국 립시츠 연속 함수에 해당하는 $f(x)$에 대해서 sup을 진행하면 되는 것입니다. 저 식에서 약간의 변수만 바꾸면 다음과 같이 됩니다.

기존 GAN의 object function과 식이 유사함을 볼 수 있습니다.

이제 위의 방정식에서 최대화 문제를 해결하는 함수 f을 찾아야합니다. 이를 근사시키기 위해서,

f는 립시츠 연속 함수가 되고, w만 파라미터로 가집니다. compact space에 w라는 파라미터를 가지기 위해서 gradient를 업데이트한 후에 W라는 compact space에 제약을 겁니다.

하지만 립시츠 연속함수에 제약을 거는 Weight clipping은 매우 치명적일 수 있습니다. 왜냐하면, f라는 함수는 립시츠 연속 함수인데 clipping parameter가 너무 크다면, sup의 하한선에 도달하는데 까지 시간이 매우 오래걸리기 때문입니다. 또한 clipping parameter가 너무 작으면 layer의 수가 많을 시, gradient vanishing 문제가 발생할 수도 있고 배치 정규화가 사용되지 않기 때문입니다.

GAN과 WGAN의 object function을 비교한 그림입니다. 자세한 설명은 여기를 누르시면 됩니다.

Algorithm

WGAN의 알고리즘 입니다.

• $P_r$과 $p(z)$를 배치 크기만큼 샘플링 한 후,
• loss function을 사용하여 parameter w를 업데이트 시킵니다.

[EM distance vs Others]

EM distance는 연속적이며 어디서나 미분가능하므로 critic을 optimal point까지 훈련시킬 수 있습니다. 또한 우리가 critic을 더욱 훈련시킬수록 더 reliable한 Wasserstein의 기울기를 얻을 수 있습니다. 다른 JS, KL,TV 같은 경우에는 gradient가 0이 나오거나 특정값으로 수렴되기 때문에 gradient vanishing 문제가 발생할 수 있습니다.

또한 GAN의 discriminator은 가짜 이미지에 대해서 매우 잘 구별하지만, reliable gradient를 얻을 수는 없습니다.

WGAN의 critic 경우, 특정 값에 saturation되지 않기 때문에 어디서나 gradient를 구할 수 있습니다.

Empirical Results

이미지를 생성하는데 있어서 WGAN과 GAN의 알고리즘을 비교 및 분석하였습니다.

Experimental Procedure

• Upper left : MLP with 4 hidden layers and 512 units at each layer
• Upper right : DCGAN without the sigmoid
• Lower half : both generator and discriminator are MLPs

위의 두 그림은 discriminator 대신 critic을 사용하여 Wasserstein 거리를 구한 것인데, 확연히 sample의 결과가 좋아 진 것을 확인할 수 있는 반면에, 아래의 discriminator를 사용한 것은 sample의 결과를 알아보기 힘들고 Loss가 일정한 것으로 보아 훈련에 실패했다고 볼 수 있다. (자신들이 아는 바로는 GAN이 나온 이래 수렴하는 특성을 보인 경우가 처음이라고 합니다..)

Meaningful loss metric

• Upper left : JS estimates for an MLP generator
• Upper right : Standard DCGAN
• Bottom : MLP with both generator and discriminator

EM distance 대신 JS distance를 사용한 것인데, sample qulity는 생성되지만, JS 거리가 constant를 유지하거나, 줄어들지 않고 증가하는 것을 볼 수 있습니다.

Why they use RMSProp for optimizer instead of ADAM?

Adam과 같은 momentum 기반의 optimizer를 사용했을 때나 높은 learning rate를 사용했을 경우, WGAN에서 불안정한 결과를 봤습니다. 그래서 momentum을 비정상적인 결과에 대한 잠재적인 원인으로 생각했습니다. (여기서 momentum은 optimize할 때 관성을 의미한다고 생각하시면 될 것 같습니다.)

왜냐하면 Adam step(관성 방향)과 gradient가 이루는 각의 cosine값이 음수가 되기 때문입니다. 이 부분이 불안정성을 야기한다고 합니다.

그래서 optimizer로 RMSProp을 사용하여 이러한 문제점을 해결한 것입니다.

Improved stability

이 부분에서는 DCGAN generator를 이용하고 다양한 실험 환경을 변경하면서 실험을 한 사진입니다.

결국 저자가 하고 싶은 말은 WGAN을 사용하면 mode collapse 문제를 해결할 수 있다는 것입니다.

Conclusion

WGAN은 기존의 GAN을 대체할 수 있다는 것을 확인하였고, 훈련의 안정성을 높일 수 있는 것을 확인하였습니다. 또한 mode collapse를 해결하기도 했습니다.