[Paper review] Encoding in Style: a StyleGAN Encoder for Image-to-Image Translation

요즘 연구실에서 StyleGAN을 활용한 latent inversion에 관련된 연구를 계획하고 있어서, 관련 논문을 쭉 읽던 중에 논문 몇 개를 포스팅으로 남기고자 리뷰를 쓴다. 

이번 포스팅에서는 소개할 논문은 psp 모델에 다루는 논문이다. psp은 latent inversion에서 끝나는 게 아니라 image translation까지 할 수 있는 모델이다. 

 

논문의 핵심을 두가지를 뽑으라면 다음과 같다.

 

1. encoder를 pretrain된 StyleGAN generator 앞에 두어, 기존 latent space에 벗어난 이미지를 input으로 받아도 이미지에 대한 latent vector를 뽑을 수 있음

2. invert first, edit later라는 기존의 latent editing과정을 벗어나 한번에 처리

3. 약간의 변형(훈련방식, 모델)을 통해, Face Frontalization, Face From Sketch, segmentation map 등 다양한 task에  활용.

 

3은 기본적으로 논문에서 각각의 task에 대해서 방법론을 자세하게 설명해놓았기에, 이번 포스팅에서는 1,2를 어떤식으로 논문에서 구현했는지 살펴볼 것이다. 

 

논문에서 제시한 모델인, psp의 구조와 방법론은 다음과 같다.

 

 

Encoder에서 feature pyramid 구조를 활용해 서로 다른 level에서 서로 다른 style vector들을 뽑는다는 것을 확인할 수 있다. 이는 StyleGAN 논문에서 언급한 서로 다른 Style input들은 다른 level의 detail을 담당한다는 점을 반영한 것이다. 이와 더불어 style vector 개수만큼 map2Style 모듈을 사용했다는 것을 알 수 있다. 

사실 구조자체는 이게 다이다. 

psp가 만들어내는 output은 input image가 encoder를 거치면서 도출된 latent vector와 기존 pretrained된 StyleGAN의 latent vector 평균과 더해준 뒤, generator를 거친 이미지이다. 식으로 표현하면 다음과 같다.

 

x는 input image임.

위 구조를 보면 알겠지만, input으로 어떤 이미지가 들어오면, output으로 target으로 하는 이미지를 뱉는 구조이다. 즉, psp는 latent를 inversion과 editing을 한번에 처리하는 모델이다. 

 

다음으로 psp가 사용하는 loss를 알아보자.

 

psp가 사용하는 모든 loss

위에서부터, l2, perceptual, regularization, identity loss이다.

먼저 l2와 perceptual loss는 psp가 만들어낸 이미지와 target 이미지 간의 차이를 구해주는 loss이다. 가령 우리가 특정 사람의 segmentation map을 input으로 받아서, 이러한 형상을 가진 사람의 얼굴(target)을 만들기 위해서는 최종적으로 psp가 만들어 내는 이미지와 target간의 차이를 줄일 수 있도록 encoder를 학습해야 한다. 

 

다음으로, regularization loss는 encoder를 거친 latent vector가 pretrained 된 StyleGAN의 latent vector의 평균과 가까워지도록하는 loss이다. 이러한 loss가 필요한 이유는 psp 모델에서도 결국 pretrained된 StyleGAN이 이미지를 생성하기 때문에, encoder가 도출하는 latent vector가 이러한 기존의 latent space 상에서 많이 벗어나버리면 real한 이미지를 만들어 내기 어렵다. 따라서 이를 보정하기 위해 필요한 것이다. 

 

마지막으로 identity loss는 input image와 psp가 만들어낸 image가 비슷하도록 유도하는 loss이다. 이러한 loss는 상식적으로 반드시 필요한데, 우리가 image translation을 하는 목적이 input 이미지로 부터의 변화를 추구하는 것이지, 전혀 새로운 이미지를 생성하는 것은 아니기 때문이다. 

 

지금까지 psp의 구조와 논문에서 제시하는 방법론을 모두 알아봤다. 논문에서는 이러한 방식을 통해 다양한 domain의 이미지를 input으로 받아서, 이를 image manipulation까지 진행할 수 있었다. 논문에서 제시하는 여러가지 결과 중에서 흥미로운 결과 몇 개를 보면 다음과 같다. 

input으로 스케치와 segmentation을 넣어줘도 관련된 실제 사람의 얼굴을 output으로 뽑아낼 수 있다.

논문에서는 pix2pixHD, DeepFace 등 다른 모델과 비교했을때도 더 좋은 결과를 낸다는 식으로 보여주고 있다. 하지만 개인적으로 위 사진에서 pix2pixHD가 psp보다 더 input 이미지와 align하다는 느낌이 든다.

 

사실 이번 논문은 예전에 관련 주제로 학회 프로젝트를 했기에, 더욱 재밌게 읽을 수 있었다. (포스팅도 있음) 당시에는 latent vector를 직접 optimization하는 방식으로 inversion을 진행했는데, 이러한 방식을 활용하다보니 inversion하는 것도 시간이 걸렸고, inversion한 latent vector를 활용해 manipulation하는 과정도 시간이 걸렸다. 하지만 psp 모델을 활용하면 초기에 훈련만 잘 시켜놓으면, 시간적으로 꽤 많이 단축될 것 같다는 느낌이 들었다. 

 

https://arxiv.org/abs/2008.00951

Encoding in Style: a StyleGAN Encoder for Image-to-Image Translation