Emerging Properties in Self-Supervised Vision Transformers

Inc Lomin's avatar
Oct 29, 2021
Emerging Properties in Self-Supervised Vision Transformers
notion image

개요

Vision Transformer(ViT)에 새로운 self-supervised learning 을 적용하였을 때, Convolution 기반 네트워크와 어떤 차이가 있는지를 연구한 논문이다.
저자가 관찰한 주요한 특징은 다음과 같다.
  1. self-supervised ViT의 feature에서는 supervised ViT 와 convnet 에서는 나타나지 않았던 segmentation feature가 포함되는 것을 확인
  1. k-NN classifier에 적용했을 때에도 괜찮은 수준의 성능을 보여줌
 
실험을 통해 다음과 같은 요소들이 ViT에서 중요한 역할을 하는 것을 확인
  • Momentum Encoder
  • Multi-Crop Training
  • Use of small patches with ViTs
 
저자들은 이러한 관찰의 결과로부터 DINO(self-distillation with no labels) 라는 self-supervised method를 제안하였고 실험을 통해 DINO와 ViT가 좋은 시너지를 보여주는 것을 확인했다고 한다.

Motivation

NLP 분야에서 self-supervised pretraining이 좋은 성능을 보여주는 것을 확인 (BERT, Language Modeling)
⇒ 문장 내에 있는 단어가 더 풍부한 학습 시그널을 만들어내는 것에 도움을 줌
⇒ 이미지에서는 object가 사전 정의된 몇 천개의 카테고리로 제한되기 때문에 visual information이 손실
⇒ 이미지에 대해서도 마찬가지로 이미지 내의 visual information이 좋은 시그널을 줄 수 있을 것
⇒ convnet 에서 이미지에 대한 self-supervised method를 적용했을 때 잠재성이 있다는 연구가 있었음
⇒ ViT 에 적용했을 때에는 어떤 결과가 나올까?

접근방식

Self-Supervised Learning(SSL) with Knowledge Distillation

notion image
notion image
 
  1. augmentation을 통해 서로 다른 2 개의 global view를 생성.
    1. (global view : 전체 면적의 50% 이상의 크기를 가지는 이미지)
  1. teacher network 에서는 global view만을 입력으로 사용
  1. student network 에서는 global view와 추가로 이미지를 crop 한 local view를 입력으로 사용
    1. (local view : 전체 면적의 50% 미만의 크기를 가지는 이미지)
      ⇒ local-to-global correspondence 를 목표로 함
  1. teacher와 student 의 feature representation의 cross entropy loss가 최소가 되도록 학습
notion image
notion image
 
Teacher Network
일반적인 Distillation 에서는 Student 모델을 학습시키기 전에 먼저 학습을 진행하지만, 여기서는 학습을 하지 않고 과거의 Student Network 로부터 만들어내는 방식이 적용되었다.
  • 이전 epoch 에서의 student network 를 teacher 로 사용
  • 이전 iteration 에서의 student network 를 teacher 로 사용
  • student network의 copy를 teacher 로 사용
  • momentum encoder 를 teacher 로 사용
본 논문에서 가장 좋은 성능을 보여주었던 방식은 exponential moving average(EMA)을 사용하는, momentum encoder 라고 한다. momentum encoder를 적용한 teacher network 의 파라미터 업데이트 식은 다음과 같다.
θt ← λθt + (1 − λ)θs
즉 새로 학습한 student의 weight와 이전에 사용한 teacher의 weight를 이동평균한 것이 새로운 teacher의 weight가 된다. 기존 momentum encoder와는 다르게 contrastive loss를 계산하거나 negative sample를 저장하는 queue 같은 것을 사용하지 않고 순수하게 이동평균한 weight만을 사용한다.
이러한 방식은 model ensemble과 비슷한 역할을 하여 teacher network가 student 보다 항상 좋은 성능을 가지도록 만들어 학습이 지속적으로 이루어질 수 있도록 만들어준다.
[Reference] Momentum Contrast for Unsupervised Visual Representation Learning
notion image
notion image
 
