[딥러닝 논문 리뷰] DOMAIN GENERALIZATION WITH MIXSTYLE

이번에는 도메인 generalization에 관한 논문을 소개해 드리도록 하겠습니다. ''MixStyle"이라는 논문으로 여러 style을 섞어서 잘 학습시켜보자 라는 컨셉을 가지고 있습니다. 현재 2021 ICLR에 submit되어 있는 상태이고 아직 decision은 나오지 않은 상태입니다. openreview.net/pdf?id=6xHJ37MVxxp

 

CNN은 이미지 분류에서 매우 뛰어난 성능을 보이지만, 보지 못한 도메인에 대해서는 성능이 좋지 못합니다. 그 이유는, 학습데이터와 테스트 데이터가 같은 분포라는 가정이 있기 때문이죠. 아래 그림을 보시면 네개의 그림이 각기 다른 style을 가지고 있지만 모두 '강아지'라는 것을 알 수 있습니다. 여기서 한 style의 강아지를 학습 시키고 다른 style의 강아지를 예측하려고 하면 잘 되지 않습니다. 여러 도메인에 걸처셔 잘 학습시키는 것을 Domain Generalization (DG)라고 부릅니다.

 

 

DG를 위한 가장 좋은 방법은 다양한 source domain을 학습시키는 것이나 현실적으로 이는 불가능 하다라는 것을 알 수 있습니다. 논문의 저자들은 '시각적인 도메인'은 image의 style과 연관성이 높다라는 것에서 motivation 을 갖습니다. 다양한 도메인/스타일 데이터는 DG에 있어서 매우 효과적일 것이다 라는 것이죠 (사실 이는 당연한 이야기 입니다)

 

 

Mixstyle은 랜덤하게 다른 도메인의 두개의 instacne를 추출하고 CNN layer단에서 probabilistic convex combination 을 진행합니다. 매우 어려운 말이지만, 이후 자세히 설명 토록 하겠습니다. 더불어 MixStyle은 Plug-and-Play방식으로 implement하기 매우 쉽다라고 주장하고 있습니다. 

 

 

MixStyle를 알기 전에 알아야할 Back ground가 있습니다. 우선 ,Instance Normalization (IN) 입니다. IN은 다음과 같이 layer를 normalization 하는 것은 얼핏 보면 Batch Normalization (BN) 과 매우 유사하다라는 것을 알 수 있습니다. 

 

그림으로 보면 직관적으로 잘 이해 할 수 있습니다. BN은 이미지가 들어왔을 때 batch 단으로 normalization 하는 것이고 Layer normalization 은 Channel 별로 nomalization 하는 것이고, IN은 이미지(R,G,B) 하나마다 normalization 하는 것입니다. (Group normalization 은 채널의 그룹을 정해서 normalization 하는 것 같네요)

IN이 BN에 비하여 학습이 잘되는 것을 보실 수 있습니다. 

 

 

그 다음으로 Adaptive Instance Normalization (ADaIN) 에 대해서 알아야 하는데, ADIN는 다음 수식과 같이 normalization 합니다. ADAIN 은 위 normalization 기법과는 다른 특징을 가지고 있습니다. 아래 수식을 보면 x와 y를 input으로 받는데, x의 contents 정보에 대하여 y의 style정보를 가지고 normalization 합니다. 즉, style y의 분포로 x의 분포를 scaling 하고 shift 시키겠다라는 것이죠. ADIN는 Style transfer쪽에서 굉장히 많이 사용되고 있습니다. 

 

ADaIN의 논문을 보시면 본인들이 제안한 ADaIN이 다른 방법들에 비하여 style를 잘 transfer한다는 것을 보여주고 있습니다.

 

 

아래와 같이 StyleGAN에도 사용되었습니다.

 

 

