Post

[논문 리뷰] DINOv2: Learning Robust Visual Features without Supervision

Posted
By Luna Lee
20 min read



Meta AI Research arXiv 2023

Introduction

[기존의 방법, 문제 제기]

Task-agnostic한 pre-train representation은 자연어 처리(NLP)의 표준이 되었다. NLP에서의 영향으로 computer vision에서도 이러한 “foundation” 모델이 나타나고 있다. 가장 유망한 방법은 text-guided pre-training, 즉 feature를 학습하기 위해 textual supervision을 이용하는 방식에 중점을 두고 있다.

하지만 이러한 방식의 학습은 caption이 이미지의 풍부한 정보를 근사할 뿐, 복잡한 pixel-level의 정보는 표현하기 어렵다. 또한 이렇게 학습된 이미지 Encoder에는 text-image pair가 항상 필요하므로 이미지만을 사용하는 유연성이 없다.

Text-guided pretraining의 대안은 image만으로 feature를 학습하는 self-supervised learning이다. 이러한 방법은 개념적으로도 pretext task에 가까울 뿐 아니라 image를 pixel-level에서 정보를 획득할 수 있고 다양한 응용이 가능하다.

하지만 이러한 잠재력에도 불구하고 이전의 작업들은 너무 작은 데이터셋에 대해 학습했거나(ImageNet-1k), 품질이 저하된 데이터를 사용하였다.

DINOv2_1.png

[논문의 방법]

본 논문에서는 self-supervised learning이 대량의 선별된 데이터에 대해서 pre-train된 경우 범용적인 visual feature를 학습할 수 있는지에 대해 탐구했다. 기존에 iBOT과 같이 image 와 patch level 모두에서 feature를 학습하는 self-supervised 방법들을 재검토하고, 더 큰 데이터셋에 대해 모델의 일부 설계에 대해 검토한다. 본 논문에서 기여하고자하는 바의 대부분은 모델 및 데이터의 크기가 확장된 상황에서 discriminative self-supervised learning을 안정화하고 가속화하는 것이라고 한다.

Pre-training을 위해 광범위한 이미지 컬렉션에서 데이터를 필터링하고 재조정하는 자동 pipeline을 구축했다. 이 방법은 데이터 similarity를 이용하여 manual annotation이 필요하지 않은 NLP 파이프라인에서 영향을 받았다고 한다.

마지막으로 다양한 ViT 구조로 pre-train된 visual model DINOv2를 제안한다. 위의 그림에서 요약된 대로 image, pixel level에서 다양한 computer vision benchmark에 대해 DINOv2의 성능을 검증헀다.

DINOv2_2.png







Method

Data Processing

선별되지 않은(Uncurated) 대규모 데이터 풀에서 기존에 선별된(Curated) 데이터셋의 이미지와 유사한 이미지를 찾는 방식으로 데이터를 선별하여 LVD-142M 데이터셋을 만들었다. 이 과정은 meta 데이터나 text가 필요 없이 이미지에 직접 적용된다. 데이터 pipeline은 아래의 그림과 같다.

  • 사용된 Curated dataset: ImageNet-22k, ImageNet-1k의 train, Google Landmarks 등의 데이터셋
  • 수집한 Uncurated data source: 1.2B 규모의 크롤링한 web 이미지

DINOv2_3.png

  1. Data source: 인터넷에서 대규모의 웹 이미지를 크롤링하여 uncurated 데이터를 수집했다. 이 과정에서 안전하지 않거나 제한된 URL은 삭제했다. 수집된 이미지에 대해 후처리(PCA hash 중복 제거 등)를 진행했다.
  2. Deduplication: SSCD paper의 copy detection pipeline을 이용하여 uncurated 데이터에서 중복에 가까운 데이터를 삭제한다.
  3. Self-supervised image retrieval: uncurated data에서 curated dataset의 이미지와 가까운 이미지를 선별하여 데이터를 구축한다. 이를 위해 ImageNet-22k에 pre-train된 ViT 모델을 사용하여 image embedding을 계산하고 cosine-similarity를 이미지 간의 거리 측정으로 사용한다. 그 다음 uncurated 데이터셋에 대해 k-means clustering을 수행한다.
    여기서 curated dataset이 충분히 크다면 각각의 query 이미지(=curated data를 의미)를 기준으로 가장 가까운 N개의 데이터를 추출하고, 충분하지 않다면 가까운 cluster로부터 M개를 샘플링한다.




