Adaptive Unified Contrastive Learning for Imbalanced Classification

Cong Cong, Yixing Yang, Sidong Liu, Maurice Pagnucco, Antonio Di Ieva, Shlomo Berkovsky, Yang Song

Abstract: Medical image classifiers often suffer from the imbalanced class distribution of datasets. For example, among the 7 classes in the ISIC2018 skin lesion detection dataset, over 67% of the instances belong to melanocytic nevus while only 1% belong to dermatofibroma. Contrastive feature learning has been shown to achieve promising results in enhancing the performance for imbalanced classification tasks. However, the contrastive learning methods are either not end-to-end or require extra memory, which may lead to less compatible and sub-optimal features and classifiers. In this paper, we propose a novel unified feature and classifier learning framework for imbalanced medical image datasets. We equip our model with an adaptive unified contrastive (AduC) loss which progressively adapts model learning between feature learning and classifier learning. Furthermore, we explore the impact of different sampling methods on model training under data sparsity. The experimental results on two long-tailed medical datasets demonstrate that our methods can substantially improve the classification accuracy and F1-score over all classes without using extra memory storage. Our code is available at https://github.com/thomascong121/AdUni.