Emerging Properties in Self-Supervised Vision Transformers 리뷰

Caron, Mathilde, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. "Emerging properties in self-supervised vision transformers." In Proceedings of the IEEE/CVF international conference on computer vision, pp. 9650-9660. 2021.

 

Dino라고도 알려진 모델의 논문이다.

코드는 https://github.com/facebookresearch/dino 에서 확인 가능하다고 한다.

 

Abstract

Convnet에 비해 ViT에서의 self-supervised learning은 효과적일까? 라는 의문을 제기한다. 우리는 실험에서 다음과 같은 결과를 확인할 수 있덨다. self-supervised ViT는 이미지의 semantic segmentation에서 supervised ViT, convnet에 비해 더욱 명확한 정보를 제공한다. 또한 완벽한 k-NN 분류기로도 작용하는데 작은패치의 사용, multi-crop training, momentum encoder의 사용 역시 매우 중요하다는 것을 강조한다. 우리는 self-DIstillation with NO labels 형태의 DINO를 제안한다.

 

Self Supervised Learning with Knowledge Distillation

DINO는 teacher, student가 동일한 구조를 가지는 모델을 이용해 knowledge distillation한다. student network의 출력인 \(g_\theta s\)는 teacher 의 출력인 \(g_\theta t\)와 동일해지도록 학습된다.

(1)을 통해 얻어진 수식으로 student network는 (2)의 손실함수를 이용해 학습이 진행된다.

DINO 의 모델 구조

\(P_t, P_x\) : 이미지 x 로부터 출력된 k 차원의 확률값 \(\tau_s\) > 0 : 출력 분포를 조정하기 위한 temperature H(a, b) : \(-a \log b\)

 

하나의 이미지에 대해 우리는 서로 다른 시각을 가진 이미지 셋 V를 생성한다. V는 \(x_1^g\), \(x_2^g\)의 global view와 cropping 되어 small resolution의 local views를 포함한다. student는 global, local views 모두 입력되지만 teacher은 오직 global views만 입력된다. standard setting에 따라 global = 224x224, local = 96x96의 해상도이다.

 

Teacher network

knowledge distillation과 다르게 teacher은 사전정의된 weight를 이용하지 않고 student의 past iteration의 weight를 이용한다. Ablation study에 의하면 student의 previous iter, copy를 이용한 teacher network update는 모델이 수렴하지 않았고, ema(지수 평균 이동)과 momentum encoder을 이용한 방법이 가장 좋은 성능을 보였다.

Network Architecture

backbone으로 ViT or Resnet을 채택하였고 projection head 는 3-layer MLP로 구성, 2048차원을 입력받아 K 차원 출력. self-supervised learning이지만 ViT를 이용하는 다른 연구들과의 일관성을 위해 CLS 토큰을 이용함.

모델 붕괴를 피하기 위해 우리는 정규화 대신 momentum teacher의 출력에 centering and sharpening을 적용했다.

centering은 \(g_t(x) \gets g_t(x) + c\) 와 같이 c 항을 추가하는 것으로 c는 다음 식과 같다.

챗지피티에 물어본 centering, sharpening의 함수는 다음과 같다고 한다. 참고만 하기 바람

def apply_centering_sharpening(teacher_output, student_output, epoch):
    centering_factor = centering_lambda ** (num_epochs - epoch)  # 코사인 스케줄 적용
    sharpening_factor = 1 - centering_factor

    # 중심화 연산 적용
    centered_teacher_output = teacher_output + centering_factor * torch.mean(student_output, dim=0)

    # 날카로움 연산 적용
    sharpened_teacher_output = centered_teacher_output * sharpening_factor

    return sharpened_teacher_output

 

Ablation Study

아주 많은 ablation study가 있지만 필자가 관심있게 본 것만 포스팅 하겠다.

ViT의 patch size에 따른 결과이다. 별로 차이가 없을것 같았지만 작은 패치에서 좋은 성능이 나왔지만 느리다고 한다.

 

collapse를 방지하기 위해 centering + sharpening을 채택했다고 하는데 이를 KL divergence와 비교한 그래프이다.

배치사이즈에 따른 모델의 결과 비교. contrastive learning에서는 배치사이즈가 크고 에포크가 많을 수록 모델의 성능이 증가하는데 128과 1024의 배치사이즈에서 비슷한 성능을 보인다고 한다(2% 차이면 큰 거 아닌가?). 또한 배치사이즈를 8로설정하고 50에포크를 훈련했을때도 35.2%가 나왔는데 이는 GPU 하나에 한 이미지가 들어가도 훈련할 수 있는 잠재력을 보인다고 한다(GPU 하나에 배치사이즈 1밖에 안될 정도로 큰 이미지를 말하는 듯).