[논문 리뷰] Understanding and Improving the Role of Projection Head in Self-Supervised Learning

Kartik Gupta, Thalaiyasingam Ajanthan외 2명|arXiv (Cornell University)|2022. 12. 22.

Domain Adaptation and Few-Shot Learning인용 수 12

한 줄 요약

이 논문은 학습 가능한 프로젝션 헤드가 훈련 후 폐기되더라도 대비 자기학습(SSL)에서 필수적임을 규명한다. 이는 프로젝션 헤드를 네트워크의 일부로 보는 대신 손실 함수의 일부로 간주하는 이중 최적화 문제로 SSL를 재정의함으로써, CIFAR-10, CIFAR-100, TinyImageNet에서 표준 SimCLR에 비해 일반화 성능을 향상시킨다.

ABSTRACT

Self-supervised learning (SSL) aims to produce useful feature representations without access to any human-labeled data annotations. Due to the success of recent SSL methods based on contrastive learning, such as SimCLR, this problem has gained popularity. Most current contrastive learning approaches append a parametrized projection head to the end of some backbone network to optimize the InfoNCE objective and then discard the learned projection head after training. This raises a fundamental question: Why is a learnable projection head required if we are to discard it after training? In this work, we first perform a systematic study on the behavior of SSL training focusing on the role of the projection head layers. By formulating the projection head as a parametric component for the InfoNCE objective rather than a part of the network, we present an alternative optimization scheme for training contrastive learning based SSL frameworks. Our experimental study on multiple image classification datasets demonstrates the effectiveness of the proposed approach over alternatives in the SSL literature.

연구 동기 및 목표

• 학습 후 폐기되더라도 대비 자기학습(SSL)에서 학습 가능한 프로젝션 헤드가 필요한 이유를 이해하는 것.
• 대비 손실 최적화를 위한 부분공간 선택에서 프로젝션 헤드의 암묵적 역할을 조사하는 것.
• 학습 가능한 헤드를 넘어서 기능적 기여도를 이해하는 데의 격차를 메우는 것.
• 프로젝션 헤드를 손실 함수의 일부로 간주하는 새로운 최적화 기법을 제안하여 훈련 안정성과 일반화 성능을 향상시키는 것.
• 학습 가능한 프로젝션 헤드가 고정 또는 폐기된 헤드보다 성능 향상에 기여함을 경험적으로 검증하는 것.

제안 방법

• SSL을 이중 최적화 문제로 재정의: 내부 루프는 대비 손실을 위한 최적의 부분공간을 선택하기 위해 프로젝션 헤드를 최적화하고, 외부 루프는 백본 특징을 업데이트한다.
• 프로젝션 헤드를 고정된 네트워크 헤드가 아니라 손실 함수의 학습 가능한 구성요소로 재정의하여 동적 부분공간 선택을 가능하게 한다.
• 교차 최적화를 적용: 대비 손실에 대해 기울기 하강법으로 프로젝션 헤드를 업데이트한 후, 업데이트된 프로젝션 헤드를 사용해 백본을 업데이트한다.
• CIFAR-10, STL-10, TinyImageNet, ImageNet에서 SimCLR 및 SimSiam에 이 방법을 적용하여 행동과 성능을 분석한다.
• 제안된 방법을 표준 SimCLR 및 DirectCLR와 비교하여 선형 평가 프로토콜에서 일반화 성능을 평가한다.
• 프로젝션 헤드의 질량과 영공간을 분석하여 특징 선택에서의 암묵적 역할을 이해한다.

실험 결과

연구 질문

• RQ1 학습 후 폐기되더라도 SSL에서 학습 가능한 프로젝션 헤드가 필요한 이유는 무엇인가?
• RQ2 학습 중에 최종 투영 레이어를 넘어서 프로젝션 헤드가 수행하는 기능적 역할은 무엇인가?
• RQ3 프로젝션 헤드의 행동—특히 특징 부분공간을 암묵적으로 선택하는 능력—을 손실 함수의 일부로 형식화하고 강제화할 수 있는가?
• RQ4 프로젝션 헤드를 손실 함수의 일부로 간주하는 것이 표준 학습 방식의 학습 가능한 헤드보다 더 나은 일반화를 이끌어내는가?
• RQ5 다양한 데이터셋에서 제안된 방법의 성능은 고정 프로젝션 헤드 또는 표준 SimCLR와 비교해 어떻게 되는가?

주요 결과

• 프로젝션 헤드는 대비 손실을 적용하기 위해 특징의 부분공간을 암묵적으로 선택하는 저랭크 매핑으로 작용하여 훈련 안정성과 일반화 성능을 향상시킨다.
• 프로젝션 헤드의 영공간은 일반화에 유용하며, 이는 헤드가 관련 없는 특징 성분을 억제하도록 학습한다는 것을 시사한다.
• 표준 SSL 훈련에서 학습 가능한 프로젝션 헤드는 고정된 프로젝션 헤드로 대체될 수 없으며, 이는 성능 저하를 초래한다.
• 제안된 이중 최적화 기법은 CIFAR-10, CIFAR-100, TinyImageNet에서 표준 SimCLR보다 뛰어난 성능을 보이며 일반화 성능 향상을 입증한다.
• 초기 훈련 단계에서는 고정 프로젝션 헤드가 더 나은 성능을 보이지만, 학습 가능한 프로젝션 헤드가 결국 이를 능가하여 종단 간 최적화의 필요성을 입증한다.
• 선형 평가 벤치마크에서 최신 기술 수준의 성능을 달성하여, 프로젝션 헤드의 역할이 우연이 아니라 학습 과정에 핵심적인 요소임을 확인한다.