[Paper review] Pivotal Tuning for Latent-based Editing of Real Images

요즘 연구실에서 진행하는 연구가 3D Inversion에 관한 연구여서, Inversion 관련 논문을 많이 읽는다. 특히 이번에 리뷰할 논문은 최근에 리뷰한 Designing an encoder for stylegan image manipulation 논문에서 다뤘던 distortion-editability 사이의 trade off에 대한 문제를 그대로 가져오면서 시작한다. 이 trade off가 구체적으로 무엇인지는 저번 논문 리뷰에서 자세히 설명했기에, 핵심만 설명하면 이거다. 어떤 이미지(x)를 manipulation하기 위해 특정 latent vector(z)로 inversion할때, 이러한 z가 StyleGAN의 original latent space(W)에 가까울 수록, distortion은 커지는 대신 editability는 증가한다는 것이다. 이러한 현상은 어떻게 보면 Pretrained StyleGAN을 그대로 사용하기에 당연한 결과이다. 여기서 가질 수 있는 의문은 결국 두마리의 토끼를 다 잡을 수는 없을까? 이다. 이번에 소개하는 논문에서는 둘을 다 잡을 수 있는 어떻게 보면 간단하면서 효과적인 방법을 소개한다. 이게 바로 "Pivotal Tuning"이다.

 

Pivotal Tuning의 아이디어는 정말 간단하다. 앞에서 trade off가 발생하는 근본적인 이유는 Pretrained StyleGAN을 그대로 사용하는 것이다. 이로 인해, 한번 latent vector가 고정되면, 현재 StyleGAN이 이 latent space를 잘 만들어 낼 수 있느냐에 따라 최종 결과가 달라진다. 이를 해결하기 위해서는 결국 고정된 StyleGAN이 학습이 가능하게 만들어줘야 한다. 이게 Pivotal Tuning의 핵심 아이디어이다. 

 

이제부터 논문에서 언급한 내용을 토대로 구체적으로 어떻게 Pivotal tuning이 이루어지는지 설명해보겠다. 

논문에서는 two-step으로 이뤄진 method를 소개한다. 첫번째 step은 Inversion으로, manipulation하고자하는 이미지에 대한 latent vector를 찾아내는 과정이고, 두번째 step이 Pivotal Tuning으로 부르는 과정이다. 각 과정에 대해서 하나씩 살펴보면, Inversion은 editability를 최대로 할 수 있는 latent vector를 찾는 것이 목적이다. 이를 위해서는 리뷰의 서두에서도 말했듯이, StyleGAN의 original latent space로 Inversion이 이루어져야 한다. 논문에서는 다음과 같은 loss를 통해 Inversion을 진행했다. 

 

Inversion에서 흔히 사용하는 loss인 perceptual loss와 noise(n)을 regularize하기 위한 reguralization loss를 더한 형태이다. 여기서는 Pretrained StyleGAN의 weight는 fix한 상태에서 진행한다. 이를 통해 Inversion된 하나의 이미지에 대한 latent vector를 논문에서는 pivot code(wp)라고 말한다. (noise도 학습하는 이유는 StyleGAN의 Generator에는 latent vector들과 더불어 noise가 input으로 들어와야 하기 때문이다.)

 

이렇게 첫번째 step을 거치고 난 후, 도출된 하나의 이미지에 대한 latent vector는 editability는 확보되었지만, distortion이 부족한 상태일 것이다. 이를 극복하기 위해 대망의 두번째 step인 Pivotal Tuning을 거치게 된다. 

Pivotal Tunining은 아까 간략히 언급했지만, Pretrained StyleGAN을 고정시켜놓는 것이 아니라 finetuning 시키는 것이다. 논문에서는 StyleGAN의 Generator 부분을 finetuning하는 방식으로 진행하고, 다음과 같은 loss를 쓴다. 

 

우리가 첫번째 step에서 N장의 latent vector를 뽑아냈다고 했을때, 이를 한꺼번에 하나의 batch로 활용한다. xi는 원본 이미지이고, xp는 pivot code를 Generator에 통과 시켰을 때 나오는 이미지이다. loss는 Perceptual loss, l2 loss를 사용한 것을 확인 할 수 있다. 핵심은 위 loss를 통해 pivot code를 optimization하는 것이 아니라, StyleGAN Generator의 weight를 업데이트 하는 것이다. 다음 그림을 보자. 

 

논문 그림 참조

오른쪽이 pivotal tuning을 진행하기 전, latent space에 놓여 있는 latent vector들이 만들어 내는 이미지(A, B)이다. A가 위치한 붉은색 지역은 Pretrained StyleGAN의 original latent space에 거의 근접한 곳으로, editability는 높지만, distortion이 심한 곳이다. 반면 붉은색이 아닌 외곽의 B가 위치한 곳은 distortion은 적지만, editability가 좋지않는 곳이다. Pivotal tuning은 B를 editability가 높은 붉은 지역으로 옮겨주는 효과를 낸다. 위 그림의 오른쪽의 C를 보면, real image와 거의 흡사하면서(distortion 적음), editing 역시 잘 이뤄진다는 것을 확인할 수 있다. 

 

이게 논문에서 말하는 Pivotal Tuning의 핵심의 전부이다. 추가적으로 Locality Regularization이라는 기법을 통해, 원본이미지의 ID를 유지해주는 기법을 적용하여, 최종적으로 StyleGAN Generator의 weight를 학습한다. 이 부분은 직접 논문을 읽어보길 바란다.

 

Pivotal Tunining을 적용하여 StyleGAN Generator를 finetuning한 이후에는 다양한 manipulation기법(StyleCLIP, linear direction..etc)을 그대로 적용하여 editing을 진행하면 된다. 아래는 논문에 제시한 몇가지 결과 이미지이다. 

 

논문 사진 참조

 

위 사진에서 알수 있듯이, input image와 거의 똑같은 figure가 editing을 통해 다른 모습을 띈다는 것을 확인할 수 있다. 이를 통해 이 논문에서 제시한 방법이 distortion과 editability 모두를 잡아낼 수 있는 방법임을 짐작할 수 있다. 

 

오늘 리뷰한 논문은 2D Inversion에서 고질적인 trade off를 해결할 수 있는 방법을 제시한 논문으로, trade off를 compromise한 것이 아니라, soving한 method이다. 이 논문을 읽고서 교수님께서는 이런 논문이 정말 좋은 논문이라고 극찬을 하셨다. 즉, 어떻게 보면 새로운 아이디어라고는 얼려놓은 StyleGAN Generator를 풀어놓은 것뿐인데, 결과는 크게 향상되었기 때문이다. 이렇게 작은 아이디어로 좋은 논문이 나오는 것을 보면, 나도 이러한 논문을 써보고 싶다는 생각이 든다. 하지만 이러한 작은 아이디어를 생각하기 위해서는 많은 양의 생각과 공부가 필요하다는 것도 누구보다 알기에 오늘도 나는 논문을 읽는다..ㅎㅎ

 

아래는 오늘 리뷰한 논문의 원본 출처이다.

https://arxiv.org/abs/2106.05744

Pivotal Tuning for Latent-based Editing of Real Images