StyleGAN: Style transfer와 mapping network를 사용하여 disentanglement를 향상시킨 generative Model

 

이번 글에서는 일명 StyleGAN이라고 불리는 'A Style-Based Generator Architecture for Generative Adversarial Network' 논문에 관해 핵심 위주로 리뷰 해 보고자 한다. StyleGAN은 출판 당시에도 회자되었던 논문이고, 특히 논문에서 추가로 제공한 StyleGAN을 통해 생성한 고해상도의 이미지 데이터 셋인 FFHQ(Flickr-Faces-HQ)는 이 글을 쓰는 시점에도 많이 쓰인다고 알려져 있다. 사실 StyleGAN 논문을 처음부터 리뷰 해야겠다고 생각한 것은 아니다. GRAF(Generative Radiance Fields for 3D-Aware Image Synthesis) 논문을 공부하다가 기존에 제안된 HoloGAN의 한계를 지적하는 내용이 나와서 preliminary를 알기 위해 HoloGAN 논문을 찾아봤고, 그 논문에서 HoloGAN에 StyleGAN에서 사용한 style transfer의 방법을 적용했다고 서술한 것이 StyleGAN 논문을 공부하게 된 계기였다. HoloGAN에서도 StyleGAN과 유사하게 style에 관한 latent vector를 각 layer에 주입시킨 후, AdaIN(Adaptive Instance Normalization)을 적용하여 좀 더 disentangled representation을 학습할 수 있도록 했다는 게 특징이다. 이처럼 StyleGAN은 당시 기존의 GAN보다 전반적인 성능에서 우위를 보이는 성과를 냈다고 알려져 있으며, 특히 generator에서 style transfer를 통해 style에 관한 좀 더 linearly separated latent factor를 이미지 생성 과정에서 점차 단계적으로 적용하여 여러 style이 혼합된 새로운 이미지를 생성한다는 점이 주요 contribution 중 하나이다. 저명한 논문이므로 이미 유튜브나 블로그에서 리뷰가 많이 작성되어 있으므로, 개인적으로 좀 더 섬세한 이해를 필요로 하는 부분을 위주로 리뷰를 써 보고자 한다.

 

Pytorch 코드로 어떻게 구현될 수 있는지를 이해하기 위해 from scratch로 구현된 코드를 reference 삼아서 작성한 코드 snippet을 바탕으로 StyleGAN의 각 부분에 관해 이해해 보고자 한다.

 

[Code Reference]

https://github.com/SiskonEmilia/StyleGAN-PyTorch

GitHub - SiskonEmilia/StyleGAN-PyTorch: A simple implementation of Style-Based Generator Architecture for Generative Adversarial

 

[Code from scratch]

https://github.com/Glanceyes/ML-Paper-Review/blob/main/ComputerVision/GenerativeAdversarialNetwork/StyleGAN/StyleGAN.ipynb

GitHub - Glanceyes/ML-Paper-Review: Implementation of ML&DL models in machine learning that I have studied and written source co

 

 

 

StyleGAN 요약

 

1. Generator에서 latent vector를 sampling 할 때 단순히 input latent space인 $\mathcal{Z}$에서 뽑은 vector $z$를 그대로 사용하지 않고, 이를 mapping network에 통과시켜서 $\mathcal{W}$에서의 intermediate latent vector $w$를 사용했다.

2. 기존의 PGGAN(Progressive Growing of GANs for Improved Quality, Stability, and Variation) 모델을 baseline으로 해서 feature map의 너비와 높이는 점차 늘려가고, channel은 점점 줄이면서 output layer에서 RGB의 3개의 channel을 지니는 고해상도 이미지를 생성할 수 있도록 한다.

3. Generator의 첫 번째 layer의 input으로 latent vector $z$를 넣지 않고, 학습된 constant tensor를 넣는다. 또한 style transfer에서 착안하여 latent vector $w$를 generator의 각각의 convolution layer의 output에 주입함으로써 좀 더 구분된(disentangled) style 정보를 반영한다. 이때, layer의 output인 feature map에 style로 들어오는 정보를 가지고 AdaIN(Adaptive Instance Normalization)을 적용하여 feature map의 channel에 관하여 기존의 통계적인 정보를 stye에 맞게 변화시키며, PGGAN baseline에 따라 점차 feature map의 해상도가 커짐으로써 coarses style부터 fine style까지 서로 다른 scale의 style을 조절할 수 있도록 한다.

4. Stochastic variation을 부여할 수 있도록 noise 정보를 매번 convolution layer의 output에 적용한다. 이때, PGGAN baseline에 의해 feature map의 크기가 block을 거치면 거칠수록 점차 커지므로, feature map의 크기에 비례하여 noise를 픽셀 단위로 적용 한다.

5. 기존에는 GAN 모델의 disentanglement를 측정하는 방법으로 StyleGAN을 평가하는 건 한계가 있어서 이 논문에서는 수치적으로 이를 더 잘 측정할 수 있는 두 가지의 방법을 제안했는데, 바로 perceptual path length와 lienar separability이다. Perceptual path length는 두 개의 latent vector에 관해 interpolation 할 때 사람이 인지하기에 얼마나 자연스럽게 변하는지를 측정하는 지표이고, linear separability에서는 선형 분류기를 사용하여 서로 다른 특징들이 얼마나 잘 선형적으로 분류되는지를 측정할 수 있는 entropy 기반의 separability score를 보였다.

 

 

 

Introduction

 

Background

 

[출처] https://arxiv.org/abs/1710.10196, Progressive Growing of GANs for Improved Quality, Stability, and Variation

 

