[X:AI] DINO 논문 리뷰

 

논문 원본 : https://arxiv.org/abs/2104.14294

Emerging Properties in Self-Supervised Vision Transformers

 

1. Introduction 

• ViT는 CNN과 경쟁력이 있지만, 아직까지 CNN 대비 명확한 이점을 제공하지 못하고 있음
• 본 논문에서는 비전 분야에서 Transformer의 성과가 제한적인 이유가 감독 학습 기반의 사전 학습 방식 때문이 아닌지 의문을 제기

• NLP에서 Transformer의 성공 요인 중 하나는 BERT 마스킹 방식이나 GPT의 언어 모델링 방식과 같은 자기지도 사전학습을 사용했다는 점
• 이러한 자기지도 방식은 문장 내 단어들을 활용해 풍부한 학습 신호 제공하며, 이는 단일 레이블 예측만을 요구하는 감독학습보다 더 나은 학습 유도

• 이미지 수준의 감독학습은 풍부한 시각 정보를 몇 천 개 카테고리 중 하나의 레이블로 축소시키는 문제
• 본 연구에서는 이러한 방식에 영감을 받아 ViT에서 자기지도 사전학습이 가져오는 특성을 분석

• ViT를 자기지도 학습시키면
  • 자기지도 ViT feature는 장면의 레이아웃과 객체 경계 정보를 명확하게 포함 (마지막 블록의 self-attention 모듈에서 직접 확인 가능)
  • 자기지도 ViT feature 간단한 k-NN 분류기만으로도 고성능 (ImageNet top-1 78.3%)

• 특히 자기학습과 더불어 작은 patch + momentum encoder + multi-crop augmentation 조합이 feature 품질 향상에 중요하다는 것도 확인

• 이러한 발견을 바탕으로 라벨 없이 지식 증류(Knowledge Distillation) 형태로 해석할 수 있는 간단한 자기지도 학습 기법인 DINO 설계
  • Teacher (momentum encoder 기반)의 출력을 Student가 예측 (Cross Entropy 사용)
  • Collapse 방지용으로 " centering, sharpening" 만 사용, 복잡한 loss나 구조는 필요 없음
• DINO의 큰 장점은 ConvNet과 ViT 모두에서 아키텍처 수정 없이 작동한다는 점

• ViT-Base + DINO → ImageNet 선형 분류 정확도에서 80.1%
• ResNet-50 + DINO → 기존 SOTA와 동등 수준 (CNN에서도 잘 작동)
• 8 GPU 서버 2대로 3일간 학습하여 ImageNet 선형 분류 정확도 76.1% 성능을 달성했고, 이는 CNN 기반 자기지도 학습보다 효율적

 

 

3. Approach

 

   3.1 SSL with Knowledge Distillation

• 입력 이미지 x가 주여졌을 때, 두 네트워크는 K차원의 확률 분포 Ps, Pt를 출력 (이 확률은 softmax함수로 정규화된 출력)
• 여기서 τs, τt는 출력 분포의 sharpness를 조절하는 온도 파라미터

 

• 학생 네트워크는 교사의 분포를 모방하기 위해 교사와 학생 사이의 cross-entropy loss 최소화 

 

   Multi-Crop

• 하나의 이미지로부터 다양한 변형된 crop(view)을 생성
• 2개의 global view : 전체 이미지의 50% 이상을 포함
• 여러 개의 local view : 전체의 50% 미만을 포함
• 모든 crop은 학생 네트워크에 통과되고, global view만 교사 네트워크에 입력
• 즉, 교사의 global view 출력과 학생의 모든 view 출력을 비교하여 학습

 

   Teacher Network

• 기존 지식 증류와 달리 고정된 교사 네트워크가 없음
• 대신, 학생의 이전 가중치를 이용해 교사를 동적으로 업데이트
• 방법: 지수이동평균(EMA) 사용

 

• λ은 학습 중에 0.996에서 1로 점점 증가하는 코사인 스케줄 사용
• 이 방식은 모델을 평균화하여 안정성과 성능을 높이는 Polyak-Ruppert 평균화와 유사함

   Network Architecture

• g = h ∘ f 형태로 구성
  • f: 백본 (ViT 또는 ResNet)
  • h: 프로젝션 헤드
    • 3층 MLP → L2 정규화 → 가중치 정규화된 FC 레이어 (K차원)
• ViT 기반이기 때문에 Batch Normalization을 사용하지 않음

  Avoiding collapse

• 다른 SSL 방식에서는 contrastive loss, predictor, BN 등으로 collapse 방지
• DINO에서는 교사 출력에 'centering' + 'sharpening' 조작을 통해 collapse 방지
  • Centering: 출력에 평균 c를 더함 → 치우친 분포 방지
    • c도 EMA로 업데이트
  • Sharpening: softmax 온도 τt 값을 작게 설정해 더 날카로운 출력 생성

 

 

