[논문 리뷰] Exploring The Role of Mean Teachers in Self-supervised Masked Auto-Encoders

Youngwan Lee, Jeffrey Willette외 3명|arXiv (Cornell University)|2022. 10. 05.

Domain Adaptation and Few-Shot Learning인용 수 7

한 줄 요약

이 논문은 마스크된 자동에코더(MAE)를 개선하기 위해 지수이동평균(EMA) 교사 모델을 통합한 자기지도 학습 방법인 RC-MAE를 제안한다. EMA 교사는 특성 유사도에 기반해 기울기를 조정함으로써 조건부 운동량 정규화 기법으로 작용한다. 이 방법은 MAE 및 최신 자기정규화 방법에 비해 수렴 속도가 빠르고, 메모리 사용량이 줄어들며, 최종 시각 작업에서 더 높은 강건성과 성능을 달성한다.

ABSTRACT

Masked image modeling (MIM) has become a popular strategy for self-supervised learning~(SSL) of visual representations with Vision Transformers. A representative MIM model, the masked auto-encoder (MAE), randomly masks a subset of image patches and reconstructs the masked patches given the unmasked patches. Concurrently, many recent works in self-supervised learning utilize the student/teacher paradigm which provides the student with an additional target based on the output of a teacher composed of an exponential moving average (EMA) of previous students. Although common, relatively little is known about the dynamics of the interaction between the student and teacher. Through analysis on a simple linear model, we find that the teacher conditionally removes previous gradient directions based on feature similarities which effectively acts as a conditional momentum regularizer. From this analysis, we present a simple SSL method, the Reconstruction-Consistent Masked Auto-Encoder (RC-MAE) by adding an EMA teacher to MAE. We find that RC-MAE converges faster and requires less memory usage than state-of-the-art self-distillation methods during pre-training, which may provide a way to enhance the practicality of prohibitively expensive self-supervised learning of Vision Transformer models. Additionally, we show that RC-MAE achieves more robustness and better performance compared to MAE on downstream tasks such as ImageNet-1K classification, object detection, and instance segmentation.

연구 동기 및 목표

• EMA 기반 정규화를 통한 자기지도 학습에서 학생과 교사 간의 동적 상호작용을 이해하기 위해.
• 비전 트랜스포머를 위한 픽셀 수준의 마스크된 자동에코딩(MIM)에서 평균 교사의 기능적 역할을 조사하기 위해.
• 기존 자기정규화 접근 방식에 비해 계산 및 메모리 비용을 줄이는 더 효율적이고 효과적인 자기지도 사전학습 방법을 개발하기 위해.
• 객체 검출 및 인스턴스 세그멘테이션과 같은 조밀한 예측 작업을 포함한 분류 작업에서의 최종 성능 향상을 위해.

제안 방법

• 학생의 재구성 페치에 대한 일致성 타겟을 제공하기 위해 EMA 교사를 통합한 수정된 MAE인 RC-MAE를 도입한다.
• 학생의 출력을 원본 마스크된 페치와 비교하기 위해 재구성 손실 $\mathcal{L}_{\text{r}}$ 를 사용한다.
• 학생의 예측을 보이는 페치에서 교사의 출력과 일치시키기 위해 일치 손실 $\mathcal{L}_{\text{c}}$ 를 적용한다.
• 특성 유사도에 기반해 현재 입력과 이전 입력 간의 특성 유사도에 따라 기울기를 조정하는 EMA 업데이트 규칙을 사용하여 천천히 업데이트되는 교사 네트워크를 유지한다.
• 선형 모델에서 기울기 역학을 분석하여, 특성이 유사할 경우 이전 방향을 제거함으로써 교사가 기울기를 조건부로 수정함을 보여준다.
• ViT-B 및 ViT-L 백본을 사용하여 ImageNet-1K, COCO 및 강건성 벤치마크에서 방법을 실증적으로 검증한다.

실험 결과

연구 질문

• RQ1 EMA 교사는 픽셀 수준의 마스크된 자동에코딩에서 기울기 업데이트에 어떻게 영향을 미치는가?
• RQ2 타겟을 제공하는 것 외에 교사는 학생의 최적화 과정에서 어떤 기능적 역할을 하는가?
• RQ3 MAE에 EMA 교사를 단순히 통합하는 것으로 기존 자기정규화 방법에 비해 더 빠른 수렴과 낮은 메모리 사용량을 달성할 수 있는가?
• RQ4 교사의 조건부 기울기 보정이 시각 작업에서의 강건성과 최종 성능 향상에 기여하는가?
• RQ5 교사의 효과는 조건부 운동량 정규화의 한 형태로 이론적으로 설명될 수 있는가?

주요 결과

• RC-MAE는 800 에포크 후 ImageNet-1K 분류에서 83.4%의 Top-1 정확도를 달성하여, MAE가 1600 에포크에서 달성한 성능을 동일하게 유지함으로써 더 빠른 수렴을 보였다.
• RC-MAE는 iBOT, MSN, BootMAE와 같은 최신 자기정규화 방법에 비해 메모리 사용량과 계산 비용을 줄였으며, 교사에 대해 보이는 페치만 입력으로 제공함으로써 이를 달성했다.
• ImageNet-C, ImageNet-A, ImageNet-R, ImageNet-Sketch에서 RC-MAE는 일관된 강건성 향상을 보였으며, mCE는 1.8포인트 감소(49.9 vs. 51.7)하고 IN-A 및 IN-R에서 최대 1.2포인트의 정확도 향상을 기록했다.
• RC-MAE의 어텐션 맵은 iBOT 및 MSN과 달리 쿼리 페치 근처의 관련 영역에 더 선명하게 집중하는 경향을 보였다.
• 이론적 분석을 통해 교사는 현재와 이전의 특성 표현 간의 유사도에 따라 기울기를 조정하는 조건부 운동량 정규화 기법으로 작용하는 것으로 밝혀졌다.
• 선형 모델과 깊은 ViT 기반 아키텍처에서의 실증적 검증을 통해 조건부 기울기 보정 메커니즘이 확인되었으며, 제안된 분석을 뒷받침한다.