Equilibrium contrastive learning for imbalanced image classification

TL;DR

Equilibrium Contrastive Learning (ECL) tackles imbalanced image classification by enforcing a batch-invariant geometric equilibrium across three interrelated components: intra-class feature collapse, uniform inter-class mean spacing, and alignment between classifier weights and class centers. It introduces BC-ECL to balance class-average features and prototypes in representation learning, CC-GE to align linear classifier weights with class prototypes, and logit compensation to mitigate bias due to imbalance, all trained end-to-end. Across five long-tailed benchmarks and two medical-imaging datasets, ECL achieves state-of-the-art accuracy and demonstrates consistent improvements in representation metrics (FC, MS, SD), while ablation studies confirm the necessity and synergy of its components. The work highlights that explicit geometric regularization, paired with prototype-based alignment, can significantly enhance generalization under class imbalance and offers scalable directions for broader modality applications and reduced batch-size requirements.

Abstract

Contrastive learning (CL) is a predominant technique in image classification, but they showed limited performance with an imbalanced dataset. Recently, several supervised CL methods have been proposed 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 prototypes, as additional samples to consider all classes. However, the existing CL methods suffer from two limitations. First, they do not consider the alignment between the class means/prototypes and classifiers, which could lead to poor generalization. Second, existing prototype-based methods treat 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, where class features, means, and classifiers are harmoniously balanced under data imbalance. The proposed ECL framework uses two main components. 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. We ran experiments with three long-tailed datasets, the CIFAR-10(0)-LT, ImageNet-LT, and the two imbalanced medical datasets, the ISIC 2019 and our constructed LCCT dataset. Results show that ECL outperforms existing SOTA supervised CL methods designed for imbalanced classification.

Figures (8)

• Figure 1: Illustration of Equilibrium Contrastive Learning ( ECL). ECL promotes the following three key geometric properties that collectively establish geometric equilibrium in the representation space: 1) the collapse of within-class representations (denoted by $\mdsmblkcircle$) to their class means (denoted by $\lozenge$), 2) equidistant spacing among class means, and 3) alignment between linear classifier weights (denoted by skyblue sticks) and class means. It can simplify the geometries in representation and classifier, leading to simple nearest-centroid classifications in the representation space. The simple geometries are useful for improving the generalization capability.
• Figure 2: Illustration of three key properties to simplify geometries via ECL under class imbalance (red and green denote features from different majority classes, and blue denotes features from the minority class). (a) Intra-class feature collapse: all the intra-class representations (denoted by $\mdsmblkcircle$) collapse to their corresponding class means (denoted by $\lozenge$). (b) Inter-class mean geometric equilibrium: distances and angles among class means become equal, forming a symmetric geometric structure. (c) CC-GE: Classifier–class center geometric equilibrium: the linear classifiers (denoted by skyblue sticks) and class means converge to each other with rescaling.
• Figure 3: Visualization of feature distribution on the unit sphere obtained by BCL with the LCCT dataset zooming in to the upper half for better visualization. The $\triangle$ and $\square$ symbols indicate the class means and classifier weights, while the red and green points represent samples from the majority and minority classes, respectively. The results show that the CC-GE property is not induced.
• Figure 4: The overall ECL architecture. The proposed ECL consists of three core modules, (a) contrastive representation learning with proposed $\mathcal{L}_{\text{BC-ECL}}$ in Section \ref{['main:sec:BC-ECL']}, (b) classifier-class center alignment with proposed $\mathcal{L}_{\text{CC-GE}}$ in Section \ref{['main:sec:ccge']}, and (c) classifier learning with $\mathcal{L}_{\text{LC}}$ introduced in Section \ref{['main:sec:lc']}. (a) The proposed $\mathcal{L}_{\text{BC-ECL}}$ formulation equalizes the contribution between class-average features and prototypes in promoting the intra-class feature collapse and geometric equilibrium among class means. We utilize the linearly transformed classifier weights as class prototypes, and $\ell _{2}$-normalize the representations $\{\mathbf{z}_{i}^{(2)}, \mathbf{z}_{i}^{(3)} : \forall i \}$. The symbols $\text{Sim.}$ and $\text{Avg.}$ denote a similarity computation between two vectors and an averaging operator, respectively. (b) The proposed $\mathcal{L}_{\text{CC-GE}}$ loss aligns the classifier weights and class prototypes to promote geometric equilibrium between classifier weights and class center representations. (c) For classifier learning, we use a logit-compensated variant of the cross-entropy loss, $\mathcal{L}_{\text{LC}}$.
• Figure 5: Visualization of feature distribution on the unit hypersphere for different supervised CL methods with LCCT dataset; $\{ x \in [-1, 1], y \in [-1, 1] \}$. If class-wise features are clustered (see Fig. \ref{['main:fig:properties']}(a)) to form a regular simplex configuration (see Fig. \ref{['main:fig:properties']}(b)) and their class means are aligned with the classifier weights (see Fig. \ref{['main:fig:properties']}(c)), then good representations and classifier are learned. Legends:denotes minority class samples; denotes majority class samples; denotes minority class mean; denotes majority class mean; denotes minority classifier weight; and denotes majority classifier weight.

Theorems & Definitions (2)

• Theorem 1
• proof