[Paper Review] DINO - Emerging Properties in Self-Supervised Vision Transformers (Mathilde Caron, 2021)

References

ICCV 2021. [Paper] [Github]

Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, Armand Joulin

Facebook AI Research | Inria | Sorbonne University

29 Apr 2021

 

Abstract

• Self-supervised learning이 ViT 분야에서 new property를 제공하는지? (CNN과 비교) → ViT에 Self-supervised Learning을 적용하여 그 효과를 확인함
• Self-supervised methods가 ViT에 잘 작동함. 단순히 성능이 높을뿐만 아니라 2가지 이점이 존재함.
  1. self-supervised ViT features에는 이미지의 semantic segmentation에 대한 explicit information이 포함되어 있음 (이는 supervised ViT나 convnet에서는 나타나지 않음)
  2. 이러한 feature들은 뛰어난 k-NN classifier 이기도 함 (78.3% top-1 on ImageNet with a small ViT)
• DINO란?
  • a simple self-supervised method
  • a form of self-distillation with no labels
• 다음 3가지 방법이 중요함

1. Momentum encoder
2. multi-crop training
3. the use of small patches with ViT

⇒ ViT-Base에 DINO를 결합한 방법으로 좋은 시너지를 낼 수 있음

 

1. Introduction

• Transformer가 visual recognition 분야에서 convnet (convolutional neural network)의 대안으로 떠오르고 있으며, 대부분 NLP에서 영감을 얻은 훈련 전략을 취함 (ex. pre-training on large dataset and fine-tuning on the target dataset)
• 그 결과 ViT는 어느정도 성능을 내었지만, 아직 convnet에 비해 명확한 이점을 제시하지 못했음 (반면, 계산 비용이 더 요구되며, 더 많은 학습 데이터가 필요하고, feature들이 unique property를 나타내지 않음)
• ⇒ Transformer 모델이 이미지 분류에서 성능이 좋지 않은 이유에 대해서 탐구함 : pre-training에서 supervision 방식을 사용하기 때문

본 논문에서는 NLP에서 Transformer가 좋은 성능을 보인 주요한 요인 중 하나는 ***‘self-supervised pre-training’***을 사용했다는 점이라고 말함. (더 풍부한 Learning Signal 전달)

⇒ Computer Vision 분야에서도 self-supervised pre-training을 사용하여 성능을 향상시킬 수 있는지 연구

• Self-supervised pre-training을 사용하여 ViT 모델 학습시켜, Self-supervised Pretraining이 ViT feature에 미치는 영향을 분석함

  

  Figure 1: Self-attention from a Vision Transformer with 8 × 8 patches trained with no supervision. We look at the self-attention of the [CLS] token on the heads of the last layer. This token is not attached to any label nor supervision. These maps show that the model automatically learns class-specific features leading to unsupervised object segmentations

   
  • Self-supervised ViT feature 들은 위의 이미지와 같이 장면 레이아웃(scene layout), 특히 객체 경계(object boundaries)를 명시적으로 포함함. 이 정보는 마지막 블록의 Self-attention 모듈에서 직접 접근할 수 있음Figure 1: Self-attention from a Vision Transformer with 8 × 8 patches trained with no supervision. We look at the self-attention of the [CLS] token on the heads of the last layer. This token is not attached to any label nor supervision. These maps show that the model automatically learns class-specific features leading to unsupervised object segmentations
  • Self-supervised ViT feature들은 finetuning, linear classifier, data augmentation 없이 기본 k-NN에서 특히 잘 수행되어 ImageNet에서 78.3%의 top-1 accuracy를 달성하였음

 

본 연구에서 발견한 이러한 특징은, 레이블이 없는 Knowledge Distillation의 한 형태로 해석될 수 있는 간단한 Self-supervised 방법의 아키텍처를 설계하도록 하였고, 이를 DINO라고 함.

표준 Cross Entropy Loss를 활용하여 Momentum Encoder로 만들어진 Teacher Network의 출력을 직접 예측하도록 Self-supervised training을 단순화함

 

※ Knowledge Distillation

: Pretrained 된 Teacher model을 활용하여, student model에서도 Teacher model과 비슷한 성능을 낼 수 있도록 지식을 전달하는 방법

