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

Meta에서 발표한 DINO의 후속 버전 DINOv2를 리뷰해 보도록 하겠다. attention과 트랜스포머 모델, DINO에 대해 알고 있는 편이 이해하기 쉬울 것 같다. 특별히 새로운 기법을 도입했다기보단 기존 다양한 연구들을 섞어서 재조합한 느낌~

 

다음 웹페이지에서 실제로 사용해볼 수 있다. 진짜 잘 작동됨!

https://dinov2.metademolab.com/

Preliminary

Transformer model

RNN, CNN 대신 Attention만 사용해 입출력 시퀀스의 연관성 처리하기 (Attention is all you need)

Attention: 한 문장(이미지) 들 간의 요소들의 연관성을 알아내는 것. 주목할 요소 찾기! 입력 벡터는 학습 가능한 Weight를 이용해서 쿼리, 키, 값 (Q,K,V) 벡터로 변환됨.

* Query: 물어보는 주체(단어 벡터), Key는 Query와 비교당하는 주체, Value는 데이터의 값

출처: https://secundo.tistory.com/96

1. Q와 모든 Key 벡터를 합친 행렬 K 간 내적 계산 -> Q와 모든 K 간의 유사도인 S를 얻을 수 있음
2. S에 softmax 취해 가중치 확률 P로 변환
3. 각 단어의 값 벡터 V에 가중치 P를 곱하고 합을 구해 단어별 최종 점수 O를 구한다.
   * 하나의 attention을 사용하는 것보다 여러 개의 attention function들을 만드는 것이 더 효율적임 -> multi-head attention

• Feed forwarding network: FC layer. * $FFN(x)=\max⁡(0,xW_1+b_1 ) W_2+b_2$
• Positional encoding: Input, output 데이터에 위치 정보를 담기 위해 추가한다.
• Pretrain: 대규모 데이터를 이용해서 transformer 모델을 pretext task를 이용해 사전학습.
• Fine-tuning: 원하는 downstream task(번역, 요약…)에 따라 pretrained 모델 업데이트해서 사용

ViT (Vision Transformer)

NLP의 트랜스포머 알고리즘을 비전 문제에 적용함.

 

• 이미지를 텍스트 시퀀스처럼 사용하기 위해 전체 이미지를 작은 패치로 나누고 순서대로 embedding한다.
• 트랜스포머 인코더에 클래스 정보를 담은 클래스 토큰도 함께 넣는다.
• 인코더는 NLP와 비슷하게 [ Multi-head self-attention layer + MLP ] 여러 층으로 구성된다.
  • 여러 패치들끼리 상호작용하며 이미지의 맥락을 이해하고, Self-attention을 통해 서로 떨어져 있는 패턴들 간의 관계도 잘 포착함.
• 인코더가 출력한 이미지 표현을 MLP head에 통과시켜서 예측된 class를 얻는다.
• 대규모 데이터로 Pretrain하고 Fine-tuning해서 사용

DINO

(Distillation of Non-contrastive Image representations)

ViT + Self-supervised learning

• 입력: 이미지를 패치로 나눠 위치 정보와 함께 임베딩해서 입력으로 사용한다.
• 구조: 같은 ViT 구조의 Teacher Network, Student Network로 구성되고, 동일 파라미터로 시작한다.
• Knowledge Distillation: Student 네트워크가 Teacher 네트워크의 출력을 따라가며 학습한다.

학습 과정

1.  입력 이미지를 random view로 크롭하고 패치로 분할한다. → student, teacher 네트워크에 각각 임베딩 
   * Teacher는 224x224의 global view image 2개, student는 global view를 포함한 96x96의 local view image n개가 통과 가능함
2. Student와 Teacher 네트워크가 각각 Self-attention을 통해 패치 간 연관성을 학습하고 클래스 토큰(이미지 특징)추출.
3. 클래스 토큰을 각각의 MLP DINO 헤드에 전달해서 score vector 계산.
4. Teacher의 output vector를 Student가 따라가도록 loss를 계산하고, Student 네트워크를 학습시킴.
   * Teacher는 student model의 가중치의 지수 이동 평균을 사용한다 (momentum)

momentum teacher model 정확도 그래프

 

DINOv2 Introduction

Problem

