Optimal Representations for Generalized Contrastive Learning with Imbalanced Datasets
Provides theoretical understanding of contrastive learning dynamics in imbalanced settings, which is crucial for practitioners dealing with real-world skewed datasets.
This paper characterizes the geometry of optimal representations in contrastive learning under class imbalance, showing that representations collapse to class means with an angular symmetry structure determined by class proportions. For extreme imbalance, minority classes collapse to a single vector when imbalance exceeds a threshold dependent on loss properties and negative sample count.
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.