sy00n / DL_paper_review

issues에 논문 요약
0 stars 0 forks source link

[14] Unsupervised Learning of Visual Features by Contrasting Cluster Assignments #16

Open sy00n opened 1 year ago

sy00n commented 1 year ago

(Facebook AI Research에서 작성한 논문임.)

개인적인 생각으로는, self-supervised learning, representation learning 분야가 발전하면서 대부분 instance discrimination을 기반으로 한 instance-wise contrastive learning 방향으로 연구가 되어왔는데, 이는 사실상 안정적인 학습을 하려면 큰 배치가 필요해서 리소스에 따라 학습이 어려울 수 있으므로 현실적인 활용 측면에서의 치명적인 단점이 존재했지만, SwAV는 feature를 직접적으로 비교하지 않고 cluster간 비교를 통해 (cluster assignments간의 consistency를 강화하는 방식으로 학습하기 때문에) pair-wise comparison 없이도 online clustering based 방식으로 self-sup을 수행할 수 있다는 점에서 큰 의의가 있다고 생각함.

Abstract

기존 contrastive learning 방식들(MoCo, SimCLR)은 배치사이즈 크기에 영향을 많이 받는다. SwAV는같은 이미지로부터 생성된 representation vector가 같은 prototype class(cluster로 할당)를 예측하도록 학습하고 서로의 cluster를 맞바꾸어 예측한다. 그리고 개별 이미지의 positivie/negative pair를 직접적으로 비교하는 것이 아니라 같은 이미지로부터 agumentation 시킨 이미지에 cluster를 할당하고 이를 서로 바꿔서 맞추도록 한다. 또 multi-crop augmentation 방법을 계산효율적으로 적용했다.

Introduction

SSL에서 모든 instance간의 contrastive learning을 수행하는 것은 현실적으로 불가능하기에 batch끼리 비교하거나 clustering을 통해 code끼리 비교한다. (ex. DeepCluster에서는 prorotype matrix를 encoder로부터 분리해서 생성하고 역전파가 따로 흘러간다. ) 이 방식은 결국 모든 dataset image를 사용해야 하기 때문에 scalable하지 않다.

Method

일반적인 clustering은 전체 data에 대한 image feature에 대해 cluster를 할당해주는 방식으로 학습되는 offilne 방식이다. 반면에 SwAV의 목표는 어떠한 supervision 없이 online으로 visual features를 학습하는 것을 목표로 한다. SwAV는 같은 이미지에서 파생된 augmented image code(군집)의 consistency를 강화하는 방향으로 학습을 진행한다. 즉 타겟 자체가 code인 것이 아니라 같은 이미지로부터 파생된 이미지는 같은 code에 할당되도록 일관화하자는 것이다. 그리고 multiple image view에 대 feature간의 비교가 아닌 cluster간의 비교를 수행하는 것이다. image fig 1을 보면, 한 이미지에 다른 augmentation을 적용해서 encoder를 통과시킨 representation을 z1, z2이라고 하면, z를 trainable한 prototype matrix C를 통해 각각의 code(cluster) q1, q2를 생성한다. 이 때 다른 augmentation을 적용했더라도 동일한 이미지에서 나온 code이기 때문에 z1로 q2를 예측할 수 있어야 하고 z2로 q1을 예측할 수 있어야 한다. 따라서 consistency loss를 다음과 같이 설정했다. image

image z는 non-linear mapping(encoder)를 통해 생성되고 각 feature는 normalize되어 unit sphere상에 projection된다. 그리고 z를 prototype vector C를 활용해서 q로 만든다. 이 때 q는 z와 c를 이어주는 그래프 형태이다. c는 k개의 대표 prototype vector를 가지는 matrix와 z들을 비교했을 때, q1은 z1이 k개 중 몇 번째 prototype vector와 가장 가까운지를 계산해서 어떤 prototype vector에 할당하면 될지를 안내해주는 일종의 path, 혹은 label 정도로 생각하면 된다.

Swapped predicion problem

image

즉, 주어진 prototype matrix C와 z_s를 통해 어떤 prototype과 가까운지에 대한 pseudo label을 생성하고 z_t와 C prototype vector와의 상대적인 거리를 구한 p_t를 생성하여 qs와 ct간의 CE를 구한다. 그리고 그 반대인 p_s, q_t 사이에서도 CE loss를 구해서 더하면 최종 loss 가된다.

Computing codes online

Q. 그럼 이제 code는 어떻게 구해지는 것일까? onlie으로 학습하기 위해서는 전체 이미지 feature가 아니라 batch내의 이미지 feature만을 사용해서 code를 계산해야 한다. 따라서 prototypes C가 서로 다른 배치에서 공유되어 batch내의 sample에 code를 할당한다. 이 때 모든 이미지가 같은 code로 매핑되는trival solution을 방지하기 위해 서로 다른 이미지들이 C에 의해 서로 다른 code로 균등배분 될 수 있게 code를 계산해야 한다.

배치 사이즈 B개의 feature vector Z와 C로 매핑된 code 벡터 Q가 있다고 할 때, 먼저 Z^TC로부터 계산된 code와 Q의 유사도가 높아야 한다. 따라서 내적 <Q, Z^TC>를 최대화시킨다. 또 배치 안 샘플들이 서로 다른 code로 균등분배되기 위해 엔트로피 항 H(Q)(정규화로 작용함)를 추가한다. 크기가 같은 행렬 A,B에 대해 <A,B> = Tr(A^TB)인 것을 이용해서 아래와 같이 code vector Q를 유도할 수 있다. image

입실론을 너무 크게 주면 엔트로피 항이 너무 정규화되어 모든 Qij가 같아버리는 trivial solution이 발생할 수 있다. 또 Q에 대해 optimal transport를 이용할 수 있도록 행렬의 합이 일정하도록 제약을 부여한다. image

이런 최적화 문제는 optimal transport문제이다. 서로 다른 분포를 가지는 데이터셋이 있을 때 하나의 데이터셋을 다른 데이터셋으로 옮기는 것으로, 다음 그림에서 빨간색을 파란색으로 옮기기 위한 최적의 couping matrix P를 구하는 문제이다. image image 이를 SwAV에 적용한 것이 아래와 같은 식의 형태가 되고 sinkhorn-knopp 알고리즘을 통해 3번의 반복만으로도 충분히 좋은 결과를 얻었다고 한다 . image

Experiments

image linear probing, semi-supervised learning 방법의 결과를 보면 feature를 frozen 했을 때 MoCov2보다 정확도가 4%만큼 증가했고 supervised와 2%밖에 차이나지 않는다.

참고: