[논문] Contrastive Learning for Unpaired Image-to-Image Translation

https://arxiv.org/pdf/2007.15651.pdf

해당 논문을 보고 작성했습니다.

 

Abstract

image-to-image translation에서, output의 patch는 domain의 상관없이 input의 patch를 반영해야 합니다. 이 논문은 contrastive learning을 토대로 하는 framework를 상용하여 둘 사이의 정보를 최대화하는 방법을 제안합니다.

 

1. Introduction

위 Fig1처럼 output이 input image의 구조 또는 내용을 유지하면서 target domain의 모양을 보이기 원합니다. 유지해야 하는 내용과 변경될 내용을 분리해야 합니다. 일반적으로 target appearance는 adversarial loss를 사용하고 content는 cycle-consistency를 사용해 보존됩니다. 하지만 cycle-consistency는 두 domain이 대립을 이룬다는 관계를 가정해 사용되는데 이는 너무 제한적입니다. 그래서 이 논문은 input patch와 output patch 간의 상호 정보를 최대화하는 간단한 방법을 보여줍니다.

이 논문은 InfoNCE loss라는 contrastive loss function을 사용했는데, 이는 서로 상응하는 patch는 연관짓고 서로 상응하지 않는 patch들은 분리하면서 embedding 또는 encoder를 학습하는 것을 목표로 합니다. 그러기 위해선 encoder는 두 domain의 공통점(예를 들어 물체가 있는 부분이나 형태)을 집중하고 서로 차이나는 부분에는 불변하도록 encoder를 학습해야 합니다. generator랑 encoder는 patch가 input으로 쉽게 추적될 수 있는 이미지를 생성할 수 있도록 해야 합니다. input image 내부에서 negative를 drawing 하는 것이 dataset에 있는 다른 image로부터 negative를 drawing하는 것보다 더 input content를 잘 보존한다고 합니다. 이 논문에서 제시한 모델은 이전의 모델들보다 더 좋은 성능을 보이고 단일 이미지로 훈련할 수 있다고 합니다.

 

2. Related Work

Image translation and cycle-consistency, Relationship preservation,  Emergent perceptual similarity in deep network embeddings, Contrastive representation learning이 있는데 Contrastive representation learning을 보겠습니다.

Contrastive representation learning

이전의 일반적인 비지도 학습은 효율적으로 input을 재구성하는 압축된 코드를 학습하려 했습니다. 데이터의 부족, loss function의 문제가 있었습니다. 그래서 최근에는 mutual information을 극대화하는 방식들이 등장하고 있습니다. 이러한 방식들은 noise contrastive estimation을 사용해 dataset에서 비슷한 신호들은 함께 모으는 임베딩을 학습합니다. 연관있는 신호들은 그 자체가 여러 종류의 이미지가 될 수 있습니다. 그리고 이 방법은 negative의 수, sample방식, hyperparameter 설정, 데이터 증강과 같은 InfoNCE loss를 설정하는 방식이 큰 영향을 준다고 합니다.

 

3. Methods

논문 저자는 input domain x에 있는 image가 output domain y에 있는 이미지와 같이 translation되는 것을 목표했습니다. 이 논문은 upaired 한 X, Y dataset을 사용했습니다. 이 논문에서 제시하는 모델은 X와 Y에 이미지가 각각 1장만 존재해도 동작할 수 있다고 합니다. 그리고 한 방향으로만 mapping 하는 걸 학습하고 inverse generator와 discriminator를 사용하지 않는다고 합니다. 이를 통해 학습 과정이 매우 간단해지고 train time이 줄어듭니다. 이들은 generator function G를 encoder G_enc와 decoder G_dec로 나누었는데, 이를 사용해 yˆ = G(z) = G_dec(G_enc(x))가 되고 output image를 얻을 수 있습니다. 

 

위 식은 adversarial loss입니다. 이런 식을 사용해 output이 target domain의 image와 시각적으로 유사하도록 만들었습니다.

이 논문은 input과 output의 상호 정보를 최대화하기 위해 noise contrastive estimation framework를 사용했습니다. contrastive learning의 아이디어는 positive는 거리를 좁히고 negative는 거리를 멀게 하는 것입니다. query와 postive 그리고 negative N은 k-dimensional vectors v를 통해 v+ ∈ R^K,  v− ∈ R^(N×K) 처럼 mapping 됩니다. v-_n ∈ R^k는 n번째 negative를 나타냅니다. 그리고 공간이 collapse되거나 expanding 되는 것을 막기 위해 unit sphere로 벡터를 정규화합니다. 과정을 거치면 N+1 way classification 문제로 설정되는데 query와 다른 example사이의 거리는 τ = 0.07와 logits을 통해 정해집니다.

 

이 논문은 위 식을 이용해 cross entropy loss를 계산하고 이는 negative보다 positive가 선택될 확률을 나타냅니다.

 

Multilayer, patchwise contrastive learning.

비지도 학습에서 contrastive learning은 image와 patch level에서 사용되어집니다. 저자들은 전체 이미지가 content를 공유할 뿐만 아니라 input과 output image 사이의 상응하는 patch도 공유해야 한다는 점에 주목했습니다. 예를 들어, output 얼룩말의 다리를 보여주는 patch가 주어졌을 때, input 말의 다른 부분을 보이는 patch가 아니라 다리에 상응하는 부분에 더 강하게 연관되어야 합니다. pixel level에서도, 얼룩말의 몸 색이 풀처럼 배경 색이 아니라 말의 몸통 부분 색에 더 강하게 연관되어야 합니다. 그래서 multilayer patch-based learning을 사용합니다.

