Emerging Properties in Self-Supervised Vision Transformers

Inc Lomin

Oct 29, 2021

Emerging Properties in Self-Supervised Vision Transformers

Contents

개요 Motivation 접근방식 실험 Ablation Study of DINO

개요

Vision Transformer(ViT)에 새로운 self-supervised learning 을 적용하였을 때, Convolution 기반 네트워크와 어떤 차이가 있는지를 연구한 논문이다.

저자가 관찰한 주요한 특징은 다음과 같다.

self-supervised ViT의 feature에서는 supervised ViT 와 convnet 에서는 나타나지 않았던 segmentation feature가 포함되는 것을 확인

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

augmentation을 통해 서로 다른 2 개의 global view를 생성.

(global view : 전체 면적의 50% 이상의 크기를 가지는 이미지)

teacher network 에서는 global view만을 입력으로 사용

student network 에서는 global view와 추가로 이미지를 crop 한 local view를 입력으로 사용

(local view : 전체 면적의 50% 미만의 크기를 가지는 이미지)

⇒ local-to-global correspondence 를 목표로 함

teacher와 student 의 feature representation의 cross entropy loss가 최소가 되도록 학습

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

Avoiding Collapse

본 논문에서 제안한 framework를 수행하는 과정에서 representation collapse가 발생하여 다양한 normalization 방법을 적용하였는데, centering + sharpening + momentum teacher 를 사용했을 때 비로소 안정적인 학습이 되었다고 한다.

sharpening

sharpening을 하게되면 feature에서 noise를 제거하고 원하는 feature만을 강조하여 teacher의 signal 극대화할 수 있다. 하지만 softmax가 지수함수이기 때문에 값이 커질수록 빠르게 증가하는데, 이는 수치적인 오류를 발생시킬 여지가 커지는 것을 의미한다. (특정 차원에서 지나치게 큰 신호가 발생할 가능성)

centering

centering은 일반적으로 많이 사용하는 정규화 방식으로, 평균을 빼서 원점을 중심으로 가지도록 만들어주는 역할을 한다. sharpening 과는 반대로 수치적인 안정성을 가지게 되지만 지수함수에서 값이 작아지게 되면 기울기가 0에 가까워 지면서 값 사이에 큰 차이가 없게 될 수 있다. (uniform distribution에 대한 가능성)

각각의 정규화 방법은 서로 반대되는 효과를 보여주는 것을 확인할 수 있는데, 두 정규화 방법을 동시에 적용하여 서로의 단점을 보완하도록 만들었다.

Architecture

구현은 Data-efficient image Transformers(DeiT) 구조를 그대로 사용하였다고 한다. 여기서 distillation token은 사용되지 않았기 때문에 base ViT와 head만 다른, 거의 동일한 구조를 가진다고 보면 된다.

[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