Cross-Supervision Similarity Network for Medical Image Classification on Imbalanced Small Datasets.
Authors
Abstract
Imbalanced small datasets are common scenarios in the field of machine learning for medical imaging, especially in real-world clinical applications. Many existing works focus on synthesize new images via data generation. However, generative methods cannot ensure reliability for medical images where categories cannot be easily distinguished, such as the pathologic complete response (pCR) evaluation via MRIs in cancer prognosis. Meanwhile, few-shot learning can deal with training with small datasets, but it depends on a balanced data distribution and a large number of image categories. In this paper, we propose an image similarity comparison classification network, referred to as Cross-Supervision Similarity Network (CSSN), using cross-supervision between class features and patch features. CSSN transforms the classification task into comparison task by calculating similarity scores at both patch and class scales, effectively training on imbalanced small datasets with limited categories. To balance the training difficulty of the two similarity branches, soft logarithmic supervision is used to construct soft labels between them. Through experiments on PCR-ISD, we observe significant performance improvements of 15% in F1 score, 6% in accuracy and 9 % in balanced accuracy over existing methods, indicating the superiority of our method in identifying minority classes and enhancing classification capabilities. Extensive experiments on three datasets and ablation experiments confirm the effectiveness and generalization ability of the proposed method. The source code is available at https://github.com/lxy-146/CSSN_TMI.