PCCT: Progressive Class-Center Triplet Loss for Imbalanced Medical Image Classification

TL;DR

This paper introduces PCCT, a two-stage triplet loss framework that effectively addresses class imbalance in medical image classification, especially for rare diseases, by combining class-balanced sampling and class-center based loss.

Contribution

The paper proposes a novel two-stage triplet loss framework with class-center involvement, extending to ranking and quadruplet losses, to improve imbalanced medical image classification.

Findings

Achieves state-of-the-art F1 scores on multiple datasets.

Effectively improves classification of rare disease classes.

Outperforms existing methods for imbalanced data handling.

Abstract

Imbalanced training data is a significant challenge for medical image classification. In this study, we propose a novel Progressive Class-Center Triplet (PCCT) framework to alleviate the class imbalance issue particularly for diagnosis of rare diseases, mainly by carefully designing the triplet sampling strategy and the triplet loss formation. Specifically, the PCCT framework includes two successive stages. In the first stage, PCCT trains the diagnosis system via a class-balanced triplet loss to coarsely separate distributions of different classes. In the second stage, the PCCT framework further improves the diagnosis system via a class-center involved triplet loss to cause a more compact distribution for each class. For the class-balanced triplet loss, triplets are sampled equally for each class at each training iteration, thus alleviating the imbalanced data issue. For the class-center involved triplet loss, the positive and negative samples in each triplet are replaced by their corresponding class centers, which enforces data representations of the same class closer to the class center. Furthermore, the class-center involved triplet loss is extended to the pair-wise ranking loss and the quadruplet loss, which demonstrates the generalization of the proposed framework. Extensive experiments support that the PCCT framework works effectively for medical image classification with imbalanced training images. On two skin image datasets and one chest X-ray dataset, the proposed approach respectively obtains the mean F1 score 86.2, 65.2, and 90.66 over all classes and 81.4, 63.87, and 81.92 for rare classes, achieving state-of-the-art performance and outperforming the widely used methods for the class imbalance issue.