Discriminative Self-supervised Pre-training

SwAV를 중심으로, DINO와 iBOT loss의 조합으로 이루어진 discriminative self-supervised 방법으로 학습을 진행한다. Knowledge distillation 학습 방법으로 볼 수 있다. 또한 regularizer와 짧은 high-resolution 학습 단계를 추가했다.

Knowledge Distillation(Teacher-Student Learning)

Knowledge Distillation이란 말 그대로 “지식(Knowledge) + 증류(Distillation)” 의 합성어로, 학습된 큰 네트워크(T)의 지식을 추출하여 작은 모델(S)로 전달하는 것을 의미한다.

DINOv2_4.png

크기가 크고 제한된 환경에서 사용하기 힘든 모델의 성능을 비교적 작고 배포에 맞는 작은 모델에서 비슷하게 학습하고자 하는 용도로 사용된다.

학습 순서는 먼저, Teacher Network를 학습시킨다. 그 다음으로 Student Netwotk를 Teacher Network를 이용하여 학습하게 되는데, 이 때 Teacher Network는 frozen한다. Loss function은 아래와 같이 구성된다.

DINOv2_5.png

Loss는 (1) Student Network의 output $f_{\theta_S}(x)$가 정답 y와 같아져야한다는 Cross Entropy Loss, (2) Student 모델의 분포 $f_{\theta_S}$가 Teacher 모델의 분포 $f_{\theta_T}$ 와 같아지기 위한 KL Divergence Loss로 구성되어있다.



Image-level objective (→DINO📄)

Student network와 Teacher network에서 추출된 feature간의 Cross-entropy loss를 사용한다. 두 feature 모두 동일한 이미지에서, 서로 다른 부분을 crop하여 얻은 ViT의 class token에서 구한다. DINO Student head(MLP로 구현)를 통과한 결과(score vector)에는 softmax를 적용하여 $p_s$를 얻고 Loss에 대한 모델 파라미터 최적화를 진행한다. DINO Teacher head를 통해 나온 결과는 softmax를 적용한 뒤 exponential moving average(EMA)를 적용하여 $p_t$를 얻고 teacher head를 업데이트 한다(기존 DINO 논문과 동일).


Patch-level objective (→iBOT📄)

Student model의 입력 이미지 patch 중 일부를 랜덤하게 making한다(Teacher model에 대해서는 masking 수행하지 않음). 그리고 각 masking patch 위치에 대해 student, teacher model feature를 구하고(patch-level) cross-entropy를 구한다. 마찬가지로 Loss를 이용 student 학습, teacher head는 EMA로 업데이트한다.


Untying head weights between both objectives

DINO(Image-level)와 iBOT(Patch-level)의 head 간에 parameter sharing(같은 head사용)을 했을 때 성능이 향상되는 이전의 연구가 있었지만, 본 논문에서는 이러한 방법을 사용하면 image-level에서 overfitting되고 patch-level에서 underfitting되는 문제를 발견했다. 따라서 두 head를 별도로 사용하였다.


Sinkhorn-Knopp centering (→ SwAV📄)

앞서 EMA(softmax-centering)를 수행했던 Teacher 모델의 업데이트 방법을 SwAV 논문의 Sinkhorn-Knopp (SK) batch normalization으로 대체한다. Sinkhorn-Knopp 알고리즘을 3 iteration 반복하고, Student에는 softmax normalization을 적용한다.

SwAV(Unsupervised Learning of Visual Features by Contrasting Cluster Assignments)
SwAV는 Clustering 기반 Self-supervised learning 방법을 제안한 논문이다.
방법은 다음과 같다. 이미지 $x$에 대해 서로 다른 Augmentation을 적용하고 각각 Encoder를 통해 feature를 추출한다. 두 개의 이미지 feature $z_t, z_s$가 주어지면 두 feature에 대한 code $q_t, q_s$를 계산하는데, 이 때 $K$개의 prototype {}(일종의 codebook과 유사한 개념)에서 이 feature를 매칭한다. 이 때, Loss는 다음과 같다.