다시 MixStyle로 돌아와서, MixStyle은 이 ADaIN 의 방법을 약간 변형한 것 이고, CNN을 학습하는 과정 중에 source domain의 style 정보를 perturbing하면서 regulariznig하는 것이라고 볼 수 있습니다. Batch를 추출할 때, 아래와 같이 data를 shuffle합니다. 도메인 label을 알고 있을 때와 모를 떄를 구분해서 제안을 하였는데, X는 original data이고 X~는 shuffle된 X데이터를 의미합니다. ADaIN는 데이터 하나하나 normalization하는 것이기 때문에 X와 X~의 순서가 매우 중요한데, 아래 그리을 보시면 X를 shuffle했기 때문에, X와 X~를 데이터 하나하나 비교해보면 서로 다른 도메인일 확률이 높습니다. 물론, domain label이 주어진다면, 완전히 서로 다른 도메인을 갖도록 할 수 있겠죠.  

 

 

그리고 X와 X~에 대하여 다음과 같이 mixed feature statistic 를 뽑 습니다. lambda는 hyper parameter로 얼마나 두메인을 섞느냐 라고 보시면 될것 같습니다. Gamma는 두 도메인의 feature statistics의 표준편차, Beta는 평균인것을 알 수 있습니다. 그리고 최종적으로 MixStyle은 아래 수식 7 과 같이 계산됩니다. 

 

ADaIN 과 매우 유사한 형태라는 것을 알 수 있습니다. 다만, mix된 도메인의 sytle transfer시키는 개념으로 볼 수 있습니다. 논문에서는 PyTorch 코드도 제공 하고 있는데, 매우 쉽게 implement 할 수 있을 거란 생각이 드네요.

 

 

실험의 경우 PACS dataset (Art Painting, Cartoon, Photo and Sketch, with totally 9,991 images of 7 classes) 을 사용하였습니다.

 

세개의 도메인을 학습시키고 나머지 한 도메인을 test 했다고 합니다. 타 알고리즘 대비 성능이 뛰어난 것을 볼 수 있습니다. 아래 표에서 ResNet-18이 그냥 세도메인 한 데이터로 합쳐서 학습시킨 것이라고 보시면 됩니다. 확실히, 그냥 학습시키는 것보다는 훨씬 더 좋은 효과를 보이고 있습니다. 이 실험외에도 사람의 포즈를 일치시키는 task와 강화학습에도 적용을 했는데 꽤 좋은 성능이 나온것을 확인 할 수 있었습니다. 

 

그리고 저자들은 어느 layer에 MixStyle을 적용하는 것이 좋은지에 대해도 연구 하였습니다. 

위 표를 보시면, 2번째 residual block부터 네번째 가지 모두 적용했을 때 가장 성능이 좋은 것을 알 수 있지만, 마지막 block까지 적용하면 성능이 떨어지는 것을 볼 수 있습니다. 이에 대해서 저자들은, 마지막 block은 prediction layer와 가장 가까운 block이기 때문에 style 보다 content/label-senstivie information 을 capture 하려는 경향이 있기 때문에 합리적인 결과라고 주장하고 있습니다. 

 

아래 그림을 보시면 2~4번째 block까지는 feature statistic들이 style 별로 군집을 이루고 있는 반면, 마지막 block에서는 feature들이 label별로 군집을 이루고 있는 것을 확인 할 수 있습니다.

 

 

더불어 저자들은 AdaIN은 하나의 style를 다른 style로 replace하는 대신 MixStyle은 두개의 style을 convex combination으로 섞어버리는 차이가 있으며 저자들은 Mixing하는게 확실히 효과가 있다고 주장합니다.

 

또한, 같은 도메인내에서도 효과가 있다고 주장 하며 논문을 마무리 하고 있습니다.

아직 ICLR2021에 under review중이지만, 이정도면 accept이 되지 않을 까 생각이 됩니다. Style transfer에서 사용하는 ADaIN을 DG에 활용하는 아이디어가 꽤 괜찮은 논문이 아닌가 싶습니다.