땡글이LAB

[논문리뷰]pSp : Encoding in Style: a StyleGAN Encoder for Image-to-Image Translation 본문

AI/논문리뷰

[논문리뷰]pSp : Encoding in Style: a StyleGAN Encoder for Image-to-Image Translation

땡글이B 2022. 6. 23. 16:01

0. Abstract

본 논문에서 주장하는 pSp 프레임워크는 사전 훈련된 Style-GAN 생성기에 공급되어 확장된 W+ 잠재 공간을 형성하는 일련의 스타일 벡터를 직접 생성하는 새로운 인코더 네트워크를 기반으로 한다.

 

그리고 본 논문에서는 인코더가 추가 최적화 없이 실제 이미지를 W+ 에 직접 내장할 수 있음을 보여준다. 그리고 얼굴 영역에서의 Translation 과정을 통해 pSp 프레임워크의 활용성을 입증하고 얼굴영역에 그치지 않고, 다른 영역으로 확장도 가능함을 보여준다.

 

1. Instruction

 

최근, 수많은 논문들이 GAN의 latent space W 즉, latent vector w를 조작하는 것을 보여준다. 이런 방법들이 대부분 "invert first, edit later" 방식을 고수한다. 하지만, 실제 이미지들을 512차원의 w vector 로 매핑한다고 해서, 그것들이 정확한 reconstruction이 되는 것은 아니라고 밝혀졌다. 이 정보에 입각해 W+ space를 한 번의 인코딩을 통해 만들려고 하는 것이 일반적인 접근 방식이 되었다. 그렇게 하더라도 이러한 방식은 이미지 하나 당 W+ space를 통해 몇 분이 걸린다. 

 

 그래서, 이 최적화 과정을 빠르게 구현하기 위해, 일부 방법은 추가 최적화가 필요한 좋은 초기 포인트 역할을 하는 W+ 의 근사 벡터를 유추하도록 인코더를 훈련시켰다. 하지만 그래도 실제 이미지를 빠르고 정확하게 W+ space로 인코딩시키는 것은 어려운 과제였다.

 

 본 논문에선, 실제 이미지를 토대로 W+ space를 만드는 새로운 인코더 아키텍처를 소개한다. 이 아키텍처는 Feature Pyramid Network를 토대로 구성되어 있다고 한다. 본 논문에서 주장한 이 인코더는 실제 입력 이미지를 직접 재구성할 수 있어 시간이 많이 걸리는 최적화 없이 Latent space Manipulation 을 수행할 수 있음을 보여준다.

 

다만, Latent space Manipulation 을 하기 위해선, input image가 invertible (가역적)이여야 한다는 한계점이 존재한다. 즉, latent code가 image를 재구성할 수 있어야 한다는 점이다. 이 한계점을 극복하기 위해, 본 논문에서는 인코더와 사전학습된 StyleGAN을 같이 사용한다. 이런 접근 방법으로, input image를 원하는 latent code로 인코딩하고, latent code를 통해, StyleGAN을 통해 새로운 이미지를 생성한다. 

 

본 논문에서 주장하는 pSp 아키텍처는 다양한 문제에도 적용 가능하고 해결할 수 있는 아키텍처라고 한다. 기존에는 단일 문제를 해결하기 위한 전용 아키텍처를 사용하는 방식이었다.

 본 논문의 대략적인 아키텍처는 pixel2style2pixel이다. 즉, 입력으로 들어온 모든 이미지가 인코딩되고 이후 다시 이미지로 바뀌는 것이다. 그래서  pSp 라고 불리우는 것이다.

 

본 논문에서 주목할만한 점은 2가지이다.

  • 실제 이미지를 latent space W+로 인코딩할 수 있는 StyleGAN Encoder.
  • 이미지간 변환 작업을 해결하기 위한 새로운 방법

 

 

2. Related Work

 

GAN Inversion

일반적으로 Inversion method에는 주어진 이미지에 대한 오류를 최소화하기 위해, latent vector를 직접 최적화하는 방법과 주어진 이미지를 latent space에 매핑하도록 인코더를 훈련시키는 방법을 모두 사용한 하이브리드 접근 방식을 사용한다. 

 

일반적으로 최적화를 수행하는 방식은 reconstruction method에서는 인코더 매핑보다 우수한 품질을 제공하지만, 많은 시간이 걸린다. 하지만, 본 논문에서 주장하는 인코더 매핑은 추가적인 최적화 과정이 없다는 것이 특징이다.

 

 

Latent space Manipulation

최근 Latent space Manipulation을 하기 위해, 많은 방법들이 사용된다. 대부분의 많은 방식들이 "invert first, edit later" 방식을 고수한다. 즉, 우선 image를 embedding 해서 latent space를 생성하고, latent space로 의미있는 방식으로 edit해주는 방식이다. 

본 논문에서 주장하는 방식과는 차이가 있다. 본 논문에서는 input image를 바로 output latent로 한 번에 인코딩해준다는 것이 차이점이다. 

 

 

Image-to-Image

지금까지의 이미지 translation 작업은 전용 아키텍처에서 진행되었다. 하지만 본 논문에서는 하나의 기능만이 아니라 범위를 넓혀서 여러 기능을 하도록 하는 범용적인 아키텍처를 제안했다.

 

 

3. The pSp Framework

