[NIPS2017] SoftmaxGAN

이 논문은 수식적으로도 어려운 것이 있지만 구현에서 상당히 애를 먹었다. 다른 논문과 달리 Generator를 훈련시키기 위해 real logit을 사용한다. 따라서 그래프를 이전과 달리 D를 두번 써서 그래프를 새롭게 만들었어야 했다. 또한 깃헙에 버전이 달라서 그런지 잘못 구현이 되어있어서 애를 먹은것 같다. 덕분에 backward 원리를 많이 공부한듯...

Abstract & Introduction

이 논문도 WGAN과 마찬가지로 GAN학습의 안정성을 추구한다. **Softmax distribution이 target distribution과 같아지지 않으면 gradient가 0이 아니라는 것을 이용한다. **

Related Works

WGAN을 상당히 많이 인용하면서 자신들도 non-vanishing gradient를 목표로 한다고 한다.

Softmax GAN

notation부터 정리하자면
$B_+$ : training data
$B_-$ : generated data
$B = B_+ + B_-$ : Union
$\mu^{\theta}(x)$ : \sum_{x \in B}e^{-\mu^{\theta}(x)}, partition function

키 아이디어는 분류기 학습할 때 맨날 쓰는 크로스 엔트로피를 사용하는 것이다. 즉, discriminator를 학습시키기 위해서는 $\mu^{\theta}(x)$를 소프트 맥스와 유사하게 자연상수의 지수를 하여 양수로 만들고, x가 training dataset, 즉 real data이면 1, latent z 에서 나오면 0으로 취급한다. 그 다음 크로스 엔트로피 손실함수의 정의는 이 계산값을 전체 데이터 갯수로 나누어 주는 것이지만, 논문에서는 레이블을 $\frac{1}{|B_+|}$로 잡아서 real data 갯수 평균과 동일한 효과를 냈다. 즉, real data만 있다고 고려하고, gene data에 대해서는 고려하지 않는다.
따라서 Discriminator의 손실함수는 아래와 같다.
$$t_D(x) = \frac{1}{|B_+|} if x \in B_+ else 0$$
$$L_D = -\sum_{x \in B} t_D (x) ln(\frac{e^{-\mu^{\theta}(x)}}{Z_B}) = -\sum_{x \in B_+} ln(\frac{e^{-\mu^{\theta}(x)}}{Z_B}) + \sum_{x' \in B_-} 0 ln(\frac{e^{-\mu^{\theta}(x')}}{Z_B}) $$
$$ = \sum_{x \in B_+} \mu^{\theta}(x) + ln(Z_B)$$

Generator의 경우에는 레이블모두를 같은 확률로 본다. 즉, 위에서는 가중치를 real data에 1, gene data에 0을 주었다면 (확률로도 real data갯수로 나누었음), 이번에는 그냥 전부다 1을 주고 전체 갯수로 나눈다. Generator의 손실함수는 아래와 같다.
$$t_G(x) = \frac{1}{|B|}$$
$$L_G = -\sum_{x \in B} t_{G}(x) ln(\frac{e^{\mu^{\theta}(x)}}{Z_B}) = -\sum_{x \in B_+}\frac{1}{|B|} ln(\frac{e^{-\mu^{\theta}(x)}}{Z_B}) - \sum_{x' \in B_-} \frac{1}{|B|} ln(\frac{e^{-\mu^{\theta}(x)}}{Z_B})$$
$$\sum_{x \in B_+} \frac{1}{|B|} \mu^{\theta}(x) + \sum_{x' \in B_-}\frac{1}{|B|}\mu^{\theta}(x') + ln(Z_B)$$

중요한 점은 맨 처음식이 소프트맥스와 거의 동일하지만 앞서 언급하였듯이 가중치 확률이 다르다. 이걸 왜 다르게 했는지는 뒤에 증명에서 나옴.

Relationship to Importance Sampling

수식이 복잡하지만 그냥 앞의 식을 일반화 하여 증명하는것에 지나지 않으므로 나중에 필요한 일이 있으면 그냥 논문을 참고하자.

결국 하고자 하는 얘기는 $L_G = KL(\frac{p_D + p_G}{2} || p_D) + KL(p_D || \frac{p_D + p_G}{2})$ 로, Jensen-Shannon divergence를 sampling만 다르게 한것과 같다는데 이건 아직 이해하지 못했다.

구현하면서 배운점

이 논문은 제대로 된 구현체가 없기 때문에 정답을 알 수가 없다.
나는 DC-GAN을 베이스로 구현하였는데 구현하면서 모델 그래프를 다시 한 번 공부 할 수 있었다.
아래 그림은 SoftGAN의 그래프이다. 이대로 구현 하면 되는데 이전의 GAN들과 다른 점은 Generator를 훈련시키기 위해 real logit도 쓴다는 점이다.
**순서에 주의하자. 먼저 D를 업데이트 시키고 업데이트 된 D에서 logit을 추출하여 G_loss를 만든다. 또한 D를 훈련시킬때에는 optimizer_D를 사용하고 G는 optimizer_G를 사용하므로 그래프가 연결되어 있어도 업데이트는 각각 될 수 밖에없다. (다시말하면, optimizer_D.step()을 하면 D의 parameter만 업데이트 된다.) **

아래 코드가 안되는 이유는 D를 훈련시키고 retain_graph=True로 하여 그래프를 남긴다. 그 뒤에 G_loss를 계산하기 위해 처음에 계산한 real_logit과 gene_logit을 사용한다. (Z_B도 마찬가지.)그런데 문제는 D가 훈련되면서 파이토치에서 parameter가 업데이트 되는 원리는 inplace연산으로, (메모리를 아끼기 위해 이렇게 한다.) 모델이 가만히 있는 것 처럼 보인다. x += 2를 생각하자. 이 때 원래 계산되었던 real_logit과 gene_logit을 사용하면서 optimizer_G.step()이 실행되는데 D는 업데이트가 되지 않지만 inplace연산을 이미 거친 그래프부분을 이용했기 때문에 계산이 안된다.

이제 이 문제를 해결 해보자.

첫번째 방법은 아래 코드처럼 그냥 gene_logit과 real_logit을 두개 만들면 된다.

이 때, D를 훈련시키는 코드에서는 gene_logit을 생성하는 부분에서 detach를 사용하였기 때문에 그래프에 G가 연결되지 않는다. 따라서 D_loss.backward()에 매개변수로 retain_graph=True를 주지 않아도 된다.

두번째 방법은 아래 코드처럼 retain graph를 활용하여 detach를 없앨 수 있다. 원리는 원래 backward를 실행하고 나면 리프부터 루트까지 연결된 부분을 없앤다. 하지만 retain graph를 True로 하면 사라지지 않아 G를 남길 수 있다. * optimizer_D.step()을 하여도 G는 하나도 바뀌지 않는다. 대신 그래프가 안 사라지게 할 수는 있는것. 여기에서 위에처럼 D를 하나 더 만들어서 업데이트 시킨다.

세번째 방법이 특이한데, 먼저 그냥 D_loss와 G_loss를 구한다. 문제점의 코드와 다른점은 먼제 다 계산하고 업데이트를 시킨다. 코드를 확인하자.
원래는 아래처럼 G를 2배 훈련시키지 않아도 돌아가는데 D의 힘이 더 센지 그냥

if idx % 2 ==0:
  ...

부분만 살려서 쓰면 생성이 안되고 그냥 검정색 사진만 나온다. 그래서 한번은 G만 훈련하면 좀 나아진다.

전체적인 결과는 깃헙에서 확인.