Towards the Sparseness of Projection Head in Self-Supervised Learning

TL;DR

This paper investigates the internal mechanisms of the projection head in self-supervised contrastive learning, revealing that a sparse projection head can improve generalization and representation quality, and introduces SparseHead regularization.

Contribution

It provides theoretical and empirical evidence that a sparse projection head enhances SSL performance and introduces SparseHead, a regularization method to enforce sparsity in the projection head.

Findings

Sparse projection head improves contrastive learning performance.

SparseHead regularization effectively constrains the projection head.

Theoretical analysis links sparsity with better generalization.

Abstract

In recent years, self-supervised learning (SSL) has emerged as a promising approach for extracting valuable representations from unlabeled data. One successful SSL method is contrastive learning, which aims to bring positive examples closer while pushing negative examples apart. Many current contrastive learning approaches utilize a parameterized projection head. Through a combination of empirical analysis and theoretical investigation, we provide insights into the internal mechanisms of the projection head and its relationship with the phenomenon of dimensional collapse. Our findings demonstrate that the projection head enhances the quality of representations by performing contrastive loss in a projected subspace. Therefore, we propose an assumption that only a subset of features is necessary when minimizing the contrastive loss of a mini-batch of data. Theoretical analysis further suggests that a sparse projection head can enhance generalization, leading us to introduce SparseHead - a regularization term that effectively constrains the sparsity of the projection head, and can be seamlessly integrated with any self-supervised learning (SSL) approaches. Our experimental results validate the effectiveness of SparseHead, demonstrating its ability to improve the performance of existing contrastive methods.