Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations
Carole H. Sudre, 1122 Wenqi Li 11 Tom Vercauteren 11 Sebastien Ourselin, 1122 M. Jorge Cardoso , 1122
Abstract
Deep-learning has proved in recent years to be a powerful tool for image analysis and is now widely used to segment both 2D and 3D medical images. Deep-learning segmentation frameworks rely not only on the choice of network architecture but also on the choice of loss function. When the segmentation process targets rare observations, a severe class imbalance is likely to occur between candidate labels, thus resulting in sub-optimal performance. In order to mitigate this issue, strategies such as the weighted cross-entropy function, the sensitivity function or the Dice loss function, have been proposed. In this work, we investigate the behavior of these loss functions and their sensitivity to learning rate tuning in the presence of different rates of label imbalance across 2D and 3D segmentation tasks. We also propose to use the class re-balancing properties of the Generalized Dice overlap, a known metric for segmentation assessment, as a robust and accurate deep-learning loss function for unbalanced tasks.
中文速览
医学图像分割（medical image segmentation）中经常需要找出只占图像极小比例的病变区域，比如脑肿瘤或白质高信号，这种"类别极不平衡"的情况会让深度学习模型训练不稳定、效果变差。研究者系统比较了四种专为不平衡问题设计的损失函数（加权交叉熵、Dice损失、敏感性-特异性损失，以及他们新提出的广义Dice损失，Generalized Dice Loss），在两种2D网络和两种3D网络上，用脑肿瘤和白质病变两个任务，测试了不同学习率与采样策略下各函数的训练表现。结果表明，轻度不平衡时几种方法差异不大，但在高度不平衡的3D场景下，广义Dice损失（GDL）对超参数变化最为鲁棒，其他方法则容易训练失败或性能大幅下滑。这项工作为从事病变检测的研究者提供了明确的损失函数选择依据，对未来挑战更极端不平衡比例（如腔隙和血管周围间隙，前景占比约百万分之一）的任务具有重要参考价值。
原文 arXiv:1707.03237;中英对照 + 大白话阅读 https://aha.fim.ai/paper/1707.03237v3