[논문 리뷰] The Mechanism of Prediction Head in Non-contrastive Self-supervised Learning

Zixin Wen, Yuanzhi Li|arXiv (Cornell University)|2022. 05. 12.

Stochastic Gradient Optimization Techniques인용 수 6

한 줄 요약

이 논문은 비대비(self-supervised) 학습에서 학습 가능한 예측 헤드가 차원 붕괴를 방지하는 메커니즘을 규명한다. 경험적 및 이론적 분석을 통해 예측 헤드가 더 강한 특징이 더 약한 특징을 대체할 수 있도록 하고, 이로 인해 학습 속도가 빨라지는 두 가지 핵심 효과—대체 및 가속—를 밝혀낸다. 이러한 효과 덕분에 손실 함수의 잠재적 최적해로 존재하는 자명한 붕괴 해를 피할 수 있다.

ABSTRACT

Recently the surprising discovery of the Bootstrap Your Own Latent (BYOL) method by Grill et al. shows the negative term in contrastive loss can be removed if we add the so-called prediction head to the network. This initiated the research of non-contrastive self-supervised learning. It is mysterious why even when there exist trivial collapsed global optimal solutions, neural networks trained by (stochastic) gradient descent can still learn competitive representations. This phenomenon is a typical example of implicit bias in deep learning and remains little understood. In this work, we present our empirical and theoretical discoveries on non-contrastive self-supervised learning. Empirically, we find that when the prediction head is initialized as an identity matrix with only its off-diagonal entries being trainable, the network can learn competitive representations even though the trivial optima still exist in the training objective. Theoretically, we present a framework to understand the behavior of the trainable, but identity-initialized prediction head. Under a simple setting, we characterized the substitution effect and acceleration effect of the prediction head. The substitution effect happens when learning the stronger features in some neurons can substitute for learning these features in other neurons through updating the prediction head. And the acceleration effect happens when the substituted features can accelerate the learning of other weaker features to prevent them from being ignored. These two effects enable the neural networks to learn all the features rather than focus only on learning the stronger features, which is likely the cause of the dimensional collapse phenomenon. To the best of our knowledge, this is also the first end-to-end optimization guarantee for non-contrastive methods using nonlinear neural networks with a trainable prediction head and normalization.

연구 동기 및 목표

• 스토케스틱 그래디언트 디센트(SGD)로 훈련된 신경망이 비대비 자기학습에서 자명한 붕괴 해가 전역 최적해임에도 불구하고 이를 피하는 이유를 이해하는 것.
• 예측 헤드가 대칭성을 깨고 다양한 비붕괴 표현을 학습하는 데 기여하는 역할을 조사하는 것.
• 비대비 대비 유사 학습에서 예측 헤드가 유도하는 암묵적 편향을 설명하는 이론적 프레임워크를 제공하는 것.
• 특히 비대칭 항목이 포함된 항등행렬로 초기화된 예측 헤드가 존재할 때의 특징 학습 동역학을 특성화하는 것, 특히 비대칭 항목의 상승과 하강 과정.

제안 방법

• 예측 헤드를 항등행렬로 초기화하고 비대칭 항목만 갱신하는 방식으로 훈련 동역학을 경험적으로 연구한다.
• 학습 과정을 네 단계로 모델링: 더 강한 특징 학습, 대체 효과, 가속 효과, 수렴.
• 대체 효과를 정의하고 분석한다. 이는 한 뉴런의 학습이 다른 뉴런의 학습을 예측 헤드 갱신을 통해 대체할 수 있음을 의미한다.
• 가속 효과를 정의하고 분석한다. 이는 대체된 특징이 더 약한 특징의 학습을 가속화하여 무시당하는 것을 방지함을 의미한다.
• 비선형 신경망에 대해 학습 가능한 예측 헤드와 배치 정규화를 갖는 최적화 보장을 유도하기 위해 단순화된 선형 모델을 사용한다.
• 텐서 거듭제곱 방법의 경계를 활용하여 다양한 학습률과 초기화 조건 하에서 기울기 갱신과 수렴성을 분석한다.

실험 결과

연구 질문

• RQ1 스토케스틱 그래디언트 디센트(SGD)로 훈련된 신경망이 비대비 자기학습에서 붕괴 해가 전역 최적해임에도 불구하고 이를 피하는 이유는 무엇인가?
• RQ2 예측 헤드가 신경망이 항등 매핑을 학습하여 무력화될 수 있음에도 불구하고 다양한 특징을 학습하는 데 어떻게 기여하는가?
• RQ3 예측 헤드의 비대칭 항목은 훈련 중에 어떤 역할을 하는가? 왜 이 항목들은 일관된 상승과 하강 경로를 따르는가?
• RQ4 예측 헤드의 대체 및 가속 효과가 함께 작용하여 특징 무시를 방지하고 전체 표현 학습을 가능하게 하는가?
• RQ5 학습 가능한 예측 헤드와 정규화를 갖는 비대비 자기학습에 대해 종단 간 최적화 보장을 제공할 수 있는가?

주요 결과

• 예측 헤드를 항등행렬로 초기화하고 비대칭 항목만 학습하는 경우에도 네트워크는 경쟁 가능한 표현을 학습한다. 이는 헤드의 구조—단지 매개변수만이 아니라—가 중요하다는 것을 보여준다.
• 훈련 과정에서 예측 헤드의 비대칭 항목에 일관된 상승과 하강 경로가 관찰되며, 이는 동적 특징 대체를 시사한다.
• 대체 효과는 한 뉴런의 더 강한 특징이 예측 헤드 갱신을 통해 다른 뉴런의 학습을 대체할 수 있을 때 발생한다.
• 가속 효과는 대체된 특징이 더 약한 특징의 학습을 가속화하여 무시당하는 것을 방지함으로써 발생한다.
• 대체 및 가속 효과가 함께 작용함으로써 네트워크는 가장 강한 특징만이 아니라 모든 특징을 학습할 수 있으며, 이는 차원 붕괴를 피하는 이유를 설명한다.
• 본 논문은 비대비 자기학습에서 학습 가능한 예측 헤드와 정규화를 갖는 비선형 신경망에 대해 종단 간 최적화 보장을 제공하는 최초의 논문이다.