딥러닝/Paper Review

[Paper Review] Jigsaw Clustering for Unsupervised Visual Representation Learning

CoGam 2025. 1. 18. 04:59
728x90

이 논문은 2021년 CVPR에 소개된 논문이다. Self-supervised Learning 중 pretext task분야에 속하는 Jigsaw puzzle을 조금 더 발전시키기 위해 노력한 논문이라고 생각되었다.

 

논문이해를 위해 아래 영상을 참고하였다.

https://www.youtube.com/watch?v=zAnc1kT2CJc

 

논문을 읽으면서 조금 헷갈렸던 개념이 있었는데 위의 영상에 언급되어있었다.

 

1. Unsupervised learning과 Self-supervised learning은 같은 개념인가?

나는 label을 전혀 가지지 않는 unsupervised learning과 직접 label을 정의하여 학습을 진행하는 self-supervised learning은 다른 개념이라고 이해하고 있었다. 위의 사진을 보면 LeCun 교수님 또한 self-supervised learning은 조금 도 supervised한 방식이기 때문에 unsupervised learning과 혼동해서는 안된다고 설명한다. 하지만 이 논문에서는 이 두 개념을 같다고 가정하고 논문을 시작하고 있기 때문에 주의해야할 것 같다.

 

2. Self-supervised learning의 세부분야는 어떻게 나뉘는가?

사실 이것은 명확한 정의가 있는 것이 아니라, 보는 관점에 따라서 조금씩 달라질 것이라고 생각한다. 나는 크게 pretext task, contrastive learning, masked modeling으로 나눠서 이해하고 있었다. 이와 달리 본 논문에서는 pretext task라는 큰 범주 안에서 크게 intra-image prediction과 inter-image prediction으로 나누고 있다는 것을 먼저 이해해야 편할 것 같다.

 

여기서 intra, inter image prediction의 의미는 아래와 같다.

Intra-image prediction: single image(논문에서는 single batch라고 표현)로 학습 진행 -> Jigsaw puzzle, Colorization

Inter-image prediction: 서로 다른 images(논문에서는 dual-batches라고 표현)를 가지고 학습 진행 -> SimCLR, MoCo

 

 

 

Related Work

본 논문의 원조가 되는 논문은 Jigsaw Puzzle이다. 따라서 해당 논문이 어떤 내용을 다루는지 간단하게 살펴보고 넘어가면 좋을 것 같다.

Jigsaw puzzle은 self-supervised learning 중 pretext task의 한 예시이다.

 

하나의 이미지를 mxm 크기의 patch로 나누고, 이것의 위치를 puzzle처럼 섞었다가 다시 원래 자리로 되돌리는 것을 학습시키는 알고리즘이다. 본 논문에서는 해당 방법을 하나의 이미지만 사용한다라는 점에서 intra-image prediction으로 분류한다.

 

 

문제 정의

이 논문에서는 기반이 되는 intra-image prediction(Jigsaw puzzle)의 단점을 보완하고자 한다. 기존 intra-image prediction 방법은 single batch를 학습에 사용하기 때문에 cost는 상대적으로 작지만 feature의 추출에서 한계점이 존재했다.

 

반면 inter-image prediction의 경우 서로 다른 이미지를 비교하며 학습하기 때문에 비교적 feature를 잘 추출할 수 있지만, dual batches를 사용한다는 점에서 더 큰 cost를 유발한다는 문제점이 있다.

 

그래서 논문의 저자는 두 방법을 적절히 혼합하여 intra-image prediction의 낮은 cost와 inter-image prediction의 높은 성능을 유지하고자 했다. 그 연구의 결과로 Jigsaw clustering을 설명한다.

 

 

 

해결 방법 및 원리

-방법론

두 방법론을 합친다는 간단한 생각에 착안하여 저자가 소개하는 방식은 위의 사진과 같다. 먼저 Jigsaw puzzle이 mxm개의 patch로 각 이미지를 나누는 것을 그대로 수행한다. 이후 inter-image prediction 방법의 특성을 부여하기 위해 서로 다른 이미지에서 나온 patch를 적절히 섞어준다. 즉 montage(몽타주) 이미지를 만드는 것이다.

 

논문의 저자는 이 몽타주 이미지를 만든 것을 주요한 특징으로 설명하고 있다. 몽타주 이미지를 만드는 것이 아니라 모든 single patch들을 가지고 학습을 진행하게 되면 이후 오직 global information만을 얻을 수 있다고 한다. 또한 small single patch가 다양한 응용 상황에서(pretext task와 downstream task에서 일치하지 않는 문제점) 흔하지 않기 때문에 문제를 유발할 수 있다고 설명한다. 때문에 몽타주 이미지를 적용하여 이러한 문제점을 해결하게 되었다.

 