Avoiding Collapse
본 논문에서 제안한 framework를 수행하는 과정에서 representation collapse가 발생하여 다양한 normalization 방법을 적용하였는데, centering + sharpening + momentum teacher 를 사용했을 때 비로소 안정적인 학습이 되었다고 한다.
notion image
sharpening
notion image
sharpening을 하게되면 feature에서 noise를 제거하고 원하는 feature만을 강조하여 teacher의 signal 극대화할 수 있다. 하지만 softmax가 지수함수이기 때문에 값이 커질수록 빠르게 증가하는데, 이는 수치적인 오류를 발생시킬 여지가 커지는 것을 의미한다. (특정 차원에서 지나치게 큰 신호가 발생할 가능성)
centering
notion image
centering은 일반적으로 많이 사용하는 정규화 방식으로, 평균을 빼서 원점을 중심으로 가지도록 만들어주는 역할을 한다. sharpening 과는 반대로 수치적인 안정성을 가지게 되지만 지수함수에서 값이 작아지게 되면 기울기가 0에 가까워 지면서 값 사이에 큰 차이가 없게 될 수 있다. (uniform distribution에 대한 가능성)
 
notion image
각각의 정규화 방법은 서로 반대되는 효과를 보여주는 것을 확인할 수 있는데, 두 정규화 방법을 동시에 적용하여 서로의 단점을 보완하도록 만들었다.

Architecture

구현은 Data-efficient image Transformers(DeiT) 구조를 그대로 사용하였다고 한다. 여기서 distillation token은 사용되지 않았기 때문에 base ViT와 head만 다른, 거의 동일한 구조를 가진다고 보면 된다.
notion image
 
[Reference] DeiT vs DINO
def forward_features(self, x): # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py # with slight modifications to add the dist_token B = x.shape[0] x = self.patch_embed(x) cls_tokens = self.cls_token.expand(B, -1, -1) # stole cls_tokens impl from Phil Wang, thanks dist_token = self.dist_token.expand(B, -1, -1) x = torch.cat((cls_tokens, dist_token, x), dim=1) x = x + self.pos_embed x = self.pos_drop(x) for blk in self.blocks: x = blk(x) x = self.norm(x) return x[:, 0], x[:, 1] def forward(self, x): x, x_dist = self.forward_features(x) x = self.head(x) x_dist = self.head_dist(x_dist) if self.training: return x, x_dist else: # during inference, return the average of both classifier predictions return (x + x_dist) / 2
def forward_features(self, x): B = x.shape[0] x = self.patch_embed(x) cls_tokens = self.cls_token.expand(B, -1, -1) x = torch.cat((cls_tokens, x), dim=1) pos_embed = self.interpolate_pos_encoding(x, self.pos_embed) x = x + pos_embed x = self.pos_drop(x) for blk in self.blocks: x = blk(x) if self.norm is not None: x = self.norm(x) return x[:, 0] def forward(self, x): x = self.mlp(x) x = nn.functional.normalize(x, dim=-1, p=2) x = self.last_layer(x) return x
[Reference] Vanilla ViT
notion image

실험

Network Configuration
notion image
  • DeiT-S, ViT-B 는 이름만 다를 뿐 동일한 구조인데 hidden dimension과 head의 수만 다르다.
  • DeiT-S/16 은 ResNet50 과의 비교할만한 비슷한 파라미터수, 처리량을 가지는 모델이다.

Comparing with SSL frameworks on ImageNet

notion image
Reference
Architecture
  • ResNet 50 (RN50)
  • DeiT-S : ResNet 50 과 비슷한 파라미터의 수(Param), 처리량(im/s), 정확도를 가지는 ViT
  • ViT-B : base vision transformer (DeiT 구조에서 distillation token이 없는 구조와 동일)
Methods
  • BYOL : negative sample 없이 positive sample 만으로 self-supervised learning을 수행
    • notion image
  • MoCov2 : key를 계산할 때 momentum encoder를 사용. queue를 사용하여 negative sample의 일관성 유지 학습이 불안정한 이슈 해결
    • notion image
  • SwAV : feature를 z, code를 q라고 하고, (z1, q1), (z2, q2) 가 주어졌을 때 z1으로 q2를, z2로 q1을 예측하도록 swap 하여 학습하는 방법
    • notion image
  • SCLR : 라벨 없이 contrastive learning 으로 학습하는 방법
    • notion image
      notion image
Comparing with the same architecture
  • ViT 아키텍쳐에 적용했을 때 다른 방법보다 좋은 성능을 확인
  • k-NN classifier 사용하여 분류를 수행했을 때에도 linear classifier 와 비슷한 성능