• NLP에서 task-agnostic representations는 보편화되어 있음.
  • fine-tuning 없이도 다양한 downstream task에서 우수한 결과를 낸다. (e.g., ChatGPT)
• 컴퓨터 비전 분야에서 이런 foundation model을 설계하기 위해 많이 쓰이는 텍스트-이미지 기반 학습은 단점이 있다
  • 텍스트는 이미지의 복잡한 정보를 충분히 표현하기 어렵고, 많은 양의 라벨링 작업이 필요함.

→ DINO / DINOv2의 접근:  self-supervised learning.

 

Main work

• 자기지도학습 (SSL) 이 충분히 다양하고 정제된 데이터에서 학습되면 task-agnostic한 visual representation을 만들 수 있음을 증명함.
• Discriminative self-supervised 방법이 대규모 데이터와 모델 크기를 다룰 수 있도록 개선.
  • Discriminative self-supervised approaches: 이미지 패치간 차이를 구별하는 것이 목표인 자기지도학습

 

Contribution

• 자기지도학습을 통해 이미지와 픽셀 레벨 둘 다 성능을 보장하는 시각적 특성 학습.
• 이미지 분류와 같은 다양한 작업에 fine-tuning 없이 그대로 사용할 수 있는 사전 학습된 시각적 특성을 제공 (frozen feature)

fig.1

• Fig.1: PCA로 세 개의 주요 특성을 추출하고, component를 서로 다른 색상으로 매핑함
  • 포즈와 객체가 다른 이미지끼리도 비슷하게 매핑됨! E.g., a) 새와 비행기의 날개
  • 객체의 위치나 스타일 변화에 관계없이 일관된 특성을 추출함

fig.2

• Fig.2 파라미터 크기(scale)에 따른 성능 변화
  • DINOv2 (파란색 선): 파라미터 크기가 커질수록 성능이 높아지며 대부분 최고 성능
  • Weakly-supervised (분홍색 삼각형): SSL보다 높음
  • 기존 Self-supervised (주황색 삼각형): DINOv2, WS에 비해 성능이 낮음

 

Related Work

1. Intra-image Self-supervised Training
   • Pretext task: 이미지의 특정 부분(패치)만으로 나머지 정보를 예측하도록 함. 이를 통해 모델이 이미지의 패턴과 구조를 이해하도록 유도할 수 있다.
   • re-colorizing (색상), transformations (구조), inpainting (손상된 부분 복원; 맥락), re-ordering (공간) … etc
   • 기존 연구 한계: downstream task에서 성능을 극대화하려면 fine-tuning이 필요하다.
2. Discriminative Self-supervised Learning
   • Pretext task: 이미지나 이미지 그룹 간의 차이를 구분하면서 모델이 특성을 학습할 수 있다.
   • 기존 연구 한계: ImageNet과 같은 표준 벤치마크에서 frozen features로서 좋은 성능을 보였지만, 대규모 모델로 확장하기 어려웠다.
   • → DINOv2는 discriminative 방식을 확장해 대규모 데이터셋에서도 안정적인 성능을 보이도록 함.
3. Scaling Self-supervised Pretraining
   • 레이블이 없는 비정제 데이터(uncurated data)로 학습 시, 모델의 크기와 데이터 크기를 키우면 성능이 향상된다.
   • Goyal et al. (2021)은 충분히 많은 데이터가 있을 때, 모델 크기를 키우는 것이 자기지도학습에 유리하다는 것을 입증했다.
   • 한계: 데이터의 품질이 낮아지면 fine-tuning이 필요하다. 따라서 다양한 uncurated data를 다루는 데 한계가 있으며, 자기지도학습이 모든 종류의 데이터에서 잘 작동하는지는 알 수 없다.
   • DINOv2: 기존의 정제된 데이터셋을 확장해서 사용함
4. Automatic Data Curation
   • DINOv2의 데이터셋 construction은 image retrieval 분야에서 아이디어를 가져왔다.
   • 기존 연구: uncurated data 필터링에 이미지의 메타데이터를 활용하거나, 사전 학습된 비전 인코더를 사용했다.
   • DINOv2: 다른 정보 필요 없이, 이미지 간의 시각적 유사성을 통해 데이터를 필터링하고 선택한다. ( -> 3 Data Processing) 이 아이디어는 Wikipedia 데이터를 활용하여 텍스트를 필터링하는 텍스트 정제 파이프라인에서 차용했다. (Wenzek et al., 2020).

 

