[2025-1] 김유현 - Wasserstein GAN

https://arxiv.org/abs/1701.07875

Wasserstein GAN

 

 

 

1. Introduction

GAN의 목적은 데이터 $x$의 분포 $P(x)$를 직접 학습하는 것이 목적이다. $P(x)$를 매개변수 θ를 사용하여 아래와 같이 표현이 가능하고 이를 학습할 수 있다. 

그러나 $P(x)$ 자체를 직접 모델링하는 것은 어렵기 때문에 GAN에서는 $x$를 결정하는 latent variable $z$의 분포 $P(z)$를 가정한 후, 이를 generator의 입력으로 사용한다. 이후, discriminator와 generator간의 관계를 학습시켜 generator의 분포를 $P(x)$에 가깝게 학습하고자 한다.

 

1) 기존 GAN의 문제점

• Discriminator와 Generator간의 균형을 유지하며 학습하기 어렵다.
  • 학습 과정에서 discriminator가 지나치게 강해지거나 너무 약해지면 generator가 제대로 학습되지 않는다. 
• Mode collapse 현상이 발생한다. 
  • 학습이 완료된 이후에도 generator가 분포의 일부 모드만을 생성하는 문제가 발생한다. 
  • 이는 discriminator가 가짜 데이터를 명확하게 분류하지 못하여, 모델이 최적점까지 학습되지 못했기 때문이다. 

 

2) Wasserstein GAN의 해결 방법

• Discriminator 대신 새롭게 정의된 critic을 사용
  • 기존 discriminator는 sigmoid 함수를 사용하여 입력의 진위여부를 판별하고 출력값은 확률로 해석한다. 
  • 하지만 WGAN의 critic은 단순히 real과 fake 데이터를 구분하는 것이 아니라, 두 분포 간의 거리를 측정하는 방식으로 동작한다. 
• Critic은 Lipschitz* 조건을 만족 
  • Lipschitz 연속성을 보장하기 위해, gradient clipping** 또는 weight clipping***을 적용하여 critic의 변화를 제한한다.
  • 이후 논문에서는 gradient penalty (WGAN-GP) 를 도입하여 보다 효과적인 Lipschitz 조건을 만족시키는 방법을 제안한다.
• Wasserstein Distance 활용
  
  • 기존 GAN은 Jensen-Shannon (JS) divergence 를 사용하지만, WGAN은 Wasserstein-1 distance (EM distance) 를 사용하여 두 분포 간의 차이를 보다 안정적으로 측정한다.
  • EM distance는 두 확률 분포 간의 물리적 이동 비용을 기반으로 정의되며, gradient가 끊기지 않고 매끄럽게 변화하도록 한다.
• KL divergence는 학습이 어려운 경우가 많다. 
  • 기존 GAN이 사용한 KL divergence와 JS divergence는 연속적이지 않거나 gradient vanishing 문제가 발생할 가능성이 높아 학습이 어려운 경우가 있다.
  • Wasserstein distance는 gradient가 더욱 안정적으로 작동하여, 효과적인 학습이 가능하다.

 

Gradient Clipping

• Gradient의 크기가 너무 커지는 것을 방지하는 방법
• 특정 임계값 $c$를 설정하고, gradient의 크기가 $c$를 넘으면 줄여주는 방식을 통해 기울기 폭발 현상을 방지한다. 

Weight Clipping

• 네트워크의 가중치를 특정 범위 [$-c, c$]로 강제로 제한하는 방법
• Lipschitz 연속성을 유지

Gradient Penalty

• Critic의 gradient가 1에 가까워지도록 학습하는 추가 손실 항을 적용

Lipschitz

• 함수의 변화율이 특정 값 이상으로 급격히 변하지 않도록 보장하는 수학적 제약 조건
• 입력값이 조금 변할 때, 함수 출력값도 제한된 범위 내에서만 변하도록 강제하는 조건

 

3) WGAN의 장점

• 기존 GAN보다 더 안정적인 학습이 가능하다. 
• Mode collapse 문제를 완화하여, generator가 더 다양한 샘플을 생성할 수 있다. 
• Discriminator와 generator 간의 균형을 엄격하게 조정하지 않아도 학습이 가능하다.

WGAN에서 사용되는 EM distance 기반 critic을 통해 WGAN은 기존의 GAN보다 더 안정적인 convergence와 샘플 품질 향상을 가능하게 한다. 

 

2. Different Distances

