[CVPR2021] The Spatially-Correlative Loss for Various Image Translation Tasks

오랜만에… Image2Image Translation 논문을 가져왔습니다. Depth 논문을 계속해서 읽어도 모자른 판에 갑자기 I2IT 관련 논문을? 하실 수 있겠지만 요새 드는 생각이 어처피 I2IT나, Segmentation, Depth Estimation, Optical Flow 등등 Encoder-Decoder 방식에 Pixel level prediction 논문들은 모두 입력 정보의 컨텐츠를 손실하는 문제를 직면하고 있습니다.

깊이 추정의 결과에서도 역시나 정확히 입력 영상의 구조를 맵핑하지 못하는 현상이 자주 발생하는데, 다른 분야에 테스크에서 (특히 Real한 결과를 원하는 I2IT 분야에서) 이러한 문제를 해결하려고 노력하지 않을까? 생각하여 찾아보던 중 찾게된 논문입니다.

Introduction

I2IT 분야는 제가 예전에 자주 설명했기에 굳이 설명할 필요는 없지만 간략한 예시를 들면 일반 말을 얼룩말로 바꿔보려하는 것입니다. 이때 얼룩 모양을 만드는데 집중하느라 말의 형태가 깨진다거나, 주변 배경도 얼룩 무늬를 칠해버리면 매우 곤란하겠죠?

그래서 옛부터 I2IT는 입력 영상의 컨텐츠(일반 말)는 잘 살리면서 동시에 우리가 원하는 스타일(얼룩말의 얼룩무늬)을 잘 입히도록 많은 노력을 했습니다. 대부분은 컨텐츠가 잘 살아나지 않기 때문에, 컨텐츠와 관련된 loss들이 많이 설계되어왔는데 크게 1) pixel-level image reconsturction loss(cycle-consistency loss), 2)feature-level perceptual loss로 나눌 수 있습니다.

그림1. Unpaired한 I2IT 결과를 비교하기 위한 다양한 Content loss. 크게 Cycle-consistency loss에 해당하는 two-sided 방식과 곧바로 비교하는 on-sided 방식이 존재.

pixel-level loss는 무엇이냐면 저희가 일반 말에서 얼룩말을 만들었다고 해봅시다. 컨텐츠가 잘 살았는지 확인하기 위해서 입력 영상과 생성된 영상을 단순히 L1 or SSIM loss로 비교해버리면 loss가 상당히 크게 나올 것입니다. 즉 영상 속 구조와 외형(color, texture 등)을 pixel-level에서는 구분하기 어렵죠.

그래서 CycleGAN과 같은 two-sided 방법론들은 만들어진 realA를 통해 만든 fakeB를 다시 네트워크를 통해 fakeA로 만들어서 realA와 fakeA끼리 비교합니다. 하지만 이러한 cycle-consistency loss는 추가적인 generator가 필요하다는 점과, fakeB로 fakeA를 만들어야 하기에 생성된 영상이 많이 불안정하고 컨텐츠 요소가 많이 손실되는 가능성이 높습니다.

Feature-level loss의 경우 pixel-level loss보다는 도움이 되지만, 도메인에 특정한 구조와 외형 속성을 완전히 분리시키지는 못한 체 융합되는 문제가 여전히 존재합니다. 그리고 대부분의 feature-level loss들은 이미지넷으로 학습된 VGG16(or 19) 네트워크를 사용하기에 특정 도메인(segmentation, thermal 등)에서는 전혀 사용을 하지 못하게 됩니다.

그림2. 제안하는 방법론을 통해 서로 다른 도메인에서도 유사한 scene structure를 인코딩하는 모습.

그래서 저자는 어떠한 도메인이더라도 영상 속 structure를 외형과 분리하여 표현할 수 있는 학습 방식 및 loss를 제안합니다. 저자는 동일한 카테고리 내의 영상들은 모든 영역에서 유사한 structure를 보여줄 것이라고 생각하였는데, 그림2의 말과 얼룩말이 외형은 다르게 생겼지만 동일한 structure를 가지는 모습입니다.

