Rethinking Dice Loss for Medical Image Segmentation

Deep learning has proved to be a powerful tool for medical image analysis in recent years. Data imbalance is a common problem in medical images. Dice Loss is widely used in medical image segmentation tasks to address the data imbalance problem. However, it only addresses the imbalance problem between foreground and background yet overlooks another imbalance between easy and hard examples that also severely affects the training process of a learning model. Empirically speaking, an easy example generally contributes less to the overall loss than a hard example. However, in practice, compared with hard examples, a large number of easy examples will be generated from a medical image and will dominate the training model, resulting in sub-optimal training or worse. To tackle this problem, we propose a novel Focal Dice Loss to alleviate the imbalance between hard examples and easy examples. Focal Dice Loss is able to reduce the contribution from easy examples and make the model focus on hard examples through our proposed novel balanced sampling strategy during the training process. Furthermore, to evaluate the effectiveness of our proposed loss functions, we conduct extensive experiments on two real-world medical image datasets with 2D and 3D convolutional neural networks. The experimental results show that our proposed Focal Dice Loss brings a significant improvement in segmentation performance compared to Dice Loss. Moreover, we find that our proposed Focal Dice Loss can effectively alleviate the over-fitting problem.

I. Introduction

Models based on deep learning have developed rapidly in recent years, and have achieved impressive performance in many fields such as automatic driving [1]–[3], assisted medical diagnosis [4]–[6] and so on. Medical image aided diagnosis system based on artificial intelligence can assist clinical diagnosis and reduce the probability of misdiagnosis. Through the recognition of artificial intelligence engine, the suspicious lesions in the image can be identified and then the doctor can read the slice, so as to improve the efficiency of the doctor. The common tasks of medical image assisted diagnosis include the detection and segmentation of lesion areas. Medical image segmentation, a fundamental step in the computer-aided diagnosis, aims to segment Region of interest (RoI) as to assisting doctors to make objective decisions. Training a medical image segmentation model is different from the case of a natural image. This is because of an issue inherent in medical images, which is the problem of data imbalance. The data imbalance in medical images can be roughly divided into two categories i.e. the imbalance between foreground and background examples, and that between easy and hard examples. Foreground and background examples usually refer to diseased and non-diseased examples in medical images. The latter is distinguished based on the classification difficulty of a single example. When the data is seriously unbalanced, even a very powerful network may lead to a relatively poor predicted result, so it is particularly important to solve the above two data imbalance problems.