Fighting Class Imbalance with Contrastive Learning

Abstract: Medical image datasets are hard to collect, expensive to label, and often highly imbalanced. The last issue is underestimated, as typical average metrics hardly reveal that the often very important minority classes have a very low accuracy. In this paper, we address this problem by a feature embedding that balances the classes using contrastive learning as an alternative to the common cross-entropy loss. The approach is largely orthogonal to existing sampling methods and can be easily combined with those. We show on the challenging ISIC2018 and APTOS2019 datasets that the approach improves especially the accuracy of minority classes without negatively affecting the majority ones.