저자는 크게 2가지 loss 방식을 제안하는데, 하나는 Fixed Self-Similarity of features(FSeSim)으로 이미지넷과 같이 사전학습된 모델을로부터 feature를 추출하여 spatially-correlative map을 뽑는 것입니다. 이 방식도 기존에 제안된 다양한 loss들과 경쟁력 있거나 더 좋은 성능을 보이기는 하지만, 이미지 넷으로 사전 학습된 모델을 사용하는 것이기에 추출된 Feature들은 photorealistic한 image에 편향되어 있습니다.(즉 segmentation같은 input에서는 사용하기 어렵다는 것이죠.)

그래서 저자는 보다 general한 spatially-correlative map을 구하기 위해, Learned Self-Similarity(LSeSim)이라는 학습 기반에 Self-Similarity를 계산하는 네트워크를 제안합니다. 사실 네트워크에는 큰 특징이 있는 것이 아닌, 학습 방식에 큰 의미가 있는 것으로 보이는데 Contrastive Learning을 통하여 외형과는 상관없이 유사한 구조를 가지는 패치는 더욱 가깝도록 학습하여 그들의 스타일이 매우 다르다 하더라도 구조의 유사성을 판단할 수 있도록 하였습니다.

Method

그럼 본격적으로 방법론에 대해서 알아봅시다. 저자는 입력 영상과 translated image 사이에 구조적 유사성을 곧바로 평가하려는 기존의 방법론들과는 달리, 각 이미지의 deep feature의 self-similarity를 계산하고, 그 다음 이미지들 사이에 self-similarity pattern을 비교하였습니다.

Fixed self-similarity(FseSim)은 이미지넷으로 사전 학습된 네트워크를 freeze 한 후 곧바로 self-similarity pattern을 비교하였으며, Learned self-similarity(LSeSim)의 경우 contrastive infoNCE loss를 사용하여 어떠한 supervision label없이 네트워크를 학습시켰다고 합니다.

Fixed Self-Similarity (FseSim)

그럼 FseSim부터 살펴보시죠. 여기서는 딱히 네트워크의 구조나 학습 방식이 따로 있는 것은 아니니, 어떻게 하면 structure를 비교하는지에 대해서 알아보면 좋을 듯 합니다. x, \hat{y}, f_{x}, f_{\hat{y}}는 각각 입력 영상, translated image, vgg16과 같은 네트워크를 통해 추출된 입력 영상의 feature, 네트워크를 통해 추출된 translated image의 feature를 의미합니다.

저자는 || f_{x} - f_{\hat{y}}||_{p}와 같이 feature-level에서 곧바로 비교하는 방식 대신, map 형식으로 self-similarity를 계산하였는데 이러한 spatially-correlative map은 아래와 같이 표현할 수 있습니다.

f^{T}_{x_{i}}, f_{x_{*}}는 각각 C 채널을 가지는 query point x_{i}의 feature와, N_{p}개의 패치와 관련있는 feature를 의미합니다. 그리고 S_{x_{i}}는 패치 내에서 query point와 다른 포인트들 사이에 spatial correlation입니다.

그림3. Self-similarity map으로부터 spatially-correlative loss를 계산하는 예시.

그림3은 spatially-correlative map을 구하는 과정을 나타낸 것으로 Similarity map을 히트맵으로 표현하였습니다. 저자는 기존 인코더를 타고 나온 feature는 색감, 조도, 텍스쳐 등과 같은 도메인에 특정한 속성들이 존재하지만, similarity map은 오직 구조적 유사성만을 지니고 있다고 강조합니다.

아무튼 이러한 Similarity map은 패치 영역에 해당하므로 전체 영상의 구조를 표현하기 위해서 multiple spatially-correlative map : S_{x} = [S_{x_{1}}, S_{x_{2}},..., S_{x_{s}}]으로 나타내게 됩니다. 여기서 s는 샘플링된 패치의 개수를 의미합니다.