당시 GAN 모델을 통해 생성되는 이미지의 해상도와 품질에서 급격한 향상을 보였지만, 정작 generator가 어떻게 새로운 이미지를 생성하는지와 관련해서는 black box에 가까운 면이 없지 않아 있었고, 그러한 이미지 생성에 있어서 확률적으로 생성되는 feature의 기원이 무엇인지 등 image synthesis 과정의 다양한 면에 관한 이해가 부족했다고 저자는 설명한다. 이는 StyleGAN이 베이스라인으로 하는 PGGAN을 고려하면 이해가 가는데, PGGAN은 너비와 높이를 증가시키고 채널의 크기는 줄인 레이어를 점차 쌓아 나가면서 고해상도와 높은 품질의 새로운 이미지를 생성할 수 있는 generative model이다. 그러나 하나의 latent vector로부터 여러 개 쌓은 convolution layer를 통과하면서 style을 구성하는 각각의 요소를 구분하여 이들을 별도로 control 하기가 어렵고, 확률적으로 생성되는 stochastic feature가 궁극적으로 어디서부터 나오는 것인지를 분석하기가 힘들다는 단점이 있다.

또한 보통 GAN은 generator에서 latent distribution에서 latent vector를 샘플링한 후 이를 가지고 새로운 이미지를 생성하는데, 그러한 latent space에 관한 이해가 당시에 부족했다고 서술한다. 그리고 당시 자주 쓰인 latent space interpolation만으로는 서로 다른 generator들을 수치적으로 비교할 수 없는 한계가 존재했다. 즉, 어떤 한 generative model의 latent space에서 각각의 subspace가 얼마나 linear 하게 잘 분리되는지를 측정하는 데 있어서 수치로 보이는 데 어려움이 있었다는 것이다.

 

 

 

StyleGAN의 Contribution

 

[출처] https://arxiv.org/abs/1812.04948, A Style-Based Generator Architecture for Generative Adversarial Networks

 

저자는 이전에 AdaIN을 사용하여 구현한 style transfer 논문의 방법을 착안하여 image synthesis process를 control 할 수 있는 새로운 방법을 제시하고자 generator를 재설계했다고 언급한다. StyleGAN은 학습된 constant input부터 시작하여 각각의 convolution layer 마다 latent code에서 나오는 style을 적용함으로써 각기 다른 scale에서 image feature의 strength를 직접 조절 가능하게 한다. 또한 네트워크의 각 layer에 직접 noise를 주입시키는 방법을 결합해서 자동화되고, 어떠한 속성이 정확히 무엇인지를 직접 관장하지 않고 비지도적으로 high-level 속성을 stochastic variation으로부터 분리할 수 있게 한다. 쉽게 말하자면, 모델이 스스로 pose, identity 등 high-level 속성과 noise를 구분하여 학습할 수 있다는 의미이다. 

본 논문의 가장 중요한 contribution 중 하나는 바로 기존의 연구들에서 흔히 진행되었던 것처럼 discriminator나 loss function을 수정하지 않고 단지 generator 디자인 자체에 집중했다는 점이다. 우리가 흔히 GAN 모델에서 generator와 discriminator 두 모델을 모두 학습하지만, 실제로 새로운 이미지를 생성할 때는 generator만을 사용한다는 점을 비추어 볼 때, 이제까지 제안된 여러 loss function에서 추가로 새로운 loss function을 제안하기 보다는 generator를 재설계하는 게 더 합리적일 수 있다.

 

StyleGAN의 generator는 input latent code인 $z$를 intermediate latent space인 $\mathcal{W}$로 보낸 latent code인 $w$를 사용했다는 점이 주목할 만한데, 이것이 생각보다 네트워크에서 variation의 factor가 어떻게 표현되는지에 심오한 영향을 끼친다고 말한다. 대개 input latent space의 분포는 학습 데이터의 확률밀도 함수를 따라가는 경우가 많은데, 일반적인 GAN처럼 latent space에서 Gaussian distribution으로 random 하게 latent code를 샘플링하면 그 latent 요소들이 entangled 될 수 있다고 한다. 즉, 우리는 일반적으로 latent space의 분포를 모르므로 종종 Gaussian distribution에서 latent code를 샘플링하는데, 실제 latent space 분포는 학습 데이터의 확률밀도 함수와 유사하므로 과연 latent space를 Gaussian distribution로 가정하여 직접 latent code를 뽑는 게 바람직하냐는 것이다.

그래서 저자는 input latent code를 intermediate latent code로 mapping 하는 network를 추가하여 intermediate latent space로 임베딩하는 것이 앞서 설명한 제한에서 자유로우면서 latent space에서 어느 정도 disentanglement도 허용할 수 있지 않느냐는 것이다. 그래서 저자는 StyleGAN의 intermediate latent space에서 얼마나 disentanglement가 잘 되는지를 측정하여 이를 증명하고자 했는데, 기존에 제안된 측정 방법으로는 한계가 있어서(왜 한계가 있는지는 후술할 예정) 논문에서 두 가지의 자동화된 새로운 측정 지표를 제안했다. 이는 perceptual path length와 linear separability인데, 이에 관해서는 자세히 후술할 것이다.

 

마지막으로 저자는 높은 품질을 지닌 고해상도의 사람 얼굴 이미지 데이터셋인 FFHQ(Flickr-Faces-HQ)를 제안했는데, 이는 기존에 존재하던 high-resolution dataset 보다 더 넓은 다양성을 지닌다고 설명하면서 기학습된 StyleGAN 모델과 source code와 함께 무료로 공개하였다. 

 

 

 

Style-based Generator

 

PGGAN baseline 응용

 

 