encoder G_enc는 이미지 translation을 생성하기 위해 계산되므로, feature stack을 쉽게 사용할 수 있고 이 점을 사용합니다. feature stack 내의 각 layer와 공간적 위치는 input image의 patch를 나타내며 더 깊은 layer는 더 큰 patch와 연관됩니다.

layer L을 선택하고 feature map은 2계층짜리 작은 MLP network H_l(SimCLR에서 사용되었던)을 통해 전달됩니다. {z_l}_L = {H_l(G^l_enc(x))}_L 식을 통해 feature stack이 생성되고 이 식에서 G^l_enc는 l번째 layer의 output을 의미합니다. layer는 l ∈ {1, 2, ..., L}로 indexing 되고 s ∈ {1, ..., S_l}를 나타내는데 여기서 S_l는 각 layer에서 spatial location의 숫자를 나타냅니다. 저자들은 correspoind feature를 z^s_l ∈ R^(C_l)로 표현하고 다른 feature들은 z^(S\s)_l ∈ R^((S_l−1) ×C_l)로 표현하는데 여기서 C_l은 각 layer의 channel 수를 의미합니다. 이런 방식으로 output image y^는 {zˆ_l}_L = {H_l(G^l_enc(G(x)))}_L로 encoding되는 걸 알 수 있습니다.

저자들은 특정 위치에서 input output patch를 일치시키는 걸 목표로 합니다. 그리고 input 내 다른 patch들을 negative로 활용합니다. 예를 들어 얼룩말 다리는 input image인 말의 다리와 더 연관되어야 하고 다른 patch들(예를 들어 말의 다른 부분, 배경, 하늘 등)은 덜 연관되어야 합니.

 

위 식을 PatchNCE loss라는 이름으로 사용합니다.

하나의 input image의 다른 patch들을 사용하는 방식이 아니라 dataset 중 다른 이미지들의 patch를 사용하는 방식도 있습니다. dataset x^에서 얻은 negative image를 encoding 하고 external NCE loss를 사용하면 됩니다. 

 

위 식이 external NCE loss입니다. 이 방식에서는 MoCo에 나온 moving average encoder를 사용해 크고 일관된 negative dictionary를 유지할 수 있습니다. MoCo를 사용하면 오래된 기록에서 negative를 sampling 할 수 있고, end-to-end update와 memory bank보다 더 효과적입니다. 위 식에서 보면 z^_l은 negatvie dataset을 의미하는데 이는 moving averaged encoder H^_l과 moving averaged MLP H^를 사용해 계산된 data들의 domain으로 생성된 external dictionary Z- 로부터 sample 된 negative data입니다.

이와 같은 방식으로 모델에서 제시한 encoder G_enc는 동물의 몸, 풀, 하늘과 같은 변하지 않는 개념을 캡처하는 방법을 배우고 decoder G_dec는 얼룩말 줄무늬와 같은 특징을 합성하는 방법을 배웁니다. 이 논문은 내부 patch만을 사용하는 게 외부 patch를 사용하는 것보다 성능이 더 우수하다고 합니다.

 

generator가 불필요한 변경을 하지 않도록 domain Y의 image에 PatchNCE loss(G, H, Y)(identity loss)를 사용합니다. 이 논문은 최종적으로 위 식을 사용하고 λ_Y = 1인 경우 λ_X = 1을 사용했고, λ_Y = 0인 경우 λ_X = 10으로 설정해 regularizer가 없는 것에 대해 해결했습니다. 이 논문에서 CUT(Contrastive Unpaired Translation)라는 이름의 전자(λ_Y = 1인 경우 λ_X = 1)가 기존 방법보다 우수한 성능을 보이고, FastCUT이라는 후자(λ_Y = 0인 경우 λ_X = 10)는 CycleGAN보다 더 빠르고 가벼운 버전으로 보면 된다고 합니다.

이 논문은 adversarial loss를 개선하기보다는 content consistency를 강화하기 위해 contrastive learning을 했습니다. 두 분포의 유사도를 측정하기 위해, Contextual Loss를 사용했고 이는 pretrain 된 network에서 뽑아낸 feature의 consine distance에 대한 softmax를 사용했습니다. 대조적으로 NCE loss를 사용해 encoder가 동일한 위치에 있는 input, output patch를 연관 지을 수 있도록 학습했습니다.

 

4. Result

 

 

위 그림은 CUT, FastCUT를 다른 방식들과 비교한 그림입니다. 논문에서 제시한 CUT는 말에 얼룩말의 무늬를 제대로 적용하는 모습을 볼 수 있습니다. FastCUT 역시 적은 train 비용으로 결과를 보여줍니다. 아래에는 실패한 예시가 있는데 말의 포즈가 다르다 보니 말이 아닌 배경에 얼룩무늬가 적용된 모습이고 두 번째는 존재하지 않는 혀를 만들어낸 모습입니다.

더 자세한 실험 결과들은 논문에 작성되어 있고 저자들의 website에도 정리되어 있습니다.

 

내 생각: cycleGAN은 generator와 discriminator가 2개씩 존재하기 때문에 학습이 오래걸리는 단점이 있습니다. 여기서 제시한 CUT 모델처럼 patch를 뽑아 negative와 postive로 나누어 cycleGAN에 적용한다면 좀 더 좋은 성능이 나오지 않을까 싶습니다. 물론 학습 속도는 더 느려지겠지만 성능은 조금 더 좋아질 것 같습니다.