Jihye Back

Jihye Back

꿈꾸는 꼬마 개발자 👩🏻‍💻 (Interested In GenAI)

[Paper Review] StyleGAN2-ADA #01: Training Generative Adversarial Networks with Limited Data 논문 분석

업데이트:

View On GitHub

✍🏻 최근에는 이미 있는 모델(pretrained model)을 잘 fine tuning하여 의미있는 결과를 내는 연구가 대세이다. (FreezeD, GANSpace, StyleCLIP 등등)

StyleGAN2-ADA도 이러한 흐름에서 나온 연구로, loss function이나 network의 architecture를 건들이지 않고 이미 학습이 된 GAN을 finetuning하며 학습을 한다. 또한, 적은 데이터로 학습을 해도 discriminator가 overfitting 되지 않도록 Adaptive Discriminator Augmentation Mechanism을 제안하였다.

1. Introduction

Prevent the augmentations from leaking

GAN은 small dataset으로 학습을 하면 Discriminator가 overfitting 된다는 문제점이 있다.

  • GAN은 Discriminator와 Generator가 경쟁하며 학습하는 모델이기 때문에 D가 overfit된다면 G도 이상하게 학습된다.(발산)
  • Overfitting 문제를 해결하기 위한 연구로 Dataset Augumentation과 같은 연구도 있다. (rotation, noise 등을 추가하여 classifier가 더 잘 학습되도록 함) 그러나 이러한 방식으로 학습을 하면, augmented distribution(leaking)가 생겨 학습이 이상하게 된다 - ex. a noise augmentation leads to noisy results, even if there is none in the dataset

StyleGAN2-ADA는 augmentation range를 조절하여 D가 overfitting되는 것을 막았다.

2. Overfitting in GANs

적은 data로 학습할수록 FID score도 높고 D도 overfitting된다.

2.1 Stochastic Discriminator Augmentation

bCR-GAN은 Discriminator를 학습시킬 때 real와 generated image 모두에 대해 consistency regularization term을 추가하였지만, Generator를 학습시킬 때에는 augmentation을 하지 않았다. 이와 같이 학습을 한다면, G가 아무런 penalty없이 자유롭게 이미지를 생성할 수 있기 때문에 augmentation이 leaking 될 수 있다.

stochastic discriminator augmentation

따라서 StyleGAN-ADA에서는 오직 augmented image만을 사용하여 Discriminator를 학습하며, Generator 역시 augment한 이미지를 바탕으로 학습을 한다. 또한, augmentation probability $p$를 통해 augmentation의 정도를 조절한다.

  • $p$가 0이면, augmentation X
  • $p$가 1에 가까우면, augmentation을 많이

2.2 (Non-leaking) Invertible Augmentation

Ambient-GAN

AmbientGAN은 문제가 있는 dataset(corruption)에 대해 학습을 할 때도 올바른 distribution을 찾을 수 있음을 밝힌 모델이다.

  • Generator가 깨끗한 이미지 $X_g$를 만들면 $f_θ$가 이를 corrupt하는 noise function을 학습해서 $Y_g$를 생성하고, Discriminator가 corrupt된 real data $Y_r$와 $Y_g$를 비교하게 하는 구조

    \[\min _{G} \max _{D} \mathbb{E}_{Y^{r} \sim p_{y}^{r}}\left[q\left(D\left(Y^{r}\right)\right)\right]+\mathbb{E}_{Z \sim p_{z}, \Theta \sim p_{\theta}}\left[q\left(1-D\left(f_{\Theta}(G(Z))\right)\right)\right]\]

  • 단, 이때 corruption process는 data space의 확률 분포에 대해 invertible transformation이 가능해야한다.

    • invertible transformation이 가능하다는 것은 다음 그림처럼 $T(x)$에서 $x$로 mapping할 수 있다는 것이다.
    • 만약에 우리가 augmentation을 uniform하게 했다면 invertible transformation이 어려울 것이다. : Augmentation 후에 {0, 90, 180, 270} 모두 25%씩 있다면 몇도에서 augmentation을 시작했는지 모름
    • 반면, $p$를 작게하여 안전하게 augmentation을 했다면 invertible transform이 가능할 것이다. : 90도 rotation Augmentation 후에 {0, 90, 180, 270}가 각각 20%, 50%, 15%, 15% 있다면, original distribution은 {0, 90, 180, 270} - 50%, 15%, 15%, 20% 일 것

StyleGAN-ada

ADA에서도 Ambient-GAN과 마찬가지로 $p$를 높게 설정하면 어떤 이미지에서 augment를 했는지 알 수 없다. 따라서 0.85 이하의 확률로 augmentation을 하도록 design하였다.

Designing non-leaking augmentations

  • (a) Isotropic image scaling : $p$에 관계 없이 non-leaking safe augmentation
  • (b) Random 90 rotation : $p$가 높으면 불안정, 0.8이하의 $p$를 사용하면 non-leaking

즉, augmentation을 모든 데이터에 대해 적용하지 않고 stochastic하게 적용한다면 leaking이 되지 않고 잘 학습된다.