GAN의 학습 과정에서는 실제 데이터 분포와 generator의 샘플 분포 사이의 거리를 측정해야 한다. 

 

1) Total Variation (TV) distance

• 두 확률 분포 간의 최대 차이를 측정하는 거리이다. 즉, 가능한 측정 값들 중 차이가 가장 큰 값으로 정의된다. 
• 두 분포가 완전히 겹치면 0, 겹치지 않으면 1이 된다.
• 불연속적인 변화가 발생하여 gradient descent 기반의 학습이 어렵다. 

 

2) The Kullback-Leibler (KL) divergence

• 두 확률 분포 $P$와 $Q$ 간의 상대 엔트로피를 측정한다. 
• 한 분포가 다른 분포를 얼마나 잘 설명하는지를 나타내는 지표이다. 
• 분포가 겹치지 않으면 KL divergence는 무한대가 된다. 

 

3) The Jensen-Shannon (JS) divergence

• KL divergence를 대칭적으로 변형한 거리이다.
• 두 분포가 완전히 동일하면 0, 겹치지 않으면 1이 된다. 
• 분포가 조금만 차이가 나도 gradient가 거의 0이 되어 학습에 어려움이 있다. 

 

4) The Earth-Mover (EM) distance or Wasserstein-1

• 두 확률 분포 간의 차이를 측정하는 방식으로, 한 확률 분포에서 다른 확률 분포로 이동하는 최소 비용을 의미한다.
• 한 확률 분포를 다른 확률 분포로 변환하는 데 필요한 최소한의 작업량을 기반으로 거리 계산이 이루어진다.
• Gradient가 연속적인 값을 가지므로 gradient descent를 통한 학습이 가능하다.
• 기존 JS/KL divergence보다 더 부드러운 거리를 제공하여 학습 안정성을 향상시켜 분포 간의 차이를 보다 정확하게 반영한다.

 

A. Example 1

 

논문에서는 Wasserstein distance와 다른 거리의 척도를 비교하기 위해 임의의 uniform distribution을 갖는 $Z$를 기반으로 $P_0$ = $(0, z)$를 구성하고, 다양한 거리에 대해 두 분포 간의 수렴 여부를 비교한다.   

• Wasserstein distance의 경우, θ 값에 관계없이 일정한 수식을 가지며 일정한 변화를 보인다.
• KL divergence, JS divergence, TV distance는 두 분포가 서로 겹치는 경우에는 0, 겹치지 않을 경우에는 무한대 또는 상수로 극단적인 거리 값을 나타낸다.
• Figure 1에서 Wasserstein distance를 제외한 다른 모든 함수는 불연속적이며, θ값이 0에 가까워지면 발산하게 된다.
• 이를 통해 KL/JS/TV distance를 사용하는 모델은 학습이 어렵거나 불안정해질 수 있지만, Wasserstein distance만이 Low-dimensional manifold에서 gradient descent를 통해 안정적인 학습이 가능하다. 

 

 

B. Theorem 1

 

Theorem 1은 Wasserstein distance $W(P_r, P_θ)$ 의 연속성을 보이며 안정적인 학습을 위해 어떤 조건이 필요한 지 명시한다.

 

1) 정의 및 수식

• $P_r$: 학습하고자 하는 목표 distribution
• $P_θ$: 학습시키고 있는 현재의 distribution
• $Z$: latent space - 입력 노이즈 분포
• $g$: latent variable $z$를 $X$로 mapping하는 함수
  • $g_θ(z)$의 분포가 $P_θ$

2) 연속성 조건

• $g$가 $θ$에 대해 연속한 경우, $W(P_r, P_θ)$도 연속적이다. 
  • 즉, Generator가 부드럽게 변하면 Wasserstein Distance도 부드럽게 변화한다. 
• $g$가 Lipschitz 조건을 만족하는 경우, $W(P_r, P_θ)$은 연속적이고 거의 모든 곳에서 미분 가능하다.
  • Lipschitz 조건은 두 점 사이의 거리를 일정 비율 이상으로 증가시키지 않는 함수를 말한다. WGAN에서는 weight clipping을 수행하는데 이는 최적의 방법이 아니다. 

 

C. Theorem 2

 

Theorem 2는 KL, JS, TV, EM distance가 각각 얼마나 강한 distance metric인지를 비교하며 강한 metric일수록 학습이 어려워진다는 점을 보여준다. Distance 결과 KL > JS > TV > EM 순으로 강한 것을 보여준다.

 