\[L (\mathbf{z}_t, \mathbf{z}_s) = \ell(\mathbf{z}_t, \mathbf{q}_s) + \ell(\mathbf{z}_s, \mathbf{q}_t)\]

DINOv2_5_2.png
이 때 두 code는 Augmentation을 다르게 적용했지만 동일한 이미지에서 생성된 것이기 때문에, 같은 정보를 캡쳐해야 한다는 것이 본문의 아이디어이다.
여기서 code를 계산하기 위한 prototype을 online으로 학습하는 방법을 제시한다. 이 때 batch내의 모든 feature가 동일한 코드로 매핑되는 trivial solution을 막기 위해, batch 내에서 매핑되는 code가 구별되도록 하는 constraint를 넣어주게 된다. 이 방법이 Sinkhorn-Knopp centering이다.




KoLeo regularizer

KoLeo regularizer는 batch내의 feature들 사이의 거리를 최대한 균등하게 만들어주는 역할을 한다. n개의 vector 집합(x_1, … , x_n)이 주어지면 아래와 같이 정의된다. KoLeo regularizer를 수행하기 전 l2-normalize는 적용된다.

\[\mathcal{L}_{koleo} = -\dfrac1n \sum_{i=1}^n \log (d_{n, i})\]

Adapting the resolution

Low-resolution에서는 작은 물체들은 사라지게 되므로, 이미지 resolution을 높이는 것은 downstream task를 위해서도 중요한 문제이다. 하지만 high-resolution 학습은 많은 시간과 메모리가 필요하므로 논문에선 pretraining의 마지막에 짧게 518×518 이미지에 대한 학습을 진행했다.




Efficient implementation

저자는 모델을 더 큰 규모로 학습하기 위한 몇가지 개선사항을 고려했다. 이 부분은 핵심만 간단하게 요약했다.

  • Fast and memory-efficient attention: Self-Attention 레이어의 메모리 사용량과 속도를 개선하기 위해 자체 버전의 FlashAttention을 구현했다.

  • Sequence packing: DINO 알고리즘에서 이미지 resolution이 다른 경우 서로 다른 patch 개수가 생성되고, token sequences 길이가 다르다. 이러한 경우에도 forward 할 수 있도록, transformer를 통해 전달해야 하는 시퀀스를 하나의 긴 시퀀스로 연결하는 trick을 적용했다.

  • Efficient stochastic depth: 저자는 결과를 masking하는 기존의 방법 대신 일부 residual path의 계산을 drop하는 방식으로 stochastic depth를 개선한 버전을 적용했다. 40% 정도의 높은 삭제율을 적용하면 컴퓨팅 효율성과 메모리 사용량이 크게 향상된다고 한다.

  • Fully-Sharded Data Parallel (FSDP): AdamW로 모델을 최적화하기 위해서는 Student, teacher, optimizer 1/2 moments 총 4개의 모델 복제본이 필요하다. 이러한 큰 모델 사용을 위해 FSDP의 PyTorch 구현을 이용하여 GPU 분할을 사용했다(data parallel). 이 방법은 기존의 DDP(DistributedDataParallel) 방식과 비교하여 통신 비용이 약 50% 감소한다.

  • Model distillation: 앞서 언급했듯 논문에서는 distillation 방법을 사용하여 큰 모델의 지식을 작은 모델에 전달하는 방법을 사용한다. 작은 모델은 처음부터 학습하지 않고 큰 모델인 ViT-g에서 distill 한다. 논문은 몇가지 예외를 제외하고 기존의 distillation 학습과 동일한 학습 루프를 사용했다.




< ablation study 부분은 다양한 component에 대한 성능 검증부분이라 생략했다. 자세한 부분은 논문을 참조해보자 >

Results

1. ImageNet Classification

첫 번째 평가로 ImageNet-1k classification 데이터셋에 대해 모델의 이미지 representation에 대해 조사했다. Frozen backbone에 대해 간단한 classifier를 학습하여 평가를 진행했다.