기존에는 latent space에서 random 하게 샘플링 한 latent code가 generator의 input layer로 들어가서 feed-forward를 통해 이후 연속되는 layer를 거치면서 latent factor가 반영된 새로운 이미지를 생성했다. 그러나 이러한 방법은 앞서 말한 것처럼 representation이 entangled 되어서 어떤 특정 요소를 개별적으로 조절하기가 어려워진다는 단점이 있다. 그래서 저자는 style에 대응되는 샘플링 된 latent code를 layer 마다 적용하여 좀 더 disentangled representation을 학습할 수 있도록 하며, 대신 latent code를 layer 마다 주입하므로 굳이 기존처럼 처음에 latent code를 input layer에 넣어서 통과시키는 과정이 필요가 없어졌다. 왜냐하면 layer를 통과하면서 단계별로 style을 계속 주입하므로 결과적으로 여러 style 정보가 반영된 복합적인 이미지를 생성하는 데 문제가 없어서다. 그래서 StyleGAN의 generator에서는 latent code가 generator의 input layer에 들어가는 대신에 input layer 자체를 generator에서 제외하고, 4×4×512 차원의 학습된 constant tensor가 바로 generator에 들어가게 했다. 다시 말해, 학습된 하나의 상수 이미지가 input으로 들어간 후, style transfer에 의해 latent code의 style이 적용되도록 한 것이다. 

 

 

 

StyleGAN 모델 구조

 

 

Latent code인 $z$를 latent space인 $\mathcal{Z}$에서 샘플링한 후, non-linear mapping network인 $f$를 통해 latent space인 $\mathcal{Z}$를 $\mathcal{W}$로 mapping 하여 intermediate latent code인 $w \in \mathcal{W}$를 생성한다. 논문에서는 $\mathcal{Z}$와 $\mathcal{W}$의 차원을 512로 동일하게 설정했고, $f$는 8개의 MLP로 구성했다. 이후 intermediate latent code인 $w$가 18개로 복제되고, 이 18개의 512 차원의 vector가 synthesis network인 $g$의 각 convolution layer 뒤에 연결되어 style transfer가 수행된다. 이 synthesis network $g$는 기존의 PGGAN 모델의 디자인과 유사하게 feature map의 너비와 높이는 점점 커지고 채널의 크기는 줄어들게끔 여러 block이 쌓인 형태로 구성되며, 논문에서는 9개의 block으로 구성했다.

첫 block을 제외한 각각의 모든 block의 앞에 bilinear upsampling을 수행하는 layer를 추가하여 feature map의 너비와 높이가 점차 커지게 구성했으며, 각각의 block은 두 개의 3×3 convolution layer로 구성된다. 논문에서는 feature map의 너비와 높이가 4×4에서 1024×1024로 커지게 구성했다. Convolution layer의 kernel 크기는 3×3으로 동일한 상태에서 이미지 feature map의 크기는 점차 커지므로, 앞쪽 block에서는 style transfer 수행 시 좀 더 coarse style이 반영되고 뒤쪽 block으로 갈수록 fine style이 반영될 것이다.

이 부분의 이해가 어려울 수 있는데, 쉽게 생각해서 그림판으로 256×256 크기의 그림을 가지고 그 위에 3×3 크기의 marquee tool을 그려서 그 영역 안의 그림 색상만 변경한다고 가정해 보자. 그림을 4×4로 축소해서 3×3 크기의 marquee tool 영역 안의 색상을 변경하는 것과 그림을 1024×1024로 확대해서 적용하는 것 중에서 어떠한 것이 그림의 전반적인 색상을 변화시키는 방법일까? 바로 전자이다. 그러면 그림의 세부적인 부분의 색상을 변화시키는 방법은 후자일 것이다. 이를 scale과 연관지어서 설명할 수 있는데, 저자는 이 부분을 가리켜 "scale-specific"이라고 표현한다. StyleGAN에서는 18개의 style 중에서 앞쪽의 4개의 style을 coarse style, 중간의 4개의 style을 middle style, 뒤쪽의 10개의 style을 fine style로 정의했다. 컴퓨터비전에서는 scale을 비롯하여 coarse와 fine에 관한 개념이 자주 등장하므로 이를 직관으로라도 한 번 정리해보고 넘어가는 것이 유용하다. 

마지막 레이어에서는 각 픽셀의 RGB 값을 예측해야 하므로, 채널을 3으로 만드는 1×1 convolution layer를 통과하도록 한다. 정리하자면, StyleGAN의 generator에서는 latent code로 뽑은 style의 정보를 각기 다른 scale을 관여하는 block에 적용하여 coarse style부터 fine style를 구분하여 학습하는 동시에 결과적으로 이 모든 style이 혼합된 이미지를 생성한다. 

 

 

Style Transfer와 AdaIN

 

이제 StyleGAN에서 style transfer를 수행하는 방법에 관해 구체적으로 분석해 보자. StyleGAN에서 사용한 style transfer는 무언가 새로운 것은 아니고 이전에 AdaIN(Adaptive Instance Normalization)을 통해 style transfer를 구현한 논문의 방법을 차용한 것이다. Style을 정의하는 방법은 gram matrix 등 여러 가지가 존재하지만, 그중에서 어떤 한 이미지에 다른 이미지의 스타일을 적용하는 과정을 다음과 같이 AdaIN으로 정의할 수 있다.

 

 

이 식을 해석하기 전에 각각의 기호의 의미를 살펴보자. $\mathbf{x}_i$는 기존의 어떤 한 이미지 feature map에서 $i$번째 channel의 feature vector이고, $\mathbf{y}_{s, i}$와 $\mathbf{y}_{b,i}$는 각각 style의 scale과 bias 요소이다. 이를 해석하면 기존의 feature map 데이터를 각각의 channel에 관하여 정규화시키고, 여기서 각 채널마다 대응되는 style의 scale을 곱하고 bias를 더하여 각각이 채널마다 기존의 feature map의 분포를 정규화하여 style의 분포로 맞추겠다는 것이다.

 

 

 

