[Paper Reivew] Emerging Properties in Self-Supervised Vision Transformers (DINO)

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

TL;DR

Vision Transformer(ViT)가 Convnet의 대안으로 떠오르기는 했지만, Convnet을 뛰어넘는 이점은 없었다.

저자들은 이와 같은 문제를 self-Supervised learning으로 해결할 수 있을 것이라 주장한다.(NLP Trasnformer와 비슷하게)

저자들은 self distillation with no label(DINO)라는 방법을 제시하여 supervised 방식으로 학습한 ViT를 훨씬 능가하는 성능을 보였다

1. Introduction.

NLP에서 사용되는 Transformer 구조를 Vision 분야에 적용(ViT)하면서 convnet의 대안으로 떠오르고 있다. 하지만, ViT는 계산량이 더 많고, 많은 학습 데이터가 필요하며, feature들이 unique하지 않다는 문제점이 있다.

본 논문에서 저자들은 pretraining 과정에서 self-supervision을 통해 ViT가 성공적으로 convnet을 뛰어넘을 수 있을지 확인한다.

self-supervised 방식은 BERT, GPT등 NLP에서 주로 사용되는 방식이였다. 하지만 convnet을 self-supervised 방식을 활용해 많은 잠재성을 보인 연구가 진행중이였고, 저자들은 이를 transformer에도 적용 가능한지 확인하였다.

또한 저자들은 ViT에서 작은 patch를 사용하는 것이 성능을 더욱 향상시킬 수 있음을 확인했다고 한다.

이러한 방법론을 저자들은 DINO(distillation with no labels.)라 한다. DINO는 momentum encoder로 구성된 teacher network로부터 출력을 직접 예측함으로써, self-supervised learning을 단순화 하였다.

또한 흥미롭게도, DINO는 centering과 sharpening만으로도 mode collapse를 피하면서 동작할 수 있으며, 유연한 구조로 인해, ViT와 convnet에서도 잘 동작할 수 있다고 한다.

2. Related work.

2.1. Self-supervised learning.

Image를 discriminating 하지 않고 unsupervised feature를 학습할 수 있는 방법이 연구되고 있다. 특히 Grill et al.이 제안한 BYOL은 momentum encoder로부터 얻은 representation을 적절히 매칭하도록 학습한다.

저자들도 BYOL에 영감을 받았다고 한다. 하지만 저자들은 BYOL과 다른 similarity matching loss를 사용하였으며, student와 teacher 모델이 정확히 똑같다는 점이 다르다고 한다.

2.2. Self-training and knowledge distillation.

Self-training 방식은 처음 작은 annotation set을 unlabeled instance set에 전달하면서 품질을 올리는 것을 목표로 한다. 이 방식은 hard assignment 혹은 soft assignment로 구분할 수 있는데, soft label을 사용한 경우 이 방식을 knowledge distillation이라 한다. (작은 student 모델이 커다란 teacher 모델을 흉내내도록 학습된다)

저자들은 DINO는 codistillation과 관련이 있다고 한다. Student와 teacher이 정확히 같은 모델을 사용하며, teacher도 student로부터 distillation 된다고 한다.

3. Approach.

3.1. SSL with Knowledge Distillation.

DINO는 최근 self-supervised 접근법과 대부분 유사하지만, knowledge distillation 관점에서도 이를 제시한다.

Self-distillation with no labels.

모델구조와 알고리즘을 간단히 설명하면 다음과 같다.

1. \(x\)라는 sample을 \(x_1, x_2\)로 augmentation한다.
2. 이를 student, teacher model에 태워 output \(s1, s2, t1, t2\)를 얻는다.
3. cross-entropy loss를 계산하여 student model만 back-prop한다.
4. teacher model은 student model로부터 knowledge를 dsitillation 받는다.

모델의 구성성분을 자세히 보자. 먼저 input image \(x\)가 주어졌을 때, model의 output probability \(P\)는 네트워크 \(g\)의 출력을 다음과 같이 softmax 취하여 구한다.

\[\begin{equation} 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)}, \quad P_t (x)^{(i)} = \frac{\exp(g_{\theta_t} (x)^{(i)} / \tau_t)}{\sum_{k=1}^K \exp(g_{\theta_t} (x)^{(k)} / \tau_t)} \end{equation}\]

이때, \(\tau (> 0)\)는 출력분포의 sharpness를 결정하는 temperature parameter이다.

고정된 teacher 모델 \(g_{\theta_t}\)이 주어졌을 때, student model은 cross-entropy를 최소화하도록 학습된다.

