이번에 소개해 드릴 논문은 A Style-Based Generator Architecture for Generative Adversarial Networks으로 StyleGAN으로 자연스러운 고해상도 이미지를 만들면서 많이 알려지게 된 논문입니다. 해당 논문은 NVDIA 팀이 CVPR 2019에서 발표한 논문입니다. 이 글은 나동빈 님의 유튜브 영상과 루닛의 블로그 글을 참고하여 작성했습니다.
Introduction
최근 GAN을 기반으로 한 이미지 합성 기술은 PGGAN 등을 포함하여 지속적으로 발전하고 있습니다. 그러나 Generator를 통한 이미지 합성 과정은 여전히 block box로 여겨지며, 이로 인해 합성되는 이미지의 attribute (성별, 연령, 헤어스타일 등) 을 조절하기가 매우 어렵다는 한계가 있습니다. 또한 생성되는 이미지 퀄리티가 불안정하여 보이시는 그림과 같이 부자연스러운 이미지도 다수 생성되어왔습니다.
StyleGAN은 이러한 문제를 해결하기위해 Style Transfer에 기반한 새로운 generator 구조를 제안했습니다. StyleGAN은 이미지를 style의 조합으로 보고, generator의 각 layer 마다 style 정보를 입히는 방식으로 이미지를 합성합니다. 이 때 각 layer에서 추가되는 style은 이미지의 coarse feature (성별, 포즈 등) 부터 fine detail (머리색, 피부톤 등) 까지 각기 다른 level의 visual attribute를 조절할 수 있습니다. 뿐만 아니라 StyleGAN은 기존의 방법들보다 훨씬 안정적이고 높은 퀄리티의 이미지를 생성할 수 있게 된 것이 Contribution이라고 볼 수 있습니다.
Background
GAN
논문에 대해 자세히 설명해 드리기 전에 제가 GAN에 대해 처음 접했기 때문에 아시는 분들도 있겠지만 배경 지식으로 설명해 드리겠습니다.
생성자와 판별자라는 두 개의 네트워크를 활용하는 것이 특징입니다. 이 때 판별자는 생성자를 학습시기키 위한 네트워크 입니다. 이런 생성자와 판별자가 서로 경쟁 관계에 있어서 동일한 목적함수를 이용해서 학습을 진행하는 과정에서 생성자가 퀄리티 높은 이미지를 만들어낼 수 있게 됩니다. 목적 함수는 다음과 같습니다. min max 게임이라고 해서 하나의 V가 있을 때, Discriminator는 커지는 방향으로 업데이트를 하고, 반면에 Generator는 이러한 식이 감소하는 형태로 네트워크를 업데이트 합니다.
이는 앞에 보이시는 Min과 Max가 의미하는 바 입니다. 동일한 함수에 대해서 한 쪽은 값을 낮추고 한 쪽은 값을 높이는 방식으로 업데이트를 진행합니다. 식에 대해 더 자세히 설명을 드리면, 거기서 매번 이미지 x를 Sampling해서 [logD(x)]라는 식에 넣습니다. 또한 z를 sampling하기 위한 distribution은 일반적으로 uniform이나 gaussian 분포를 사용합니다. 마찬가지로 z를 매번 sampling한 뒤에 식(log(1-D(G(z))에 넣어서 거기에 대한 기댓값(E) 즉 mean값을 구합니다. discriminator는 그러한 이미지인 x가 들어왔을 때 이때 x는 원본데이터의 분포를 따르기 때문에 이런 원본데이터에 대해서는 1을 도출할 수 있는 형태로 학습을 진행합니다.
결론적으로 G는 진짜 이미지를 진짜라고 분류하고 D는 가짜 이미지를 가짜 이미지로 분류 하도록 학습을 진행합니다. 그 과정에서 Generator가 만드는 이미지는 실제 이미지와 비슷하게 보이도록 학습이 되며 학습이 다 진행이 되면 이 discriminator는 더이상 실제이미지와 가짜이미지를 구분할 수 없게 되면서 output값으 0.5로 수렴하게 됩니다. Generator에서 sampling되는 image들은 이러한 real image distribution을 따를 수 있도록 수렴합니다.
그러나 GAN의 초창기 모습은 그렇게 성능이 뛰어나지 않았고, 네트워크가 수렴하는 것도 불안정했다고 합니다.
WGAN-GP
이러한 초창기 GAN모델이 가진 수렴의 불안정성을 해결하고자 다양한 Loss term들이 제안되었습니다. SoTA 중 하나인 WGAN-GP는 1-Lipshichtz(립시츠) 조건을 만족하도록 하여 안정적인 학습을 유도했습니다. StyleGAN Architecture에서도 WGAN-GP Loss를 사용했다고 합니다.
PGGAN
StyleGAN이 나오게 된 배경인 Progressive Growing of GANs 즉 PGGAN에 대해 설명해 드리겠습니다. PGGAN의 특징은 한번에 전체 네트워크를 다 구상하는 것이 아니라 학습 과정에서 레이어를 추가하면서 해상도를 높인다는 점입니다. 이러한 방법으로 한번에 네트워크를 만드는 것보다 훨씬 안정적으로 architecture를 만들 수 있었습니다. 위 소요 시간을 보시면 며칠이 걸리는 training 과정을 거쳐야 하지만, 고해상도 이미지를 만들 수 있었습니다. 그러나 PGGAN은 각 특징들이 분리되어 있지 않기 때문에 다른 특징까지 개입되는 문제가 발생하곤 했습니다. 이러한 문제를 StyleGAN 개선했다고 할 수 있습니다.
StyleGAN은 PGGAN을 baseline으로 삼고 있습니다. PGGAN은 점진적으로 네트워크를 붙여나가는 방식으로 진행했다는 것이 특징입니다.
생성자 부분을 보시면, 4x4의 레이어부터 출발해서 1024x1024라는 고해상도 이미지로 발전해 나갑니다. 판별자도 마찬가지로 생성자와 대칭으로 점진적으로 진행이 됩니다. 이와 함께 WGAN-GP Loss를 사용해서 안정적인 학습또한 가능하게 했습니다.
Propesed Method
Mapping Network
StyleGAN의 가장 큰 특징은 input vector Z로부터 직접 이미지를 생성하는 것이 아니라, mapping network를 거쳐 intermediate vector W로 먼저 변환한 후 이미지를 생성한다는 것입니다. 기존의 방법에서처럼 input vector로부터 이미지를 직접 생성할 경우, (a) 그림을 보시면 고정된 input distribution에 학습 이미지의 distribution 을 맞춰야 한다는 한계가 있습니다. 이로 인해 (b) 그림처럼 visual attribute가 input space에 non-linear하게 mapping되고, input vector로 visual attribute를 조절하기가 매우 어려워집니다. 예를 들어, 학습 이미지에 검은 머리의 사람이 대부분일 경우, input space의 대부분의 영역이 검은 머리를 표현하게되고, input vector를 조절하여 머리 색을 변경하기가 어려워집니다. 그러나 그(c) 그림을 보시면 mapping network를 사용할 경우 w는 고정된 distribution을 따를 필요가 없어지기 때문에, 학습 데이터를 훨씬 유동적인 공간 (intermediate latent space)에 mapping할 수 있고 w를 이용하여 visual attribute를 조절하기가 훨씬 용이해집니다 . 이러한 특징을 disentanglement라고 합니다.
Style Modules
ADaIN이라는 Layer를 활용해 style 정보를 입힐 수 있도록 만들었습니다. 하나의 이미지를 생성할 때 여러개의 스타일 정보가 레이어를 거칠 때마다 입혀질 수 있도록 하는 방식입니다. ADaIN은 학습시킬 파라미터가 필요하지 않습니다. (𝛾와 𝛽사용하지 않음) feed-forward 방식의 style transfer 네트워크에서 사용되어 좋은 성능을 보인 방식입니다. 위 그림 중 Instance Normalization은 하나의 이미지에 대해서 정규화를 수행하기 때문에 Batch size가 N이라고 채널 단위로 정규화를 수행하는 것을 알 수 있습니다. 이렇게 정규화를 수행한 뒤 별도의 스타일 정보를 입력 받아서 feature 상의 statistics를 바꾸는 방식으로 style을 입히는 것이 Adaptive Instance Normalization이라고 보시면 됩니다.
앞서 설명한 Style Module 즉 AdaIN Layer은 위 같은 방식으로 들어갑니다. 이 매핑 네트워크를 통해 Latent vector인 w를 뽑은 다음에 synthesis network가 이미지를 생성하는 과정에서 style을 입히는 데 사용됩니다. 이때 Latent vector인 W는 affine transformation을 거쳐서 AdaIN에 들어가는 스타일 정보를 만듭니다. 생성 네트워크는 4x4x512 constant tensor를 convolution, upsampling을 통해 1024x1024x3 이미지로 변환시킵니다. 이 때 각 convolution layer 이후마다 AdaIN을 통해 style이 입혀집니다.
이러한 구조는 Synthesis network의 매 layer마다 AdaIN을 통해 style을 normalize 한 후 새로운 style을 입히게 되므로, 특정 layer에서 입혀진 style은 바로 다음 convolutional layer에만 영향을 끼칩니다. 따라서 각 layer의 style이 특정한 visual attribute만 담당하는 것이 쉬워집니다.
Noise
Stocastic Varidation도 처리할 수 있도록 네트워크를 구성했습니다. 하나의 이미지에 포함될 수 있는 다양한 확률적인 측면을 컨트롤할 수 있도록 만드는 것입니다. 같은 사람에 대한 이미지라 할지라도 머리카락, 수염, 주름 등 stochastic 하다고 볼 수 있는 요소가 많이 존재하며, 이러한 요소는 이미지를 더욱 디테일하고 사실적으로 보이게 만들어줍니다. StyleGAN에서는 이를 표현하기 위해 synthesis network의 각 layer마다 random noise를 추가하였습니다. 이렇게 stochastic한 정보를 따로 추가해주면 더욱 사실적인 이미지를 생성하게 될 뿐 아니라, input latent vector는 이미지의 중요한 정보(성별, 인종, 헤어스타일 등)를 표현하는 데에만 집중할 수 있게 되고 이를 조절하는 것도 더욱 쉬워집니다.
스타일은 high-level global attributes를 컨트롤 하며, noise는 stochastic variation을 control할 수 있게 합니다. (a)는 모든 레이어에 노이즈를 적용한 것이고, (b)는 stochastic한 정보를 적용하지 않은 것이고, (c)는 뒷쪽에 layer에만 noise를 넣어준 것이며 (d)는 앞쪽에 noise 를 넣어주어서 더 큰 규모의 stochastic한 정보들이 표현된 것을 알 수 있습니다.
본 논문은 입력되는 layer의 위치에 따라서 해당 style에 미치는 영향력의 규모가 달라질 수 있다는 점을 언급합니다. coarse, middle, fine style로 나뉩니다. Latent Vector는 18 x 512 차원의 행렬 벡터입니다. 앞쪽 4개의 latent vector가 coarse style이며 그 다음이 middle이며 그 다음이 fine style을 담당합니다.
Experiments
FID
이러한 StyleGAN은 두 데이터셋 (CelebA-HQ, Flickr-Faces-HQ (FFHQ))에서 state-of-the-art 성능 (Frèchet inception distance (FID) score)을 기록했습니다. 여기서 사용된 FID는 GAN을 사용해 생성된 이미지의 집합과 실제 생성하고자 하는 클래스 데이터의 분포의 거리를 계산한 것이며, 거리가 가까울 수록 즉 숫자가 낮을 수록 좋습니다. 논문에서 제시한 A부터 F가지의 요소들 (style mapping, constant input, adding noise, style mixing)이 성능 향상에 도움이 됨을 보였습니다. (F) Mixing Regularization은 Style Mixing이라고 하는데 이는 아래에서 더 자세히 설명해 드리겠습니다.
그래서 각 데이터 셋에 적용한 결과이며 A는 Baseline PGGAN이며, B에서는 interpolation technique을 적용할 수 있었습니다. C에서는 mapping network와 AdaIN layer를 이용해서 성능을 더 올릴 수가 있었고 D 에서는 초기 input layer로서 4x4x 512 차원의 상수 텐서값을 사용했을 때 성능이 올라갑니다. E와 F에서 Noise input을 넣고 Mixing regulariation를 추가했을 때 더 성능이 올라간 것을 확인할 수 있습니다.
Style Mixing
이전에 설명드린 Mapping network으로부터 나온 intermediate vector w는 synthesis network의 모든 layer에 대한 style을 표현하도록 학습되는데, 이는 여러 layer에 대한 style이 correlate 되는 위험을 발생시킬 수 있습니다. 논문에서는 이러한 문제를 해결하기 위해 다음과 같은 style mixing 기법을 제시하였습니다. StyleGAN을 학습 할 때, 우선 두개의 input vector로부터 w_1, w_2를 만들어냅니다. 그리고 synthesis network의 초반에는 w_1로 학습하다가 특정 layer 이후부터는 w_2를 적용하여 학습합니다. 이 때 style이 교체되는 layer를 매번 random하게 결정해줌으로써 연속된 두 layer간의 style의 correlate 되는 현상을 방지할 수 있습니다 (dropout과 비슷한 원리). 이렇게 학습된 모델은 아래 그림과 같이 각 layer가 담당하는 style이 뚜렷하게 구분되는 것을 확인할 수 있습니다.
Disentanglement studies
또한, 본 논문은 disentanglement 측면에서 2가지 성능 지표를 제안했습니다. Path Length는 두 벡터를 보간(interpolation)할 때 얼마나 급격하게 이미지 특징이 바뀌는지를 측정하는 지표입니다. 이미지가 잘 disentanglement되어 있다고 하면 image가 급격하게 바뀌지 않고 서서히 자연스러운 이미지로 바바뀐다고 합니다. 그래서 이 path length를 숫자가 낮을 수록 좋은 성능이라고 할 수 있습니다. 그리고 linear serability도 측정할 수 있습니다. Separability는 latent space에서 attributes가 얼마나 선형적으로 분류될 수 있는지 평가하는 지표입니다. 그래서 테이블을 보시면 w공간으로 먼저 mapping을 한 다음에 이 w공간에서 interpolation을 하든 다양한 스타일 정보를 바꾸는 것이 이 z 공간에서 latent vector를 뽑은 다음에 control하는 것보다 훨씬 이상적인 결과를 낼 수 있었다고 합니다. 결론적으로 다양한 지표를 통해서 w공간이 z공간보다 이상적인 성질을 가지고 있음을 알 수 있었습니다.
Conclusion
이 논문은 이미지 합성을 위해 style에 기반한 새로운 generator 구조를 제시하였고 높은 성능 향상을 보여주었습니다. 합성되는 이미지의 구성요소를 조정하는 것이 가능하며, 나아가 CNN에서 다양한 visual attribute를 어떻게 encoding하는지에 대한 이해도 증가에도 기여했습니다. generator 구조에만 집중하였기 때문에 dicsriminator 구조나 loss function등의 변경을 통해 더욱 개선될 여지도 보여진다고 결론 지을 수 있습니다.
댓글