# Investigating the Benefits of Projection Head for Representation Learning

arXiv: https://arxiv.org/abs/2403.11391 · HTML: https://mlsanity.com/paper/2403.11391/

## Key facts

- Cited by 20 papers in the map (data through 14 Oct 2026).
- Builds on 56 of its 67 references that are in the map.
- First posted on arXiv 18 Mar 2024.

## Builds on (most cited first, 30 of 56)

- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Supervised Contrastive Learning](https://mlsanity.com/paper/2004.11362/) (2020)
- [The Many Faces of Robustness: A Critical Analysis of Out-of-Distribution Generalization](https://mlsanity.com/paper/2006.16241/) (2020)
- [Barlow Twins: Self-Supervised Learning via Redundancy Reduction](https://mlsanity.com/paper/2103.03230/) (2021)
- [Do ImageNet Classifiers Generalize to ImageNet?](https://mlsanity.com/paper/1902.10811/) (2019)
- [ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness](https://mlsanity.com/paper/1811.12231/) (2018)
- [Natural Adversarial Examples](https://mlsanity.com/paper/1907.07174/) (2019)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Adversarial Examples Are Not Bugs, They Are Features](https://mlsanity.com/paper/1905.02175/) (2019)
- [Learning Representations by Maximizing Mutual Information Across Views](https://mlsanity.com/paper/1906.00910/) (2019)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [Measuring Robustness to Natural Distribution Shifts in Image Classification](https://mlsanity.com/paper/2007.00644/) (2020)
- [Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations](https://mlsanity.com/paper/2204.02937/) (2022)
- [Understanding Dimensional Collapse in Contrastive Self-supervised Learning](https://mlsanity.com/paper/2110.09348/) (2021)
- [Using Pre-Training Can Improve Model Robustness and Uncertainty](https://mlsanity.com/paper/1901.09960/) (2019)
- [Noise or Signal: The Role of Image Backgrounds in Object Recognition](https://mlsanity.com/paper/2006.09994/) (2020)
- [In Search of the Real Inductive Bias: On the Role of Implicit Regularization in Deep Learning](https://mlsanity.com/paper/1412.6614/) (2014)
- [Implicit Regularization in Matrix Factorization](https://mlsanity.com/paper/1705.09280/) (2017)
- [Provable Guarantees for Self-Supervised Deep Learning with Spectral Contrastive Loss](https://mlsanity.com/paper/2106.04156/) (2021)
- [An Investigation of Why Overparameterization Exacerbates Spurious Correlations](https://mlsanity.com/paper/2005.04345/) (2020)
- [On the Optimization of Deep Networks: Implicit Acceleration by Overparameterization](https://mlsanity.com/paper/1802.06509/) (2018)
- [Whitening for Self-Supervised Representation Learning](https://mlsanity.com/paper/2007.06346/) (2020)
- [Understanding self-supervised Learning Dynamics without Contrastive Pairs](https://mlsanity.com/paper/2102.06810/) (2021)
- [Revisiting Self-Supervised Visual Representation Learning](https://mlsanity.com/paper/1901.09005/) (2019)
- [Unsupervised Embedding Learning via Invariant and Spreading Instance Feature](https://mlsanity.com/paper/1904.03436/) (2019)

## Built on by (most cited first, 20 of 20)

- [Explaining and Mitigating the Modality Gap in Contrastive Multimodal Learning](https://mlsanity.com/paper/2412.07909/) (2024)
- [INTRA: Interaction Relationship-aware Weakly Supervised Affordance Grounding](https://mlsanity.com/paper/2409.06210/) (2024)
- [Projection Head is Secretly an Information Bottleneck](https://mlsanity.com/paper/2503.00507/) (2025)
- [Position: An Empirically Grounded Identifiability Theory Will Accelerate Self-Supervised Learning Research](https://mlsanity.com/paper/2504.13101/) (2025)
- [Self-Supervised Contrastive Learning is Approximately Supervised Contrastive Learning](https://mlsanity.com/paper/2506.04411/) (2025)
- [DUNIA: Pixel-Sized Embeddings via Cross-Modal Alignment for Earth Observation Applications](https://mlsanity.com/paper/2502.17066/) (2025)
- [Directional Neural Collapse Explains Few-Shot Transfer in Self-Supervised Learning](https://mlsanity.com/paper/2603.03530/) (2026)
- [DualFed: Enjoying both Generalization and Personalization in Federated Learning via Hierachical Representations](https://mlsanity.com/paper/2407.17754/) (2024)
- [On the Similarities of Embeddings in Contrastive Learning](https://mlsanity.com/paper/2506.09781/) (2025)
- [On the Alignment Between Supervised and Self-Supervised Contrastive Learning](https://mlsanity.com/paper/2510.08852/) (2025)
- [scSSL-Bench: Benchmarking Self-Supervised Learning for Single-Cell Data](https://mlsanity.com/paper/2506.10031/) (2025)
- [The Geometry of Projection Heads: Conditioning, Invariance, and Collapse](https://mlsanity.com/paper/2605.17180/) (2026)
- [ProFeAT: Projected Feature Adversarial Training for Self-Supervised Learning of Robust Representations](https://mlsanity.com/paper/2406.05796/) (2024)
- [An Empirical Study on Improving SimCLR's Nonlinear Projection Head using Pretrained Autoencoder Embeddings](https://mlsanity.com/paper/2408.14514/) (2024)
- [Representation learning with a transformer by contrastive learning for money laundering detection](https://mlsanity.com/paper/2507.08835/) (2025)
- [Soft Separation and Distillation: Toward Global Uniformity in Federated Unsupervised Learning](https://mlsanity.com/paper/2508.01251/) (2025)
- [ANCHOR: Integrating Adversarial Training with Hard-mined Supervised Contrastive Learning for Robust Representation Learning](https://mlsanity.com/paper/2510.27599/) (2025)
- [Beyond Routing: Characterising Expert Tuning and Representation in Vision Mixture-of-Experts](https://mlsanity.com/paper/2605.20610/) (2026)
- [Radial Basis Function Networks as Projection Heads in Self-Supervised Learning](https://mlsanity.com/paper/2606.21590/) (2026)
- [On the Role of the Projector in Contrastive Self-Supervised Learning: Last-Layer Rank Dynamics Drive Representation Quality](https://mlsanity.com/paper/2609.26334/) (2026)

## Abstract (arXiv)

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.