아무튼 이렇게 입력 영상 x와, 변환된 영상 \hat{y}에 대하여 multiple structure similarity map을 계산한 다음 이 둘을 비교하는 식으로 content loss를 계산합니다.

두 similarity map를 비교하는 것은 단순히 거리를 통해서 비교하게 되는데, 저자는 크게 L1과 cosine distance 중 하나를 선택했다고 합니다. 전자는 패치 내 모든 포인트들이 동일해지도록 강조하는 효과가 있으며, 후자는 크기와는 상관없이 패턴의 상관관계?를 좀 더 집중한다고 합니다.

Learned Self-Similarity(LseSim)

자 이번에는 LseSim에 대해서 알아봅시다. LseSim역시 FseSim과 동일하게 structure similarity map을 계산할테니 해당 섹션에서는 어떻게하면 고정된 네트워크를 사용하지 않고 추가적인 비용이 발생하는 label 없이 모델을 학습하는지에 대해서 집중하면 좋을 듯 합니다.

일단 학습 방식으로는 patch level에서 contrastive loss를 사용합니다. 자꾸 contrastive라는 용어가 나오는데 contrastive에 대해서 간략히 설명하자면, self-supervised 방법론들 중 하나로 유사한 데이터 포인터끼리는 가깝게, 다른 데이터 포인터는 멀리 존재하도록하는 metric learning을 하는 것과 유사한 방식입니다.

이 때 어떤 방식으로 positive와 negative를 분리할 것인지가 중요한데, 해당 논문에서는 원래의 입력 영상에 다양한 augmentation을 적용한 후 원래 입력 영상과 유사하다는 방식으로 학습을 진행하였습니다.

그림4. Patchwise contrastive learning을 나타낸 요약도.

v = S_{x_{i}} 를 query patch의 spatially correlative map이라고 하고, v^{+} = S_{\hat{x_{i}}}, v^{-} 를 각각 패치 샘플의 positive와 negative라고 합시다.

query patch는 augmentation이 적용된 영상 x_{aug}의 동일한 위치 i와 positive한 관계를 가지고 있으며 그 외에 다른 x_{aug} 속 다른 위치의 패치 또는 target image y의 패치들과는 negative 관계를 가지고 있습니다.

negative patch의 개수는 K = 255로 설졍하였으며 contrastive loss를 아래와 같이 표현할 수 있습니다.

sim(\bold{v}, \bold{v^{+}}) = \bold{v^{T}v^{+}}/||\bold{v}|| ||\bold{v^{+}}||는 두 spatially-correlative map 사이에 코사인 유사도를 계산한 것이며, 타우는 temperature paramter입니다.

해당 loss를 최소화하기 위해서, 네트워크는 대응되는 패치들이 비록 다른 시각적 외형을 가진다 하더라도, 구조가 유사하면 가깝도록 학습할 것이며, 이는 image translation의 목적인 입력 영상의 컨텐츠를 보존한다와 부합하다고 할 수 있겠습니다.

여기서 중요한 점은, 해당 losss는 오직 structure representation network를 최적화하는데만 사용되며 생성 모델과 구분 모델에는 어떠한 loss 영향을 주지 않습니다.

Full Objective

거의 다 끝났습니다. 해당 논문은 I2IT 분야의 논문이기에 Generator와 Discriminator를 사용하므로 이 둘의 Loss가 사용되며, 거기에 논문에서 제안하는 structure representation network를 학습시키는 loss가 추가됩니다. 아 그리고 generator loss 뒤에 structure network에서 나온 structure 유사도 loss도 같이 사용됩니다.

Experiments

자 그럼 이제 즐거운 결과 확인 시간을 가져볼까요?

해당 표는 Cityscape와 Horse->Zebra 데이터 셋에 대하여 다양한 I2IT 방법론들과 논문에서 제안하는 방법론을 비교한 정량적 결과입니다. 해당 방법론의 Generator와 Discriminator는 CycleGAN을 사용하였는데, 이때 CycleGAN의 네트워크 중 절반만을 사용하여, Cycle-consistency loss를 제외하고 자신들이 제안하는 FseSim과 LSeSim을 사용하였습니다.

