Class-Balanced Affinity Loss for Highly Imbalanced Tissue Classification in Computational Pathology

Abstract

Early detection of cancer, and breast cancer in particular, can have a positive impact on the survival rate of cancer patients. However, visual inspection by expert pathologists of whole-slide-images is subjective and error-prone given the lack of skilled pathologists. To overcome this limitation, many researchers have proposed deep learning driven approaches to detect breast cancer from histopathology images. However, these datasets are often highly imbalanced as patches belonging to the cancerous category is minor in comparison to the healthy cells. Therefore, when trained, the classification performance of the conventional Convolutional Neural Network (CNN) models drastically decreases, particularly for the minor class, which is often the main target of detection. This paper proposes a class balanced affinity loss function which can be injected at the output layer to any deep learning classifier model to address the imbalance learning. In addition to treating the imbalance, the proposal also builds uniformly spread class prototypes to address the fine-grained classification challenge in histopathology datasets, which conventional softmax loss cannot address. We validate our loss function performance by using two publicly available datasets with different levels of imbalance, namely the Invasive Ductal Carcinoma (IDC) and Colorectal cancer (CRC) datasets. In both cases, our method results in better performance, especially for the minority. We also observe a better 2D feature projection in multi-class classification with the proposed loss function, making it more apt to handle imbalanced fine-grained classification challenges.