또한 SOTA open-source weakly supervised model과 비교를 위한 실험을 진행했다.

DINOv2_6.png



High quality frozen feature를 생성하는 모델의 능력이 특정 데이터셋에 대해 supervised 학습을 통해 fine-tuning 될 때 성능에 영향을 미치는지 조사했다.

DINOv2_7.png



생성된 feature의 generalization 성능을 조사하기 위해 domain generalization benchmark에서 linear classification head로 학습된 ImageNet-1k 모델을 평가했다. SOTA SSL 방법과 비교할 때 DINOv2 모델은 훨씬 더 나은 robustness를 보여준다.

DINOv2_8.png



2. Additional Image and Video classification Benchmarks

다음으로는 downstream classification benchmark에 대한 feature의 generalization 성능을 연구한다. 먼저 iNaturalist, Places205와 같은 크고 세분화된 데이터셋을 사용한다. 또한 SimCLR에서 12가지 이미지 classification task를 수행한다.

DINOv2_9.png



12 transfer classification benchmark에서 선택된 frozen feature를 비교했다.

DINOv2_10.png



3. Instance Recognition

Non-parametric approach를 사용하여 instance-level recognition에 대한 모델의 성능을 조사했다.

DINOv2_11.png



4. Dense Recognition Tasks

Dense downstream task에 대해 학습된 network에서 추출된 patch-level feature의 품질을 조사했다. 다양한 데이터셋에 대해 Semantic image segmentation과 Monocular depth estimation에 대한 평가를 수행했다.

DINOv2_12.png

위의 결과는 segmentation을 mIoU로 측정한 결과이다. Semantic segmention 평가를 위해 두가지 설정을 고려했다.

  • Linear: Patch token에 대한 class logit을 예측하도록 학습 된다. low-resolution logit map을 생성하는데 사용되며, 그 다음으로 segmentation map을 얻기 위해 전체 해상도로 upsampling된다. 이 절차는 간단하지만 high-resolution segmentation을 쉽게 생성하기는 어렵다.
  • +ms: Linear 설정의 향상된 버전. 마지막 4개 layer의 patch token을 concat하고, 더 큰 이미지 resolution을 사용하고, multiscale test-time augmentation을 사용한다.



DINOv2_13.png

이 실험에서는 NYUd, KITTI, NYUd에서 SUN3d로의 zero-shot transfer 이렇게 세 가지 monocular depth estimation benchmark에서 patch-level feature를 평가한 결과이다.

마찬가지로 depth estimation 평가를 위한 세가지 다른 설정을 고려했다.

  • lin. 1: transformer → bi-linearly upsample(x4배) → simple linear layer(256 depth prediction range, 256 bin으로 나눔), linear normalization
  • lin. 4: lin 1과 동일하나, transformer를 하나에서 4개로 늘림.
  • DPT: frozen 모델 위에 DPT decoder 사용



5. Qualitative Results

마지막으로 feature에 대한 정성평가를 위한 이미지이다.

DINOv2_14.png

위의 그림에서 ADE20K 데이터셋에 대한 segmentation과 KITTI, SUN RGB-D, NYUd데이터셋에 대한 depth estimation에 대한 몇 가지 정성적 결과를 보여주고 있다. DINOv2 backbone을 사용하는 linear segmentation 모델은 좋은 결과를 생성하고 동일한 평가 설정에서 OpenCLIP 모델보다 훨씬 더 잘 동작한다.



DINOv2_15.png

DINOv2 모델에 의해 추출된 patch feature에 대해 수행된 principal component analysis(주성분 분석, PCA) 결과이다. 이 과정은 이미지의 주요한 개체를 배경에서 잘 분리하는 것을 나타내고 있다.



DINOv2_16.png

마지막으로 전체 이미지에서 patch-level feature의 matching을 확인하여 어떤 유형의 정보를 포함하는지 탐색한다. 위의 그림에서 저자는 특징이 다른 물체나 동물에서 유사한 목적을 가지고 있는 semantic region에 대한 정보를 포착한다는 것을 관찰했다(비행기의 날개-새의 날개).



AI, Paper Review
Multi-modal Knowledge Distillation Self-supervised
This post is licensed under CC BY 4.0 by the author.