위의 AdaIN을 상기하면서 StyleGAN에서의 style transfer를 살펴보자. 각각의 convolution layer로 style 정보인 latent code $w$가 들어가기 전에, 학습된 affine transformation인 $\mathrm{A}$를 통해 $w$를 style을 나타내는 vector인 $\mathbf{y}$로 변환한다. $\mathbf{y}$는 $\mathbf{y}_{s, i}$와 $\mathbf{y}_{b,i}$로 구성되는데, 이들은 기존의 feature map 정보를 AdaIN을 통해 얼마나 변화시킬지를 결정짓는 scale과 bias 요소이며, 마치 두 요소가 concatenated 된 것과 같다. 즉, synthesis network에서 convolution layer를 통과하면서 나오는 output인 feature map을 각각의 channel에 관해 정규화하고, 이를 style 분포로 맞추기 위해 scaling 하고 bias를 더하여 style transfer를 완료한다. $\mathbf{y}_{s, i}$와 $\mathbf{y}_{b,i}$ 각각의 차원은 $\mathbf{x}_i$ 차원과 동일해야 하므로 이 두 요소가 concatenated 되어 있는 $\mathbf{y}$의 차원은 $\mathbf{x}_i$ 차원의 두 배여야 한다.

기존의 AdaIN과 함께 제안된 style transfer와의 차이라면, StyleGAN에서는 latent code에 대응되는 style로부터 $\mathbf{y}$를 가져온 데 반해 기존의 style transfer에서는 style로서 적용하고자 하는 어떤 예시 이미지에서 그러한 style vector를 가져온 점이다. 그래서 기존의 style transfer와 완전히 일치하는 것은 아니지만 같은 매커니즘을 사용하고 의미 상 유사하므로 $\mathbf{y}$에 style이라는 이름을 붙인 것이다.

 

먼저 $w$를 affine transform $A$를 통과시켜서 style mixing을 적용하기 위한 style distribution 정보의 scaling과 bias를 얻는 코드는 다음과 같이 구현할 수 있다.

 

# Apply affine transformation A to w.
# Input: w
# Shape of w: (n_channel)
# Output: style = [y_si, y_bi]
# Shape of style: (2 * n_channel)

class LearnedA(nn.Module):
    def __init__(self, dim_latent, n_channel):
        super().__init__()
        self.transform = ScaledLinear(dim_latent, n_channel * 2)
        self.transform.linear.bias.data[:n_channel] = 1
        self.transform.linear.bias.data[n_channel:] = 0

    def forward(self, w):
        style = self.transform(w).unsqueeze(2).unsqueeze(3)
        return style

 

$W$에서의 차원인 `dim_latent`를 linear layer를 통해 synthesis block에서의 데이터 channel 차원인 `n_channel`의 두 배로 mapping 시키는 과정을 수행하는 부분이다. 그러면 batch마다 각각 `n_channel` 채널을 지니는 style의 scaling와 style bias를 얻을 수 있다. 이는 뒤에서 `split` 또는 `chunk`를 통해 두 부분으로 나누면 된다.

 

class AdaIN(nn.Module):
    def __init__(self, n_channel):
        super().__init__()
        self.norm = nn.InstanceNorm2d(n_channel)
    
    def forward(self, image, style):
        # Split style tensor to two parts along channel dimension.
        scaling, bias = style.chunk(2, dim=1)
        image_norm = self.norm(image)
        result = scaling * image_norm + bias
        return result

 

앞서 구한 affine transformation의 적용 결과는 style의 scaling과 bias가 채널 차원에서 연결되어 있는데, 이를 `chunk`를 통해 채널 차원에서 두 부분으로 나눈 걸 각각 scaling, bias에 할당한다. 그리고 기존의 `image` 데이터를 각 차원마다 instance normalization 한다. 그러면 `InstanceNorm2d`를 통해 channel 차원에 관하여 정규화를 수행하는 instance normalization을 하면 각 채널별로 정규화된 결괄르 얻는 것이고, 여기에 scaling를 곱하고 bias를 더하면 그 정규화한 결과를 style의 분포로 맞춰주게 된다. 

 

 

Noise 추가

 

 

마지막으로, generator에서 synthesis network를 구성하는 convolution layer의 output에 관하여 AdaIN을 수행하기 전에 명시적으로 noise를 주입함으로써 stochastic detail을 부여한다. 이 noise는 서로 연관되지 않은 Gaussian noise로 구성된 single channel 이미지들이며, synthesis network에서 이러한 이미지들을 각각의 convolution layer의 output에 넣는 것이다. 단, convolution layer를 통과해서 나온 output은 다른 block에 위치해 있으므로 feature map의 크기가 달라서 서로 다른 scale 정보를 지니고 있을 것이다. 그런데 noise를 적용할 때도 이러한 feature map의 크기에 맞춰줘야 하므로 이러한 noise 이미지가 각각의 feature map 단계에 관해 학습된 scaling factor에 의해 feature map 크기로 broadcasting 되어 output에 더해진다. 

 

 

저자는 StyleGAN의 성능이 기존의 단순한 PGGAN의 베이스라인만을 사용했을 때보다 얼마나 더 우수한지를 보이기 위해 마치 ablation study처럼 StyleGAN에서 사용된 방법들을 하나씩 추가해 가는 방법으로 생성 작업을 수행했을 때의 결과를 표로 정리했다. 성능 측정 지표는 FID(Fréchet Inception Distance)라고 해서 generator가 생성한 결과의 품질을 측정하기 위한 전형적인 방법을 사용했으며, 이 지표 값이 낮을수록 모델의 성능이 좋다. 그리고 실험 데이터로는 CelebA-HQ와 FFHQ의 두 가지 데이터셋을 사용했는데, 후자는 앞서 말한 것처럼 저자가 논문과 함께 공개한 고해상도의 인물 이미지 데이터셋이다. (A)는 단순한 PGGAN의 베이스라인일 때의 결과, (B)는 upsampling 수행 시 bilinear interpolation을 사용한 것, (C)는 mapping network와 style transfer를 적용한 것, (D)는 기존 방식처럼 generator의 input을 latent code로 하는 것이 아니라 학습된 constant tensor로 교체한 것, (E)는 noise를 적용한 것, (F)는 후술할 mixing regularization 방법을 나타낸다. 확실히 E 또는 F까지 적용했을 때의 FID 값이 낮은 것을 확인할 수 있다.

 

 

