First Pass

5C

1. Category : Self-distillation 기반 Self-supervised Vision Transformer 학습(DINO)과 emergent attention property 분석
2. Context : 당시 CV에서는 CNN 기반 SSL(SimCLR, MoCo 등)이 representation learning에서 강력한 성능을 보이고 있었다. 동시에 ViT가 visual recogintion 분야에서 주목받기 시작했지만, supervised learning 환경에서는 CNN 대비 명확한 구조적 이점이 아직 불분명했다. 저자들은 NLP의 BERT/GPT 계열 self-supervised pretraining에서 영감을 받아 self-distillation 기반 SSL을 ViT에 적용하였고, 그 과정에서 ViT가 sementic object boundary를 자연스럽게 학습하는 emergent property를 보인다는 점을 발견하였다.
   1. Radford, Alec, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving Language Understanding by Generative Pre-Training. n.d. Paper
   2. Devlin, Jacob, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. “BERT: Pre-Training of Deep Bidirectional Transformers for Language Understanding.” arXiv:1810.04805. Preprint, arXiv, May 24, 2019. Paper.
3. Correctness :저자들은 기존 CNN 기반 SSL의 self-distillation 전략을 ViT에 적용하였다. 특히 contrastive negative pair 없이도 안정적인 representation learning이 가능함을 보였으며, 학습 과정에서 attention map이 semantic object boundary에 자연스럽게 집중하는 현상을 실험적으로 확인하였다. 이는 supervised label 없이도 high-level visual concept가 형성될 수 있음을 보여준다.
4. Contributions
   1. Teacher-student self-distillation 기반의 self-supervised ViT 학습 방법(DINO) 제안
   2. 학습된 ViT의 self-attention map이 별도의 supervision 없이 semantic object boundary에 집중하는 emergent property를 보인다는 점을 발견
   3. 단순 linear evaluation이나 k-nn classification에서도 강력한 representation quality를 보이며, SSL 기반 ViT feature의 일반화 가능성을 입증
   4. Vision Transformer가 CNN 대비 Self-supervised representation learning과 특히 잘 결합된다는 가능성을 제시
5. Clarity : 논문이 올바르게 작성되었는지?

Second Pass

Related Works

Self-supervised Learning

당시에 자기지도 학습에 관한 연구는 instance classification이라는 discriminative한 접근법에 초점을 맞추고 있었는데, 이는 각 이미지를 서로 다른 클래스로 간주하고 데이터 augmentation을 통해 모델을 판별해 학습시키는 방법이다.

이 방법은 모든 이미지들을 discriminate하기 위해 classifier를 학습하는건 확장성이 떨어지므로, 이를 해결하기 위해 Noise Contrastive Estimator(NCE)가 고안됨. NCE는 instance를 분류하는 대신 서로 비교함. 그러나 NCE도 동시에 많은 양의 이미지에서 뽑은 feature를 비교해야한다.

여기서 다시 고안된게 통계 학습법인 BYOL. BYOL은 momentum encoder로 뽑은 representation으로 서로 매칭하도록 feature들이 훈련되었다.

DINO의 저자들은 BYOL의 접근법에 착안하였지만, 다른 similarity matching loss를 사용하고, student와 teacher에 정확히 동일한 구조를 사용한다.

Self-training and knowledge distillation

Self-training은 label이 달린 적은 수의 데이터를 라벨링되지 않은 많은 양의 데이터에 확장시켜 feature들의 퀄리티를 향상시키는 것이 목표이다. 이 과정에서 label을 soft label을 사용하는 경우 일반적으로 knowledge distillation이라고 부르며, 작은 모델을 거대한 모델의 output을 흉내내도록 훈련시켜 모델을 압축하는 형태라고 볼 수 있다.

DINO의 저자들은 이러한 self-training과 distillation의 관계를 기반으로, knowledge distillation을 label이 전혀 없는 self-supervised learning 환경으로 확장하였다. 또한 기존 연구들이 pre-trained fixed teacher를 사용하는 것과 달리, DINO는 training 과정에서 EMA 기반으로 teacher를 동적으로 업데이트한다. 이를 통해 distillation을 단순한 post-processing 단계가 아니라 self-supervised objective 자체로 사용한다.

또한 DINO는 student와 teacher가 동일한 architecture를 사용한다는 점에서 codistillation과 유사성이 있다. 그러나 codistillation이 teacher와 student 간의 상호 distillation을 수행하는 반면, DINO의 teacher는 student의 EMA를 기반으로 갱신된다는 차이가 있다.

Approach

SSL with Knowledge Distillation

Knowledge Distillation은 student network $g_{\theta_s}$가 teacher network $g_{\theta_t}$의 출력 분포를 모방하도록 학습하는 방법이다.

입력 이미지 $x$가 주어졌을때 두 네트워크는 K차원의 output logits를 생성하며, softmax 함수를 통해 각각의 확률분포 $P_s,P_t$를 얻는다.

\[P_s(x)^{(i)} = \frac{\exp(g_{\theta_s}(x)^{(i)}/\tau_s)}{\sum_{k=1}^K\exp(g_{\theta_s}(x)^{(k)}/\tau_s)} \tag{1}\]

이 때, $\tau$는 출력 분포의 sharpness을 조절하는 temperature 파라미터이다. 작은 $\tau$를 사용할수록 분포는 더욱 sharp해진다.

기존 knowledge distillation은 고정된 teacher network의 출력 분포를 student network가 모방하도록 학습한다.

\[\min_{\theta_s} H(P_t(x),P_s(x)) \tag{2}\]

