MBNM: Multi-branch network based on memory features for long-tailed medical image recognition

Deep learning algorithms show revolutionary potential in computer-aided diagnosis. These computer-aided diagnosis techniques often rely on large-scale, balanced standard datasets. However, there are many rare diseases in real clinical scenarios, which makes the medical datasets present a highly imbalanced long-tailed distribution, leading to the poor generalization ability of deep learning models. Currently, most algorithms to solve this problem involve more complex modules and loss functions. But for complicated tasks in the medical domain, usually suffer from issues such as increased inference time and unstable performance. Therefore, it is a great challenge to develop a computer-aided diagnosis algorithm for long-tailed medical data.

We proposed the Multi-Branch Network based on Memory Features (MBNM) for Long-Tailed Medical Image Recognition. MBNM includes three branches, where each branch focuses on a different learning task: 1) the regular learning branch learns the generalizable feature representations; 2) the tail learning branch gains extra intra-class diversity for the tail classes through the feature memory module and the improved reverse sampler to improve the classification performance of the tail classes; 3) the fusion balance branch integrates various decision-making advantages and introduces an adaptive loss function to re-balance the classification performance of easy and difficult samples.

We conducted experiments on the multi-disease Ophthalmic OCT datasets with imbalance factors of 98.48 and skin image datasets Skin-7 with imbalance factors of 58.3. The Accuracy, MCR, F1-Score, Precision, and AUC of our model were significantly improved over the strong baselines in the auxiliary diagnosis scenario where the clinical medical data is extremely imbalanced. Furthermore, we demonstrated that MBNM outperforms the state-of-the-art models on the publicly available natural image datasets (CIFAR-10 and CIFAR-100).

The proposed algorithm can solve the problem of imbalanced data distribution with little added cost. In addition, the memory module does not act in the inference phase, thereby saving inference time. And it shows outstanding performance on medical images and natural images with a variety of imbalance factors.