Pixel-wise Modulated Dice Loss for Medical Image Segmentation

TL;DR

This paper introduces a simple, computationally efficient pixel-wise modulated Dice loss that improves medical image segmentation by effectively addressing both class and difficulty imbalance issues.

Contribution

The paper proposes a novel, minimal-cost modification to Dice loss using pixel-wise modulation to better handle class and difficulty imbalance in medical segmentation.

Findings

Outperforms existing methods on three medical segmentation tasks.

Effectively addresses class and difficulty imbalance.

Maintains low computational cost.

Abstract

Class imbalance and the difficulty imbalance are the two types of data imbalance that affect the performance of neural networks in medical segmentation tasks. In class imbalance the loss is dominated by the majority classes and in difficulty imbalance the loss is dominated by easy to classify pixels. This leads to an ineffective training. Dice loss, which is based on a geometrical metric, is very effective in addressing the class imbalance compared to the cross entropy (CE) loss, which is adopted directly from classification tasks. To address the difficulty imbalance, the common approach is employing a re-weighted CE loss or a modified Dice loss to focus the training on difficult to classify areas. The existing modification methods are computationally costly and with limited success. In this study we propose a simple modification to the Dice loss with minimal computational cost. With a pixel level modulating term, we take advantage of the effectiveness of Dice loss in handling the class imbalance to also handle the difficulty imbalance. Results on three commonly used medical segmentation tasks show that the proposed Pixel-wise Modulated Dice loss (PM Dice loss) outperforms other methods, which are designed to tackle the difficulty imbalance problem.