Mapping network와 affine transformation은 학습된 distribution으로부터 각각의 style에 관한 샘플을 그리는 것으로 이해할 수 있고, synthesis network에서는 이러한 style들의 집합을 바탕으로 새로운 이미지를 생성하는 역할을 한다고 볼 수 있다. 저자는 이러한 각각의 style의 효과가 네트워크에서 localized 된다고 말하는데, 이는 스타일의 특정 일부분만을 수정하는 것이 이미지를 생성하는 데 있어서 오직 일부분에만 영향을 끼칠 것으로 예상할 수 있다는 것이다. 왜 이렇게 생각할 수 있는지에 관해 저자는 AdaIN의 과정을 근거로 들고 있다. Style transfer에 의해 새롭게 바뀐 statistics(평균과 분산)가 그 이후에 계속 수행하게 될 convolution operation에 대하여 feature의 상대적인 중요성을 수정하지만, 기존의 statistics에 관해서는 이미 정규화를 했기 때문에 거기에 영향을 받지 않는다. 즉, style transfer에서 정규화 과정으로 인해 기존의 feature map의 statistics는 어느 정도 영향이 사라지고, 그 이후부터는 style에 의해 변화된 statistics가 영향을 끼치므로 결과적으로 각각의 style은 오직 하나의 convolution layer output에만 영향을 끼친다는 것이다.

 

 

 

StyleGAN Generator의 Property

 

Mixing Regularization

 

여기서 더 나아가 좀 더 style을 localize 하기 위해 저자는 mixing regularization을 사용했다고 한다. 이 방법은 잘 보면 정말 간단한데, 하나의 latent code만을 사용하지 않고 두 개의 latent code를 사용해서 StyleGAN을 통해 정해진 percentange 만큼의 이미지를 생성한다는 것이다. 상세히 설명하자면, 앞서 설명한 StyleGAN에서는 latent code인 $z$ 하나를 샘플링해서 이를 mapping network에 통과시킨 후 얻은 $w$를 style로 사용해서 synthesis network에 관해 style transfer를 적용했다. 그러나 mixing regularization은 두 개의 latent code인 $z_1$과 $z_2$를 뽑아서 이를 mapping network에 통과시킨 후, 그 결과로 나온 $w_1$, $w_2$를 두 개를 synthesis network에 style 정보로 활용하겠다는 것이다. 대신에 random으로 뽑은 임의의 한 지점을 crossover point로 지정해서 그 이전까지는 style transfer 수행 시 $w_1$만을 사용하고, 그 지점 이후부터는 $w_2$만을 사용해서 좀 더 다채로운 이미지를 생성하는 동시에 인접한 style이 서로 연관되지 않도록 구분지을 수 있으며 이를 통해 style의 localization을 달성할 수 있다고 말하고 있다. 

 

 

위 그림처럼 논문에서 나온 Figure 3는 StyleGAN 모델의 style이 얼마나 잘 localized 되어서 style 요소들이 잘 disentangled 되었는지를 확인할 수 있도록 한눈에 쉽게 정리한 것이다. Source A에서의 style 정보와 source B에서의 style 정보를 두 개의 latent code로 하여 새로운 이미지를 생성할 때 그 결과가 어떻게 되는지를 확인할 수 있는데, 기본적으로 source A로부터 이미지를 생성하되 style transfer를 위해 stlye의 일부분을 source B의 latent code로 사용했다. 대신에 그 latent code를 변경하는 구간을 여러 실험마다 다르게 했는데, 위의 그림에서 'Coarse styles from source B'는 앞쪽의 4개의 style을, 'Middle styles from source B'는 그 이후의 4개의 style을, 'Fine from B'는 뒤의 10개의 style에 source B의 latent code로 적용했다는 의미이다. (앞서 coarse와 fine을 설명할 때 저자가 총 18개의 style에 대응되는 latent code에서 앞쪽 4개의 style은 coarse, 중간 4개는 middle, 뒤쪽 10개는 fine style로 정의했다고 언급했다.)

결과를 보면 상당히 흥미롭다. Source B의 style을 coarse로 적용했을 때, source B에서 안경을 쓴 모습이 source A를 기반으로 생성된 이미지에 반영되었고, source B가 어린 아이일 경우 source A 기반으로 생성된 이미지에서 인물의 모습이 어린 아이의 얼굴형과 흡사해졌다. 즉, coarse style을 source B의 latent code로 대체하면 얼굴형, 안경 착용 여부, pose 등 전반적인 특징이 style로 반영된다.

Source B의 style이 middle에 적용되면, coarse style 보다는 좀 더 세부적인 정보들이 반영되었음을 볼 수 있다. Source B의 인물이 바라보는 눈동자의 방향이 생성 이미지에 반영되거나 눈을 찌푸린 모습, 입을 벌렸는지 여부 등 구체적인 모습들에서 source B의 특징이 나타나는 것을 확인할 수 있다.

Fine style의 경우 이보다 더 세밀한 정보들이 style로서 반영된 것을 볼 수 있는데, 특히 머리 색상, 피부 색, 배경색 등 매우 디테일한 정보들에 source B의 특징이 녹아 있음을 엿볼 수 있다.

 

 

위의 표는 mixing regularization이 얼마나 모델의 성능을 높이는지를 앞서 언급한 FID 지표를 바탕으로 보인 것이다. 가로 축은 Latent code를 하나만 샘플링 한 것 뿐만이 아니라 두 개, 세 개, 네 개를 뽑은 경우를 모두 비교한 것이고, 세로 축은 학습 데이터에서 얼마만큼 비율의 이미지에 mixing regularization을 적용한 실험인지를 나타낸 것이다. 대체로 학습 과정에서 mixing regularization을 아예 적용하지 않았을 때보다 학습 데이터에 90% 또는 100% 모두 적용했을 때 generator의 생성 품질이 좋음을 확인할 수 있다.

 

 