1) Total Variation Distance와 JS Divergence는 동일한 수렴성을 가진다. 

• $P_n -> P$일 때 TV Distance가 0으로 수렴하면 JS Divergence도 0으로 수렴한다. 
• TV Distance와 JS Distance는 Discrete Setting에서 분포가 살짝만 차이나도 0이나 무한대 같은 극단적인 값을 가질 수 있어 학습이 어렵다. 

2) Wasserstein Distance가 0으로 수렴하는 것은 $P_n$이 $P$로 확률 분포의 개념에서 수렴하는 것과 같다. 

• Wasserstein Distance를 사용하면 확률 분포가 점진적으로 변하는 형태로 학습이 가능하다. 

3) KL Divergence가 0으로 수렴하면 TV distance와 JS distance도 0으로 수렴한다. 

• KL Divergence는 분포가 겹치지 않으면 무한대가 되어 학습이 끊길 수 있다. 

4) TV/JS Distance의 수렴은 Wasserstein Distance의 수렴을 보장하지 않는다. 

• 기존 GAN이 사용한 JS Distance와 KL Distance는 WGAN에서 사용하는 Wasserstein Distance보다 학습이 불안정할 수 있다. 

 

D. Corollary 1

 

 

Corollary 1은 Wasserstein distance를 최소화해 학습하는 것이 neural network 학습에서 타당하게 작동한다는 것을 보여준다. 즉, Neural Network 기반 Generator를 사용하여 Wasserstein Distance를 최소화하는 것이 이론적으로 타당하다.

실제 데이터는 low-dimensional manifold 상에서 분포하는 경우가 많은데 KL, JS, TV distance는 이런 manifold 구조에 민감하지 않아 학습이 더 어렵다. 하지만, EM distance는 manifold 구조에서도 민감한 반응을 보이므로 더 좋은 cost function이 될 수 있다. 실제 실험에서도 WGAN이 기존 GAN보다 안정적으로 학습하며 mode collapse가 적게 발생한다. 이는 GAN 연구에서 adversarial training을 개선하는 중요한 방향을 제시한다.

 

3. Wasserstein GAN

Wasserstein GAN은 Wasserstein distance를 최적화하는 것이 목표이며 기존 GAN의 학습 문제를 해결한다.

 

A. Wasserstein GAN

Kantorovich-Rubinstein duality를 이용하여 $W$를 위와 같이 설정한다. 여기서 sup 아래의 의미는 $f$가 1-Lipschitz 함수라는 것을 나타낸다. 만약 $f$가 1-Lipschitz가 아니라 어떤 상수 $k$에 대해 K-Lipschitz 조건을 만족하는 경우 아래와 같이 변형된다. 즉, $f$가 엄격한 1-Lipschitz가 아니더라도, 적절한 K-Lipschitz 조건을 만족하면 Wasserstein distance를 최적화할 수 있다. 

 

WGAN에서는 최대화 문제를 푸는 함수를 신경망을 이용해 근사하는데 $f$를 학습하는 critic network를 정의하고, 이 critic의 weight $w$를 최적화하여 Wasserstein distance를 최대화한다.

• 신경망의 가중치 $w$가 compact space $W$에 놓여 있다고 가정한다.
• 최적화 과정에서 $E_{z~p(z)}[▽_θf_w(g_θ(z))]$를 이용해 backpropagation을 수행한다.

이 과정에서 기존 GAN의 훈련 과정과 유사하지만 차이점이 있는데, WGAN에서는 critic이 제한 없는 scalar 값을 출력하여 Wasserstein distance를 근사한다. Lipschitz 조건을 만족시키기 위해, weight clipping 또는 gradient penalty를 사용한다. 

 

B. Theorem 3

 

 

Theorem 3에서는 Wasserstein distance를 최적화할 때의 수학적 정당성을 보인다.

첫 번째 식의 상한이 어떤 최적의 $w$ ∈ $W$에 대해 달성된다면, 이 최적화 과정이 $W(P_r, P_θ)$의 값을 multiplicative constant 내에서 근사함을 보인다. 또한 $W(P_r, P_θ)$의 미분을  $E_{z~p(z)}[▽_θf_w(g_θ(z))]$ 를 추정할 수 있음을 증명한다. 즉, Wasserstein distance를 gradient-based optimization을 통해 학습 가능함을 보인다. 

 

