Handling Imbalanced Data: Uncertainty-Guided Virtual Adversarial Training With Batch Nuclear-Norm Optimization for Semi-Supervised Medical Image Classification.

Abstract

In manyclinical settings, a lot of medical image datasets suffer from imbalance problems, which makes predictions of trained models to be biased toward majority classes. Semi-supervised Learning (SSL) algorithms trained with such imbalanced datasets become more problematic since pseudo-supervision of unlabeled data are generated from the model's biased predictions. To address these issues, in this work, we propose a novel semi-supervised deep learning method, i.e., uncertainty-guided virtual adversarial training (VAT) with batch nuclear-norm (BNN) optimization, for large-scale medical image classification. To effectively exploit useful information from both labeled and unlabeled data, we leverage VAT and BNN optimization to harness the underlying knowledge, which helps to improve discriminability, diversity and generalization of the trained models. More concretely, our network is trained by minimizing a combination of four types of losses, including a supervised cross-entropy loss, a BNN loss defined on the output matrix of labeled data batch (lBNN loss), a negative BNN loss defined on the output matrix of unlabeled data batch (uBNN loss), and a VAT loss on both labeled and unlabeled data. We additionally propose to use uncertainty estimation to filter out unlabeled samples near the decision boundary when computing the VAT loss. We conduct comprehensive experiments to evaluate the performance of our method on two publicly available datasets and one in-house collected dataset. The experimental results demonstrated that our method achieved better results than state-of-the-art SSL methods.