Data Processing

• LVD-142M dataset: DINOv2 모델을 학습하기 위해 만든 대용량 고품질 데이터셋
• Data Sources
  • Curated Data: ImageNet-22k, ImageNet-1k, Google Landmarks 등 다양한 기존의 고품질 데이터셋
  • Uncurated Data: 웹에서 크롤링한 1.2억 개의 이미지. (안전하지 않은 URL이나 NSFW 콘텐츠, 식별 가능한 얼굴 등을 필터링했다.)
• Deduplication: Uncurated data의 중복 이미지를 제거하기 위해 코사인 유사도가 높은 이미지 쌍을 제거한다.
• Self-supervised Image Retrieval: Uncurated에서 curated와 유사한 이미지를 선택해서 데이터셋 확장.
  • Curated data로 모델을 Pretrain -> 모델로 이미지의 고유한 특징을 벡터로 표현해서 각 이미지의 임베딩을 만든다.
  • Uncurated를 클러스터링하고, curated 이미지와 비슷한 이미지 4개를 검색해서 curated 데이터에 추가한다.
    • 클러스터에서 쿼리 이미지와 가장 가까운 4개 이웃을 코사인 유사도 기준으로 선택한다.

 

Discriminative Self-supervised Pre-training

Image-level objective (이미지 단위): DINO

교사는 학생보다 천천히 변화하며, 학생이 안정된 방향으로 학습할 수 있도록 일관된 목표(출력)를 제공함. 학생 네트워크와 교사 네트워크가 서로 비교하며 학습한다. 

 

• 각각 동일한 이미지의 서로 다른 부분에서 클래스 토큰(이미지의 대표 특징) 추출.
• 학생의 클래스 토큰을 학생 DINO 헤드에 통과시켜 프로토타입 점수 g_θs 계산 → softmax를 적용하여 확률 분포 p_s 를 얻는다.
• 교사의 클래스 토큰을 교사 DINO 헤드에 통과시켜 프로토타입 점수 g_θt 계산 → 지수이동평균으로 centering 하고 softmax를 적용해서 p_t 를 얻는다.
  • Centering: 출력이 특정 방향으로 편향되지 않도록 조정하는 것.
    • momentum
    • $g_t (x)←g_t (x)+c$
    • $c←mc+(1-m) \frac {1}{B} ∑g_{θt} (xi)$
• Loss: $p_s, p_t$ 의 cross-entropy loss를 사용한다.
  • $L_{DINO}=-∑ p_t \log⁡ p_s$
• Loss를 활용해서 θs를 학습시키고, 교사 파라미터 θt는 지수이동평균을 이용해서 업데이트.
  • $θ_t ← λθ_t + (1-λ) θ_s$

 

Patch-level objective: iBOT

학생이 일부 가려져 있는 패치(masked patch)를 교사와 유사하게 복원하도록 학습한다.

MAE (ViT)처럼 학생은 마스킹을 복원하면서 교사와 비교한다.

• 학생에게 입력할 이미지의 패치들을 무작위로 마스킹해서 보지 못하도록 하고, 교사 네트워크에는 모든 패치 정보를 포함한 이미지를 입력한다.
• 학생 네트워크에서는 mask patch token에 대해 iBOT 헤드를, 교사 네트워크에서는 해당 위치의 visible patch token에 대해 iBOT 헤드를 적용한다.
• 각 패치에 softmax를 적용하여 확률 분포를 계산한다. 교사의 출력은 centering해서 특정 방향으로 치우치지 않게 조정한다.
• 학생이 마스킹된 패치에 대해 교사 네트워크와 유사한 특징을 학습하도록 loss을 계산한다
  • $L_{iBOT}=-∑_i p_{ti} log⁡ p_{si}$
  • i: 학생 네트워크에서 마스킹된 패치의 인덱스.
• 이미지 단위 손실에서처럼, 학생 네트워크의 파라미터는 loss을 통해 학습되고 교사 네트워크는 이 값을 일정 비율로 반영하여 조금씩 변화한다.

 

Untying head weights between both objectives

