# The Mechanism of Prediction Head in Non-contrastive Self-supervised Learning

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

## Key facts

- Cited by 38 papers in the map (data through 14 Oct 2026).
- Builds on 76 of its 91 references that are in the map.
- First posted on arXiv 12 May 2022.

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

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [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)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Exploring Simple Siamese Representation Learning](https://mlsanity.com/paper/2011.10566/) (2020)
- [Unsupervised Learning of Visual Features by Contrasting Cluster Assignments](https://mlsanity.com/paper/2006.09882/) (2020)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [BEiT: BERT Pre-Training of Image Transformers](https://mlsanity.com/paper/2106.08254/) (2021)
- [Improved Baselines with Momentum Contrastive Learning](https://mlsanity.com/paper/2003.04297/) (2020)
- [SimCSE: Simple Contrastive Learning of Sentence Embeddings](https://mlsanity.com/paper/2104.08821/) (2021)
- [Barlow Twins: Self-Supervised Learning via Redundancy Reduction](https://mlsanity.com/paper/2103.03230/) (2021)
- [An Empirical Study of Training Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.02057/) (2021)
- [Contrastive Multiview Coding](https://mlsanity.com/paper/1906.05849/) (2019)
- [Learning deep representations by mutual information estimation and maximization](https://mlsanity.com/paper/1808.06670/) (2018)
- [VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning](https://mlsanity.com/paper/2105.04906/) (2021)
- [What Makes for Good Views for Contrastive Learning?](https://mlsanity.com/paper/2005.10243/) (2020)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks](https://mlsanity.com/paper/1901.08584/) (2019)
- [A Theoretical Analysis of Contrastive Unsupervised Representation Learning](https://mlsanity.com/paper/1902.09229/) (2019)
- [Implicit Regularization in Deep Matrix Factorization](https://mlsanity.com/paper/1905.13655/) (2019)
- [Learning and Generalization in Overparameterized Neural Networks, Going Beyond Two Layers](https://mlsanity.com/paper/1811.04918/) (2018)
- [Gradient Descent Maximizes the Margin of Homogeneous Neural Networks](https://mlsanity.com/paper/1906.05890/) (2019)

## Built on by (most cited first, 30 of 38)

- [BYOL-Explore: Exploration by Bootstrapped Prediction](https://mlsanity.com/paper/2206.08332/) (2022)
- [Understanding Self-Predictive Learning for Reinforcement Learning](https://mlsanity.com/paper/2212.03319/) (2022)
- [A Theoretical Study of Inductive Biases in Contrastive Learning](https://mlsanity.com/paper/2211.14699/) (2022)
- [On the Power of Foundation Models](https://mlsanity.com/paper/2211.16327/) (2022)
- [InfoNCE Loss Provably Learns Cluster-Preserving Representations](https://mlsanity.com/paper/2302.07920/) (2023)
- [Joint Embedding vs Reconstruction: Provable Benefits of Latent Space Prediction for Self Supervised Learning](https://mlsanity.com/paper/2505.12477/) (2025)
- [Towards a Unified Theoretical Understanding of Non-contrastive Learning via Rank Differential Mechanism](https://mlsanity.com/paper/2303.02387/) (2023)
- [Matrix Information Theory for Self-Supervised Learning](https://mlsanity.com/paper/2305.17326/) (2023)
- [Understanding Augmentation-based Self-Supervised Representation Learning via RKHS Approximation and Regression](https://mlsanity.com/paper/2306.00788/) (2023)
- [Investigating the Benefits of Projection Head for Representation Learning](https://mlsanity.com/paper/2403.11391/) (2024)
- [Information Flow in Self-Supervised Learning](https://mlsanity.com/paper/2309.17281/) (2023)
- [What Improves the Generalization of Graph Transformers? A Theoretical Dive into the Self-attention and Positional Encoding](https://mlsanity.com/paper/2406.01977/) (2024)
- [Identifiable Causal Representation Learning: Unsupervised, Multi-View, and Multi-Environment](https://mlsanity.com/paper/2406.13371/) (2024)
- [Link Prediction with Non-Contrastive Learning](https://mlsanity.com/paper/2211.14394/) (2022)
- [More Synergy, Less Redundancy: Exploiting Joint Mutual Information for Self-Supervised Learning](https://mlsanity.com/paper/2307.00651/) (2023)
- [Projection Head is Secretly an Information Bottleneck](https://mlsanity.com/paper/2503.00507/) (2025)
- [iQRL -- Implicitly Quantized Representations for Sample-efficient Reinforcement Learning](https://mlsanity.com/paper/2406.02696/) (2024)
- [Rethinking Invariance Regularization in Adversarial Training to Improve Robustness-Accuracy Trade-off](https://mlsanity.com/paper/2402.14648/) (2024)
- [Collapse-Proof Non-Contrastive Self-Supervised Learning](https://mlsanity.com/paper/2410.04959/) (2024)
- [GEDI: GEnerative and DIscriminative Training for Self-Supervised Learning](https://mlsanity.com/paper/2212.13425/) (2022)
- [Unraveling Projection Heads in Contrastive Learning: Insights from Expansion and Shrinkage](https://mlsanity.com/paper/2306.03335/) (2023)
- [Understanding Self-Supervised Learning via Gaussian Mixture Models](https://mlsanity.com/paper/2411.03517/) (2024)
- [Deciphering the Projection Head: Representation Evaluation Self-supervised Learning](https://mlsanity.com/paper/2301.12189/) (2023)
- [Towards Understanding Why FixMatch Generalizes Better Than Supervised Learning](https://mlsanity.com/paper/2410.11206/) (2024)
- [Towards the Sparseness of Projection Head in Self-Supervised Learning](https://mlsanity.com/paper/2307.08913/) (2023)
- [Dual Perspectives on Non-Contrastive Self-Supervised Learning](https://mlsanity.com/paper/2507.01028/) (2025)
- [Augmentation Invariant Manifold Learning](https://mlsanity.com/paper/2211.00460/) (2022)
- [Diversified Ensemble of Independent Sub-Networks for Robust Self-Supervised Representation Learning](https://mlsanity.com/paper/2308.14705/) (2023)
- [Feature Normalization Prevents Collapse of Non-contrastive Learning Dynamics](https://mlsanity.com/paper/2309.16109/) (2023)
- [Unifying Self-Supervised Clustering and Energy-Based Models](https://mlsanity.com/paper/2401.00873/) (2023)

## Abstract (arXiv)

Recently the surprising discovery of the Bootstrap Your Own Latent (BYOL) method by Grill et al. shows the negative term in contrastive loss can be removed if we add the so-called prediction head to the network. This initiated the research of non-contrastive self-supervised learning. It is mysterious why even when there exist trivial collapsed global optimal solutions, neural networks trained by (stochastic) gradient descent can still learn competitive representations. This phenomenon is a typical example of implicit bias in deep learning and remains little understood. In this work, we present our empirical and theoretical discoveries on non-contrastive self-supervised learning. Empirically, we find that when the prediction head is initialized as an identity matrix with only its off-diagonal entries being trainable, the network can learn competitive representations even though the trivial optima still exist in the training objective. Theoretically, we present a framework to understand the behavior of the trainable, but identity-initialized prediction head. Under a simple setting, we characterized the substitution effect and acceleration effect of the prediction head. The substitution effect happens when learning the stronger features in some neurons can substitute for learning these features in other neurons through updating the prediction head. And the acceleration effect happens when the substituted features can accelerate the learning of other weaker features to prevent them from being ignored. These two effects enable the neural networks to learn all the features rather than focus only on learning the stronger features, which is likely the cause of the dimensional collapse phenomenon. To the best of our knowledge, this is also the first end-to-end optimization guarantee for non-contrastive methods using nonlinear neural networks with a trainable prediction head and normalization.