Stochastic Variation

 

인물의 사진을 생성하는 데 있어서 사람의 주근깨, 여드름, 머리의 배치, 땀 구멍 등 다양한 확률적인 요소가 가미될 수 있는데, 저자는 기존의 GAN에서는 noise가 오직 네트워크에서의 첫 input에만 의존적이고 이를 여러 네트워크에 통괴시키면서 network capacity를 차지함과 동시에, 생성된 signal의 주기성을 숨기기가 어려워지므로 이러한 stochastic feature를 직접 control 하기가 어려웠다고 지적한다. 또한 이러한 stochastic feature들은 적절한 distribution을 따른다면 이미지에 관한 전반적인 인식에 영향을 주지 않아야 한다고 덧붙인다. 그래서 StyleGAN에서는 사진 생성의 stochastic variation을 부여하기 위해 noise를 synthesis network의 convolution layer output 마다 픽셀 단위로 직접 더했다.

 

 

위의 그림에서 (a)는 모든 레이어에 noise를 준 것, (b)는 noise를 주지 않은 것, (c)는 오직 fine layer에 noise를 주어서 머리카락이 매우 곱슬거림을 보여주고, (d)는 오직 coarse layer에 noise를 주어서 머리카락이 (c)보다 덜 곱슬거리는 것을 확인할 수 있다. 이 그림에서 알 수 있는 것은 noise가 이미지의 identity나 pose 같이 전반적인 compotisition이 아니라 오직 stochastic aspect에만 영향을 끼치며, coarse부터 fine까지 각각 layer에 직접 다른 noise를 주입함으로써 앞서 설명한 style 뿐만이 아니라 noise 자체도 localized 되는 효과를 가져온다는 점이다.

 

짚고 넘어가야 할 점은 pose를 바꾸거나 identity를 변화시키는 global effect가 발생해도, noise는 오직 중요하지 않은 부분에 관해서만 stochastic variation을 준다는 것이다. 즉, 공간적으로 invariant 한 statistics는 style을 나타내는 반면에, 공간적으로 변화하는 feature는 매우 특정한 instance에 변화가 있는 noise처럼 나타난다는 것이다. 좀 더 구체적으로 보자면, style은 앞서 설명한 style transfer에 의해 이미지의 feature map의 전반적인 분포가 style 분포로 맞춰지면서 일관성을 보이지만, noise는 픽셀 단위에 독립적으로 적용되면서 확률적인 variation을 구현하는 데 더 이상적이다. 만약에 synthesis network가 noise를 가지고 pose를 조절하려고 하면, 이는 style로 조절하는 것보다는 비일관적이게 변할 수 있기에 discriminator가 금방 generator가 생성한 이미지인지 찾아낼 수 있어서 결과적으로 어떠한 가이드 없이 스스로 최적화에 의해 style에 의해 변화되는 정보와 noise에 의해 생기는 정보를 분리하여 생성하는 것이다.

 

 

위의 그림에서 생성된 이미지인 (a)에서 stochastic variation이 어떻게 나타날 수 있는지를 (b)에 정리했다. 인물의 머리카락의 배치가 noise에 따라서 일부 변한다는 걸 엿볼 수 있다. 또한 그러한 noise가 인물의 이미지에서 대개 어떠한 부위에서 나타나는지 그 변화에 대한 표준편차를 시각화하여 표현한 것이 바로 (c)이다. 사람의 얼굴형 등 전반적인 정보보다는 머리카락, 옷 등 전체 content의 인식에는 큰 변화를 주지 않는 디테일한 부분에 주로 noise가 추가되었다는 사실을 확인할 수 있다. 또 주목할 만한 점은 인물의 눈동자도 noise가 자주 반영되는 대상이 되었다는 점이다.

 

 

 

StyleGAN과 Disentanglement

 

 

Disentanglement를 정의하는 방법은 다양하지만, 공통된 내용으로 정리하자면 좀 더 linear한 subspace로 구성된 latent space를 만드는 것이 disentanglement의 목표이다. 이러한 각각의 linear한 subspace는 하나의 factor에 관한 variation을 control 하는 것으로 볼 수 있다. 그러나 $\mathcal{Z}$에서 factor들의 combination에 관해 샘플링하는 확률 분포가 학습 데이터의 확률밀도 함수를 따라가야 하는데, 우리는 원래 학습 데이터의 분포를 알 수가 없으므로 보통 uniform distribution이나 Gaussian distribution 같은 고정된 분포를 가정하여 latent code를 샘플링한다.

 

 

그러나 논문에서는 이러한 방법이 disentanglement에 있어서 좋지 않다고 말하고 있다. 예를 들어, 위의 그림을 살펴보자. (a)는 학습 데이터의 분포인데, 세로 축과 가로축이 각각 남성다움(masculinity)과 머리 길이를 의미한다. 세로 축에서 아랫쪽은 남성다움이 낮고, 위쪽은 남성다움이 높으며, 가로축의 왼쪽은 머리 길이가 길고, 오른쪽은 머리 길이가 짧은 것을 의미한다. 이 학습 데이터의 분포에서 남성성이 높으면서 머리 길이가 긴 데이터는 아예 배제되어 있다. 이러한 상태에서 latent space인 $\mathcal{Z}$에서 Gaussian distribution으로 feature를 샘플링하면 문제가 발생할 수 있다. 예를 들어, 남성성이 낮으면서 머리 길이가 짧은 특징을 지니는 latent code를 뽑았는데, 위에서 파란색 칸이 원래 데이터의 분포와 비교했을 때 오히려 남성성이 높고 머리 길이가 긴 특징으로 치우쳐 있음을 볼 수 있다. 이를 논문에서는 "warping"이 발생했다고 말한다.

