StarGAN 리뷰

GAN을 이용한 image translation에 관한 논문이다.

 

Contribution

1. Multi-domain image translation을 하나의 generator와 discriminator로 가능하게 했다.

2. 그 과정에서 서로 다른 데이터셋을 같이 사용할 수 있게 했다.

3. 성능 또한 기존 모델보다 좋다.

Multi-Domain Image Translation

  StarGAN을 CelebA 데이터셋에 적용한 결과물이다. Input image에 여러 attribute를 합성하였는데, 놀랍게도 하나의 generator가 수행한 것이다. 또한, 오른쪽 3개의 열에서 Angry, Happy, Fearful과 같은 label은 심지어 CelebA에는 존재하지 않는 label이다. RaFD라는 facial expression 데이터셋에서 얻은 지식을 CelebA에 적용한 것이다. 즉, 여러 데이터셋을 함께 사용하여 하나의 G와 D를 학습시켰다.

 

CycleGAN

기존의 image translation에서는 G가 하나의 domain에서 다른 domain으로의 mapping을 학습했다. 위의 그림에서 G는 말이라는 domain에서 얼룩말이라는 domain으로의 mapping을 하고 있다. 그런데 이 말을 백마로 바꾸고 싶다면 어떻게 해야 할까? 간단하다. 백마 domain의 데이터셋을 구하고, 새로운 G와 D를 다시 백지에서 학습시키면 된다. 즉, mapping을 만들고 싶은 domain의 매 쌍 마다 새로운 G와 D를 학습해야 하는 단점이 있다.

 

하지만 StarGAN은 하나의 G와 D만으로 여러 domain으로의 mapping을 만들 수 있다는 것을 보여줬다. 그렇다면 말에서 얼룩말로 바꿀지, 백마로 바꿀지 어떻게 결정할까? Input과 같이 domain information을 넣어주면 된다. Domain information은 binary 혹은 one-hot vector인데, 가령 [0,1,0,0,0]을 같이 넣어주면 G가 2번 domain으로 mapping 해주는 식이다. 여기서 같이 넣어준다는 것은 결국 cnn에 넣을 것이므로 shape을 맞춰서 concatenate 하던지 했을텐데, 구체적인 방법은 나오지 않는다.

 

이런 방식을 사용하면 단순히 G와 D를 하나만 만들어도 된다는 것 외에 모델 자체의 성능이 올라가는데, 왜냐하면 G와 D를 학습시킬 때 모든 domain의 데이터를 다 쓸 수 있기 때문이다. 심지어 label을 다르게 한 여러 데이터셋을 전부 이용할 수 있다.

 

Overview

D에 가짜와 진짜 이미지가 들어가서 real/fake를 판단하여 확률값을 뱉는 것은 기존의 GAN과 같다. 다만 여기서는 추가적으로 domain classification까지 수행한다. Real image를 넣어서 나온 domain classification은 D를 학습하는데 사용되고, fake image를 넣어서 나온 domain classification은 G를 학습하는 데에 사용된다. G에는 이미지와 domain information으로 one hot vector가 들어간다.

 

Loss

Multi-task loss를 사용하는데, 우선 당연히도 adversarial loss가 들어간다.

실제로는 W-GAN loss가 사용되었다.

 

그런데 D 밑에 붙은 src는 무슨 뜻일까?

StarGAN에서는 D의 output이 두 개라는 것을 잊지 말자.

 

Domain Classification Loss

위의 overview에서 real image의 domain classification 값은 D를 학습하는 데에 사용된다고 했다. 여기서 x는 real image이고 c'는 x의 domain이다. 즉, D가 진짜 이미지 x의 domain을 잘 맞추게끔 학습시킨다.

 

이번엔 fake image의 domain classification 값으로 G를 학습시킨다. 여기서 c는 target domain이다. G에 x와 target domain c를 넣어서 나온 가짜 이미지를 D가 c로 생각하게끔 G가 속이도록 학습한다.

 

그런데, 이게 꼭 필요한가 하는 의문이 든다. 직관적으로 봤을 때는 이 loss가 없어도 원하는 domain으로 G가 잘 만들어야 할 것 같기 때문이다. 왜냐하면 엉뚱한 domain으로 보냈을 경우 그 domain의 real image랑 달라서 D를 속일 수 없기 때문이다. 당연히 성능이 좋아서 들어간 것일 테지만 ablation study를 해줬다면 좋았을 것 같다.

 

Reconstruction Loss

