[논문 리뷰] Momentum^2 Teacher: Momentum Teacher with Momentum Statistics for Self-Supervised Learning

Zeming Li, Songtao Liu외 1명|arXiv (Cornell University)|2021. 01. 19.

Music and Audio Processing참고 문헌 53인용 수 11

한 줄 요약

이 논문은 배치 정규화(BN) 통계에 모멘텀 업데이트를 적용함으로써 교사 네트워크의 안정성을 향상시키는 자기지도 학습 방법인 Momentum² Teacher를 제안한다. 이는 학생 네트워크 가중치뿐 아니라 BN 통계에 대해서도 모멘텀 업데이트를 적용함으로써 이루어지며, '모멘텀 BN'—이전 BN 통계의 모멘텀 기반 업데이트—를 도입함으로써, 소규모 배치(128) 학습만으로도 TPU나 동기화 BN을 요구하지 않고 ImageNet에서 선구적 74.5%의 top-1 정확도를 달성한다.

ABSTRACT

In this paper, we present a novel approach, Momentum$^2$ Teacher, for student-teacher based self-supervised learning. The approach performs momentum update on both network weights and batch normalization (BN) statistics. The teacher's weight is a momentum update of the student, and the teacher's BN statistics is a momentum update of those in history. The Momentum$^2$ Teacher is simple and efficient. It can achieve the state of the art results (74.5\%) under ImageNet linear evaluation protocol using small-batch size(\eg, 128), without requiring large-batch training on special hardware like TPU or inefficient across GPU operation (\eg, shuffling BN, synced BN). Our implementation and pre-trained models will be given on GitHub\footnote{https://github.com/zengarden/momentum2-teacher}.

연구 동기 및 목표

• 소규모 배치 자기지도 학습에서 배치 정규화의 불안정성을 해결하여 학생-교사 프레임워크에서의 성능을 제한하는 문제를 해결한다.
• BYOL과 같은 선구적 방법들이 대규모 배치 학습이나 전용 하드웨어(예: TPU)에 의존하는 문제를 해결하고, 비효율적인 동기화 BN 연산을 피한다.
• 학습 네트워크 가중치 외에도 배치 정규화 통계에 모멘텀 업데이트를 적용함으로써 교사 네트워크의 안정성을 향상시킨다.
• 일반적인 GPU 하드웨어에서 표준적이고 효율적인 소규모 배치 학습을 통해 고성능 자기지도 표현 학습을 가능하게 한다.

제안 방법

• 학습 단계 동안 현재 미니배치 통계에 의존하는 대신, 배치 통계(평균 및 분산)의 누적 평균을 유지하는 새로운 연산인 '모멘텀 BN'을 제안한다.
• 기본적인 학생-교사 프레임워크에서와 같이 학생 네트워크 가중치와 교사 네트워크 가중치에 모두 모멘텀 업데이트를 적용한다.
• 교사의 배치 정규화 통계에도 모멘텀 업데이트를 확장하여 이전 배치 통계의 이동 평균을 사용함으로써 정규화의 안정성을 향상시킨다.
• 교사 네트워크에 역전파를 수행하지 않아, 배치 정규화 통계의 역사가 비가역적이기 때문에 발생하는 기울기 문제를 피한다.
• 효율성을 위해 학생 네트워크는 표준 소규모 배치 BN을 사용하고, 교사 네트워크는 안정적이고 정확한 특징 정규화를 위해 모멘텀 BN을 활용한다.
• 학습은 대조 또는 부트스트랩 학습 목표함수를 통해 수행되며, 교사가 타겟 표현을 제공한다.

실험 결과

연구 질문

• RQ1 모멘텀 기반의 배치 정규화 통계 업데이트가 학생-교사 자기지도 학습의 안정성과 성능을 향상시킬 수 있는가?
• RQ2 대규모 배치나 TPU 기반 학습 없이도 소규모 배치 학습(예: 128)이 자기지도 표현 학습에서 선구적 성능을 달성할 수 있는가?
• RQ3 정확도와 훈련 효율성 측면에서 표준 BN, 셔플링 BN, 또는 동기화 BN에 비해 모멘텀 BN이 우월한가?
• RQ4 제안된 Momentum² Teacher 프레임워크가 다양한 최종 작업 및 데이터셋에 일반화 가능한가?
• RQ5 BYOL 및 MoCo와 같은 선구적 자기지도 모델과 비교했을 때, 선형 평가 및 파인튜닝 벤치마크에서 성능이 어떻게 비교되는가?

주요 결과

• Momentum² Teacher 방법은 소규모 배치 크기 128만으로도 선형 평가 프로토콜 하에서 ImageNet에서 선구적 74.5%의 top-1 정확도를 달성한다.
• 이 방법은 대규모 배치 학습이나 TPU 하드웨어에 의존하지 않음에도 불구하고, 셔플링 BN을 사용하는 MoCo를 능가하고 BYOL과 유사한 성능을 기록한다.
• COCO 객체 검출 및 인스턴스 세그멘테이션에서, 이 방법은 지도 학습 기반 사전 훈련 베이스라인을 뛰어넘고 MoCo를 상당한 격차로 능가한다.
• 장수분포 카테고리가 포함된 LVIS 인스턴스 세그멘테이션에서, 이 방법은 지도 학습 사전 훈련 대비 mAP를 1.7%p 향상시켜 저자원 카테고리로의 일반화 능력이 뛰어나다는 것을 보여준다.
• ImageNet 레이블의 1% 및 10%만을 사용하는 준지도 학습 환경에서, 이 방법은 각각 57.7% 및 62.3%의 top-1 정확도를 기록하여 이전의 모든 자기지도 및 준지도 학습 방법을 뛰어넘는다.
• 이 방법은 검출 및 세그멘테이션 작업 전반에 걸쳐 강력한 전이 능력을 보이며, 특히 장시간 훈련 스케줄에서 RetinaNet 및 FCOS와 같은 밀도 높은 객체 검출기에서 두각을 나타낸다.