Research map / Label-efficient and robust learning / Self-supervised and contrastive
Investigating the Benefits of Projection Head for Representation Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Theoretically explains why pre-projection representations outperform post-projection ones via layer-wise progressive feature weighting in linear models and feature selection via non-linearity, validated on MNIST-on-CIFAR-10, CIFAR-100, and UrbanCars, and proposes a fixed reweighting head as an alternative.
Problem
Why are pre-projection representations (before a projection head) better for downstream tasks than post-projection ones in contrastive and supervised learning, when the pre-projection features are not directly optimized by the loss?
Approach
Shows that gradient-based training (min-norm or gradient flow) yields balanced layer weights W1W1^T = W2^T W2, so the pre-projection encoder weights features more equally; with non-linearity, lower layers retain features that are zeroed after the head, and a fixed reweighting head mimics this effect.
Key results
- On MNIST-on-CIFAR-10 with SimCLR, the fixed reweighting head reaches 97.3% digit accuracy, matching the trainable projection head (97.3%) and well above vanilla (77.0%).
- On CIFAR-100 fine-grained linear evaluation with SCL, pre-projection reaches 36.0% vs 25.7% post-projection.
- On UrbanCars few-shot adaptation, the reweighting head achieves 87.0% vs 86.1% for projection head and 82.2% vanilla.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use pre-projection instead of post-projection representations | SCL · Synthetic coarse-to-fine | fine accuracy: $36.0_{\pm 0.1}$ vs $25.7_{\pm 0.7}$ | Helped |
| Use pre-projection instead of post-projection representations | SL · Synthetic coarse-to-fine | fine accuracy: $33.7_{\pm 0.3}$ vs $29.7_{\pm 0.3}$ | Helped |
| Use pre-projection instead of post-projection representations | SL · Synthetic coarse-to-fine | coarse accuracy: $53.1_{\pm 0.3}$ vs $53.2_{\pm 0.7}$ | No clear effect |
| Use pre-projection instead of no-head vanilla representations | SimCLR ResNet-18 · MNIST-on-CIFAR-10 | digit classification accuracy: 77.0 -> 97.3 | Helped |
| Use pre-projection instead of no-head vanilla representations | SCL ResNet-18 · CIFAR-100 fine-grained classification | fine-grained accuracy: 21.8 -> 36.0 | Helped |
| Use pre-projection instead of no-head vanilla representations | SL ResNet-18 · CIFAR-100 fine-grained classification | fine-grained accuracy: 31.44 -> 33.7 | Helped |
| Use pre-projection instead of no-head vanilla representations | SL · UrbanCars few-shot distribution shift | few-shot adaptation accuracy: 82.2 -> 86.1 | Helped |
| Replace the trainable projection head with a fixed reweighting head | SimCLR ResNet-18 · MNIST-on-CIFAR-10 | digit classification accuracy: 77.0 -> 97.3 | Helped |
| Replace the trainable projection head with a fixed reweighting head | SCL ResNet-18 · CIFAR-100 fine-grained classification | fine-grained accuracy: 21.8 -> 30.2 | Helped |
| Replace the trainable projection head with a fixed reweighting head | SL ResNet-18 · CIFAR-100 fine-grained classification | fine-grained accuracy: 31.44 -> 32.2 | Helped |
| Replace the trainable projection head with a fixed reweighting head | SL · UrbanCars few-shot distribution shift | few-shot adaptation accuracy: 82.2 -> 87.0 | Helped |
| Increase weight decay (ablation) | SimCLR ResNet-18 with projection head · MNIST-on-CIFAR-10 with p_drop=1 | digit classification accuracy | Hurt |
Surprising findings
- Having a strong downstream-relevant feature in pretraining can still make pre-projection representations preferable, because the model moderates overly strong features.
- Large weight decay can eliminate and even reverse the benefit of pre-projection representations.
Limitations
- Theory is restricted to 2-layer linear and diagonal non-linear networks with simplified independent-coordinate data distributions.
- Optimal layer for a given downstream task is hard to determine in practice.
- Reweighting head kappa is a fixed hyperparameter tuned per experiment.
Experimental setup
Datasets: MNIST CIFAR-10 CIFAR-100 UrbanCars ImageNet
Benchmarks: MNIST CIFAR-100 coarse/fine UrbanCars ImageNet shifted variants
Metrics: linear probe accuracy few-shot adaptation accuracy sample complexity indicator
Baselines: post-projection representations no-projection vanilla training without head trainable projection head
Builds on: SimCLR supervised contrastive learning neural collapse
Closest papers by meaning
- Projection Head is Secretly an Information Bottleneck (ICLR 2025 · similarity 0.90)
- When Does Contrastive Visual Representation Learning Work? (CVPR 2022 · similarity 0.87)
- Perceptual Inductive Bias Is What You Need Before Contrastive Learning (CVPR 2025 · similarity 0.85)
- Harnessing small projectors and multiple views for efficient vision pretraining (NeurIPS 2024 · similarity 0.85)
- Tracing Representation Progression: Analyzing and Enhancing Layer-Wise Similarity (ICLR 2025 · similarity 0.84)
- $\mathbb{X}$-Sample Contrastive Loss: Improving Contrastive Learning with Sample Similarity Graphs (ICLR 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.