• Distillation Loss
  • Teacher model이 가지고 있는 Knowledge를 Student model에게 Distillation을 하기 위한 Loss
  • Tempearture hyperparameter로 softning한 각각의 Teacher와 Student의 soft target을 KL divergence Loss를 이용하여 계산
  • KL divergence로 두 확률분포의 차이를 계산할 수 있기 때문에 Distillation Loss에서 KL divergence를 사용함으로서 Teacher와 Student의 확률분포가 비슷해지도록 학습을 유도
• Student Loss
  • Student Loss를 계산하기 위해 Cross Entropy Loss를 사용하고 soft targets이 아닌 T = 1인 Hard targets을 사용
  • Student model는 Distillation Loss를 통해 Teacher model의 knowledge를 전수받고, Student Loss를 통해 더 정확한 값을 학습할 수 있음

또한, DINO Framework는 유연하기 때문에 아키텍처를 수정하거나 내부 정규화를 수정할 필요 없이, Convolution Network와 ViT 모두 작동함

ImageNet Benchmark를 통해 DINO와 ViT간의 시너지를 검증하였고, Convolution Network와도 함께 작동한다는 사실을 확인함

+기존의 Self-supervised 방법에서는 segmentation mask 의 출현이 공통 속성으로 나타나는 것으로 보임

+하지만 k-NN에서 우수한 성능은 momentum encoder와 multi-crop augmentation과 같은 특정 구성 요소를 결합할 때만 나타남

+또한 ViT와 resulting features의 퀄리티를 개선하기 위해선 smaller patch를 사용하는 것이 중요하다는 것을 발견함

 

2. Related work

Self Supervised Learning

Supervision 학습 방법과 달리, unlabeled data를 사용하여 모델을 학습시키는 방법

데이터에서 유용한 feature를 추출, 이후 다양한 task를 수행

본 논문에서는 BYOL의 방법론을 사용하지만, Loss가 다르고 동일한 구조의 Teacher, Student Network를 사용한다는 차이가 있음

Self-training and knowledge distillation

이전에 Self-Supervised Learning과 Knowledge Distillation 방법을 결합한 연구가 많지만, 본 논문에서는 Pre-trained Teacher Network를 사용하지 않는다는 차이가 있음

 

3. Approach

• DINO (Distillation of Non-contrastive Image representations) 제안
• 라벨이 없는 knowledge distillation (with no label)의 한 형태로 해석될 수 있는 간단한 self-supervised 방법
• Contrastive learning이 아닌, distillation 기반의 self-supervised learning 방법으로, 이미지의 feature를 추출하는 teacher 모델과 student 모델을 사용
• Contrastive learning (CRL)
  : 입력 샘플 간의 비교를 통해 학습을 하는 것. self-supervised learning에 사용되는 접근법 중 하나. 사전에 정답 데이터를 구축하지 않는(no label) 판별 모델.
• DINO 학습 순서
  1. Teacher 모델을 먼저 학습시키고, 이를 통해 이미지의 feature를 추출
  2. Student 모델 학습. 이 때 teacher 모델이 추출한 특징을 distillation을 통해 student 모델에 전달 이를 통해, student 모델은 teacher 모델이 추출한 특징을 학습함

 

3.1. SSL with Knowledge Distillation

• DINO 는 SSL(Self Supervised Learning) 방법을 사용하고 있으며, Knowledge distillation 형태를 띔
• 위 그림에서 볼 수 있듯이 student network $g_s$ 를 학습할 때, teacher network $g_t$와 일치시키는 방법으로 학습

 

1. 이미지 x가 주어졌을 때, Teacher, Student network 모두 $K$ 차원의 확률 분포 Ps와 Pt 출력함. 이때, 확률분포 P는 네트워크 g의 출력을 softmax function 으로 정규화하여 얻어짐 (이때, tau > 0는 temperature parameter로 출력 분포의 첨도, 즉 뾰족한 정도를 조절)

 

2. 고정된 teacher network gt가 주어졌을 때, student network의 parameter에 대한 cross-entropy loss 를 최소화하여, 두 신경망의 출력 분포를 일치시킴

여기서,

 

