DINO 논문 리뷰

Emerging Properties in Self-Supervised Vision Transformers

---

Summary

DINO는 Self-Supervised Learning 환경에서 ViT를 활용하여 레이블 없이 의미 있는 표현을 학습하는 방법을 제시한다. DINO는 학생(Student) 네트워크와 교사(Teacher) 네트워크 구조를 가지고 있으며, student 네트워크가 teacher 네트워크의 출력을 모방하도록 학습한다. 또한, 본 논문은 (1) Softmax + Temperature Sharpening, (2) Centering, (3) multi-crop augmentation을 사용하여 Self-supervised learning의 collapse 위험성을 해결한다.

이러한 접근 방식은 ViT가 객체 분할과 같은 고수준 시각 표현을 자발적으로 학습하며, 자연스러운 군집화와 지역적 특징 분리를 하는 능력을 가지게 한다.

---

Methods

Self-Distillation

보통 Distillation이라 함은, 우수한 성능을 가진 teacher model과 학습하고자 하는 student 모델의 출력 분포가 유사하도록 훈련하는 방식이다. 그러나 self-distillation은 다음과 같은 학습을 통해 진행된다:

1. 초기 image $x$는 서로 다른 augmentation을 통해서 각각 $v_1$과 $v_2$로 변환되고 student와 teacher 모델의 input으로 들어간다.
2. 각 모델의 layer을 거치면서 만들어지는 representation $v_1 '$과 $v_2 '$의 표현의 차이를 loss로 계산하여, student 모델은 gradient update 방식으로 학습되고, teacher model은 EMA를 통해서 파라미터가 갱신된다.

이러한 방법은 student 모델이 서로 다르게 augment된 input들을 정렬시키도록 학습되기 때문에, 별다른 지도 없이도 좋은 표현을 학습하도록 유도할 수 있다.

 

방법 면에서는 BYOL과 거의 같지만, (1) predictor을 사용하지 않는 점과 (2) loss를 MSE가 아닌 Cross entropy를 사용한다는 점, 그리고 (3) Teacher와 student에 적용하는 augmentation이 다르다는 점이 핵심 차이이다.

여기서 바로 의문이 들 수 있는 것은, model collapse에 너무 취약하다는 문제점이 있지 않느냐는 것이다. 따라서 DINO는 이를 어떻게 해결하는지 알아보자.

 

Multi-Crop Augmentation

출처: https://www.keep-current.dev/unsupervised-learning-of-visual-features-by-contrasting-cluster-assignments/

위의 그림처럼, teacher model에는 image crop을 적용하되, 전체적인 모습을 볼 수 있도록 global crop을 적용한다. Student 모델의 경우, 작은 영역을 여러 개 만드는, local crop을 적용한다. 이때 global crop는 보통 2개 정도로 적게, local crop는 8개~9개 정도로 많이 적용하는데, 그 이유는 다음과 같다:

• 계산 효율성: Global crop의 해상도는 local crop에 비해서 해상도가 높아서, 연산 측면에서 local crop을 더 많이 가져가는 것이 유리하다.
• Representation 균형: Global은 큰 그림을 담기 때문에 몇 개면 충분하지만, local은 세밀한 패턴과 부분적 특징을 다양하게 많이 보게 해야 representation의 consistency를 형성할 수 있다.
• Collapse 방지: Global view끼리는 consistency를 맞추는 것이 간단하지만, local view와의 consistency는 훨씬 어려워서, 모델이 단순히 하나의 벡터로 수렴해버리기 어렵다.

이렇게 $N_1$개의 global representation에 대한 확률분포와 $N_2$개의 local representation이 만나서, $N_1 × N_2$만큼의 cross entropy 연산이 일어나며 loss가 업데이트된다.

추가적으로, Collapse는 두 가지 방향으로 일어난다. 출력 확률 분포가 모든 클래스에 대해 uniform으로 수렴하는 uniform collapse와 특정 방향으로 수렴하는 mode collapse가 있다. 이번에는 이를 해결하기 위한 방법인 centering과 sharpening을 알아보도록 하자.

Sharpening

Teacher의 softmax에 낮은 temperature을 적용하여 분포를 날카롭게 만들 수 있다. 이는 분포가 uniform으로 퍼지는 것을 막아 uniform collapse를 방지할 수 있다.

Centering

Teacher의 출력 분포에서 특정 클래스 차원으로 쏠리지 않게 하기 위하여, center 벡터 $c$를 EMA로 업데이트한다.

각 batch마다 teacher 출력의 평균을 추적하여 중심을 맞춰주면서, 동시에 EMA 계수 $m$으로 점진적인 보정이 일어나도록 적용할 수 있다.

---

Experiments

성능이 뛰어난 것을 알 수 있다. 무엇보다 주목할 수 있는 것은 image segmentation인데, experiment 결과를 확인해보자:

Object segmentation에서 우수한 성능을 내는 것을 확인할 수 있으며, 실제로 attention map을 통해 보았을 때

attention map에서 object를 잘 감지하는 것을 확인할 수 있다. 특히 supervised learning보다 더욱 성능이 높은 것이 인상적이다.

마지막으로 DINO에 적용된 방법론이 모델의 성능에 얼마나 기여하는지 확인해보자.

여기서 Momentum Encoder가 적용되지 않은 경우 model이 collapse되는 것을 확인할 수 있는데, 안정된 target distribution을 가지기 위해서는 EMA알고리즘이 반드시 필요함을 확인할 수 있다. 또한, Predictor이 따로 없더라도 DINO에서 적용된 방법으로도 collapse를 막고 좋은 성능을 낼 수 있도록 할 수 있는 것을 알 수 있다.