[논문 리뷰] Investigating the Benefits of Projection Head for Representation Learning

Yihao Xue, Eric Gan외 3명|arXiv (Cornell University)|2024. 03. 18.

Educational Games and GamificationPsychology인용 수 3

한 줄 요약

이 논문은 대비 학습에서 투영 헤드가 표현 학습을 향상시키는 이유에 대한 이론적 분석을 제공하며, 깊이 있는 레이어에서의 특징 전문화를 완화하고 분포 이탈 상황에서의 강건성을 향상시킨다고 보여준다. 이는 사전 투영 표현이 암묵적인 계층별 특징 가중치화 덕분에 더 일반화 가능하다는 것을 밝히며, 고정된 재가중 헤드가 투영 헤드 성능을 따라잡을 수 있음을 보여주어 더 해석 가능한 대안을 제공한다.

ABSTRACT

An effective technique for obtaining high-quality representations is adding a projection head on top of the encoder during training, then discarding it and using the pre-projection representations. Despite its proven practical effectiveness, the reason behind the success of this technique is poorly understood. The pre-projection representations are not directly optimized by the loss function, raising the question: what makes them better? In this work, we provide a rigorous theoretical answer to this question. We start by examining linear models trained with self-supervised contrastive loss. We reveal that the implicit bias of training algorithms leads to layer-wise progressive feature weighting, where features become increasingly unequal as we go deeper into the layers. Consequently, lower layers tend to have more normalized and less specialized representations. We theoretically characterize scenarios where such representations are more beneficial, highlighting the intricate interplay between data augmentation and input features. Additionally, we demonstrate that introducing non-linearity into the network allows lower layers to learn features that are completely absent in higher layers. Finally, we show how this mechanism improves the robustness in supervised contrastive learning and supervised learning. We empirically validate our results through various experiments on CIFAR-10/100, UrbanCars and shifted versions of ImageNet. We also introduce a potential alternative to projection head, which offers a more interpretable and controllable design.

연구 동기 및 목표

• 대비 자기지도 학습에서의 투영 헤드의 경험적 성공 배경에 있는 이론적 메커니즘을 이해하기 위해.
• 직접 최적화되지 않은 사전 투영 표현이 왜 더 나은 최종 성능을 달성하는지 조사하기 위해.
• 데이터 증강과 네트워크 비선형성의 특징 분포와 일반화에 미치는 영향을 분석하기 위해.
• 결과를 감독 대비 학습 및 감독 학습으로 확장하여 강건성 향상과 클래스 붕괴 감소를 입증하기 위해.
• 표준 투영 헤드의 더 해석 가능한 대안으로서의 고정 재가중 헤드를 제안하고 검증하기 위해.

제안 방법

• 대비 손실로 훈련된 선형 모델의 이론적 분석을 통해 최적화 편향으로 인한 계층별 점진적 특징 가중치화를 밝혀내었다.
• 사전 투영 표현이 더 유익한 조건을 유도하는 조건을 유도하였다. 특히 데이터 증강이 유용한 특징을 왜곡할 경우에 해당한다.
• 비선형 네트워크로의 분석 확장을 통해, 하위 레이어가 상위 레이어나 투영 헤드 출력에 존재하지 않는 특징을 학습함을 보여주었다.
• CIFAR-10/100, UrbanCars, 그리고 이격된 ImageNet 변형에서의 실증적 검증을 통해 사전 및 사후 투영 표현을 비교하였다.
• 학습 가능한 투영 헤드의 대안으로서의 고정 재가중 헤드를 도입하고 성능을 평가하였으며, 유사한 성능을 기록하였다.
• 반-합성 데이터셋(MNIST on CIFAR-10)과 OOD 벤치마크를 활용한 ImageNet에서의 미세조정을 통해 강건성을 테스트하였다.

Figure 1: Weights of features in a two-layer fully connected ReLU network trained with CL. Left : With all features having equal strength, those that are more disrupted by augmentation have smaller/zero weights. Right : With augmentation treating all features equally. features with the largest/small

실험 결과

연구 질문

• RQ1 직접 최적화되지 않은 투영 헤드 이전의 표현이 왜 이후 표현보다 더 잘 일반화되는가?
• RQ2 최적화의 암묵적 편향이 대비 학습에서 계층 간 점진적 특징 가중치화를 어떻게 이끌어내는가?
• RQ3 데이터 증강이 유용한 특징을 손상시키거나 특징이 너무 강하거나 약할 경우, 투영 헤드가 어떤 상황에서 강건성을 향상시키는가?
• RQ4 하위 레이어가 투영 헤드 출력에 존재하지 않는 특징을 학습할 수 있으며, 이는 일반화에 어떤 영향을 미치는가?
• RQ5 고정 재가중 헤드가 학습 가능한 투영 헤드와 유사한 성능을 달성할 수 있으며, 더 해석 가능한 대안이 될 수 있는가?

주요 결과

• 데이터 증강이 유용한 특징을 왜곡하거나 훈련 중 특징이 너무 강하거나 약할 경우, 사전 투영 표현이 사후 투영 표현보다 더 강건하고 일반화 능력이 뛰어나다.
• 최적화의 암묵적 편향은 계층별 점진적 특징 가중치화를 유도하며, 이는 깊이 있는 레이어가 사전 훈련 목표에 더 전문화되게 하고, 하위 레이어를 더 정규화되고 덜 전문화된 상태로 만든다.
• 비선형 활성화는 특징 전문화를 강화하여 하위 레이어가 투영 헤드 출력에 존재하지 않는 특징을 학습할 수 있게 하며, 이는 일반화 능력 향상에 기여한다.
• 이격된 ImageNet 벤치마크에서 사전 투영 표현은 사후 투영 표현 및 원본 모델보다 더 높은 OOD 정확도를 달성하여 배경 편향에 대한 강건성이 향상됨을 보여준다.
• 고정 재가중 헤드는 표준 투영 헤드와 유사한 성능을 기록하며, 이론적 통찰인 특징 재가중이 핵심 메커니즘이라는 것을 검증하고 더 해석 가능한 설계 대안을 제안한다.
• 감독 대비 학습 및 감독 학습에서 사전 투영 표현은 클래스 붕괴를 감소시키고 분포 이탈에 대한 소수의 샘플 적응 능력을 향상시키며, 특히 UrbanCars 및 ImageNet에서 두드러진다.

Figure 2: Left: Weights of features in a two-layer fully connected ReLU network trained using SCL. The subclass feature is not represented post-projection but is represented pre-projection. Right: As a result, the four subclasses are only separable in pre-projection representations.