Class-Specific Distribution Alignment for Semi-Supervised Medical Image Classification

TL;DR

This paper introduces CSDA, a semi-supervised learning framework that effectively handles class imbalance in medical image classification by aligning class-specific distributions and maintaining balanced unlabeled sample pools.

Contribution

The paper proposes a novel distribution alignment method considering class-specific predictions and introduces VCQ for balanced unlabeled data sampling in semi-supervised medical imaging.

Findings

CSDA achieves competitive results on three public datasets.

The method effectively addresses class imbalance issues.

Experimental results demonstrate improved classification performance.

Abstract

Despite the success of deep neural networks in medical image classification, the problem remains challenging as data annotation is time-consuming, and the class distribution is imbalanced due to the relative scarcity of diseases. To address this problem, we propose Class-Specific Distribution Alignment (CSDA), a semi-supervised learning framework based on self-training that is suitable to learn from highly imbalanced datasets. Specifically, we first provide a new perspective to distribution alignment by considering the process as a change of basis in the vector space spanned by marginal predictions, and then derive CSDA to capture class-dependent marginal predictions on both labeled and unlabeled data, in order to avoid the bias towards majority classes. Furthermore, we propose a Variable Condition Queue (VCQ) module to maintain a proportionately balanced number of unlabeled samples for each class. Experiments on three public datasets HAM10000, CheXpert and Kvasir show that our method provides competitive performance on semi-supervised skin disease, thoracic disease, and endoscopic image classification tasks.