과정을 조금 더 구체적으로 살펴보면 아래와 같다.

하나의 batch X 안에는 총 n개의 서로 다른 이미지가 존재 = {x1, x2, ..., xn}

각 이미지 xi를 mxm개의 patches로 나누기 = 하나의 batch에는 총 nxmxm개의 patches 존재

서로 다른 이미지에서 나온 patch들을 섞어서 새로운 batch X' 만들기 = {x1', x2', ..., xn'}

 

최종적으로 network는 총 2개의 문제를 해결하게 된다.

1. 각 patch가 서로 같은 이미지를 나타내고 있는가 (inter-image prediction)  -> supervised clustering으로 해결 가능

2. 각 patch의 위치가 올바른가 (intra-image prediction)  -> classification problem으로 해결 가능

 

 

- 장점

이런 혼합된 방법이 어떠한 장점을 가져다주는 것일까.

먼저 기존의 contrastive learning은 dual-batches를 사용하는 반면 이 방법은 single-batch를 사용해서 절반의 training batches를 사용한다는 장점이 있다. 따라서 cost도 감소하는 것이다.

 

두 번째로 서로 다른 이미지를 조합한 몽타주 이미지에서 이것들을 구분해내는 것을 학습하는 과정에서 instance-level information을 학습할 수 있다. (다른 contrastive learning에서는 일반적으로 얻을 수 없는 정보이다)

 

세 번째로 구분한 각 이미지들을 같은 이미지끼리 clustering하는 과정에서는 image-level information을 학습할 수 있게 된다.

 

 

다음은 실험을 하며 저자가 고려한 사항이다. Patch를 나눌 때 어떠한 overlap도 없이 disjoint하게 patch를 나누게 되면 이미지에 따라서 전혀 같은 이미지 같지 않은 patch들도 존재할 수 있다. 따라서 어느정도의 overlap을 주어서 network로 하여금 더욱 쉽게 이해할 수 있도록 만들었다고 설명한다. 그 비율은 실험 결과 0.3이 최적이었다고 한다.

 

 

-네트워크 구조

백본으로는 ResNet-50을 적용하였다.

 

이후 각 이미지를 임베딩하여 벡터로 표현한다.

 

이때 특이한 decouple module이다. 이 부분은 임베딩된 각 이미지의 벡터의 가로, 세로가 m의 정수배가 아닌 경우 patch를 만들 때 정보 손실이 발생하는 것을 보완하고자 interpolation을 진행하는 부분이다. Inerpolation과 Average pooling만을 진행하기 때문에 파라미터가 존재하지 않는다는 특징을 가진다.

 

이후 하나의 batch를 nxmxm개의 patch를 가지는 몽타주 이미지들로 재구성한 뒤 각각 clustering branch와 location branch로 통과시킨다. 이 두 branch는 앞서 살펴본 해결해야하는 2가지 문제를 위한 branch라고 이해할 수 있다. Clustering branch는 N개의 class로 분류하기 위해 1개의 FC layer와 1개의 Normalization layer 총 2개의 layer로 구성되어 있고, location branch는 각 patch의 본래 위치를 학습하기 위해 1개의 FC layer로 구성되어 있다.

 

 

 

Loss function

Clustering branch의 경우, 식이 복잡했지만 조금 찾아보며 이해를 해보았다. 먼저 Ci는 i번째 patch와 같은 class에 속하는 patch들의 index를 담고 있다. 그 index들을 j로 표현하고 i번째 patch와 j번째 patch의 cosine similarity를 계산하고 이것을 log softmax에 대입한다. 그리고 이것의 평균을 구한다. 최종적으로 이것을 모든 patch 각각에 대해서 반복하고 평균을 구하게된다.

 

이전에 살펴봤던 contrastive learning의 loss function과 유사하다고 생각된다. 서로 같은 class를 가지는 patch끼리는 거리를 줄이고, 서로 다른 class를 가지는 patch 사이의 거리를 늘려주는 loss function인 것이다.

 

 

Location branch는 앞서 언급했듯이 classification task로 바라볼 수 있다. 따라서 cross entropy loss를 적용한 것을 알 수 있다.

 

이 두 가지 branch의 loss를 최종적으로 조합하여 loss를 계산한다. 실험 결과로는 알파, 베타 모두 1인 것이 좋은 결과를 만들었다고한다.

 

 

 

결과

 

결과를 살펴보면 당연하게도 Supervised learning이 가장 좋은 성능을 나타내는 것을 알 수 있다. 표가 3등분 되어있는데 가운데 부분은 intra-image prediction, 마지막 부분은 inter-image prediction이다. 본 논문의 정확도를 살펴보면 66.4로  intra-image prediction 방법 중에서는 가장 뛰어나며 inter-image의 SOTA model인 MoCo v2와 1.3 정도의 차이만을 보여주는 것을 알 수 있다. 즉 single-batch를 사용하여 cost를 줄였음에도 성능은 dual-batches method와 비슷하거나 뛰어난 것을 알 수 있다.

 

