StyleGAN

StyleGAN#

오늘 알아볼 모델은 StyleGAN입니다. 기존에 다뤘던 GAN과 같이 이미지를 생성하는 모델입니다. generator 구조를 변경함으로써 성능을 올리고 feature의 control이 가능하게 했습니다. loss나 discriminator 구조 개선에 관한 논문은 아닙니다. 먼저 결과를 보도록 하죠.

이 논문의 contribution은 다음과 같습니다.

새로운 구조를 제안하여 성능을 높이면서 feature의 control이 가능해졌습니다.
새로운 데이터셋을 제안했습니다. (FFHQ)

이 중에서 첫 번째 contribution을 자세히 보도록 하겠습니다. 논문의 abstract에는 다음과 같은 문장이 있습니다.

The new architecture leads to an automatically learned, unsupervised separation of high-level attributes (e.g., pose and identity when trained on human faces) and stochastic variation in the generated images (e.g., freckles, hair), and it enables intuitive, scale-specific control of the synthesis.

논문에서 제안한 새로운 generator 구조가 할 수 있는 일을 설명하는 부분입니다. 여기서 보시면 high level attribute의 separation이 가능하다고 얘기하고 있습니다. 저는 개인적으로 이 부분이 StyleGAN의 가장 중요한 특징이라고 생각합니다.

생성 모델로 이미지를 생성하고자 할 때, 사용자는 어떠한 목적을 가지고 자신이 원하는 이미지를 만들고자 할 것입니다. 이미지의 품질이 좋더라도 모델이 사용자의 의도와 상관없는 랜덤한 이미지를 내뱉어준다면 그 모델의 실용성이 좋다고 할 수 없을 것입니다. 근래에 Text-to-Image 모델들이 인기를 얻었던 이유도 누구나 쉽게 텍스트를 통해서 생성되는 이미지를 조절할 수 있다는 점도 한몫했다고 생각합니다. StyleGAN은 그런 controllability를 어느 정도 가능하게 한 모델이라는 측면에서 의미있다고 생각합니다.

StyleGAN의 구조는 아래 그림과 같습니다. synthesis network는 해상도를 4x4에서 시작해서 1024x1024까지 높여줍니다. 최종적으로 1024x1024 해상도를 가지는 이미지를 갖게됩니다. 아래 구조를 보면 기존 GAN하고 비교해서 특이한 점이 세 가지 있습니다.

z를 input으로 받는 mapping network
style과 AdaIN
noise와 B (stochastic variation)

이 각각에 대해서 알아보도록 합시다.

stylegan_02 — Fig. 63 Structure of StyleGAN#

Mapping Network#

stylegan_03 — Fig. 64 Mappings with \(w\) and without \(w\)#

기존 GAN을 생각해보면 z를 input으로 받아서 generator를 거쳐서 이미지를 생성하는 구조입니다. 이 z는 보통 Gaussian distribution에서 샘플링으로 얻습니다. GAN은 학습을 통해 Gaussian distribution을 data distribution으로 보내는 방법을 배우게 될 것이고, 이 분포는 (b)처럼 생기게 될 것입니다. 그런데 데이터가 (a)처럼 주어져서 특정한 데이터가 없거나 적을 수도 있을 것입니다. 예를 들어, 데이터에 피부가 희면서 머리가 긴 샘플들이 없다고 해봅시다. 그러면 피부색과 머리 길이라는 두 feature는 서로 얽히게(entangled)되어, 하나를 바꿀 때 다른 하나도 같이 바뀌는 현상이 일어나게 됩니다. 이런 현상을 완화하기 위해 논문에서는 Gaussian에서 뽑은 z를 바로 사용하는 것이 아니라 mapping network를 통해 learnable distribution에서 뽑은 w를 사용합니다.

Style and AdaIN#

instance normalization은 샘플 하나의 각 채널마다 정규화를 취해주는 방법입니다.

stylegan_04 — Fig. 65 Normalization methods#

adaptive instance normalization (AdaIN) 은 instance normalization에 scale을 곱해주고 bias를 더해주는 형태입니다. 그런데 이 scale과 bias가 style vector의 linear transformation으로 주어지는 형태입니다. linear layer를 통해서 w는 \(\mathbf{y}=(\mathbf{y}_{s},\mathbf{y}_{b})\)로 보내지게 됩니다. AdaIN의 수식은 아래와 같습니다.

\[ AdaIN(\mathbf{x}_{i},\mathbf{y})=\mathbf{y}_{s,i}\frac{\mathbf{x}_{i}-\mu(\mathbf{x}_{i})}{\sigma(\mathbf{x}_{i})}+\mathbf{y}_{b,i} \]

AdaIN은 각 블록마다 두 개씩 들어가서 style은 총 열여덟 번 AdaIN을 통해 generator에 들어가게 됩니다. AdaIN은 localization이라는 특징과도 연관이 있습니다. 여기서 말하는 localization이란 열여덟 개의 style 중에서 일부를 바꿈으로써 이미지의 일부 특징들을 바꿀 수 있다는 의미입니다. AdaIN은 각 convolution layer 다음에 적용이 됩니다. 이 때 feature map들은 normalization되고 style에 의해 새로운 statistics를 가지게 됩니다. style은 하나의 convolution에 적용되고, 다음 convolution에서 다시 normalization이 수행되기 때문에 이전 layer에 적용된 style과 다음 layer에 적용된 style이 분리되게 학습될 수 있습니다.

