Equilibrium contrastive learning for imbalanced image classification

The data imbalance issue is critical in image classification. Contrastive learning (CL) is a predominant technique in image classification, but they showed limited classification performance when the class distribution in a training dataset is skewed (i.e., an imbalanced dataset). Recently, several supervised CL methods have been proposed particularly to promote an ideal regular simplex geometric configuration in the representation space—characterized by intra-class feature collapse and uniform inter-class mean spacing—especially for imbalanced datasets. In particular, existing prototype-based methods include class representatives, i.e., prototypes, as additional samples to achieve a more balanced treatment of all classes. However, the existing supervised CL methods for imbalanced datasets still suffer from two major limitations. First, they do not consider the alignment between the class means/prototypes and linear classifiers, which could lead to poor generalization. Second, existing prototype-based methods treat class prototypes as only one additional sample per class, making their influence depend on the number of class instances in a batch and causing unbalanced contributions across classes. To address these limitations, we propose Equilibrium Contrastive Learning (ECL), a supervised CL framework designed to promote geometric equilibrium in the representation space, where class features, means, and classifiers are harmoniously balanced under data imbalance. The proposed ECL framework uses two main strategies. First, ECL promotes the representation geometric equilibrium (i.e., a regular simplex geometry characterized by collapsed class samples and uniformly distributed class means), while balancing the contributions of class-average features and class prototypes. Second, ECL establishes a classifier-class center geometric equilibrium by aligning classifier weights and class prototypes in the representation space. We ran experiments with three long-tailed natural image datasets, the CIFAR-10-LT, CIFAR-100-LT and ImageNet-LT benchmark datasets, and the two imbalanced medical image classification datasets, the ISIC 2019 benchmark and our constructed lung cancer chest computed tomography dataset. Results with the five imbalanced datasets show that ECL outperforms existing state-of-the-art supervised CL methods designed for imbalanced/long-tailed classification. Codes are available at this link.