2.3 Our augmentation pipeline

  • 본 논문은 18개의 Transformation pipeline을 사용하였다. (자세한 내용은 Appendix B 참고)

Effectiveness of stochastic discriminator augmentation by performing exhaustive sweeps over p for different augmentation categories and dataset sizes

Impact of stochasticity for augmentations

dataset의 크기나 augmentation categories에 따라 최적의 $p$가 다르다. 따라서 ADA에서는 Discriminator의 output에 따라 $p$를 Adaptive하게 바꿔주었다.

3. Adaptive Discriminator Augmentation

매번 augmentation probability $p$ 를 조정하는 것은 어렵다. 따라서 본 논문에서는 overfitting의 정도에 의거하여 동적으로 $p$를 조절한다.

  • Overfitting이 되면 validation set(green)은 generated images(orange)처럼 행동한다.

Overfitting Heuristics

\[r_{v}=\frac{\mathbb{E}\left[D_{\text {train }}\right]-\mathbb{E}\left[D_{\text {validation }}\right]}{\mathbb{E}\left[D_{\text {train }}\right]-\mathbb{E}\left[D_{\text {generated }}\right]} \quad r_{t}=\mathbb{E}\left[\operatorname{sign}\left(D_{\text {train }}\right)\right]\]
  • heuristic이 1이면 complete overfitting / 0이면 overfitting X
  • $r_{v}$ : validation set의 결과가 training set와 비슷한지, generated image와 비슷한지를 검사
  • $r_{t}$ : training data에서 D가 real이라고 훈련하는 데이터가 얼마나 있는지를 검사
  • $r_v$의 성능이 좋을 것 같지만 실제로 $r_t$의 성능이 더 좋음

본 논문에서는 처음에 $p$($=r_t or r_h$)를 0으로 초기화한 후에 4번의 minibatch마다 이 값을 조절한다. 따라서 Adaptive Discriminator Augumentation (ADA) 라고 불린다.

Control the Augmentation Strength $p$

  • If the heuristic (e.g. $r_v$ > 0.6 ) indicates too much overfitting, the augmentation probability $p$ is increased by a fixed amount.
  • If the heuristic (e.g. $r_v$ < 0.6 ) indicates too underfitting, the augmentation probability $p$ is decreased by a fixed amount
  • 논문에서는 비교하기 위한 $p$로 0.6의 고정된 값을 사용하였다. (hyper-parameter)

**Impact of target value in heuristics **

  • (a) : 2k의 data에서는 optimal $r_v = 0.7$, 10k에서는 optimal $r_v = 0.4 pr 0.5$, 50k에서의 optimal $r_v = 0.3$
    • data의 수에 따라 optimal한 augmentation probability가 다르다. (training의 개수가 클수록 overfitting이 덜되다보니 augmentation을 적게 해도 괜찮다.)
  • (b) : $r_t$를 사용하면 데이터의 크기와 상관없이 0.6이 best heuristic!
    • $r_t$가 $r_v$보다 안전하게 실험이 되므로 본 실험에서는 $r_t$를 heuristic으로 사용하였고 overfitting/underfitting을 확인하는 target value로는 0.6을 사용하였다.
  • (c) : heuristic $r_t$를 사용하여 training을 하면 $p$는 점점 증가한다. 이때 dataset이 적을수록 augmentation을 많이 하는게 효과가 좋으므로 $p$가 빠르게 증가한다.
  • (d) : ADA를 쓰지 않고 고정된 $p$값을 바탕으로 훈련을 하면 시간이 흐름에 따라 heuristic이 발산한다.(점선) 반면, ADA를 사용하여 학습을 하면(heuristic이 0.6보다 크면 augmentation이 적게 하고, heuristic이 0.6보다 작으면 aug를 많이 하는 식으로) heuristic이 (b)의 고정된 $p=0.6$로 수렴하게 된다.

Effect of adaptive discriminator augmentation

Figure1는 ADA를 적용하지 않았을 때의 결과값이고, Figure6은 ADA를 적용했을 때의 FID score이다.

Augmentation을 하면 D가 overfitting이 안되도록 학습을 할 수 있기 때문에 G가 더 다채로운 이미지를 만들 수 있다. 또한, ADA를 사용하면 loss function을 더 강력하게 학습할 수 있으므로 gradient field를 좀더 detail하게 유지할 수 있다.

4. Evaluation

4.1 Training from scratch

  • (a), (b) : 적은 dataset을 사용하더라도 ADA를 사용하면 FID 수치가 향상된다 !
  • (d) : 그림 (d)에서는 bCR-GAN을 사용했을 때 생성된 이미지(blurring)의 분포가 leaking이 됨을 보여준다. ADA는 생성이미지가 real image와 비슷한 것을 보아 leaking이 안된다.
  • (c) : ADA와 bCR-GAN을 같이 사용하면 좋다.

다른 augmentation methods(PA-GAN, WGAN-GP, zCR 등등)과 ADA를 비교해봐도 ADA가 좋다고 함

4.2 Transfer Learning

  • Freeze-D와 ADA를 같이 쓰면 Transfer-Learning이 잘됨

4.3 Small Datawets

  • ADA는 small dataset에도 효과적

댓글남기기

참고