그렇지만 이를 앞서 설명한 mapping network에 통과시켜서 intermediate latent code를 얻게 되면 그 latent code가 위치한 space $\mathcal{W}$가 좀 덜 "warp"되어 subspace가 어느 정도 linear 하게 구성되므로 굳이 샘플링을 위해 Gaussian처럼 어떤 한 특정 분포를 정해서 따를 필요가 없어지게 된다. 즉, subspace에 대응되는 latent factor들이 좀 덜 entangled 된다는 의미이다. 

 

저자는 실제로 mapping network를 거친 intermediate latent space $\mathcal{W}$가 얼마나 잘 disentangle 되었는지를 보이고자 했는데, 불행히도 이를 증명할 수 있는 당시에 사용되는 방법으로는 StyleGAN에 사용하기가 어려웠다고 한다. 왜냐면 input image를 latent code로 보낼 수 있는 encoder network가 필요한데, StyleGAN의 generator에서는 encoder를 사용하지 않고 단지 latent space에서 뽑은 latent code를 가지고 이미지를 생성한다. 다시 말해, encoder network를 StyleGAN에 추가하지 않는 이상 이를 측정하기가 어렵다. 그래서 저자는 disentanglement를 수치적으로 보일 수 있는 두 가지의 새로운 metric을 제안했는데, 바로 perceptual path length와 linear separability이다.

 

 

 

Perceptual Path Length

 

Perceptual path length는 latent space에서 뽑은 두 벡터를 interpolation 할 때 얼마나 벡터가 급격하게 또는 완만하게 바뀌는지 그 정도를 수치로 보일 수 있는 지표이다. 두 벡터를 interpolation 할 때 급격하게 벡터가 바뀌면 덜 그럴 듯한 이미지가 생성되는 것이고, 반대로 완만하게 바뀌면 더 그럴 듯한 이미지가 생성된다고 볼 수 있다. 이를 좀 더 풀어서 설명하자면, 두 벡터를 interpolate 한 어떤 중간 부분의 벡터를 구하는데, interpolation으로 구한 벡터에서 각각의 끝점에 있는 두 벡터가 나타내는 특징과는 전혀 다른 특징이 나타나면 이는 latent space가 entangled 되었다고 볼 수 있다.

 

 

 

이 부분이 이해가 안 간다면 이 예시를 읽어 보자. Latent space에서 웃는 표정을 나타내는 벡터 $w_1$과 여성을 나타내는 벡터 $w_2$를 interpolation 한 벡터를 구한다고 하자. 우리가 interpolation을 위해 차츰 보간 벡터를 $w_1$에서 $w_2$로 이동시키면서 그 벡터가 어떤 특징을 나타내는지 확인해 본다. 우리가 기대하는 바는 점점 웃는 표정의 여성을 나타내는 특징이 나오는 건데, 갑자기 interpolation 과정에서 어느 시점에 뜬끔없이 안경을 쓴 여성을 나타내는 특징이 나온다고 해 보자. 그러면 우리는 과연 이 latent space가 제대로 disentangled 되었다고 볼 수 있을까? 그렇지 않다. 갑자기 뜬끔포로 안경을 쓴 인물의 특징이 나왔으니까. 소위 말해 '갑툭튀' 해 버린 것이다. 요소들이 latent space에서 잘 구분되었다면, 기대했던 바와 같이 위의 그림처럼 결과가 나와야 한다.

 

 

이러한 아이디어를 바탕으로 perceptual path length를 보면 이해가 쉬워진다. 먼저 latent space에서 두 벡터 $\mathbf{z}_1$과 $\mathbf{z}_2$를 뽑고, 그 두 벡터의 interpolation한 $t$ 지점을 구한다. 이때 $t$는 어떠한 비율로 interpolation을 수행했을 때의 중간 지점에 대응되며, 0과 1 사이의 구간에서 uniform distribution으로 sampling 한 결과이다. 이 때 interpolation은 slerp(spherical linear interpolation)을 사용하는데, 이는 latent space $\mathcal{Z}$에서는 Gaussian distribution을 따르면서 두 latent vector가 정규화가 되었으므로 linear interpolation 대신 slerp를 적용한다. 그러면 두 벡터 $\mathbf{z}_1$과 $\mathbf{z}_2$를 보간한 새로운 벡터가 나오고, 이를 mapping network에 통과시키고 synthesis network를 통해 어떤 한 이미지를 생성한다. 이를 VGG16 모델에 넣어서 그 모델에서의 embedding을 얻는다. 이러한 전체 과정을 $t + \epsilon$ 보간 지점에 관해 한 번 더 이미지를 생성하여 이미지를 생성하고, 이를 마찬가지로 VGG16 모델에 넣어서 embedding을 구한다. 여기서 $\epsilon$은 $10^{-4}$의 매우 작은 값이며, 보간 과정에서 얼마나 이미지가 급격하게 변화하는지를 측정하기 위해 그 변화의 작은 구간을 설정한 것으로 볼 수 있다. 또한 논문에서 VGG16의 embedding을 사용하는 이유는 인간의 인지적인 유사도에 대응되는 perceptual distance를 측정하기에 좋기 때문이라고 말한다. 이 두 보간 지점으로 생성한 이미지의 VGG16 embedding의 거리$d$를 $\epsilon$의 제곱으로 나눈 것이 perceptual path length이며, 이렇게 가능한 모든 $\mathbf{z}_1$과 $\mathbf{z}_2$, 그리고 $t$에 관해 구한 값들의 기댓값이 바로 위의 식이다. 실제로 논문에서는 100,000개의 샘플에 대한 기댓값을 계산했다.

 

 

