On Neural Collapse in Contrastive Learning with Imbalanced Datasets

Neural collapse is a phenomenon in neural networks where all the samples from the same class collapse to their class mean and the class means have a specific geometric structure called Equiangular Tight Frame (ETF). Recent studies have empirically and theoretically confirmed that neural collapse solutions attain a global minimum of Contrastive Learning (CL) losses if the class distribution is uniform (balanced classes). This paper investigates the neural collapse phenomenon in Contrastive Learning for imbalanced datasets. We show that even if the classes are imbalanced, global minima of Contrastive Learning losses include solutions in which all samples from the same class collapse to their class mean. However, the geometric structure of the class means is, in general, different from an ETF. In addition, we show that under certain conditions the optimal geometry of the class means can be found by solving a convex optimization problem. We provide a program based on the CVX package to find these optimal class means.