\[\begin{equation} \min_{\theta_s} H (P_t (x), P_s (x)) \end{equation}\]

이때, \(H(a,b) = -a \log b\)이다.

저자들은 실험에서 다양한 distorted, crop view를 만들었는데, 두개의 global crop과 여러개의 local crop을 사용하였다. 이때 global view는 teacher 모델에만 통과할 수있으며, 이는 local-to-global 대응을 장려한다고 한다.

\[\begin{equation} \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')) \end{equation}\]

Teacher Network.

일반적인 knowledge distillation과는 다르게, teacher model(\(g_{\theta_t}\), priori)이 없으므로, student network의 이전 iteration을 이용해(EMA, i.e. momentum encoder) teacher model을 만든다.

여러가지 update rule로 실험을 한 결과 EMA(exponential moving average)를 사용하는 것이 특히 잘 동작했다고 한다.

\[\begin{equation} \theta_t \leftarrow \lambda \theta_t + (1-\lambda) \theta_s \end{equation}\]

또한 저자들은 teacher 모델이 학습 내내 student model보다 성능이 좋았으며, 따라서 높은 quality의 feature를 student에게 적절히 guide해주었는데, 이는 이전의 연구서는 발견되지 않는 특징이였다.

Network Architecture.

네트워크 \(g\)는 Backbone \(f\)와 projection head \(h\)로 구성되어있다.

\[g = h \circ f\]

또한 일반적인 convnet과는 달리, ViT 구조에서는 BatchNorm을 사용하지 않았으며, DINO system을 entirely BN-free로 만들었다.

Avoiding Collapse.

많은 self-supervised 방법은 contrastive learning, clustering constraints, predictor, batch norm… 등 다양한 방법으로 collapse를 피하려 한다. 저자들의 DINO에서도 다른 방법으로도 collapse를 피할 수 있었지만, centering과 sharpening만으로 momentum teacher의 collapse를 피할 수 있었다고 한다.

Centering과 sharpening이 반대의 효과가 있기 때문에, 적절하게 균형을 맞춘다면 momentum teacher를 유지하면서 collapse를 피할 수 있었다고 한다.

centering 연산이 오직 batch의 1차 statistic에만 의존하며 이는 teacher에 bias 항을 더하는 것으로 해석할 수 있다.

\[\begin{equation} g_t(x) \leftarrow g_t(x) + c \end{equation}\]

이때 \(c\)는 EMA방식으로 student의 knowledge를 이용하여 구성한다.

\[\begin{equation} c \leftarrow mc + (1-m) \frac{1}{B} \sum_{i=1}^B g_{\theta_t} (x_i) \end{equation}\]

3.2. Implementation and evaluation protocols.

Vision Transformer.

저자들이 사용한 model architecture는 다음과 같다.

Networks configuration

ViT는 \(N \times N\)의 이미지를 patch로 입력받는다. 본 논문에서는 \(N= 8, 16\)을 일반적으로 사용하였다. 이후 patch는 linear layer를 통과하고, embedding된다.

또한 저자들은 추가적인 학습 가능한 token을 추가하였는데, 이는 전체 sequence의 정보를 집계하고, projection head \(h\)에 붙였다. 이 token을 이전 연구들과 통일성을 위해 [CLS]이라 한다.

4. Main Results.

Linear and k-NN classification on ImageNet.

저자들은 ImageNet의 다른 self-supervised 방식과 비교해 DINO의 성능을 검증하였다. 이후 image retrieval, object discovery, transfer-learning과 같은 특징을 연구했다고 한다.

4.1. Comparing with SSL frameworks on ImageNet.

저자들은 크게 same architecture와 across architecture 두가지 세팅을 고려했다.

Comparing with the same architecture.

ResNet-50과 ViT-small 의 구조를 다양한 방법으로 self-supervised 학습한 결과는 위 표의 위 pannel에 나와있다고 한다.

Comparing across architecture.

이 세팅의 흥미로운 점은, 방법들을 직접 비교한 것이 아니라, DINO로 train한 ViT의 한계를 평가했다는 점이다. 저자들은 큰 사이즈의 ViT를 DINO로 학습할 때 성능이 향상되되었으며, patch size를 작게했을 때 더욱 큰 성능향상이 있었다고 한다.

4.2. Properties of ViT trained with SSL.

Nearest neighbor search, object location retaining, transferability 측면에서 DINO의 특성을 확인하였다고 한다.

4.2.1. Nearest neighbor retrieval with DINO ViT.

Image Retrieval. Oxford와 Paris 데이터셋에서 k-NN을 사용한 검색 실험 결과, DINO로 학습된 특징은 ImageNet 레이블로 학습된 특징보다 더 높은 성능을 보였다.

또한, DINO는 annotation이 필요하지 않은 SSL이기 때문에, annotation 없는 Google Landmarks v2(GLDv2) 데이터셋에서도 학습이 가능하며, 이전의 방법들보다 더 우수한 성능을 기록했다.

Copy detection. INRIA Copydays 데이터셋에서 왜곡된 이미지(blur, insertions, print, scan, etc.)를 인식하는 copy detection 작업에서도 DINO로 학습된 ViT는 높은 성능을 보여주었다. 이때, ViT의 출력 토큰을 이용한 feature를 사용하며, cosine similarity를 활용하여 copy detection을 수행한다. DINO로 학습된 ViT는 복제 탐지에서 매우 경쟁력 있는 성능이 확인되었다.

4.2.2. Discovering the semantic layout of scenes.

저자들의 Self-attention map은 image의 semantic한 정보도 가지고 있으며, 이런 특징을 이용해 standard benchmark를 측정하는 연구를 진행했다.

Video instance segmentation. DAVIS-2017 비디오 인스턴스 분할 벤치마크에서 실험한 결과, DINO로 학습된 ViT는 별도의 추가 학습이나 미세 조정 없이도 경쟁력 있는 성능을 보였다. 특히 작은 패치 크기(“/8”)를 사용하는 모델이 더 높은 성능(+9.1%)을 기록했으며, 이는 네트워크가 공간 정보를 잘 유지하고 있음을 시사한다.

Probing the self-attention map. 다양한 attention head가 이미지의 서로 다른 의미적 영역에 주목할 수 있으며, 심지어 가려지거나 작은 객체에도 잘 반응하는 것을 시각화로 확인했다고 한다. 또한, DINO로 학습된 ViT가 객체의 혼잡한 환경에서도 supervised ViT보다 더 나은 객체 인식을 보여주었으며, Jaccard 유사도를 기준으로 DINO 모델이 더 우수한 성능을 나타냈습니다.

4.2.3. Transfer learning on downstream tasks.

저자들은 DINO로 사전 학습된 ViT 모델의 특징을 ImageNet에서 감독 학습된 동일한 아키텍처의 모델과 비교했다. Downstream tasks에서 DINO는 supervised learning의 특징보다 더 좋은 성능을 보였으며, 이는 이전의 convnet 연구에서 관찰된 결과와 일치한다.

5. Ablation Study of DINO.

ViT에 DINO를 적용했을 때 ablation study를 진행했다고 한다. (meta 정도 되어야 이정도 ablation을 할 수 있나….)

5.1. Importance of the Different Components.

다양한 요소를 추가/제거 하면서 실험한 결과이다 .

주요 결과는 다음과 같다.

1. momentum 없이는 DINO가 동작하지 않는다.
2. 3번, 9번 row를 비교하면 Performance에서 momentum encoder가 얼마나 중요한 지 알 수 있다.
3. 4번, 5번 row를 비교하면 multi-crop training과 cross-entropy loss가 중요하다는 것을 알 수 있다.
4. 6번 row를 보면 DINO에서 predictor는 크게 중요하지 않았으며, 이는 BYOL과 대비되는 결과.

Importance of the patch size. 저자들은 ViT-S 모델이 서로 다른 패치 크기(16×16, 8×8, 5×5)로 학습되었을 때의 성능을 비교하였다. 모델은 모두 300 epoch 동안 학습되었으며, 패치 크기를 줄이면 성능이 크게 향상됨을 관찰하였다. 흥미로운 점은 추가적인 파라미터 없이도 성능이 향상된다는 점이다. 그러나 패치 크기를 작게 사용할수록 처리 속도(throughput)는 감소하는 단점이 발생합니다.

5.2. Impact of the choice of Teacher Network

이 절에서는 DINO에서 사용된 Teacher Network의 역할을 이해하기 위해 다양한 teacher 네트워크를 실험하였다.

5.2.1. Building different teachers from the student.

저자들은 momentum teacher외에 다양한 전략을 실험했다. Momentum 방식이 가장 성능이 좋았지만(당연하게도,,) previous epoch의 student를 teacher로 사용하는 방법은 collapse없이 MoCo-v2나 BYOL과 같은 기존 방식과 경쟁력 있는 성능을 보여주었다.

5.2.2. Analyzing the training dynamic.

Figure 6(왼쪽)에 따르면, momentum teacher는 훈련 중 지속적으로 student 보다 우수한 성능을 보였는데, 이 현상은 ResNet-50을 훈련할 때도 동일하게 관찰되었다.(Appendix D 참고)

이런 현상은 다른 momentum 방식이나, previous epoch을 사용하는 방식에서도 관측되지 않았는데, 저자들은 DINO가 Polyak-Ruppert 평균화의 일종으로 해석할 수 있다고 제안한다.

Polyak-Ruppert 평균화는 종종 훈련 종료 시 네트워크의 성능을 향상시키기 위해 모델 앙상블링을 시뮬레이션하는 데 사용됩니다. 저자들의 방법은 훈련 과정에서 지속적으로 우수한 성능의 모델 앙상블을 구축하여 학생 네트워크의 훈련을 이끄는 방식으로 Polyak-Ruppert 평균화를 적용하는 것으로 해석할 수 있습니다.

5.3. Avoiding collapse.

이 부분에서 저자들은 모델 붕괴(collapse)를 방지하기 위한 Centering과 Target Sharpening의 상호 보완적인 역할을 확인하였다.

붕괴에는 두 가지 형태가 있습니다:

1. 입력과 상관없이 모델 출력이 모든 차원에서 일정
2. 특정 차원이 지배적인 경우

Centering은 특정 차원이 지배적인 붕괴(1)를 방지하지만, 출력이 일정(2)해지는 경향을 촉진합니다. 반면, Sharpening은 반대 효과를 유도한다.

저자들은 상호 보완적인 효과를 설명하기 위해, cross entropy(\(H\) )를 entropy(\(h\))와 Kullback-Leibler(KL) 다이버전스(DKL)로 분해해 분석 하였다. (만약 KL 값이 0에 수렴하면 출력이 일정해지고, 이는 붕괴를 나타낸다)

\[H(P_t,P_s)=h(P_t)+D_{KL}(P_t \rm∥P_s)\]

위 그래프에서는 Centering 및 Sharpening을 적용한 경우와 적용하지 않은 경우 훈련 중의 엔트로피와 KL의 변화를 보여준다. 둘 중 하나의 연산이 빠지면 KL이 0으로 수렴하여 붕괴가 발생하지만, 그러나 엔트로피(h)는 다른 값으로 수렴한다.:

Centering이 없을 때는 0으로, Sharpening이 없을 때는 \(-\log(1 / K)\)로 수렴하는데, 이는 두 연산이 서로 다른 형태의 붕괴를 유도한다는 것을 의미하며, 두 연산을 모두 적용함으로써 이 효과들이 균형을 이루게 된다.

5.4. Compute requirements.

저자들은 multi-crop 사용이 DINO 실행의 정확도와 실행 시간의 trade-off를 개선함을 확인 하였다. 예를 들어, multi-crop 없이 46시간 훈련한 경우 성능이 72.5%였지만, \(2×224^2+10×96^2\)설정으로 multi-crop을 적용하면 24시간 만에 74.6% 성능을 달성했습니다. 이는 시간을 절반으로 줄이면서 성능이 2% 개선된 것이며, 다만 메모리 사용량은 더 높았다(15.4 GB versus 9.3 GB).

5.5. Training with small batches.

또한 저자들은 batch 크기가 DINO로 훈련된 특징에 미치는 영향을 확인하였다. 작은 batch 크기(128)로도 고성능 모델을 훈련할 수 있음을 확인했으며, 작은 배치 크기일수록 하이퍼파라미터 조정이 필요했다고 한다.

6. Conclusion.

저자들의 기여는 다음과 같다.

• 표준 ViT 모델을 Self-supervise pretraining을 통해 convnet과 비교해 경쟁력 있는 성능을 달성하였다.
• k-NN 분류에서 특징의 품질이 이미지 검색에 유망하다는 것과 장면 배치 정보가 약한 지도 학습 이미지 분할에 유리할 수 있다는 특성을 확인하였다.
• 마지막으로, 저자들의 가장 중요한 발견은 self-supervised learning이 BERT-like model을 만드는게 핵심 역할을 할 수 있는 증거를 발견했다는 것이다.

따라서 저자들은 향후, 대형 ViT 모델을 무작위로 선택한 uncurated 이미지로 pretraining하여 visual feature 한계를 확장할 가능성을 탐구할 계획이라고 한다.

Reference

JiYeop Kim’s blog를 참고하여 작성하였습니다.