이전에 계속 공부해오던 StyleGAN을 효과적으로 architecture을 재구성하는 방법이 소개되어 오랜만에 글을 쓰게 되었습니다 :)
몇개월 전 Nvidia에서 StyleGAN3를 발표했었는데, 1, 2, 3 시리즈를 계속 봐오면서 현재 나타나는 문제점을 파악하고 분석하며 이에 대한 효과적인 해결책을 제시하기까지의 과정이 정말 정교하고 체계적이라고 느꼈습니다. 그래서 StyleGAN 이라는 시리즈에 더 관심을 갖게 되네요 ㅎㅎ.. 단순히 논문에서 소개하는 지식 외에도 느끼는 점이나 얻어가는 점이 굉장히 많다고 생각합니다. (물론 이 논문이 이 시리즈에 속한다고 하기는 좀 그렇지만..)
Abstract
StyleGAN 모델은 이미지넷 데이터셋과 같이 크고 구조화되지 않은 데이터셋에서는 잘 작동하지 않습니다. 이에 대한 문제점으로 논문에서는 StyleGAN의 구조에 대해 언급합니다. (Progressive Growth Method를 사용하지 않는 StyleGAN3의 Generator, Pretrained Feature Extractor의 부재 등..)
그래서 이 논문에서는 StyleGAN의 구조를 다소 변경하고 굉장히 큰 네트워크로 구성하는 방법을 제시합니다.
Introduction
StyleGAN은 variation의 disentangle factors나 interpolation을 좀 더 원활하고 효과적으로 제어할 수 있다는 점에서 큰 장점을 가집니다. 하지만 다른 standard한 generator network들과는 다르게 복잡하고 diverse한 데이터셋을 훈련시키는데는 다소 안좋은 성능을 보입니다.
그래서 본 논문에서는 이러한 문제를 해결하기 위해 몇가지 요소를 수정한 architecture를 제안합니다.
- Progressive Growing Strategy
- Projected GAN에서 사용된 것과 같이 pretrained feature network 사용 (CNN, ViT 둘 다)
- Classifier guidance (DIffusion Model에서 사용된 방법)
Scaling StyleGAN to ImageNet
1. Adapting Regularization and Architectures
우선 ImageNet은 여러 클래스가 존재하는 데이터셋이기 때문에 class-conditional model로 구성을 해야합니다.
그리고 Generator는 StyleGAN3-T 모델을 차용합니다.
StyleGAN-R 모델도 사용해봤는데 rotational-equivariant 한 모델이기 때문에 복잡한 데이터셋에서는 대칭적인 이미지를 생성해내서 사용하지 않았다고 합니다.
Regularization
기존 StyleGAN 에서 사용되던 Path Length Regularization과 같은 Regularization 기법은 uni-modal dataset에서는 확실한 성능 향상을 이뤄냈지만, ImageNet과 같은 multi-modal datasets에서는 치명적이라고 합니다.
하지만 StyleGAN에서 제시한 Path Length Regularization의 경우 High Quality Inversion을 가능하게 하는 너무 매력적인 방법이라 최대한 효과적으로 사용하려고 했습니다.
모델이 어느정도 충분히 훈련된 후 (첫 200k image를 훈련 후) Path Length Regularization을 사용하면 위의 문제가 어느정도 해결되었다고 합니다.
Discriminator에서는 기본적으로 GP 없이 Spectral Normalization을 사용했고 첫 200k images를 훈련시킬 때는 \(\sigma = 2\)의 하이퍼파라미터로 가우시안 블러를 적용했습니다.
이 방법은 StyleGAN3-R 모델에서 제안된 방법입니다.
Low Dimensional Latent Space
기존 연구 (The Intrinsic Dimension of Images and Its Impact on Learning) 에서는 nature image datasets의 intrinsic dimension은 낮은 dimension과 연관이 있다고 합니다. ImageNet의 경우 40정도의 dimension 이구요.
StyleGAN에서는 latent code 의 dimension은 512로 사용하는데 이는 좀 과하기 때문에 64 dimension으로 낮춰서 사용합니다. 하지만 unwarp 시킨 style code w는 그대로 512 dimension으로 사용합니다.
Pretrained Class Embeddings
앞서 언급했듯이 ImageNet Dataset은 1000개의 클래스가 존재하기 때문에 GAN 모델에서 또한 클래스에 대한 정보를 제공해야합니다.
그래서 본 논문에서는 pretrained Efficientnet-lite0으로 lowest resolution features를 가져와서 class 당 mean을 계산해서 해당 class에 대한 embedding으로 사용합니다. 이러한 과정을 거친 Embedding vector를 z의 크기와 맞춰주기 위해 (imbalance를 피하기 위해) linear projection을 한번 더 거칩니다.
Reintroducing Progressive Growing
StyleGAN3에서 Progressive Growing을 사용하지 않은 이유는 interpolation 과정을 거칠 때 머릿결과 같은 texture가 해당 위치에 고정되버리는 texture sticking problem을 해결하기 위했습니다. 이에 대한 근본적인 이유는 aliasing 이었구요.
하지만 안정적으로 high resolution에 대한 image를 생성해낼 때 Progressive Growing 방법이 큰 도움이 되기 때문에 조심히 적용하는 방법을 제안합니다.
우선 16 x 16 이미지를 11개의 layer를 통해 학습 시킵니다. 그리고 그 다음 resolution을 학습시킬 때는 마지막 2개 layer는 없애고 나머지 레이어는 freeze 시킵니다. 그런 다음 7개 레이어를 붙여서 upsampling 시킵니다. 이러한 과정을 반복하고 마지막 resolution에 대해서는 7개 레이어가 아닌 5개 레이어만을 붙입니다.
1024 x 1024 이미지의 경우 총 39개 레이어를 사용하게 됩니다.
이러한 PG 형식의 Generator를 사용하는 반면에 Discriminator의 경우는 고정된 네트워크를 사용합니다.
대신 Multi Resolution의 결과를 내놓는 Generator의 output을 pretrained feature network의 input size에 맞춰 upsampling 시킵니다. (더 작은 resolution인 경우만.)
PG 방법의 경우 fixed growing schedule을 통해 늘려나가는데, 이런 schedule을 사용하는게 아니라 FID를 측정해가며 그만 줄어들 때 까지 훈련시킵니다.
그리고 작은 resolution의 이미지의 경우 (16 ~ 64) 2048 정도의 큰 배치사이즈를 사용할 때 이점이 있었고, 더 큰 resolution의 경우 더 작은 배치사이즈를 사용했습니다.
Exploiting Multiple Feature Networks
GAN의 performance를 극대화하기 위해 최근에는 pretrained network를 통해 feature extraction 하는 방법이 많이 연구되고 있습니다. 이런 아이디어를 차용해서 Projected GAN에서 사용된 CCM, CSM을 차용합니다.
이때 단순히 하나의 feature network만을 사용하는 것이 아닌 CNN 구조와 ViT 구조 둘 다 사용하여 feature extraction 하게 됩니다. (EffNet, DeiT-M)
Classifier Guidance for GANs
Diffusion model에서 class information을 주입하기 위해 사용된 classifier guidance 방법을 차용합니다.
이런 방법을 통해 gradient에서 추가적으로 class에 대한 information을 고려할 수 있습니다.
경험적으로, 앞단에서 Class Embedding으로 class에 대한 정보를 다 줬어도 해당 방법을 차용함으로써 이점이 있었습니다.
Experimental 결과입니니다. truncation 기법은 stylegan 방법과 동일합니다.
