[평범한 학부생이 하는 논문 리뷰] StarGAN : Unified Generative Adversarial Networks for Multi_Domain image-to-image Translation

https://arxiv.org/abs/1711.09020

StarGAN: Unified Generative Adversarial Networks for Multi-Domain Image-to-Image Translation

GAN 스터디를 하면서 몇달 전에 읽었던 논문 포스팅을 미루고 미루다 이제서야 하게 되었다.

논문 리뷰 포스팅은 논문을 읽음과 동시에 해야겠다...ㅎㅎㅎㅎ

논문 리뷰 포스팅에 사용된 이미지 중 출처가 따로 써져있지 않은 이미지들은 논문 원본에서 가져온 이미지입니다!

0. Introduction & Abstract

기존 GAN의 한계

기존의 CycleGAN, PixToPix와 같은 Image-to-Image translation 작업을 위한 GAN모델들은 두개의 도메인의 이미지를 통해 Image-to-Image translation작업을 해왔다. 하지만 여러 도메인($k$개의 도메인)으로 이러한 작업을 하기 위해서는 $k(k-1)$개의 generator를 학습해야한다. 또한 공통으로 존재하는 global한 특징이 존재하지만 각 generator에서 각각 따로 훈련되어야 하기 때문에 엄청나게 비효율적이다.

(a)기존의 GAN모델을 가지고 multi-domain image-to-image translation을 할 경우, (b) StarGAN을 통해 multi-domain image-to-image translation을 할 경우

Contribution

• 단일 generator와 discriminator만을 사용해 여러 도메인 간의 매핑을 학습하고 모든 도메인의 이미지로부터 효과적으로 학습하는 새로운 GAN인 StarGAN을 제안.
• StarGAN이 사용 가능한 모든 도메인 레이블을 제어할 수 있는 마스크 벡터 방법을 활용하여 여러 데이터 세트 간의 멀티 도메인 image translation을 성공적으로 학습하는 방법을 보여줌.
• StarGAN을 사용한 얼굴 속성 전송 및 표정 합성 작업에 대한 정성적, 정량적 결과를 모두 제공하여 기준 모델보다 우수한 성능을 보여줌.

Related Work

• Generative Adversarial Net
• conditional GAN
• Image-to-Image Translation

1.  Star Generative Adversarial Networks

이제 StarGAN에 대해 자세하게 알아보자.

1.1. Multi-Domain Image-to-Image Translation

아래의 이미지는 본 논문에서 제안하는 접근법의 훈련 과정을 보여준다.

전체적인 훈련 과정을 나타내는 이미지

StarGAN의 목표는 여러 도메인 간의 매핑을 학습하는 단일 generator $G$를 훈련하는 것이다. 이를 위해 $G$가 입력 이미지 $x$를 목표 도메인 레이블 $c$에 따라 조건부 출력 이미지 $y$로 변환하도록 훈련된다. 또한 하나의 discriminator가 여러 도메인을 제어할 수 있는 보조 classifier를 도입한다.

StarGAN에는 3가지의 loss 함수가 사용되는데 각 loss 함수에 대해 알아보자.

 

Adversarial Loss