보시면은 CycleGAN과 비교하였을 때 상당히 큰 성능 차이를 보이고 있습니다. 또한 CUT과 비교하였을 때도 우수한 성능을 보이고 있는데, CycleGAN과 CUT의 저자가 동일하고 코드 셋팅도 동일해서.. 저자가 이 두 녀석과의 성능 차이를 강하게 강조하고 있는 것으로 보입니다.

또한 저어기 MeM은 학습할 때 사용되는 메모리 코스트라고 얘기하는데, 이를 추가한 이유는 content loss를 계산하는데 있어 one-side 방식으로 간단하기에 two side 방식보다 훨씬 빠르다는 것을 보여주는 듯 합니다.

그림5. 여러 방법론들에 대한 정성적 결과 비교.

위에는 horse2zebra, 아래는 cityscape의 정성적 결과입니다. 일단 뭐 DistanceGAN이나 MUNIT의 경우 컨텐츠가 그냥 작살나는 것을 확인할 수 있지만, LSeSim과 FSeSim은 상당히 안정적인 모습이네요.

그리고 city scape의 경우 잘 안보이실 수도 있으신데 확대해서 보시면 차량과 나무의 선명함 및 전체적인 배경이 보다 뚜렷한 것을 확인할 수 있습니다.

마지막으로 정말로 LSeSim이 RGB 도메인뿐만 아니라 다른 도메인에서도 잘 적용되는지, structure를 잘 보고 있는 것이 맞는지에 대한 결과를 보고 글을 마무리 짓도록 하겠습니다.

그림6. Error map visualization

보시면 입력 영상으로 사용할 x와 이에 대응되는 target 영상 y가 있다고 해 봅시다. 이때 y의 경우 절반은 x와 align이 맞지만, 나머지 절반은 align이 안맞는 모습입니다.

이때 단순히 pixel level에서 비교하는 l1 loss의 경우 단순히 픽셀의 값만을 비교하는 것이기에, structure와 appearance가 구분되지 못하여 align이 맞든 안맞든 상관없이 에러가 큰 모습입니다. Perceptual loss의 경우 Feature level에서 비교한 것이다 보니 조금 나아지는 듯 해보이지만, 그래도 여전히 align과 unalign 모두 큰 error를 가지는 모습입니다.

PatchNCE Loss는 이러한 문제를 해결하고자 feature들끼리의 코사인 유사도를 계산하였지만, 이 역시도 loss map이 여전히 highlight 되는 모습입니다. 이는 feature를 추출할 때 완벽하게 appearance 속성을 분리하지 못하였기 때문입니다.

하지만 LSeSim loss의 경우 appearance 속성을 제대로 무시하였기에, align이 맞는 왼쪽 절반 영역에는 에러가 낮은 반면, unalign한 오른쪽 영역은 error 값이 크게 나타나고 있으며 이를 통해 도메인이 크게 다른 영상에서도 공통적인 structure를 잘 보고 있음을 확인할 수 있습니다.

논문에서는 이 외에도 다양한 실험 결과 및 Ablation study도 보이고 있으니 관심있으신분들은 한번쯤 보셔도 좋을 듯 하네용.

결론

해당 논문은 비록 I2IT만에서 평가하였지만, Segmentation이나 Depth Estimation 등과 같은 Encoder-Decoder 모델을 사용하는 테스크에서도 충분히 컨텐츠를 강제하도록 사용할 수 있지 않을까 생각이 듭니다. 논문에서 augmentation을 어떤 것을 적용했는지에 대해 설명이 전혀 없었는데, augmentation의 종류를 새롭게 늘리거나 또는 structure를 보는 방식을 조금 바꿔봄으로써 성능을 향상시킬 수도 있을 것 같습니다.

Author: 신 정민

답글 남기기

이메일 주소는 공개되지 않습니다. 필수 항목은 *(으)로 표시합니다