C. Algorithm 1

 

<WGAN의 학습 과정>

1) Critic 학습 단계

• 목표 분포 $P_r$와 latent 분포 $P(z)$를 각각 미니배치 크기만큼 샘플링.
• Critic의 loss function을 이용하여 weight $w$ 를 업데이트.
• Weight clipping 수행: Lipschitz 조건을 강제하기 위해 $w$를 [$−c, c$] 범위로 제한.

2) Generator 학습 단계

• Generator는 critic이 제공하는 Wasserstein distance를 줄이는 방향으로 업데이트
• Critic이 최적화될수록 더 안정적인 gradient 정보를 제공하여 generator 학습을 지원

 

<Weight Clipping의 한계점 및 문제점>

1) $c$가 너무 큰 경우 

• clipping 범위가 넓은 경우
• Lipschitz constraint가 약해지면서 최적의 critic을 찾는 데 오랜 시간이 걸린다.
• 즉, training이 비효율적이 될 수 있다.

2) $c$가 너무 작은 경우

• clipping 범위가 좁은 경우
• Critic의 weight들이 작은 값에 제한되어 gradient vanishing 현상이 발생.
• 즉, gradient가 너무 작아져 학습이 제대로 진행되지 않는다. 

 

WGAN의 weight clipping은 Lipschitz 조건을 강제하는 방식이지만, 한계가 존재한다. 이 문제를 해결하기 위해 WGAN-GP가 제안되는데 weight clipping 대신 gradient penalty를 추가하여 Lipschitz 조건을 만족하도록 개선한다. 

 

D. Figure 2

Fig 2는 GAN discriminator와 WGAN critic(discriminator 역할)을 최적화했을 때의 차이점을 보여준 결과이다.

 

 

 

1) GAN의 Discriminator

 

• Fake/Real을 잘 구분할 수 있도록 학습되지만, reliable gradient information을 제공하지 못함.
• Discriminator가 너무 강해지면 generator의 gradient가 0에 가까워져 학습이 멈추는 문제인 gradient vanishing가 발생.
• 기존 GAN의 cost function은 mode collapse 문제 발생 가능.

 

2) WGAN의 Critic

 

• 기존 GAN과 달리 sigmoid를 사용하지 않고, raw score를 출력하여 Wasserstein distance를 근사.
• 결과적으로 loss function이 linear하게 변하며 clean gradient를 유지할 수 있음.
• Discriminator가 너무 강해지더라도 gradient 소실 문제가 발생하지 않고 안정적인 학습이 가능.

 

F. GAN vs WGAN

1) GAN

 

 

2) WGAN

 

3) GAN vs WGAN

• Discriminator/Critic과 Generator의 역할
  • GAN: Discriminator $D$는 실제 데이터 $x$와 Generator $G$가 생성한 가짜 데이터 $G(z)$를 구분하는 역할을 한다. 
  • WGAN: Critic $f$는 Wasserstein distance를 근사하도록 학습된다. 실제 데이터 분포와 생성된 데이터 분포 간의 거리를 측정하는 함수이다. 
• GAN의 Loss Function
  • Discriminator: Binary cross-entropy loss를 사용하며, 목적은 진짜 데이터 $x$에 대해서 $D(x)$를 1에 가깝게, 가짜 데이터 $G(z)$에 대해서 $D(G(z))$를 0에 가깝게 만든다. 이때 Discriminator는 gradient ascent (최대화)를 수행한다. 
  • Generator: Generator는 Discriminator가 생성한 샘플을 진짜로 믿도록 만드는 방향으로 학습된다. Discriminator의 출력을 1에 가깝게 만드는 방향으로 $θ_g$에 대해 gradient descent (최소화)를 수행한다. 
• WGAN의 Loss Function
  • Critic: WGAN에서는 Discriminator 대신 critic $f$를 사용하며, Binary cross-entropy loss가 아닌 Wasserstein distance를 학습하도록 설계되었다. Critic은 실제 데이터에 대해 높은 값을 주고, 가짜 데이터에 대해 낮은 값을 주는 방향으로 학습된다. Wasserstein distance의 특성상, gradient ascent (최대화)를 수행하여 $f$가 최적의 함수를 찾도록 한다. 
  • Generator: Generator는 Critic이 출력한 값 $f(G(z))$를 최대화하도록 학습한다. Critic이 가짜 데이터에 대해서 높은 점수를 주도록 만드는 방향으로 $θ$에 대해 gradient descent (최소화)를 수행한다. 

 