pSp framework는 사전학습된 StyleGAN genreator와 W+ latent space를 사용해서 구축되었다. StyleGAN generator를 이용해서 생성하기 위해선 input image를 정확한 latent space로 인코딩 시키는 강력한 인코더가 필요하다. 가장 간단한 방법으로는 인코더 네트워크의 마지막 레이어에서 얻은 512차원의 vector를 통해 주어진 입력 이미지를 직접 W+ space로 인코딩해서 18가지 style vector를 학습시키는 것이다.

 

 하지만, 앞서 말한 방식은 주어진 input image의 finer한 detail을 잃어버린다는 한계점이 존재한다. 그렇기에 reconstruction의 퀄리티에도 한계가 있다. 

 

 StyleGAN에서는 3가지의 스타일로 latent space가 구분되는데 그것은 coarse, medium, fine 이다. pSp에선 feature pyramid를 활용한 backbone을 통해 encoder를 확장시켰고, 중간에 간단한 intermediate network 로 "map2style" 라는 네트워크를 추가해줬다. 구조는 다음과 같다. 

 
Figure 2

 계층적 표현과 정렬된 스타일은 출력 이미지를 생성하기 위해 스케일에 따라 생성기로 공급되며, 따라서 중간 스타일 표현을 통해 입력 픽셀에서 출력 픽셀로의 변환을 완료한다.

 또한, StyleGAN에서와 같이 본 논문에서는 사전 훈련된 Generator의 평균 style vector를 추가로 정의한다. 입력 이미지 x가 주어지면, 본 논문의 모델의 출력은 다음과 같다.

 

pSp 출력 값

 위의 공식에서 본 논문의 인코더는 평균 스타일 벡터와 관련해서 latent code를 학습하는 것을 목표로 한다. 우리는 이것이 더 나은 초기화를 초래한다는 것을 발견했다. 

  • E( ) 는 인코더를 의미하고, G( )는 Generator를 의미한다.

 

3.1 Loss Functions

 본 논문의 인코더는 여러 가지 Loss Function을 사용한다. 첫 번째로는 L2 loss function이다.

L2 loss function

두 번째로는 LPIPS loss function을 사용하는데, 이는 이미지의 퀄리티를 더 좋게 해주게끔 보장한다. (F( ) : perceptual feature extractor 의미한다)

LPIPS loss function

 또한 인코더가 평균 latent vector 에 더 가까운 latent style vector를 출력하도록 장려하기 위해, 본 논문에서는 다음과 같은 regularization loss(정규화 손실)을 추가로 정의한다. 그리고 본 논문에서는 인코더 훈련에 이 정규화를 추가하면 출력의 fidelity(충실도)를 손상시키지 않고 이미지 품질이 향상된다는 것을 발견했다.

regularization loss

마지막으로, 얼굴 이미지 인코딩의 특정 작업을 처리할 때 input identity의 보존 문제가 존재하는데, 이를 해결하기 위해 출력 이미지와 소스 사이의 코사인 유사성을 측정하는 recognition loss 를 사용한다.

  • R( ) : 사전학습된 ArcFace network

recognition loss

 

즉, 위의 loss 들을 사용해서 전체적인 loss 는 아래와 같다. 앞의 계수들은 loss 가중치에 의해 정해지는 상수들이다. 

 

3.2 The Benefits of The StyleGAN Domain

 스타일 도메인(domain)을 통한 이미지 간 변환은 픽셀 간 대응 없이 지역적(local)으로 동작하는 것이 아닌 아닌 전체적(global)으로 모델을 작동시키기 때문에 많은 표준 이미지 간 변환 프레임워크와 pSp를 차별화 시켜준다. 

Figure 3

 지금까지의 연구결과들로 StyleGAN에서의 계층구조 덕분에 latent space는 disentangle하게 작동할 수 있었다. 그래서 본 논문에서도 위의 그림과 같이 StyleGAN의 계층구조를 활용해서 StlyeGAN에서와 같이 세세한 style을 변형할 수 있도록 모델을 설계했다. 또한, StyleGAN의 계층구조 덕분에 multi-modal synthesis가 가능할 수 있었다. 

 일종의 이미지 변환 작업들은 애매하기에 하나의 input image에 대해서 여러 결과가 나올 수 있다. 본 논문에서는 style vector를 랜덤하게 샘플링함으로써 output image에 대응되는 W+ space를 만들 수 있었다. 

 그런 다음 Style Mixing은 latent space의 선택된 layer를 랜덤하게 생성된 latent space의 layer로 대체함으로써 수행되었으며, 두 스타일 간의 블렌딩을 위한 매개 변수 a로도 가능하다.

 

4. Applications and Experiments

Figure 4

 Figure 4에서는 StyleGAN을 기반으로 한 Encoder 관련 논문들(ALAE, IDInvert)과 본 논문(pSp)의 Inversion 결과를 비교한다. 즉, input image에 대해 변형 없이 output latent vector가 얼마나 잘 찾아지는지를 의미한다. 아래의 표는 각 모델들의 loss 값 비교와 실행 시간을 나타낸다. pSp를 사용해서 인코딩을 했을 경우에 실행시간이 낮은편에 속하며 Similarity가 다른 모델들에 비해 높은 것을 보여준다. Karras 는 similarity는 높지만 실행시간이 너무 길어 활용하기 어려운 부분이 있을 것으로 보인다.

 

Comments