Comparing across architectures
  • 더 작은 patch size를 사용했을 때 성능이 더 좋은 것을 확인 (ViT-B/8 vs ViT-B/16)
  • 적은 파라미터만을 가지고도 좋은 성능을 보여주는 것을 확인 (DeiT-S/8, ViT-B/8)
  • k-NN 의 경우 월등한 성능
    • ⇒ 좋은 Feature를 뽑아낸다
 

Properties of ViT trained with SSL

Nearest neighbor retrieval with DINO ViT
Image Retrieval
notion image
  • DINO method 로 학습시키는 경우, convnet 보다 ViT로 학습시킨 것이 더 좋은 성능
  • DINO 의 장점은 라벨이 없어도 어떤 데이터셋이든 학습이 가능
    • ⇒ ImageNet이 아닌 GLDv2 로 pretrain 이 가능, 더 많은 데이터셋을 학습에 사용하여 더 좋은 성능을 이끌어내는 것을 확인
 
Discovering the semantic layout of scenes
Video instance segmentation
notion image
  • dense task 에 대해 학습하지 않아도 경쟁력 있는 결과 (DeiT-S/16, ViT-B/16)
  • dense task 에 대해 학습하는 경우 상당한 성능 향상 확인 (DeiT-S/8, ViT-B/8)
  • self-supervised 방식을 적용한 결과를 비교했을 때 가장 좋은 성능
Probing the self-attention map
notion image
notion image
 
 
  • 서로 다른 head가 서로 다른 sementic region을 주목하는 결과가 나타남
  • 무언가에 가려져 있거나 작아도 잘 찾아내는 것을 확인할 수 있음
  • supervised ViT 는 object 에 attention이 생기지 않지만, DINO 로 학습한 모델은 확실한 차이를 보임
  • self-supervised convnet도 segmentation 정보를 포함하고 있지만 특별한 처리가 필요
Transfer learning on downstream tasks
notion image
  • 같은 architecture 에서도 self-supervised learning 을 적용한 것이 더 성능이 좋은 것을 확인

Ablation Study of DINO

Importance of the Different Components
notion image
  • Momentum 없이는 동작하지 않음
  • Sinkhorn-Knopp 와 같은 collapse 회피 기법은 momentum 과 같이 사용할 때에는 효과가 별로 없음
  • 3, 9 를 비교해보면 momentum의 유무가 성능에 크게 영향을 미친 것을 확인
  • 1, 4 를 비교해보면 multi-crop이 성능에 큰 영향을 미친 것을 확인
  • 1, 5 를 비교해보면 cross entropy loss 를 사용하는 것이 적합하는 것을 확인
  • student network의 predictor 의 유무는 큰 영향이 없음 (BYOL 에서는 critical 한 요소)
 
notion image
  • patch size가 작을수록 좋은 성능을 보여줌
  • patch size가 작을수록 처리량이 감소
 
Impact of the choice of Teacher Network
notion image
  • 현재 혹은 이전 iteration 에서의 student weight 를 가져오는 경우에는 학습이 되지 않는 현상
    • ⇒ 모델이 자체에 큰 차이가 없어서 learning signal 이 너무 약해진 것으로 생각됨
  • 이전 epoch의 weight를 사용하는 경우에는 collapse 되지 않는 것을 확인
 
Momentum 방식으로 teacher의 weight를 설정했을 때 항상 student보다 좋은 성능을 보여줌
  • 이 부분은 convnet을 사용했을 때에도 비슷한 결과
  • momentum 방식 자체가 Polyak-Ruppert averaging 을 의도하고 만든 것인데, 이 방법을 통해 model ensemble 과 같은 효과를 얻을 수 있었기 때문에 이러한 결과가 나타난 것으로 보인다.
notion image
Avoiding collapse
notion image
  • 둘 중 하나만 적용한 경우 KL divergence가 0으로 수렴
    • (=teacher와 student가 항상 동일한 representation을 출력)
  • entropy의 경우 sharpening과 centering이 서로 다른 값으로 수렴
    • (= 서로 다른 형태의 collapse, 0: sharpening, - log 1/K: centering)
 
Compute requirement
notion image
  • BYOL : 4096 batch 512 x TPU cores
  • SwAV : 4096 batch 64 x V100 16GB GPU ( 2 x 224 + 2 x 160 + 4 x 96 = 88064 )
  • DINO : 1024 batch 16 x V100 16GB GPU ( 2 x 224 + 10 x 96 = 92160)
 
Training with small batches
notion image
  • Batch size가 작아도 잘 학습되는 것을 확인
 
Share article