4. Empirical Results

아래의 실험은 LSUN-Bedrooms dataset을 학습하게 했고 DCGAN을 baseline으로 설정했다. 

A. Experimental Procedure

 

DCGAN과 MLP에서 EM distance의 추정을 loss function으로 설정해 학습한 결과를 나타낸다. 

왼쪽 그림은 generator로 Multi Layer Perceptron, 오른쪽 그림은 DCGAN을 이용한 결과이다. Sigmoid 활성 함수를 사용하지 않아 Wasserstein 거리가 점차 감소하며, 샘플의 결과 또한 좋아진 것을 확인할 수 있다.

아래 그림은 discriminator와 generator 모두 MLP를 사용한 결과인데, 샘플 이미지는 품질이 낮아 식별하기 어려우며 각 샘플에 대해 Wasserstein distance를 계산한 결과 값이 상수에 가깝게 유지되고 있다. 

이는 다른 모델과 성능을 비교하기 위해 생성된 샘플을 확인할 필요가 없기 때문에 adversarial network를 연구할 때 매우 유용하다. 

 

B. Meaningful Loss Metric

위의 그래프는 Jensen-Shannon(JS) distance를 사용하여 실험한 결과이며, 곡선이 샘플 품질과 명확한 상관관계를 가지지 않음을 보여준다. A에서 사용한 critic+MLP, critic+DCGAN, MLP+MLP 모델 구조를 동일하게 적용한 후, generator의 각 iteration 마다 JS distance를 측정하여 그래프를 확인한 결과를 보인다. 결과적으로 샘플 품질이 향상되어도 JS distance는 증가하거나 상수 값을 유지하는 것을 볼 수 있다.

논문에서는 WGAN이 Adam과 같은 momentum에 기반한 optimizer를 사용하면 학습이 불안정해지므로 비정상(nonstationary) 문제에도 잘 작동하는 RMSProp을 사용해 학습을 진행했다. 

 

C. Improved Stability

DCGAN의 generator를 사용하여 다양한 변화를 주며 실험을 진행한 후, discriminator와 critic의 성능을 비교한다. 실험 결과, WGAN을 사용하면 discriminator와 critic간의 균형을 맞추는 문제를 더 이상 신경 쓸 필요가 없으며, mode collapse 현상이 발생하지 않는다는 것이 확인된다.  

• Figure 5
  • 일반적인 GAN과 WGAN을 비교한 결과, WGAN은 더 높은 품질의 샘플의 이미지를 생성한 결과를 확인 할 수 있다.
• Figure 6
  • Batch normalization을 제거하고, generator 내 DCGAN의 filter 개수를 고정하여 전체적인 parameter 수를 줄인 결과이다.
  • WGAN에서는 위 설정에서 잘 작동하지만, 일반적인 GAN에서는 성능 저하가 발생하는 것을 확인할 수 있다. 
• Figure 7
  • Generator를 MLP+ReLU로 변형하여 실험한 결과이다.
  • 왼쪽은 WGAN, 오른쪽은 일반적인 GAN의 결과이며, DCGAN을 사용했을 때보다 생성된 이미지의 품질이 떨어졌지만 mode collapse 현상에 대해서 비교할 수 있다.
  • 오른쪽의 경우, 비슷한 이미지가 반복적으로 생성되며 특정 이미지에 대해서는 전혀 생성하지 못하는 현상이 발생한다.
  • 왼쪽의 경우, 다양한 이미지를 비교적 균형 있게 생성하는 것을 확인할 수 있다. 

 

5. Conclusion

WGAN은 기존 GAN 훈련 방식의 대안으로, 학습 안정성을 향상시키고 mode collapse 문제를 해결하며, 디버깅과 하이퍼파라미터 탐색에 유용한 학습 곡선을 제공한다. 또한, 이 최적화 문제가 이론적으로 타당함을 입증하고, 다양한 확률 분포 간 거리와의 깊은 연관성을 이론적으로 분석했다. 

 

Reference

---

https://ahjeong.tistory.com/7

[논문 읽기] Wasserstein GAN

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

GAN — Wasserstein GAN & WGAN-GP

https://velog.io/@pabiya/Wasserstein-GAN

Wasserstein GAN

https://cumulu-s.tistory.com/31

5. Wasserstein GAN (WGAN) - paper review