# Perfectly Balanced: Improving Transfer and Robustness of Supervised Contrastive Learning

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

## Key facts

- Cited by 40 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 57 references that are in the map.
- First posted on arXiv 15 Apr 2022.

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

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (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)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Unsupervised Learning of Visual Features by Contrasting Cluster Assignments](https://mlsanity.com/paper/2006.09882/) (2020)
- [Supervised Contrastive Learning](https://mlsanity.com/paper/2004.11362/) (2020)
- [Improved Baselines with Momentum Contrastive Learning](https://mlsanity.com/paper/2003.04297/) (2020)
- [Understanding deep learning requires rethinking generalization](https://mlsanity.com/paper/1611.03530/) (2016)
- [Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere](https://mlsanity.com/paper/2005.10242/) (2020)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [What Makes for Good Views for Contrastive Learning?](https://mlsanity.com/paper/2005.10243/) (2020)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [Skin Lesion Analysis Toward Melanoma Detection 2018: A Challenge Hosted by the International Skin Imaging Collaboration (ISIC)](https://mlsanity.com/paper/1902.03368/) (2019)
- [Contrastive Learning with Hard Negative Samples](https://mlsanity.com/paper/2010.04592/) (2020)
- [A Theoretical Analysis of Contrastive Unsupervised Representation Learning](https://mlsanity.com/paper/1902.09229/) (2019)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [Understanding Dimensional Collapse in Contrastive Self-supervised Learning](https://mlsanity.com/paper/2110.09348/) (2021)
- [On Mutual Information Maximization for Representation Learning](https://mlsanity.com/paper/1907.13625/) (2019)
- [Provable Guarantees for Self-Supervised Deep Learning with Spectral Contrastive Loss](https://mlsanity.com/paper/2106.04156/) (2021)
- [Contrastive Learning Inverts the Data Generating Process](https://mlsanity.com/paper/2102.08850/) (2021)
- [Hidden Stratification Causes Clinically Meaningful Failures in Machine Learning for Medical Imaging](https://mlsanity.com/paper/1909.12475/) (2019)
- [No Subclass Left Behind: Fine-Grained Robustness in Coarse-Grained Classification Problems](https://mlsanity.com/paper/2011.12945/) (2020)
- [Self-supervised Pretraining of Visual Features in the Wild](https://mlsanity.com/paper/2103.01988/) (2021)
- [Neural Collapse Under MSE Loss: Proximity to and Dynamics on the Central Path](https://mlsanity.com/paper/2106.02073/) (2021)
- [A Kernel Theory of Modern Data Augmentation](https://mlsanity.com/paper/1803.06084/) (2018)
- [On the Role of Neural Collapse in Transfer Learning](https://mlsanity.com/paper/2112.15121/) (2021)
- [Towards the Generalization of Contrastive Self-Supervised Learning](https://mlsanity.com/paper/2111.00743/) (2021)
- [Model Patching: Closing the Subgroup Performance Gap with Data Augmentation](https://mlsanity.com/paper/2008.06775/) (2020)
- [A Broad Study on the Transferability of Visual Representations with Contrastive Learning](https://mlsanity.com/paper/2103.13517/) (2021)

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

- [M-FLAG: Medical Vision-Language Pre-training with Frozen Language Models and Latent Space Geometry Optimization](https://mlsanity.com/paper/2307.08347/) (2023)
- [Learning Generalized Zero-Shot Learners for Open-Domain Image Geolocalization](https://mlsanity.com/paper/2302.00275/) (2023)
- [Demystifying Structural Disparity in Graph Neural Networks: Can One Size Fit All?](https://mlsanity.com/paper/2306.01323/) (2023)
- [InfoNCE Loss Provably Learns Cluster-Preserving Representations](https://mlsanity.com/paper/2302.07920/) (2023)
- [Which Features are Learnt by Contrastive Learning? On the Role of Simplicity Bias in Class Collapse and Feature Suppression](https://mlsanity.com/paper/2305.16536/) (2023)
- [OrCo: Towards Better Generalization via Orthogonality and Contrast for Few-Shot Class-Incremental Learning](https://mlsanity.com/paper/2403.18550/) (2024)
- [Understanding and Improving Transfer Learning of Deep Models via Neural Collapse](https://mlsanity.com/paper/2212.12206/) (2022)
- [A Survey on the Robustness of Computer Vision Models against Common Corruptions](https://mlsanity.com/paper/2305.06024/) (2023)
- [Benchmarking and Building Long-Context Retrieval Models with LoCo and M2-BERT](https://mlsanity.com/paper/2402.07440/) (2024)
- [Investigating the Benefits of Projection Head for Representation Learning](https://mlsanity.com/paper/2403.11391/) (2024)
- [Rethinking Weak Supervision in Helping Contrastive Learning](https://mlsanity.com/paper/2306.04160/) (2023)
- [Contrastive Tuning: A Little Help to Make Masked Autoencoders Forget](https://mlsanity.com/paper/2304.10520/) (2023)
- [Neural Collapse in Multi-label Learning with Pick-all-label Loss](https://mlsanity.com/paper/2310.15903/) (2023)
- [FALCON: Fairness Learning via Contrastive Attention Approach to Continual Semantic Scene Understanding](https://mlsanity.com/paper/2311.15965/) (2023)
- [CLOSER: Towards Better Representation Learning for Few-Shot Class-Incremental Learning](https://mlsanity.com/paper/2410.05627/) (2024)
- [Neural Collapse is Globally Optimal in Deep Regularized ResNets and Transformers](https://mlsanity.com/paper/2505.15239/) (2025)
- [Quantifying the Variability Collapse of Neural Networks](https://mlsanity.com/paper/2306.03440/) (2023)
- [Cell-ontology guided transcriptome foundation model](https://mlsanity.com/paper/2408.12373/) (2024)
- [Collapse-Proof Non-Contrastive Self-Supervised Learning](https://mlsanity.com/paper/2410.04959/) (2024)
- [SC-Block: Supervised Contrastive Blocking within Entity Resolution Pipelines](https://mlsanity.com/paper/2303.03132/) (2023)
- [A Tale of Two Classes: Adapting Supervised Contrastive Learning to Binary Imbalanced Datasets](https://mlsanity.com/paper/2503.17024/) (2025)
- [Self-Supervised Representation Learning with Cross-Context Learning between Global and Hypercolumn Features](https://mlsanity.com/paper/2308.13392/) (2023)
- [SINCERE: Supervised Information Noise-Contrastive Estimation REvisited](https://mlsanity.com/paper/2309.14277/) (2023)
- [On the Effectiveness of Supervision in Asymmetric Non-Contrastive Learning](https://mlsanity.com/paper/2406.10815/) (2024)
- [Continual Contrastive Finetuning Improves Low-Resource Relation Extraction](https://mlsanity.com/paper/2212.10823/) (2022)
- [DualTalker: A Cross-Modal Dual Learning Approach for Speech-Driven 3D Facial Animation](https://mlsanity.com/paper/2311.04766/) (2023)
- [A Generalized Theory of Mixup for Structure-Preserving Synthetic Data](https://mlsanity.com/paper/2503.02645/) (2025)
- [A Theoretical Framework for Preventing Class Collapse in Supervised Contrastive Learning](https://mlsanity.com/paper/2503.08203/) (2025)
- [Escaping Plato's Cave: JAM for Aligning Independently Trained Vision and Language Models](https://mlsanity.com/paper/2507.01201/) (2025)
- [Balanced Supervised Contrastive Learning for Few-Shot Class-Incremental Learning](https://mlsanity.com/paper/2305.16687/) (2023)

## Abstract (arXiv)

An ideal learned representation should display transferability and robustness. Supervised contrastive learning (SupCon) is a promising method for training accurate models, but produces representations that do not capture these properties due to class collapse -- when all points in a class map to the same representation. Recent work suggests that "spreading out" these representations improves them, but the precise mechanism is poorly understood. We argue that creating spread alone is insufficient for better representations, since spread is invariant to permutations within classes. Instead, both the correct degree of spread and a mechanism for breaking this invariance are necessary. We first prove that adding a weighted class-conditional InfoNCE loss to SupCon controls the degree of spread. Next, we study three mechanisms to break permutation invariance: using a constrained encoder, adding a class-conditional autoencoder, and using data augmentation. We show that the latter two encourage clustering of latent subclasses under more realistic conditions than the former. Using these insights, we show that adding a properly-weighted class-conditional InfoNCE loss and a class-conditional autoencoder to SupCon achieves 11.1 points of lift on coarse-to-fine transfer across 5 standard datasets and 4.7 points on worst-group robustness on 3 datasets, setting state-of-the-art on CelebA by 11.5 points.