Improving Class Balancing at Both Feature Extractor and Classifier Head

Kanghao Chen, Huijuan Lu, Ruixuan Wang, Wei-Shi Zheng

Abstract: Training data are often imbalanced across classes in practice, and such class imbalance issue often causes model predictions biased toward majority classes during inference. Different from existing solutions which employ various training strategies to alleviate the class imbalance issue, this study proposes a novel two-head model architecture to help alleviate the issue. One auxiliary classifier head helps the feature extractor of the classifier more fairly learn to extract features for each class, and the main classifier head learns in a more class-balanced manner by dividing each majority class into multiple clusters in advance and considering each cluster as a new class. Extensive empirical evaluations on four class-imbalanced image datasets showed that the proposed approach achieves state-of-the-art classification performance.