해당 그래프로는 Jigsaw clustering의 빠른 convergence 속도를 이해할 수 있다. 논문의 저자는 효과적은 contrastive pair의 사용으로 빠른 수렴을 가능하게 했다고 설명한다. 이는 contrastive learning에서는 positive, negative 이미지를 나눠서 학습을 하지만, jigsaw clustering의 경우 patch로 나눈 이미지에서 이러한 pair를 학습하기 때문에 더 효과적이라는 의미라고 이해하였다.

 

조금 더 작은 dataset에 대해서도 뛰어난 성능을 보여주는데, 이것은 해당 모델이 instance, image level의 information을 모두 학습할 수 있기 때문이라고 설명한다.

 

위의 결과는 ImageNet에서 pretrain하고 semi-supervised learning에 적용한 결과이다. 가운데 부분은 semi-supervised learning이고 마지막 self-supervised learning의 결과이다. UDA와 유사한 성능 혹은 조금더 뛰어난 성능을 보이는 것을 알 수 있다. 또한 supervised learning보다 뛰어난 것을 알 수 있다.

 

이것은 ImageNet에서 pretrain한 MoCo와 JigClu를 Faster-RCNN에 transfer하여 object detection을 한 결과이다. 상대적으로 더 좋은 성능을 보여주는 것을  알 수 있다.

 

위의 모델은 ImageNet에서 pretrain한 model을 classifier에서 transfer한 뒤 CIFAR 데이터에서 분류 작업을 수행한 결과이다. 여기서도 Jigsaw clustering이 우수한 성능을 보여주는 것을 알 수 있다. (linear는 백본 모델을 그대로 사용하고 마지막 classifier layer만 학습시킨 것을 의미한다)

 

 

위의 결과는 앞서 살펴봤던 montage image를 사용한 것과 small size patch, scaled up patch를 사용한 것의 차이를 보기위해 진행한 실험이다. 언급했듯이 small-size single patch를 사용하게 되면 일반 이미지들과의 차이가 생기고, downstream task를 처리할 때 그 차이로 인해 문제가 생기게 된다. 따라서 정확도가 낮아지는 것을 확인할 수 있다.

 

조금 신기했던 부분은 그 작은 이미지를 scale-up해서 사용했을 때 정확도가 증가한 부분이다. 정확한 이유는 알 수 없으나 미세하게 몽타주 방식보다 성능이 높게 나타난 것을 알 수 있다. 하지만 각 이미지를 interpolation하는 과정, 그리고 각각의 patch를 따로 처리하다보니 높은 cost와 긴 시간을 소비하는 것을 알 수 있다.

 

 

다음은 data augmentation을 언제 진행하는 것이 가장 효과적인지를 실험한 것이다. 생각보다 차이가 커서 놀랐던 결과이다. Jigsaw clustering에서는 각 patch를 나누는데 이로 인해서 언제 augmentation을 진행하는지가 결과에 영향을 미치게 되었다. 실험 결과 mxm patches로 이미지를 분할한 후에 augmentation을 수행하는 것이 가장 좋은 성능을 내는 것을 알 수 있다. 

 

split 이전에 수행하게 될 경우, clustering branch가 image feature를 배우는 것이 아니라 augmentation에 치우쳐서 학습을 진행하게 되는 문제가 생기게된다. 또한 augmentation을 진행하게 되면 original 이미지에서 정확하게 어떤 location에 두어야하는지 배우기가 어려워진다. 이로인해 3.7%로 매우 낮은 정확성을 보이는 것을 알 수 있다.

 

 

마지막은 ablation study의 결과인데, 각 branch의 중요성을 나타냈다고 생각할 수 있다. 여기서 clustering branch가 더욱 높은 정확도를 나타냄을 보아, 네트워크에서 더 높은 중요성을 가진다고 생각할 수 있다. Location branch에서는 조금 더 detail한 정보를 학습하는 것이라고 설명한다.

 

 

 

결론

Self-supervised learning에서 intra, inter image prediction의 차이와 그 특징에 집중하여 고안한 논문이라고 생각되었다. 특히 마지막 결과에서 jigsaw puzzle의 location을 맞추는 해결 능력에 기반을 두었지만 최종 branch의 중요도에서 clustering branch가 더욱 중요하게 나타났다는 점이 신기하게 느껴지기도 하였다. 기존에 존재하는 두 가지 방법의 장점을 결합하기위해서, 간단해보이는 아이디어를 생각하고 좋은 결과까지 이끌어내도록 만들었다는 것에 좋은 논문이라는 생각이 들었다.

 

반응형