• L_DINO, L_iBOT 각 손실은 학습 가능한 MLP 헤드를 갖는다.
• 기존 연구에서는 파라미터를 공유했을 때 더 좋은 성능을 낸다고 하였으나 큰 scale에서는 패치 레벨에서 underfit & 이미지 레벨에서 overfit되는 것을 발견 → 개별 헤드를 적용해서 서로 다른 파라미터를 사용하도록 함.

 

Sinkhorn-Knopp Centering

• SwAV에서 도입된 방식을 적용해서, DINO와 iBOT의 teacher의 softmax-centering 단계를 Sinkhorn-Knopp(SK) 배치 정규화로 대체했다. Student는 softmax 정규화를 적용했다.
• Sinkhorn-Knopp 알고리즘: 행렬의 각 행과 열의 합을 특정 목표값으로 맞추기 위해 반복적으로 조정 → 출력값이 균일하게 분포하도록 강제함

 

KoLeo Regularizer

• L2 정규화를 feature에 먼저 적용한 후 다음 정규화 계산.
  • $L_{koleo} =- \frac 1 n ∑_{i=1}^n \log⁡ d_{n,i}$
  • $d_{n,i}= \min_{j≠i}∥x_i-x_j∥$: 특정 벡터 xi와 배치 내 다른 벡터들 사이의 최소 거리

Kozachenko-Leonenko differential entropy estimator에 기반한 방식으로 feature 분포를 균일하게 만든다.

 

Adapting the Resolution

• Apative하게 높은 이미지 해상도에서 학습을 진행해서 작은 물체를 더 잘 인식할 수 있도록 한다.
  • Segmentation과 같은 픽셀 단위의 downstream task에서, 작은 물체는 낮은 해상도에서는 잘 인식되지 않는다.
• 높은 해상도로 학습하면 시간과 메모리 소모가 크기 때문에, 마지막 단계에서만 518×518로 높여 학습한다.

 

Efficient Implementation

대규모 모델을 학습하기 위한 효율적 구현 방법

 

Fast and Memory-Efficient Attention

• Self-attention: Attention을 자기 자신과 수행해서 한 문장(이미지) 안에서 요소들의 연관성을 알아내는 것 DINO는 한 이미지 안에서의 각기 다른 위치의 패치간의 관계를 이해하는 것이 중요하므로 self-attention 수행
  • Required) 입력 벡터(단어 임베딩 or 이미지 패치)는 쿼리, 키, 값 (Q,K,V) 벡터로 변환됨.
  1. 쿼리 벡터 Q와 모든 키 벡터 K 간 내적 계산 -> Q와 모든 K 간의 유사도인 S를 얻을 수 있음
  2. S에 softmax 취해 가중치 확률 P로 변환
  3. 각 단어의 값 벡터 V에 가중치 P를 곱하고 합을 구해 단어별 최종 점수 O를 구한다. (O는 쿼리 벡터에 얼마나 연관이 있는지를 나타냄!)
• Standard attention implementation의 문제:
  • GPU의 계산은 실제로는 HBM을 거쳐 SRAM에서 수행됨
  • GPU HBM write load가 빈번하게 발생함

 

기존 방법

 

• 해결: FlashAttention 기반의 self-attention을 도입해서 메모리 사용량과 속도를 개선함.

FlashAttention

• FlashAttention: GPU에서 연산 시 HBM access를 줄여 속도를 높인 방법
  • row 방향을 따라 Q를 T_r개, K, V를 T_c개의 블록으로 쪼갠다. 이후 블록별로 SRAM에 load하고 softmax를 계산해서 한 번에 O를 write.
• 특히 embedding dimension이 64의 배수일 때 효율성이 극대화된다는 점을 반영해, 기존 구조보다 계산 효율을 높이도록 조정함.

 

Sequnce Packing

NLP에서 사용된 기법으로, 서로 다른 길이의 시퀀스(이미지 크롭)를 하나로 결합하여 Transformer에 입력한다.

•  기존
  • 교사는 global crop(224x224), 학생은 global crop / local crop(98x98) 을 랜덤하게 입력받아서 forward 한다.
  • 한 이미지의 서로 다른 사이즈 크롭들을 한 개씩 네트워크에 넣고 forward → loss도 각 크롭별로 계산하고 backward pass. 즉 크롭 갯수만큼 forward/backward 따로 했어야 함.
