컴퓨터 비전

[컴퓨터 비전] DINO v1

🐋 Growth Whale 2025. 6. 18. 21:47
"Emerging Properties in Self-Supervised Vision Transformers"
ICCV 2021, 29 Apr 2021

 

1. Self Supervised Learning을 통한 ViT 학습

 

최근 NLP 분야에서 "Transformer" 모델을 통해서 휼롱한 성능을 보여주고, AI가 크게 주목받게 되어 있습니다. 이처럼 이미지 분야에서도 Vision Transformer(ViT)가 등장하면서 CNN(합성곱 신경망)의 대안을 제시하고 있었습니다.

 

하지만 이 논문이 나오기 전까지는 ViT를 Supervised Learning통해 학습을 진행했지만, 기존 CNN에 대비해서 놀라만한 성능을 보여주지 못했습니다. 저자는 Supervised Learning이 특정 카테고리만 분류하는 학습으로는 이미지에 대한 특징을 충분히 반영하지 못하는 문제라고 생각합니다.

 

이에 논문은 흥미로운 질문을 합니다.

 

ViT도 NLP처럼 자기지도학습(self-supervised learning)으로 더 강력한 성능을 낼 수 있을까?”

 

 

페이스북 AI 연구팀은 DINO (self-DIstillation with NO labels) 라는 자기지도 학습 방법에 대해서 제안합니다. 

그리고 라벨 없이 학습함으로 객체를 분리해서 보는 능력(세그멘테이션)이 자연스럽게 생긴다는 것과 k-NN 분류에서도 놀라운 성능을 보인다는 것을 발견합니다.

 

2. 사전 지식 : Knowledge Distillation

 

해당 논문의 자가 학습 방법 이름이 DINO(Self-Distillation without Labels) 라고 지었습니다.이건 기존의 Knowledge Distillation에서 아이디어를 가져온 겁니다. 

 

Knowledge Distillation은  쉽게 말하면 큰 모델(Teacher)의 지식을 작은 모델(Student)에게 전달하는 과정입니다.

큰 모델의 결과와 작은 모델의 결과를 동일하게 만들어 주도록 작은 모델을 학습 시키는 방식으로 학습합니다.

 

보통 이 방법은 성능이 좋은 큰 모델을 좀 더 가볍게 만들기 위해서 사용되는 방법입니다. 하지만 해당 방법은 Teacher 모델과 Student 모델이 동일한 구조를 사용해 경량화 목적보다는 자가 학습의 목적으로 사용합니다. 

 

 

3. 자가 학습 방법

DINO는 Knowledge Distillation 처럼 두 모델(Teacher, Student)을 이용해서 자가 학습합니다. Student 모델의 결과값이 Teacher 모델과 비슷하게 나오도록 Student 모델을 학습합니다.

 

이를 좀 더 수식적으로 표현해 보겠습니다. 모델의 결과를 K 차원의 벡터로 나오는데 이를 Softmax를 통해서 확률 분포로 표현할 수 있습니다. 

 

 

또한 Teacher 모델과 Student 모델을 같게 만들기 위해서 Cross Entrophy Loss를 사용해서 두 모델이 같게 나오도록 학습합니다. 

 

 

여기서 주의할 점은 기존 Knowledge Distillation과 다르게 Teacher 모델과 Student 모델은 같은 구조를 사용합니다. 또한 Teacher 모델은 학습하지 않고, Student 모델만 학습을 진행합니다. 그럼 언제 Teacher 모델을 학습할까요?

이는 추후 이 논문에서 제안하는 Momentum encoder를 알아보도록 합니다. 

 

Momentum encoder 

 

Teacher 모델은 직접 학습하지 않고, 대신 Student 모델의 가중치를 천천히 따라가도록 유지하는 방식을 사용합니다. 이런 모델을 Momentum encoder라고 부릅니다. 이 방법을 사용했을 때 안정적으로 학습을 진행했다고 합니다.

 

이를 수식으로 표현하면 아래와 같습니다. 

 

즉, Teacher 모델은 Student 모델의 일종의 지수이동평균(EMA: Exponential Moving Average)이라고 볼 수 있습니다. Teacher 모델은 급격하게 변화하지 않고, 천천히 변화하면서 학습하도록 만듭니다. 이는 아래처럼 말할 수 있습니다. 

 

“선생님이 학생의 과거 모습을 기억해서 조금씩 가르쳐주고, 학생은 계속해서 자신을 업데이트합니다.”

 

 

또한 논문에서는 해당 방식이 모델 앙상블 효과도 낸다고 합니다. 모델 앙상블은 여러 모델을 사용해 성능을 높이는 방법인데, 해당 모델도 매번 학습한 모델들의 합으로 결과를 내기 때문에 앙상블 효과를 낸다고 합니다. (자세한 내용은 더 확인해 봐야 될 것 같습니다.)

 

이를 통해 항상 Teacher 모델은 Student 모델에 비해 성능이 좋아, 안정적인 학습을 하도록 도와줍니다. 

 

