Optimal Representations for Generalized Contrastive Learning with Imbalanced Datasets

2026-05-11

AI Summary

This work investigates the optimal representation geometry of contrastive learning under class imbalance. Addressing the lack of theoretical characterization of feature configurations in imbalanced settings, we rigorously derive—via convex optimization and geometric analysis—the exact solution for generalized contrastive loss: samples from the same class collapse onto their class mean, and these class means are arranged symmetrically on the unit hypersphere, with angular symmetry dictated by class proportions. Notably, we uncover a “minority-class collapse” phenomenon in highly imbalanced regimes, wherein all minority-class samples converge to a single vector, and we establish precise threshold conditions under which this collapse occurs. Numerical experiments confirm our theoretical predictions, demonstrating strong agreement between analysis and empirical observation.

Abstract

In this paper, we provide a computable characterization of the geometry of optimal representations in Contrastive Learning (CL) when the classes are imbalanced. When classes are balanced and the representation dimension is greater than the number of classes, it is well-known that the optimal representations exhibit Neural Collapse (NC), i.e., representations from the same class collapse to their class means and the class means form an Equiangular Tight Frame (ETF). For imbalanced classes and a large, generalized family of CL losses, we prove that the optimal representations of all samples from the same class collapse to their class means and their geometry exhibits an angular symmetry structure that is determined by the relative class proportions. In general, we show that the geometry can be determined by solving a convex optimization problem. Exploiting this symmetry structure, we analytically investigate a special case where class imbalance is extreme and prove that CL exhibits a phenomenon called Minority Collapse (MC) where all samples from the minority classes (classes with small probabilities) collapse into a single vector, whenever the class imbalance exceeds a threshold, which in turn depends on the regularity properties of the CL loss used and on the number of negative samples. Numerical results are provided to illustrate these phenomena and corroborate the theoretical results. We conclude by identifying a number of open problems.

Problem

• Contrastive Learning
• Class Imbalance
• Optimal Representations
• Minority Collapse
• Neural Collapse

Innovation

• Contrastive Learning
• Class Imbalance
• Neural Collapse
• Minority Collapse
• Equiangular Tight Frame