Liu, Bingchen, et al. "Towards faster and stabilized gan training for high-fidelity few-shot image synthesis." International Conference on Learning Representations. 2020.
StyleGAN2와 같은 GAN 모델의 SOTA들을 학습하기 위해선 고성능의 GPU 들과 많은 양의 이미지들이 필요합니다. 하지만 현실 세계에서는 컴퓨팅 버짓이나 학습에 활용될 이미지들이 부족합니다. Pretrained model을 사용하는 것이 하나의 방법일 수 있지만, 생성하고자 하는 데이터에 적절한 모델이 없거나 fine-tuning한 경우 성능이 떨어질 수도 있습니다.본 논문에서는 하나의 GPU와 100개 미만의 이미지들로 몇 시간만에 고해상도(1024X1024)의 이미지들을 생성할 수 있는 모델을 제안합니다. 아래 그림은 제안하는 모델이 하나의 RTX 2080-Ti GPU, 1000여개의 이미지들로 짧은 시간 안에 생성한 이미지들입니다.
제안하는 모델의 핵심 아이디어는 skip-layer channel-wise excitation module과 self-supervised discriminator 입니다.
1. Skip-layer channel-wise excitation module (SLE)
고해상도의 이미지를 만들기 위해서 Generator는 깊어야 합니다. 깊을수록 모델의 파라미터 수가 많아지고 gradient flow가 약해지기 때문에 학습 시간이 더 소요됩니다. Skip-layer connection을 활용한 residual structure로 layer간 gradient flow를 강화할 수 있지만, 여전히 계산 비용이 많이 큽니다. 이를 개선하고자 본 논문에서는 Skip-layer를 재정의한 Skip-Layer Excitation(SLE) 모듈을 제안합니다.
위 내용을 자세히 적어보겠습니다. 기존 skip-layer connection이 GAN에 적용된 방식은 conv layer 의 출력인 feature map을 더하는 것 이었습니다. 아래 그림은 Wang et al.[1]의 연구에서 사용된 residual block의 방식을 나타냅니다. 먼저 저해상도의 이미지로 G1 네트워크를 학습하고, G1 네트워크 앞 뒤에 G2 네트워크를 붙여 고해상도의 이미지로 함께 학습합니다. 이때 G1과 G2의 같은 크기를 가지는 feature map을 element-wise addition을 하여 학습하게 됩니다.
본 연구에서 제안하는 Skip-Layer Excitation module은 element-wise addition 대신 channel-wise multiplication을 합니다. 이 방식을 통해 연산 수도 줄이고, 다른 크기의 feature map끼리 skip-layer를 적용할 수 있는 두 가지 장점이 있습니다. 아래 이미지의 왼쪽은 SLE 모듈을 설명하고 오른쪽은 본 연구의 Generator 구조를 나타냅니다.
먼저 오른쪽 이미지부터 보시면, 8^2 크기의 feature map과 128^2 크기의 feature map 간 SLE 연산을 적용하는 것을 볼 수 있습니다. 이렇게 다른 크기의 feature map 간 Skip-layer를 적용할 수 있는 것이 SLE 모듈의 장점 중 하나라고 볼 수 있습니다. 왼쪽 이미지는 64X128X128 feature map 과 512X8X8 feature map 간 SLE 연산을 나타낸 그림입니다. 먼저 512X8X8의 feature map을 64X128X128 feature map 채널 수에 맞추기 위하여 64로 맞추고, 연산을 줄이기 위하여 1X1로 spatial size를 압축합니다. 이때 pooling layer와 conv layer 등이 사용됩니다. 그 다음 64X128X128 feature map과 64X1X1 벡터 간 channel-wise multiplication을 합니다.
2. Self-supervised discriminator
적은 양의 이미지로 네트워크를 학습하는 경우 overfitting이 될 수 있습니다. 이를 방지하고 적은 양의 이미지로 성공적인 학습을 하기 위해 Self-supervised discriminator 를 도입합니다. Discriminator는 구조상 네트워크가 작아지는 모양이기 때문에 autoencoder의 encoder와 유사합니다. Discriminator의 중간 정도의 layer까지를 encoder로 보고, 다시 원본이미지보다 저해상도의 이미지로 복원하는 decoder를 붙여서 autoencoder의 reconstruction error를 계산하여 네트워크를 학습하는 아이디어입니다. 이 학습 과정에서 crop하거나 다운샘플링한 이미지를 사용하기 때문에 data augmentation의 효과가 있게 됩니다.
아래 그림은 본 연구의 discriminator 구조 입니다. 먼저 generator와 같이 SLE 모듈이 적용된 것(파란색 박스)을 알 수 있습니다. 또 GAN과 같이 generator와 discriminator가 경합적으로 학습되는 점은 동일합니다. 재밌는 점이 초록색 박스인 decoder입니다. Decoder는 입력 feature map을 16배 크기로 증가 시키게끔 되어 있고, 동일한 decoder가 discriminator의 16^2 feature map을 출력하는 layer와 8^2 feature map을 출력하는 layer에 달려 있습니다. 먼저 16^2의 feature map을 8^2로 crop한 후에 decoder로 넣어 128^2 이미지를 생성합니다. 이때 reconstruction loss 계산을 위해 실제 이미지에서 crop 한 데이터를 사용합니다. 그리고 8^2의 feature map은 그대로 decoder로 입력하여 128^2 이미지를 생성하고, 학습을 위해선 다운 샘플링한 데이터를 사용합니다.
AAE처럼 GAN에 autoencoder를 합친 것으로 보기는 어렵습니다. Generator를 decoder로 학습하지 않았고, 단순히 discriminator를 regularization하기 위하여 decoder를 활용하였기 때문입니다.
References
[1] Ting-Chun Wang, Ming-Yu Liu, Jun-Yan Zhu, Andrew Tao, Jan Kautz, and Bryan Catanzaro. High-resolution image synthesis and semantic manipulation with conditional gans. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 8798–8807, 2018.