$$\mathcal{L}_{adv} = \mathbb{E}_x [log D_{src}(x)] + \mathbb{E}_{x,c}[log (1 - D_{src}(G(x, c))], \quad (1)$$

기존의 GAN모델들에 사용됐던 Adversarial Loss Function에서 $c$라는 조건이 추가된 loss function이다.

 

Domain Classification Loss

주어진 입력 이미지 $x$와 목표 도메인 레이블 $c$에 대해 $x$를 목표 도메인 $c$로 적절하게 분류된 출력 이미지 $y$로 변환하는 것이 StarGAN의 목표이다. 이를 위해 $D$ 위에 보조 classifier를 추가하고 $D$와 $G$를 모두 최적화할 때 도메인 분류 손실을 도입한다. $D$를 최적화하는 데 사용되는 실제 이미지의 도메인 분류 손실과 $G$를 최적화하는 데 사용되는 가짜 이미지의 도메인 분류 손실이라는 두 가지 term으로 objective를 분해한다. 

 

1. 실제 이미지의 Domain Classification Loss

$$\mathcal{L}^ r_{ cls} = \mathbb{E}_{x,c^{\prime}} [- log D_{cls}(c^{\prime} |x)],\quad (2)$$

여기서 $D_{cls}(c^{\prime} |x)$ term은 $D$가 계산한 도메인 라벨의 확률 분포를 나타낸다. 이 objective를 최소화함으로써 $D$는 실제 이미지 $x$를 해당 원본 도메인 $c^{\prime}$ 로 분류하는 학습을 하게 된다. (입력 이미지와 도메인 레이블 쌍$(x, c^{\prime} )$이 훈련 데이터에 의해 주어졌다고 가정)

 

2. 가짜 이미지의 Domain Classification Loss

$$\mathcal{L}^ f_{ cls} = \mathbb{E}_{x,c}[- log D_{cls}(c|G(x, c))]\quad (3)$$

$G$는 목표 도메인 $c$로 분류할 수 있는 이미지를 생성하기 위해 이 objective를 최소화한다.

 

Reconstruction Loss

수식 (1)과 (2) loss를 최소화해도 입력 이미지의 도메인 관련 부분만 변경하면서 번역된 이미지가 입력 이미지의 내용을 보존한다고 보장할 수는 없다. 이 문제를 완화하기 위해 아래와 같은 cycle consistency loss를 generator에 적용한다. 

$$\mathcal{L}_{rec} = \mathbb{E}_{x,c,c^{\prime}} [||x - G(G(x, c), c^{\prime})||_1 ],\quad (4)$$

여기서 $G$가 번역된 이미지 $G(x, c)$와 원본 도메인 라벨 $c^{\prime}$을 입력으로 받아 원본 이미지 $x$를 재구성하려고 시도하는 것으로 정의된다. 재구성 손실로 $L_1$ norm을 채택한다. 먼저 원본 이미지를 대상 도메인의 이미지로 변환한 다음 변환된 이미지에서 원본 이미지를 재구성하기 위해 단일 생성기를 두 번 사용한다.

 

Full Objective

$$\mathcal{L}_D = -\mathcal{L}_{adv} + λ_{cls} \mathcal{L}^ r_{ cls} ,\quad (5)$$

$$\mathcal{L}_G = \mathcal{L}_{adv} + λ_{cls} \mathcal{L}^ f_{ cls} + λ_{rec} \mathcal{L}_{rec} ,\quad (6)$$

여기서 $λ_{cls}$와 $λ_{rec}$는 각각 도메인 분류 및 재구성 손실의 상대적 중요성을 적대적 손실에 비해 제어하는 하이퍼파라미터이다. 모든 실험에서 $λ_{cls} = 1, λ_{rec} = 10$을 사용했다고 한다.

1.2 Training with Multiple Datasets

StarGAN의 중요한 장점은 다양한 유형의 레이블을 포함하는 여러 데이터 세트를 동시에 통합하여 테스트 단계에서 모든 레이블을 제어할 수 있다는 것이다. 그러나 여러 데이터 세트에서 학습할 때 발생하는 문제는 각 데이터 세트에서 라벨 정보를 부분적으로만 알 수 있다는 것이다.

예) CelebA에는 머리 색과 성별에 대한 레이블은 있지만 '행복', '화난'과 같은 RaFD에 있는 속성에 대한 레이블이 없고, RaFD에는 반대의 경우를 가짐.

이는 변환된 이미지 $G(x,c)$에서 입력 이미지 $x$를 재구성할 때 라벨 벡터 $c^{\prime}$에 대한 완전한 정보가 필요하기 때문에 문제가 된다. (Recunstruction Loss $\mathcal{L}_{rec} = \mathbb{E}_{x,c,c^{\prime}} [||x - G(G(x, c), c^{\prime})||_1 ]$에서 확인 가능)

 

Mask Vector

이 문제를 완화하기 위해 StarGAN이 지정되지 않은 레이블을 무시하고 특정 데이터 세트에서 제공하는 명시적으로 알려진 레이블에 집중할 수 있는 마스크 벡터 $m$을 도입했다. 라벨의 통합 버전을 아래와 같은 벡터로 정의한다.

$$\tilde{c} = [c_1,\cdots, c_n, m],\quad (7)$$

$c_i$는 $i$번째 데이터셋의 레이블 벡터를 나타내고, $m$은 어떤 레이블을 쓸건지 정하는 one-hot 벡터이다.

 

Training Strategy

여러 데이터셋으로 StarGAN을 훈련할 때는 식 (7)에 정의된 도메인 라벨 $\tilde{c}$를 생성기에 입력으로 사용한다. 이렇게 함으로써 generator는 제로 벡터인 지정되지 않은 레이블을 무시하고 명시적으로 주어진 레이블에 집중하도록 학습한다. Generator의 구조는 입력 레이블 $\tilde{c}$의 차원을 제외하고는 단일 데이터 세트로 학습할 때와 완전히 동일하다. 반면, 판별기의 보조 classifier를 확장하여 모든 데이터 세트의 레이블에 대한 확률 분포를 생성한다. 그런 다음, 판별자가 알려진 레이블과 관련된 분류 오류만 최소화하려고 시도하는 다중 작업 학습 환경에서 모델을 훈련한다. 예를 들어, CelebA의 이미지로 훈련할 때 판별기는 CelebA 속성과 관련된 레이블에 대한 분류 오류만 최소화하고 RaFD와 관련된 얼굴 표정은 최소화하지 않는다. 이러한 설정에서 판별기는 CelebA와 RaFD를 번갈아 사용함으로써 두 데이터 세트에 대한 모든 판별 특징을 학습하고 생성기는 두 데이터 세트의 모든 레이블을 제어하는 방법을 학습한다.

Multiple dataset으로 training하는 과정을 나타내는 이미지