관련 코드

class StyleMod(nn.Module):
    def __init__(self, latent_size, channels, use_wscale):
        super(StyleMod, self).__init__()
        self.lin = EqualizedLinear(latent_size,
                                   channels * 2,
                                   gain=1.0, use_wscale=use_wscale)

    def forward(self, x, latent):
        style = self.lin(latent)  # style => [batch_size, n_channels*2]

        shape = [-1, 2, x.size(1)] + (x.dim() - 2) * [1]
        style = style.view(shape)  # [batch_size, 2, n_channels, ...]
        x = x * (style[:, 0] + 1.) + style[:, 1]
        return x

class LayerEpilogue(nn.Module):
    """Things to do at the end of each layer."""

    def __init__(self, channels, dlatent_size, use_wscale,
                 use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer):
        super().__init__()

        layers = []
        if use_noise:
            layers.append(('noise', NoiseLayer(channels)))
        layers.append(('activation', activation_layer))
        if use_pixel_norm:
            layers.append(('pixel_norm', PixelNormLayer()))
        if use_instance_norm:
            layers.append(('instance_norm', nn.InstanceNorm2d(channels)))

        self.top_epi = nn.Sequential(OrderedDict(layers))

        if use_styles:
            self.style_mod = StyleMod(dlatent_size, channels, use_wscale=use_wscale)
        else:
            self.style_mod = None

    def forward(self, x, dlatents_in_slice=None):
        x = self.top_epi(x)
        if self.style_mod is not None:
            x = self.style_mod(x, dlatents_in_slice)
        else:
            assert dlatents_in_slice is None
        return x

code from huangzh13/StyleGAN.pytorch

아래 그림은 source A의 style 중 일부를 source B의 style로 변경해서 만든 이미지들입니다. style은 총 18곳에서 사용되는데 처음 4곳 (\(4^2 - 8^2\))을 coarse, 그다음 4곳 (\(16^2-32^2\))을 middle, 마지막 10곳 (\(64^2-1024^2\))을 fine style로 정의하였습니다. 그림을 보시면 윗 부분에서는 포즈나 전체적인 머리 스타일같이 coarse style은 source B의 것을 유지하고, 아래로 갈수록 source A의 큰 틀을 유지하면서 세부적인 부분들을 B에서 가져왔음을 볼 수 있습니다.

stylegan_05 — Fig. 66 Mixing two styles#

Stochastic Variation#

한 사람의 이미지 안에는 확률적으로 바뀔 수 있는 부분이 있습니다. (주근깨, 머릿결, 피부) 이를 모델링하기 위해서 noise를 추가적인 input으로 사용하여 각 convolution layer 다음에 더해집니다. 아래 그림에서 (a)의 생성된 한 사람의 이미지 안에서도 디테일들은 (b)와 같이 달라질 수 있습니다. (c)와 같이 standard deviation을 구해봤을 때 얼굴형과 같은 attribute는 변하지않지만 noise에 의해서 머리카락과 같은 부분은 variation이 생김을 볼 수 있습니다.

stylegan_06 — Fig. 67 Examples of stochastic variation#

아래 그림에서 (a)는 모든 layer에 noise를 준 경우, (b)는 noise를 주지 않은 경우, (c)는 fine layers (\(64^2 - 1024^2\))에만 noise를 준 경우, (d)는 coarse layers (\(4^2 - 32^2\))에만 noise를 준 경우입니다. (b)를 보면 noise가 없을 경우 머리카락같은 디테일이 제대로 살아있지 않은 것을 볼 수 있습니다. (c)와 (d)를 보면 fine layers에 들어간 noise가 머리카락의 더 세밀한 부분에 영향을 끼친다는 것을 볼 수 있습니다.

stylegan_07 — Fig. 68 Effect of noise inputs at different layers#

Mixing Regularization#

논문에서는 localization이 더 잘 되게하기 위해 style mixing이라는 방법을 훈련에 사용합니다. 두 개의 style vector \(\mathbf{w}_{1},\mathbf{w}_{2}\)를 사용하여 앞 쪽 layer에는 \(\mathbf{w}_{1}\)을, 뒤 쪽 layer에는 \(\mathbf{w}_{2}\)를 사용하는 방법입니다. 이는 generator가 인접한 style끼리 correlated되어있다고 학습하는 것을 막아서 localization을 더 잘 되게 하는 목적입니다.

실험 결과#

마지막으로 저자들이 제안한 방법들이 실제로 효과가 있었는지 확인해봅시다. 아래 표와 같이 실험적으로 보았을 때 저자들이 제안한 방법들을 모두 사용한 경우 FID가 가장 우수하게 나왔습니다.

stylegan_08 — Fig. 69 FID for various generator designs#