3. Self-supervised learning에 적용되는 부분

우선, Multi-crop view strategy를 활용하여 다양한 이미지 셋 V를 만들고, 이 때 Local view와 Global view를 포함하여 최종으로 아래와 같은 Loss를 최소화함

 

< DINO 의 pseudo-code >

 

Teacher Network

DINO는 다른 knowledge distillation과는 다르게, pretrained Teacher network를 가지지 않음

따라서 Iteration의 구성을 Student network 이전으로 구성하였음

여러 업데이트 규칙을 실험해 본 결과 Student의 가중치에 EMA 방식을 사용하는 Momentum Encoder가 좋았음

 

※ EMA(Exponentially Moving Average; 지수 이동 평균)

데이터의 이동 평균을 구할 때 오래된 데이터가 미치는 영향을 지수적으로 감쇠되도록 만들어주는 방법

 

Mean Teacher

Mean Teacher 는 self-supervised learning에서 학습하는 student model의 성능을 향상시키기 위해 사용되는 기술임.

이 기술에서는 student model과 teacher model 두 개의 모델이 사용되는데, Teacher model은 student model의 가중치의 이동 평균으로, 학습 중에 student model에게 soft target을 제공하는 데 사용됨

student model은 teacher model이 제공하는 soft target과의 차이를 최소화하도록 학습함. 이를 통해 student model은 더 강한 특징(두드러지는 특징)을 학습하고, unseen data에 대해 더 일반화된 예측을 할 수 있게 됨

Mean Teacher 는 self-supervised learning 모델, 특히 ViT 모델의 성능을 향상시키는 데 효과적이라는 것이 입증됨. 본 논문에서는 Momentum Teacher 라는 Mean Teacher의 변형을 사용하는데, Momentum Teacher는 student model의 가중치의 지수적으로 감소하는 이동 평균을 teacher model로 사용함

 

이와 같은 teacher network의 파라미터는 student의 평균이 되는 데 이는 앙상블한 효과를 가짐 논문에서 실험 결과 teacher의 성능이 student보다 높다는 것을 보여줌

 

Multi-Crop Strategy

Multi-crop self-supervised learning 에서 사용되는 data augmentation 기법 중 하나

이는 이미지를 여러 개의 작은 패치로 나누어 각각을 모델에 입력시키는 방법

multi-crop을 사용하여 ViT 모델을 학습시키고, 성능을 향상시킴

두 가지 종류의 crop 사용

1. 224 x 224 크기의 큰 crop (→ Global view → teacher network에 사용)
2. 96 x 96 크기의 작은 crop (→ Local view → student network에 사용)

• 0.4~1배로 zoom

각 이미지에 대해 두 종류의 crop을 적용하여 총 10개의 crop을 얻음

⇒ 이렇게 smaller patch를 얻은 것이 성능을 향상에 도움이 됨

같은 이미지를 넣은 것이므로 분포가 유사해야 함 ⇒ Cross-Entropy Loss 계산

 

Centering 과 Sharpning

self-supervised learning 에서 ViT 모델의 안정적인 학습을 위해 centering 과 sharpening 이라는 방법을 제안함

Centering 은 teacher 네트워크가 생성한 feature 벡터의 평균을 각 벡터에서 빼는 방법. 이를 통해 어떤 차원도 지배하지 않도록 하고 feature 벡터가 더 균일하게 분포되도록 함

Sharpening 은 softmax 함수에 적용되는 temperature parameter($\Gamma$)를 조정하여 feature 벡터를 정규화하는 것을 의미함. temperature parameter의 값이 낮을수록 가장 유사한 feature 벡터에 가장 높은 확률이 할당되는 더 선명한 분포가 생성됨

이를 통해 feature 벡터를 더 구별 가능하게 만들고 학습 중에 쉽게 구분할 수 있도록 도움 (이는 ViT for small datasets에서도 사용한 기법)

Centering과 sharpening의 조합은 서로의 효과를 균형있게 조절하여 모델이 학습 중에 붕괴하는 것을 방지함

이는 논문에서 self-supervised ViT 모델의 학습을 안정화하고 downstream task에서의 성능을 향상시키는 데 효과적이라는 것이 입증되었다고 함