4. Main Results

 

   4.1 Comparing with SSl frameworks on ImageNet

 

   1) Comparing with the same architecture

• DINO는 ResNet-50 기반에서도 다른 SSL 기법과 유사한 성능
• ViT-S 기반에서는:
  • BYOL, MoCov2, SwAV보다 +3.5% 높은 선형 분류 정확도
  • k-NN 정확도는 +7.9% 높음
• 놀라운 점: 단순 k-NN 분류 성능 ≒ 선형 분류기 성능
  • ViT + DINO 조합에서만 나타나는 특성

   2) Comparing across architecture 

• 패치 크기를 줄이면 (예: /16 → /8) 성능 크게 향상
• ViT-B/8 (8×8 패치):
  • 선형 분류 80.1%, k-NN 77.4%
  • 기존 SOTA보다 10배 적은 파라미터, 1.4배 빠른 속도

 

 

   4.2 Properties of ViT trained with SSL

 

    4.2.1 Nearest neighbor retrieval with DINO ViT

• Oxford/Paris 데이터셋에서 기존 supervised 모델보다 더 나은 성능
• Google Landmarks v2로 학습된 DINO ViT → landmark 검색에서도 기존 SOTA 추월
• Copy detection (INRIA Copydays):
  • 다양한 왜곡(블러, 인쇄, 스캔)에도 견고
  • [CLS] + patch token pooling 결과를 사용한 1536차원 특징으로 cosine similarity 측정
  • 사전학습만으로도 강력한 복제 탐지 성능

 

    4.2.2 Discovering the semantic layout of scenes

• ViT의 Self-Attention map만으로 객체 위치와 형태 파악 가능
• DAVIS-2017 Video Segmentation:
  • 학습 없이 최근접 프레임 간 특징 유사도로 분할 수행
  • /8 패치가 /16보다 +9.1% 성능 향상
• Self-Attention을 Threshold해서 만든 마스크가 supervised 모델보다 뛰어남

 

 

    4.2.3 Transfer learning on downstream tasks

• 다양한 다운스트림 태스크에서 finetune 성능 측정
• DINO 사전학습된 ViT는 Supervised ViT보다 일관되게 더 좋은 성능
• 특히 ImageNet에서도 +1~2% 정확도 향상

 

 

5. Ablation Study of DINO

 

   5.1 Importance of  the Different Components

Momentum encoder 없음	학습 실패 (collapse 발생)
Multi-crop 없음	성능 급감
Cross-entropy loss 제거	성능 저하
Predictor 추가	거의 영향 없음 (BYOL과는 다름)

 

 

 

   5.2 Impact of the choice of Teacher Network

• Momentum teacher (EMA)가 가장 효과적
• 단순히 이전 epoch/iteration의 student 사용 → collapse 발생하거나 낮은 성능
• DINO의 momentum teacher는 Polyak-Ruppert averaging처럼 동작
  • 학습 도중에도 성능 좋은 앙상블 역할

 

 

   5.3 Avoiding collapse

• Collapse = 출력이 모두 같아지는 현상 (정보 손실)
• Centering: 특정 차원만 지배하는 collapse 방지
• Sharpening: 출력이 너무 퍼지는 collapse 방지
  • 두 기술이 서로 상반된 collapse를 막아주는 보완 관계

 

   5.4 Compute requirments

• multi-crop 전략: 적은 시간에 높은 성능 달성
  
  • 2×224²만 사용 → 46시간에 72.5%
  • 2×224² + 10×96² → 24시간에 74.6% (시간 절반에 성능 +2%)
• 더 많은 crop 수는 수익 감소 효과 있음 (6 → 10개에서 +0.2%)

 

 

   5.5 Training with small batches

• 기본: bs=1024 (성능 최고)
• bs=128로도 괜찮은 성능 유지 가능
• 심지어 1 GPU, bs=8으로도 50에포크에 35.2% 달성
  • DINO는 소규모 자원으로도 학습 가능성 있음

 

 

6. Conclusion

• DINO + ViT는 특별한 구조 없이도, 최고 성능의 CNN 기반 모델들과 비슷한 성능을 달성
• 실험을 통해 발견된 두 가지 중요한 특징:
  • k-NN 분류 성능이 뛰어남 → 이미지 검색(검색 기반 유사 이미지 탐색)에 활용 가능
  • 장면의 공간 구조(scene layout) 정보 보존 → 약지도(weakly-supervised) 이미지 분할에 유리
• 핵심 주장:
  • 자기지도학습(Self-supervised learning)은 ViT 기반의 BERT처럼 범용적인 비전 모델을 만드는 핵심 열쇠가 될 수 있음
• 향후 계획:
  • 라벨 없는 무작위 이미지 대량 사전학습을 통해 ViT 기반 범용 비전 모델의 한계를 확장해보고자 함