[논문 리뷰] DisCo: Remedy Self-supervised Learning on Lightweight Models with Distilled Contrastive Learning

Yuting Gao, Jiaxin Zhuang외 5명|arXiv (Cornell University)|2021. 04. 19.

Domain Adaptation and Few-Shot Learning참고 문헌 46인용 수 8

한 줄 요약

이 논문은 경량 모델에서 자기지도 학습(SSL)의 성능을 향상시키기 위해 대규모 교사 모델의 최종 대비 임베딩을 소규모 학생 모델로 MSE 손실을 통해 전달하는 지식 정복 프레임워크인 DisCo를 제안한다. 이 방법은 추론 시 추가 파라미터 없이도 선형 평가 정확도를 크게 향상시키며, ResNet-101 교사 모델을 사용할 경우 EfficientNet-B0의 ImageNet top-1 정확도를 22.1% 향상시킨다.

ABSTRACT

While self-supervised representation learning (SSL) has received widespread attention from the community, recent research argue that its performance will suffer a cliff fall when the model size decreases. The current method mainly relies on contrastive learning to train the network and in this work, we propose a simple yet effective Distilled Contrastive Learning (DisCo) to ease the issue by a large margin. Specifically, we find the final embedding obtained by the mainstream SSL methods contains the most fruitful information, and propose to distill the final embedding to maximally transmit a teacher's knowledge to a lightweight model by constraining the last embedding of the student to be consistent with that of the teacher. In addition, in the experiment, we find that there exists a phenomenon termed Distilling BottleNeck and present to enlarge the embedding dimension to alleviate this problem. Our method does not introduce any extra parameter to lightweight models during deployment. Experimental results demonstrate that our method achieves the state-of-the-art on all lightweight models. Particularly, when ResNet-101/ResNet-50 is used as teacher to teach EfficientNet-B0, the linear result of EfficientNet-B0 on ImageNet is very close to ResNet-101/ResNet-50, but the number of parameters of EfficientNet-B0 is only 9.4\%/16.3\% of ResNet-101/ResNet-50. Code is available at https://github. com/Yuting-Gao/DisCo-pytorch.

연구 동기 및 목표

• 모델 크기가 감소함에 따라 정확도가 급격히 떨어지는 경량 모델에서 자기지도 학습(SSL)의 성능 저하 문제를 해결한다.
• 기존 정복 방법이 최적의 지식(예: 소프트맥스 출력 또는 특징 맵)을 전달하는 데 한계가 있음을 초월하여, 가장 정보가 풍부한 최종 임베딩을 전달함으로써 문제를 해결한다.
• 최종 대비 임베딩을 주요 지식 소스로 삼아 경량 모델의 일반화 능력을 향상시킨다.
• 학생 모델의 MLP 헤드에 제한된 은닉 차원으로 인해 발생하는 '정복 병목 현상'—성능의 한계—를 식별하고 해결하기 위해 임베딩 차원을 확장한다.
• 실제 응용에서 효율성을 보장하기 위해 후속 배포 시 추가 파라미터를 도입하지 않는 방법을 설계한다.

제안 방법

• 학생 모델의 최종 대비 임베딩이 교사 모델의 최종 대비 임베딩과 일치하도록 하기 위해 평균 제곱오차(MSE) 손실을 사용하여 대규모 교사 모델에서 경량 학생 모델로 지식을 전달한다.
• 대조 학습에 가장 정보가 풍부한 표현을 담고 있는 교사 모델의 최종 프로젝션 헤드 출력(MLP 이후)을 주요 지식 신호로 사용한다.
• 학생의 최종 임베딩이 사전 훈련 중에 교사의 최종 임베딩과 일치하도록 정규화하는 수정된 훈련 파이프라인을 도입한다.
• 학생의 MLP 헤드의 은닉 차원을 늘림으로써 '정복 병목 현상'을 해결하여 지식 전달 능력을 향상시킨다.
• 사전 훈련 중에 사용된 추가 MLP 헤드는 미세조정 및 추론 전에 제거하여 모델 효율성을 유지한다.
• MoCo-V2, SwAV, DINO 등의 다양한 SSL 프레임워크에 DisCo를 통합하여, 다양한 아키텍처와 학습 철학에 걸쳐 호환성과 일반화 능력을 검증한다.

실험 결과

연구 질문

• RQ1 대규모 교사 모델의 최종 대비 임베딩을 경량 학생 모델로 전달함으로써 자기지도 학습에서 학생 모델의 표현 품질을 크게 향상시킬 수 있는가?
• RQ2 학생의 최종 임베딩을 MSE 손실을 통해 교사의 임베딩과 일치시키는 것이 기존 정복 방법보다 더 나은 최종 성능을 이끌어내는가?
• RQ3 소규모 모델에서의 성능 저하가 지식 전달의 병목 현상 때문인지, 그리고 MLP 헤드 차원을 확장함으로써 이 문제를 완화할 수 있는가?
• RQ4 DisCo는 아키텍처 수정 없이도 다양한 SSL 프레임워크(MoCo-V2, SwAV, DINO 등)에 효과적으로 적용될 수 있는가?
• RQ5 DisCo를 다른 정복 기법(AT, RKD, KD)과 조합하면 추가 성능 향상이 이루어지는가?

주요 결과

• ResNet-101 교사 모델을 사용할 경우 DisCo는 EfficientNet-B0의 ImageNet 선형 평가 top-1 정확도를 22.1% 향상시켜 66.5%의 정확도를 달성한다.
• ResNet-50를 교사로 사용할 경우 DisCo는 EfficientNet-B0의 top-1 정확도를 19.7% 향상시켜 66.6%의 정확도를 달성한다.
• MobileNet-v3와 ResNet-18을 포함한 모든 테스트된 경량 모델에서 SEED 및 MoCo-V2와 같은 최신 기술을 초월한다.
• 학생 모델의 MLP 헤드 차원을 늘릴수록 성능이 단조롭게 향상되며, 2048 차원에서 가장 큰 성능 향상이 관찰된다.
• DisCo를 다른 정복 기법(AT, RKD, KD)과 조합하면 추가 성능 향상이 이루어지며, 이는 그 호환성과 효과성을 확인한다.
• MoCo-V2, SwAV, DINO를 포함한 다양한 SSL 프레임워크에서 DisCo가 강력한 성능을 내며, 일반화 능력과 호환성을 입증한다.