Projection Head

A projection head is typically implemented as a multi-layer perceptron (MLP) with one to three hidden layers. A common configuration is a 2-layer MLP with a hidden dimension matching the encoder output, a ReLU activation, and a final linear layer that projects to a lower-dimensional space (e.g., 128 or 256 dimensions).

• Input: Encoder representations (e.g., 2048-d from ResNet-50)
• Hidden layer: Linear → BatchNorm → ReLU
• Output layer: Linear → L2-normalized embedding
• Depth matters: SimCLR demonstrated that a 2-layer head significantly outperforms a linear head, but deeper heads (3+ layers) yield diminishing returns

The projection head acts as a bottleneck that absorbs loss-specific invariances, preventing the backbone encoder from discarding useful information. By applying the contrastive loss in a separate, lower-dimensional space, the head shields the encoder from learning shortcut features that only optimize the pretext task.

• The head learns to be invariant to augmentations (color jitter, rotation, cropping)
• The backbone is free to retain semantically meaningful features useful for downstream tasks
• Removing the head before evaluation consistently yields superior transfer performance
• This phenomenon is known as the information bottleneck effect in self-supervised learning

A defining characteristic of the projection head is that it is discarded after pre-training. During the linear evaluation protocol, the frozen backbone encoder is used directly to extract features, and a new linear classifier is trained on top.

• Pre-training: Encoder → Projection Head → Contrastive Loss
• Evaluation: Encoder → Linear Classifier → Task Labels
• The head's representations are loss-specific and not transferable
• This pattern is consistent across SimCLR, MoCo, BYOL, and SwAV
• Fine-tuning may optionally retain a new task-specific head

The final output of the projection head is typically L2-normalized to project embeddings onto the unit hypersphere. This normalization is critical for contrastive loss functions like NT-Xent (InfoNCE), which compute cosine similarity between normalized vectors.

• L2 normalization: ||z||₂ = 1, constraining embeddings to the hypersphere
• Enables efficient cosine similarity computation via dot product
• Temperature parameter τ scales the similarity logits
• Batch normalization in hidden layers stabilizes training and prevents divergence
• The hyperspherical space encourages uniform distribution of representations

Empirical studies show that the projection head depth and width directly influence the quality of backbone representations. A non-linear projection head consistently outperforms a linear one, as measured by downstream classification accuracy.

• Linear head: 85.8% top-1 accuracy on ImageNet (SimCLR, ResNet-50)
• 2-layer MLP head: 91.6% top-1 accuracy (same setup)
• The head filters out augmentation-specific information
• Representations before the head retain more spatial and semantic structure
• This property is exploited in dense prediction tasks like segmentation

Different self-supervised methods employ projection heads with distinct architectures tailored to their objectives:

• SimCLR: 2-layer MLP with ReLU, outputs 128-d normalized embeddings
• MoCo v2: Identical MLP head on both query and momentum encoders
• BYOL: Asymmetric design — online network has a predictor (additional MLP) on top of the projector; target network has only the projector
• Barlow Twins: Projector outputs high-dimensional embeddings (8192-d) for cross-correlation computation
• VICReg: Separate expander head maps to high dimensions for variance and covariance regularization