Multi-Crop 전략

 

이 논문에서는 Global-Local 정보를 학습하기 위해서 Multi-Crop 전략을 사용합니다. Multi-Crop은 입력 이미지를 여러 해상도(크기)로 잘라서 모델에 넣는 방식입니다. 

 

Crop 방식은 한 이미지에서 2 장의 Global Crop (이미지의 50% 이상), 여러 장의 Local Crop (이미지의 50% 미만)을 생성합니다. 그리고 Teacher 모델에는 Global Crop만 넣고, Student 모델에는 모든 Crop을 다 넣어 학습을 진행합니다. 

 

이는 Student 모델이 부분만 보고 전체를 이해할 수 있게 만드는 것입니다. 이를 통해서 Global-Local 대응 관계를 자연스럽게 배우도록 합니다. 

 

Collapse 방지 : Centering과 Sharpening

 

자기 지도 학습에 가장 큰 문제는 Collapse 현상 입니다. Collapse 현상은 모델이 모든 입력에 대해 동일한 출력을 내놓는 현상입니다. 예를 들어 어떤 이미지를 넣든, 항상 [0.1, 0.1, 0.1, 0.1, 0.1, ...] 같은 결과가 나오는 현상 입니다. 

 

이 논문에서는 Collapse를 피하기 위해서 Centering과 Shapening을 적용했습니다.

 

Centering은 Teacher 모델의 출력에 평균 값을 빼주는 작업입니다. 이 과정을 통해서 모든 차원 값들이 평평해져 특정 차원으로 몰리도록 학습되는 것을 방지합니다. 이를 수식으로 표현하면 아래와 같습니다. 

 

 

Sharpening은 출력 값을 표현하는 Softmax 함수에 temperature 변수를 추가하여 분포를 좀 더 날카롭게 만드는 작업이빈다. 이 과정을 통해서 확률 분포가 평평해지는 것을 방지합니다. 이를 포함한 Softmax 함수는 아래와 같습니다. 

 

 

Centering과 Sharping은 서로 다른 역할을 합니다. Centering은 분포를 평평하게 만드는 역할을 하고, Sharpening은 분포를 날까롭게 하는 역할을 합니다. 하나만 사용하면 Collapse될 확률이 높은 데, 이를 둘 다 사용하면 학습이 잘 되도록 도와줍니다.

 

4. ViT에서 나타나는 특징

논문에서는 DINO가 놀라운 일반화 능력을 보임을 나타냈습니다. 

 

1. Image retrieval에 뛰어난 능력을 보임

 

ViT는 Patch 이외에도 [CLS] 클래스 토큰을 Transformer에 넣어 학습합니다. 그럼 [CLS]에는 이미지에 대한 글로벌한 정보를 얻을 수 있습니다. 해당 [CLS] 토큰으로 KNN 이미지 분류로 사용해도 뛰어난 결과를 얻을 수 있습니다. 

 

K-NN 이미지 분류는 저장된 이미지 벡터 중에서 비슷한 이미지를 찾고, 해당 이미지에서 가장 많은 클래스로 답을 정하는 분류 방법입니다. 여기에서 좋은 성능을 보인다는 것은 비슷한 종류의 이미지의 벡터는 비슷한 위치에 존재하여 학습이 잘 되었다는 걸 증명합니다. 

 

실제로 장소를 찾는 Landmark Retrieval과 비슷한 복사본을 찾는 Copy Detection 태스크에서 꽤 좋은 성능을 보여주었습니다. 

 

 

2. 세그먼테이션 정보를 학습

 

아래 그림은 마지막 레이어의 [CLS] 토큰의 self-attention 맵을 시각화 한 것입니다. 이를 보면 확실히 [CLS]에 이미지 중요한 객체가 어디 있는 지 알고, 객체 영역을 정확히 알고 있다는 것을 알 수 있습니다. 분명 지도 학습이 아닌데도 자가 지도 학습으로 스스로 해당 부분을 학습했다는 게 신기합니다.

 

그리고 각 헤더 Attention 맵을 다른 색깔로 시각화해보면, 서로 다른 헤더들이 서로 다른 객체를 학습했다는 것을 알 수 있습니다. 

 

 

 

5. 정리

 

DINO 논문은 Momentum Encoder를 사용해 이미지 도메인에서 Self-Supervsed Learning을 할 수 있는 방법을 알려주고, DINO 학습시 나타나는 특징에 대해서 설명했습니다.

 

ViT의 Self-Superved Learning을 위한 학습 방법 

  • Knowledge Distillation 기반의 Momentum Encoder
  • Multi-Crop 전략
  • Collapse 방지 : Centering과 Sharpening

DINO 학습시 나타나는 특징

  • KNN 분류에서 좋은 성능을 보임 (비슷한 이미지의 벡터는 동일한 위치에 존재)
  • Segmentation 정보을 담음 (객체 영역 특징을 잘 학습)