Unraveling Projection Heads in Contrastive Learning: Insights from Expansion and Shrinkage

TL;DR

This paper investigates how contrastive learning projectors influence representation quality, revealing expansion and shrinkage effects that impact downstream classification performance, supported by theoretical analysis and extensive experiments.

Contribution

It introduces the concepts of expansion and shrinkage effects caused by contrastive loss and models projector behavior with linear transformations to explain downstream accuracy.

Findings

Contrastive loss induces expansion or shrinkage in representations.

Linear projectors in shrinkage regime hinder downstream accuracy.

Theoretical analysis matches experimental results on synthetic and real data.

Abstract

We investigate the role of projection heads, also known as projectors, within the encoder-projector framework (e.g., SimCLR) used in contrastive learning. We aim to demystify the observed phenomenon where representations learned before projectors outperform those learned after -- measured using the downstream linear classification accuracy, even when the projectors themselves are linear. In this paper, we make two significant contributions towards this aim. Firstly, through empirical and theoretical analysis, we identify two crucial effects -- expansion and shrinkage -- induced by the contrastive loss on the projectors. In essence, contrastive loss either expands or shrinks the signal direction in the representations learned by an encoder, depending on factors such as the augmentation strength, the temperature used in contrastive loss, etc. Secondly, drawing inspiration from the expansion and shrinkage phenomenon, we propose a family of linear transformations to accurately model the projector's behavior. This enables us to precisely characterize the downstream linear classification accuracy in the high-dimensional asymptotic limit. Our findings reveal that linear projectors operating in the shrinkage (or expansion) regime hinder (or improve) the downstream classification accuracy. This provides the first theoretical explanation as to why (linear) projectors impact the downstream performance of learned representations. Our theoretical findings are further corroborated by extensive experiments on both synthetic data and real image data.