Cycle consistency를 이용해서 image의 concept를 유지하며 domain만 바꾸게 했다. CycleGAN에서는 G와 F 두 개가 있었지만 여기서는 모든 domain 간의 mapping을 G 혼자 하므로 저런 notation이 나왔다.

 

전체적인 loss는 아래와 같다.

 

classification과 reconstruction loss에 가중치가 붙은 모습이다.

 

Mask Vector

두 개 이상의 데이터셋을 가지고 학습을 할 때, 각각의 one-hot labeling이 독립적이면 학습에 문제가 생긴다. 가령 CelebA에서는 [1,0,0]이 'Young'인데, RaFD에서 [1,0,0]은 'Disgusted'라고 하자. G가 그 둘을 구분할 방법이 없다. 따라서 하나의 통합된 label로 만드는 과정이 필요하다. 위에서 n은 데이터셋의 수이고(여기서는 CelebA와 RaFD 두 개이므로 n=2), 각각의 c는 각 데이터셋의 label에 해당하는 one-hot vector이다. 이 경우에서 m은 [1,0] 또는 [0,1]일 수가 있는데, 가령 [1,0]이면 c1만을 사용한다는 뜻이고 [0,1]이면 c2만을 사용한다는 뜻이다. 안쓰는 데이터셋의 label들(즉 나머지 n-1개의 c들)은 0으로 만든다. 이런 기준에 따라 우리는 여러 데이터셋의 label들 중에 원하는 것을 선택하여 사용할 수 있다.

 

Baseline Models

DIAT은 face recognition 모델이다.

이런 reconstruction loss를 사용하는데, 여기서 F는 pre-train된 face feature extraction model이라고 한다. 저 F를 이용해서 가짜 이미지의 concept을 추출해서 x와 같은지 비교하는 방식인데, F가 face recognition에서 사용하던 모델이니까 다른 reconstruction 방법에 비해 얼굴 이미지의 concept 유지를 더 잘 할 것이다.

 

IcGAN에서 G는 {z,c} -> x의 mapping을 수행하는데, z는 latent vector이고 c는 conditional vector이다. 그리고 encoder 2개를 학습시켜서 각각 x -> z, x -> c의 mapping을 하게 한다. 여기서 conditional vector는 이 논문에서 말하는 domain information vector와 비슷한 개념인 것 같다. c 값만 바꿔주면 다른 domain으로 mapping이 된다고 한다. 그런데 뒤에서 실험 결과를 보면 E가 z랑 c라는 low dimensional latent vector만 x에서 되살릴 수 있으면 되기 때문에 concept 유지가 잘 안된다고 한다. 여기서 z만 latent vector가 아니라 이미지로 바꿔주면 StarGAN의 아이디어랑 흡사하게 되는 것 같다. 

 

여기에 CycleGAN까지 합쳐서 3개의 baseline이 사용된다.

 

Results

확실히 제일 선명하고 퀄리티도 좋은 모습이다. 위에서 말했듯이 IcGAN같은 경우는 여자 사진을 넣었는데 남자 사진이 나오는 등 concept 유지가 잘 안되고 있다.

 

AMT도 들어갔다. 여기서는 4개 모델의 결과를 사지선다로 고르게 했는데, CycleGAN처럼 real image랑 대결을 시켰다면 더 좋았을 것 같다. StarGAN이 압도적인 결과를 보여주고 있는데, 아래의 표처럼 여러 domain이 결합된 경우에도 잘 한다. 

 

여기서는 ResNet-18 구조로 만든 facial expression classifier를 훈련하고, 그걸로 각 모델의 결과물을 classification해서 error를 구한 것이다. 이 classifier는 real image에 대해 99.55% accuracy를 기록했다. 그리고 저기 써있는 parameter는 starGAN은 이처럼 한 쌍의 G와 D로 이 많은 domain에 translation을 할 수 있다는 뜻으로 받아들이면 될 것 같다.

 

여기서 classification error를 떠올린 김에, StarGAN에 붙은 D의 classifier로 facial expression classification task를 수행하고 error rate를 보여줬으면 어땠을까 싶다.

 

결과를 보면 단순히 multi-domain transfer를 하나의 G가 하니까 편리하다를 넘어서, 한 쌍의 domain만 가지고 훈련한 모델들보다 성능 또한 좋다. 논문에서는 G가 훈련 시에 랜덤하게 주어진 domain으로의 translation을 학습하면서 regularization의 효과가 나타나기 때문이라고 설명한다. 즉, 다양한 domain에서 general하게 적용될 수 있는 유의미한 feature를 얻어내게끔 학습된다는 것이다.