위의 이미지는 training의 전체적인 과정을 나타낸다. 위 이미지에서 위에 있는 (a),(b),(c),(d)는 CelebA 데이터셋을 가지고 훈련시키는 과정이고 아래의 4개의 이미지는 RaFD 데이터셋을 가지고 훈련시키는 과정을 나타낸다.

 

Discriminator Training

(a)와 (e)는 D(discriminator)를 훈련시키는 과정이다. Real Image가 들어오면 그 이미지가 진짜 이미지인지 판별하고 그 이미지가 해당 도메인 라벨에 속하는 지를 학습한다.

 

Generator Training

나머지는 G(generator)를 훈련시키는 과정이다. (b)와 (f)에서는 input image와 target domain label이 주어진다.  CelebA에 관한 label과 RaFD에 관한 label, Mask Vector가 주어진 것을 볼 수 있다. (b) 이 input들을 generator에 넣어주면 그 domain의 label에 맞게 이미지를 생성한다. 그렇게 생성된 이미지를 가지고 도메인 classification loss를 계산한다. (c) 그렇게 나온 output image와 원본 domain label을 다시 generator를 넣어서 나온 이미지를 원래 input 이미지와의 recunstruction loss를 계산한다. (d) (b)에서 생성된 가짜 이미지를 discriminator에 넣어서 실제 이미지로 판별되고 target domain으로 분류되도록 generator가 학습된다.

2. Implementation

Improved GAN training.

훈련 프로세스를 안정화하고 더 높은 품질의 이미지를 생성하기 위해 adversarial loss(식 (1)) 기울기 페널티가 있는 Wasserstein GAN 목표로 대체한다.

$$\mathcal{L}_{adv} = \mathbb{E}_x[D_{src}(x)] - \mathbb{E}_{x,c}[D_{src}(G(x, c))] - λ_{gp} \mathbb{E}_{\hat{x}}[(||\triangledown{\hat{x}}D_{src}(\hat{x})||_2 - 1)^2 ] ,\quad (8)$$

$\hat{x}$는 실제 이미지와 생성된 이미지 쌍 사이의 직선을 따라 균일하게 샘플링한다. 모든 실험에서 $λ_{gp} = 10$을 사용한다.

 

Network architecture.

StarGAN의 generator 네트워크는 CycleGAN을 채택해 다운샘플링을 위한 stride 크기가 2인 컨볼루션 레이어 2개, residual 블록 6개, 업샘플링을 위한 stride 크기가 2인 transposed 컨볼루션 레이어 2개로 구성되어 있다. Generator에는 instance normalization를 사용하지만 discriminator에는 정규화를 사용하지 않는다. 로컬 이미지 패치가 진짜인지 가짜인지를 분류하는 discriminator 네트워크에는 PatchGAN을 활용한다.

3. Experiments

3.1 기준 모델

논문에 실행한 실험에서 StarGAN과의 비교를 위한 기준 모델로 쓰인 모델들은 아래와 같다.

• DIAT
• CycleGAN
• IcGAN

3.2 데이터셋

실험에서 사용한 데이터셋은 CelebA와 RaFD이다.

3.3 Training Detail

• $β_1 = 0.5, β_2 = 0.999$의 Adam을 사용하여 훈련.
• 데이터 증강을 위해 0.5의 확률로 이미지를 수평으로 뒤집는다.
• 5번의 discriminator 업데이트 후 한번의 generator 업데이트를 수행.
•  배치 크기는 16으로 설정.
• CelebA에 대한 실험에서는 처음 10개의 에포크 동안 모든 모델을 0.0001의 학습률로 훈련하고 다음 10개의 에포크 동안 학습률을 0으로 선형적으로 감쇠시킴.
• 데이터 부족을 보완하기 위해 RaFD로 훈련할 때는 모든 모델을 100개의 에포크 동안 0.0001의 학습률로 훈련하고 다음 100개의 에포크 동안 동일한 감쇠 전략을 적용.

3.4 CelebA에 대한 실험 결과

 

3.5 RaFD에 대한 실험 결과

3.6 CelebA + RaFD에 대한 실험 결과

좌측은 joint training의 효과를 보여주는 실험을 통해 나온 결과이고 우측은 mask vector의 효과를 보여주는 실험의 결과이다. 좌측의 실험 결과는 두 데이터셋을 공동으로 학습할 경우 다른 데이터셋의 특징을 변환하는게 더 자연스럽다는 것이다. 우측의 실험 결과는 mask vector 값을 잘못 주었을 경우 이미지가 이상하게 생성되는 것을 확인할 수 있다. 그렇기에 mask vector이 제대로 역할을 해내고 있다는 것을 확인했다.

 

4. Conclusion

흥미로운 실험 결과들이 많이 나온 논문인 것 같다. 논문 읽은지 한참 뒤에 포스팅하지만 논문을 읽을 당시엔 엄청 흥미롭게 읽었었다. 이 논문을 읽고 같이 스터디한 스터디원들과 이 모델로 각자의 얼굴로 표정이나 머리 색을 바꾸는 것을 해봤는데 웃긴 결과가 나왔던 것으로 기억한다 ㅎㅎ.