DINO는 이를 SSL로 확장한다. 하나의 이미지로부터 생성한 다양한 augmented view를 사용하여, student network는 global/local crop 모두를 입력받고 teacher network는 global crop만 입력받는다. 이를 통해 local feature와 global semantic representation 사이의 correspondence를 학습한다.

\[\min_{\theta_s}\sum_{x\in\{x_1^g,x_2^g\}} \sum_{x'\in V, x'\ne x} H(P_t(x),P_s(x)) \tag{3}\]

두 네트워크는 동일한 arichitecture $g$를 공유하지만 서로 다른 파라미터 셋 $\theta_s, \theta_t$를 사용한다. Student parameter는 SGD로 업데이트되며, teacher parameter는 student parameter의 EMA를 통해 갱신된다.

Teacher Network

Knowledge distillation과는 다르게 DINO는 $g_{\theta_t}$를 사전에 제공하지 않고, student network의 과거 parameter들을 기반으로 동적으로 teacher network를 구성한다. 또한 저자들은 student weight에 EMA$(\theta_t \leftarrow \lambda\theta_t + (1-\lambda)\theta_s)$를 사용하는것이 DINO에 특히나 잘 작동함을 발견하였다. 이때 $\lambda$는 고정값이 아닌 cosine schedule을 활용해 $0.996\rightarrow1$로 상승시켰다.

Momentum encoder는 다른 framework와는 다르게 contrastive learning의 queue encoder 역할이 아니라, self-training의 mean teacher 역할에 더 가깝게 사용되었다.

Mean teacher : SSL에서 teacher network의 prediction을 target으로 사용하여 student network가 일관된 representation을 학습하도록 만드는 방식

최종적으로 훈련 과정 중에서 teacher network가 student network보다 더 향상된 성능을 보이며, 더 높은 품질의 target feature를 제공함으로써 student network를 효과적으로 guiding함을 관찰하였다.

Network architecture

Neural Network $g$는 backbone model $f$(ViT or ResNET)과 projection head $h:g=h\circ f$로 구성되어있다. Downstream task에서 사용되는 feature는 $f$의 출력물이며, projection head는 2048개의 은닉층을 사용하는 3-layer MLP로 구성되며, 출력 feature에 l2 norm을 적용한 뒤 K차원의 weight-normalized fully connected layer를 통과시킨다.

Student, teacher networks 모두 동일한 구조를 가지기에 predictor를 사용하지 않았으며, ViT는 기본적으로 Batch Normalization을 사용하지 않기 때문에, DINO의 projection head에도 BN을 적용하지 않아 전체 시스템을 BN-free 구조로 구성하였다.

Avoiding collapse

다른 SSL 모델은 collapse를 방지하기 위해 contrastive loss, clustering constraints, predictor, batch normalization 같은 방법들을 사용했다. DINO는 여러개의 normalization을 사용하여 모델을 안정화할 수도 있지만, momentum teacher network 출력을 centering하고 sharpening하여 model collapse를 방지할 수도 있다.

Model Collapse : 입력 이미지와 관계없이 모든 representation이 동일하거나 uniform한 값으로 수렴하는 현상. 이 경우 모델은 의미있는 semantic feature를 학습하지 못한다.

Centering은 한 개 차원이 모델을 과하게 점유하는 것을 방지하지만 지나치게 적용될 경우 overly uniform 분포로 model이 collapse 될 가능성이 있으며, Sharpening은 그 반대의 역할을 한다. 따라서 두 개의 연산을 같이 적용하여 self-distillation 과정에서 발생할 수 있는 collapse를 효과적으로 방지한다.

DINO는 BN을 기반으로 하는 안정화 기법에 대한 의존도를 줄이는 대신, Centering과 Sharpening을 사용하면 학습 안정성을 확보한다. Centering은 batch의 평균값과 같은 first-order 통계량에만 의존하고, 수학적으로 teacher network에 bias term $c$를 더하는 것으로도 해석될 수 있다$(g_t(x) \leftarrow g_t(x) + c)$. 이때 center $c$는 EMA로 업데이트 되어 다양한 배치 사이즈에서도 잘 작동한다.

\[c \leftarrow mc+(1-m)\frac{1}{B}\sum^B_{i=1}g_{\theta_t}(x_i) \tag{4}\]

결과적으로 centering과 sharpening은 output entropy를 서로 반대 방향으로 조절하며, 두 연산의 균형을 통해 stable self-distillation을 가능하게 만든다.

Main Results

Claim 1

ViT는 SSL과 특히 더 잘 결합된다.

Evidence

• ViT-Based DINO가 CNN(ResNet) based SSL보다 linear probing 성능도 높고, k-NN에서도 강한 representation quality 가짐.

Claim 2

DINO를 적용한 ViT는 패치 사이즈가 작아질수록 모델 효율이 오른다.

Evidence

• ViT의 크기가 커질수록 성능도 상승했지만, 패치 사이즈를 줄일수록 성능이 더 비약적으로 향상되었다.
• 패치 사이즈를 줄였지만 파라미터 수는 늘지 않았고, 실행 시간과 메모리 사용량의 절약으로 이어졌다.
• 또한 이전 SOTA 모델에 비해 1.4배 더 빠르게 작동하고, 10배 적은 파라미터를 가지면서 SOTA 달성

Claim 3

DINO로 학습된 ViT는 별도 supervision 없이도 semantic object boundary를 학습한다.

Evidence

• CLS token attention map 시각화
• Foreground/background 분리
• Unsupervised object segmentation 가능
• Supervised Segmentation 없이도 object mask 형성