Highly Efficient Representation and Active Learning Framework for Imbalanced Data and its Application to COVID-19 X-Ray Classification

We propose a highly data-efficient classification and active learning framework for classifying chest X-rays. It is based on (1) unsupervised representation learning of a Convolutional Neural Network and (2) the Gaussian Process method. The unsupervised representation learning employs self-supervision that does not require class labels, and the learned features are proven to achieve label-efficient classification. GP is a kernel-based Bayesian approach that also leads to data-efficient predictions with the added benefit of estimating each decision's uncertainty. Our novel framework combines these two elements in sequence to achieve highly data and label efficient classifications. Moreover, both elements are less sensitive to the prevalent and challenging class imbalance issue, thanks to the (1) feature learned without labels and (2) the Bayesian nature of GP. The GP-provided uncertainty estimates enable active learning by ranking samples based on the uncertainty and selectively labeling samples showing higher uncertainty. We apply this novel combination to the data-deficient and severely imbalanced case of COVID-19 chest X-ray classification. We demonstrate that only $\sim 10\%$ of the labeled data is needed to reach the accuracy from training all available labels. Its application to the COVID-19 data in a fully supervised classification scenario shows that our model, with a generic ResNet backbone, outperforms (COVID-19 case by 4\%) the state-of-the-art model with a highly tuned architecture. Our model architecture and proposed framework are general and straightforward to apply to a broader class of datasets, with expected success.