[论文解读] A Mixed Focal Loss Function for Handling Class Imbalanced Medical Image Segmentation.

Michael Yeung, Evis Sala+2 more|arXiv (Cornell University)|Feb 8, 2021

Radiomics and Machine Learning in Medical Imaging被引用 9

一句话总结

本文提出一种混合焦点损失函数，结合改进的焦点损失与焦点Dice损失，以提升医学图像分割中类别不平衡问题的性能。该方法在KiTS19和BraTS20数据集上的二分类与多分类分割任务中，均实现了更优的查全率-查准率平衡性与鲁棒性，优于标准损失函数及各组件损失函数。

ABSTRACT

Automatic segmentation methods are an important advancement in medical imaging analysis. Machine learning techniques, and deep neural networks in particular, are the state-of-the-art for most automated medical image segmentation tasks, ranging from the subcellular to the level of organ systems. Issues with class imbalance pose a significant challenge irrespective of scale, with organs, and especially with tumours, often occupying a considerably smaller volume relative to the background. Loss functions used in the training of segmentation algorithms differ in their robustness to class imbalance, with cross entropy-based losses being more affected than Dice-based losses. In this work, we first experiment with seven different Dice-based and cross entropy-based loss functions on the publicly available Kidney Tumour Segmentation 2019 (KiTS19) Computed Tomography dataset, and then further evaluate the top three performing loss functions on the Brain Tumour Segmentation 2020 (BraTS20) Magnetic Resonance Imaging dataset. Motivated by the results of our study, we propose a Mixed Focal loss function, a new compound loss function derived from modified variants of the Focal loss and Focal Dice loss functions. We demonstrate that our proposed loss function is associated with a better recall-precision balance, significantly outperforming the other loss functions in both binary and multi-class image segmentation. Importantly, the proposed Mixed Focal loss function is robust to significant class imbalance. Furthermore, we showed the benefit of using compound losses over their component losses, and the improvement provided by the focal variants over other variants.

研究动机与目标

• 为解决医学图像分割中类别不平衡问题，特别是小肿瘤与大面积背景区域之间的不平衡问题。
• 评估多种基于Dice和交叉熵的损失函数在真实世界医学影像数据集上的鲁棒性与性能表现。
• 识别现有损失函数在严重类别不平衡情况下的局限性，并推动开发更有效的复合损失函数。
• 提出一种新型混合焦点损失函数，结合焦点损失与焦点Dice损失的优势，以提升分割性能。

提出的方法

• 作者首先在KiTS19 CT数据集上评估七种现有损失函数（包括交叉熵与基于Dice的变体），以识别表现最佳的候选函数。
• 随后，从表现最佳的三种损失函数中选取样本，并在BraTS20 MR影像数据集上进一步评估，以验证其在不同模态间的泛化能力。
• 所提出的混合焦点损失通过结合改进的焦点损失与改进的焦点Dice损失构建，旨在降低易样本（尤其是负样本）的影响，同时集中训练注意力于难分类、误分类的样本。
• 复合损失函数在模型训练过程中端到端优化，超参数经调优以在二分类与多分类分割任务中实现查全率与查准率的平衡。
• 该方法利用焦点调制机制，降低背景主导的易样本的影响，从而提升对罕见但关键结构（如肿瘤）的学习效果。

实验结果

研究问题

• RQ1 在医学图像分割中，标准交叉熵损失与基于Dice的损失函数在严重类别不平衡条件下表现如何？
• RQ2 结合焦点损失与焦点Dice损失的复合损失函数，是否能实现优于各独立组件的查全率-查准率平衡？
• RQ3 所提出的混合焦点损失是否能在不同成像模态（如CT与MR）之间实现良好泛化？
• RQ4 与标准损失相比，焦点机制在小而罕见的结构（如肿瘤）上的性能提升程度如何？
• RQ5 在类别不平衡分割任务中，复合损失函数是否比其组成部分更鲁棒、更有效？

主要发现

• 混合焦点损失函数在KiTS19与BraTS20数据集的二分类与多分类分割任务中，显著优于所有基线损失函数。
• 所提出的损失函数实现了更优的查全率与查准率平衡，尤其提升了对小而罕见肿瘤区域的检测能力，这些区域常被标准损失函数忽略。
• 复合损失函数对严重类别不平衡具有强鲁棒性，在肿瘤区域占总体积不足1%时仍能保持高性能。
• 焦点变体损失函数始终优于其非焦点对应版本，证实聚焦于难分类样本的益处。
• 研究结果表明，使用复合损失函数可带来可测量的性能提升，优于单独使用各组件损失函数，验证了混合焦点方法设计的有效性。