• 개선
  • 다른 사이즈의 크롭에서 얻은 패치 토큰 시퀀스들을 하나로 결합해서 Transformer에 입력 → 한 번에 여러 크롭을 처리!
  • 다만 서로 다른 패치 간 영향을 주는 것을 막기 위해 Self-attention 레이어에서 block-diagonal mask를 사용함.
  • 여러 시퀀스를 한 번에 처리하면서 이전처럼 개별 처리하는 것과 같은 결과를 얻으면서 계산 비용을 크게 줄일 수 있다.

 

Efficient stochastic depth

• Stochastic depth
  • Depth가 높은 네트워크의 Layer 일부분을 무작위로 drop해서 train하고, test할땐 모든 layer 사용.
  • 기존의 방법은 drop하기로 결정한 레이어의 출력을 0으로 마스킹함. (계산에 포함)
• Efficient Stochastic Depth
  • Layer를 스킵하는 것이 아니라, 드롭 비율(d=40%)에 따라 배치 내 전체 샘플 B개 중에 랜덤하게 (1−d)×B개 샘플만을 선택해서 계산함.
    • 그러면 stochastic depth가 아니라 그냥 SGD 아닌가? → stochastic depth에서도 SGD처럼 샘플을 몇 개 선택함
  • 드롭 비율에 비례하여 계산 비용과 GPU 메모리 사용량을 절감함.

 

Fully-Sharded Data Parallel (FSDP)

• 문제
  • DINO는 AdamW Optimizer를 사용하는데 4개의 모델 복제본이 필요하다. (student, teacher, optimizer 1st moment m, 2nd moment v)
    • AdamW: Adam의 generalization 성능을 높이기 위해 weight decay 항을 명시적으로 넣어준 optimizer
  • 따라서 큰 모델 파라미터를 저장해 두어야 하므로 메모리 부담이 커진다.

• FSDP
  • 모델의 파라미터를 GPU 간에 분할(sharding)하여 저장한다. 따라서 모델 크기가 개별이 아닌 전체 GPU 메모리의 합에 의존하게 된다.
  • Mixed Precision: Optimizer가 필요로 하므로 weight는 float32로 저장되지만, GPU 간 파라미터 전송과 그라디언트 연산은 float16으로 수행됨.
  • 다른 SSL에서 사용한 float32로 모든 파라미터를 전송하는 DDP 방식에 비해 통신 비용이 50% 줄어들었음.

 

Model Distillation

• DINOv2는 여러 개의 모델을 제공하는데, 이 중 가장 큰 모델인 ViT-g를 사전 학습시킨다.
• ViT-L과 같은 작은 모델을 처음부터 학습시키는 대신, ViT-g의 출력을 따라가도록 학습시킨다.
• 기존 학습 루프를 활용하지만 몇 가지 차이가 있다.
  • Teacher model is frozen: 파라미터가 업데이트되지 않는다.
  • spare EMA of the student: student의 spare에 EMA를 수행해서 final model로 사용한다. (기존의 teacher)
  • remove the masking and stochastic depth: 안정적인 학습을 위해 제거한다.
  • iBOT loss on global crops: iBOT Loss는 global crop 두 장에만 적용된다. (masking 없이?)

 

 

 

Ablation Study와 Result를 보면 결과가 매우 좋다.

역시 좋은 거 다 넣으면 결과가 좋을 수밖에 ...

 

Patch matching

PCA

 

스타일이 완전히 달라도 유사한 이미지 featrue끼리 잘 매칭하고, PCA로 분석했을 때도 앞에서 봤던 것처럼 서로 스타일이 다른 객체끼리도 유사한 부분(머리, 꼬리)는 같은 색으로 표시된다.

 

심지어 비행기랑 독수리는 다른 클래스인데도 날개와 머리 부분을 보면 유사하게 표시된 것을 볼 수 있다.

 

이 DINOv2로 학습시킨 feature는 파인튜닝 없이 사용할 수 있는 frozen feature이다. 

즉 backbone을 freeze시키고 linear classifier만 학습시켜서 사용해도 다양한 downstream task에서 잘 작동한다.