이를 intermediate latent space인 $\mathcal{W}$에 관해 계산한 식이 바로 위의 식이다. $\mathcal{Z}$에서 $\mathbf{z}_1$과 $\mathbf{z}_2$를 샘플링하는 것은 동일한데, 대신 $\mathcal{Z}$에서 보간하지 않고 이를 mapping network $f$에 통과하여 $f(\mathbf{z}_1)$와 $ㄹ(\mathbf{z}_1)$를 얻고, 어떠한 샘플링 분포의 가정이 존재하지 않아 두 벡터가 정규화되지 않았으므로 그 $\mathcal{W}$ 공간 위에서 $t$에 관해 slerp 대신 lerp(linear interpolation)을 수행한다. 그 보간 지점에 관해 synthesis network에 입력으로 넣어서 어떤 한 이미지를 생성한다. 마찬가지로 전체 과정을 $t + \epsilon$ 보간 지점에 관해 이미지를 생성하여 그 VGG embedding 거리 차이를 구하는 것은 위에서와 동일하다.

 

 

그러면 결과는 어떻게 될 것인가? 의도한 결과는 당연히 mapping network를 통과시킨 intermediate latent space인 $\mathcal{W}$에서의 perceptual path length가 latent space인 $\mathcal{Z}$보다 더 짧아야 한다. 거리가 더 가깝다는 것은 그만큼 subspace를 구성하는 요소들이 entangled 덜 되어서 보간하고자 하는 벡터의 특징 외에 다른 특징이 '갑툭튀'하지 않고 그럴 듯한 이미지를 생성할 수 있는 것이고, 논문의 저자가 주장하고자 하는 바는 샘플링한 latent code를 이처럼 mapping network를 통과시켜서 좀 더 disentanglement가 보장된 intermediate latent space로 보내는 것이 더 낫다는 것이므로 실제로 이와 같은 결과가 나와야 한다. 논문에서는 기존의 일반적인 GAN과 StyleGAN 모두 mapping network로 보낸 intermediate latent space에서의 perceptual path length가 더 짧다는 것을 확인하여 이를 위의 표로 정리했다.

 

 

 

Linear Separability

 

만약 latent space가 충분히 disentangled 하다면, variation에 관한 각각의 독립적인 factor에 일관성 있게 대응되는 direction vector들을 찾을 수 있어야 한다. 그래서 저자는 이러한 현상이 얼마나 잘 만족되는지를 수치화하고자 latent space에 존재하는 데이터 포인트들이 hyperplane에 의해 두 개의 개별적인 set으로 잘 분리되는지 측정하고자 했다.

 

먼저 GAN 또는 StyleGAN이 생성한 이미지들을 잘 레이블링 할 수 있도록 보조의 선형 분류기 네트워크를 학습하는데, 이 분류기들은 기존의 CelebA 데이터셋으로부터 나온 CelebA-HQ라는 40개의 각각의 속성들을 지닌 데이터셋을 잘 분류할 수 있도록 했다. 주의할 점은 각각의 선형 분류기들은 오직 하나의 속성만을 분류할 수 있도록 학습했다는 점이며, 그러므로 총 40개의 분류기를 학습했다고 볼 수 있다. 각각의 속성마다 separability를 측정하기 위해 CelebA-HQ로부터 200,000개의 이미지를 기존의 GAN 또는 StyleGAN의 generator로 생성하여 앞서 학습한 선형 분류기에 넣어서 이를 분류할 수 있도록 했다. 이 선형 분류기의 분류 결과의 confidence가 낮은 100,000개 데이터는 제거하고, 나머지 100,000개의 데이터에 관해 latent-space vector(latent code)를 구해 그 레이블과 함께 latent space point 데이터로 저장했다. 이때 기존의 GAN을 사용할 때는 $z$를 latent-space point로 사용하고, (여기서 latent-space vector를 구한다는 말이 어색할 수 있는데, 어차피 GAN 또는 StyleGAN으로 해당 이미지를 생성하면서 샘플링된 $z$ 또는 $w$가 있으므로, 이를 그대로 가져와서 활용했다고 개인적으로 이해했다.) 

 

 

이후 각각의 속성마다 앞서 구한 latent-space point들에 관해 선형 SVM(Support Vector Machine) 모델을 학습하고, 이 label을 잘 구분지을 수 있는 hyperplane을 구한다. 그리고 conditional entropy인 $\mathrm{H}(Y|X)$를 구하는데, 여기서 $X$는 SVM이 해당 데이터를 예측하여 구분한 클래스이고, $Y$는 실제 선형 분류기가 분류한 레이블이다. 쉽게 말하자면, 어떠한 데이터가 존재할 때 그 데이터가 실제 레이블에 맞게 잘 분류되었는지에 관한 entropy를 구한 것이다. 그런데 entropy가 어떠한 데이터를 잘 분류하는 데 필요한 최소 질문(정보)의 개수(크기)임을 고려하면, 이 값이 크면 클수록 그만큼 데이터를 잘 분류하기 위해 필요한 feature의 수가 부족함을 의미한다. 이를 모든 속성에 관해 entropy를 구하여 그 합에 exponential을 취하면 latent space가 실제 데이터에 맞게 정말로 잘 disentangled 되었는지를 알 수 있다. 실제로 논문에서 $z$ 샘플 포인트를 사용한 GAN과 $w$를 사용한 StyleGAN의 linear separability 값을 비교했을 때 후자가 더 낮은 값을 보이는 것을 확인할 수 있다. 또한 기존 GAN을 사용하더라도 mapping network를 통해 $\mathcal{Z}$를 $\mathcal{W}$로 보냈을 때의 linear separability가 확실히 더 좋아졌음을 실험을 통해 증명했다.

 

 

출처
1. https://arxiv.org/abs/1812.04948, A Style-Based Generator Architecture for Generative Adversarial Networks, Samuli Laine
2. https://www.youtube.com/watch?v=HXgfw3Z5zRo&list=LL&index=2&t=1893s, StyleGAN: 고해상도 가상 얼굴 이미지 생성 기술 (꼼꼼한 딥러닝 논문 리뷰와 코드 실습), 나동빈