[论文解读] Unified Focal loss: Generalising Dice and cross entropy-based losses to handle class imbalanced medical image segmentation

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

Retinal Imaging and Analysis被引用 9

一句话总结

本文提出统一焦点损失（Unified Focal Loss），一种分层框架，将Dice损失和交叉熵损失泛化，以更好地处理医学图像分割中的类别不平衡问题。通过结合焦点损失原理与可微分的自适应加权机制，该方法在五个不同且类别不平衡的数据集上实现了稳定的性能提升，优于标准损失函数和复合损失函数，在二分类与多分类任务中均表现更优。

ABSTRACT

Automatic segmentation methods are an important advancement in medical image analysis. Machine learning techniques, and deep neural networks in particular, are the state-of-the-art for most medical image segmentation tasks. Issues with class imbalance pose a significant challenge in medical datasets, with lesions often occupying a considerably smaller volume relative to the background. Loss functions used in the training of deep learning algorithms differ in their robustness to class imbalance, with direct consequences for model convergence. The most commonly used loss functions for segmentation are based on either the cross entropy loss, Dice loss or a combination of the two. We propose the Unified Focal loss, a new hierarchical framework that generalises Dice and cross entropy-based losses for handling class imbalance. We evaluate our proposed loss function on five publicly available, class imbalanced medical imaging datasets: CVC-ClinicDB, Digital Retinal Images for Vessel Extraction (DRIVE), Breast Ultrasound 2017 (BUS2017), Brain Tumour Segmentation 2020 (BraTS20) and Kidney Tumour Segmentation 2019 (KiTS19). We compare our loss function performance against six Dice or cross entropy-based loss functions, across 2D binary, 3D binary and 3D multiclass segmentation tasks, demonstrating that our proposed loss function is robust to class imbalance and consistently outperforms the other loss functions. Source code is available at: https://github.com/mlyg/unified-focal-loss

研究动机与目标

• 解决医学图像分割中持续存在的类别不平衡挑战，其中病灶或器官仅占图像体积的一小部分。
• 克服标准交叉熵损失和Dice损失的局限性，这些损失对类别不平衡敏感，可能导致收敛性差或预测偏差。
• 开发一种统一的、可微分的框架，泛化基于Dice和交叉熵的损失函数，以提升在类别不平衡情况下的鲁棒性与性能。
• 在多种医学影像任务中（包括2D二分类、3D二分类和3D多分类分割）证明所提损失的一致优越性。
• 提供一种灵活、即插即用的损失函数，可轻松集成到现有U-Net及类似架构中，无需修改网络结构。

提出的方法

• 提出一种分层损失框架，通过可微分的自适应加权机制统一Dice损失和交叉熵损失。
• 引入广义焦点损失组件，对各类别应用动态加权，以降低易样本的权重并聚焦于难样本。
• 将统一焦点损失表述为改进后的Dice损失与焦点交叉熵损失的加权和，其超参数可学习或固定。
• 使用伽马（γ）参数控制聚焦效果，在多分类任务中为所有类别统一使用γ=0.5以简化实现。
• 确保损失函数可微分，且与U-Net及类似编码器-解码器架构中的标准反向传播兼容。
• 在五个公开数据集上，使用标准U-Net和3D U-Net架构，对2D与3D分割网络进行训练时应用该损失函数。

实验结果

研究问题

• RQ1 能否设计一种统一的损失函数，泛化Dice和交叉熵损失，同时在医学图像分割中对类别不平衡具有鲁棒性？
• RQ2 在多样化且类别不平衡的医学影像数据集上，统一焦点损失与标准Dice损失、交叉熵损失及复合损失（如Combo损失）相比，性能如何？
• RQ3 所提出的损失是否在高度不平衡场景（如大器官中微小肿瘤分割）下保持梯度稳定并改善收敛性？
• RQ4 统一焦点损失能否泛化到包含多个稀有类别的多分类分割任务中，例如肾肿瘤分割？
• RQ5 统一焦点损失在不同成像模态和网络架构下的性能增益是否保持一致？

主要发现

• 在所有五个数据集（CVC-ClinicDB、DRIVE、BUS2017、BraTS20、KiTS19）上，统一焦点损失均一致优于六种基线损失函数（包括Dice、交叉熵、Combo损失和Tversky损失）。
• 在高度不平衡的KiTS19数据集中，统一焦点损失在肾脏和肾肿瘤分割任务中均取得了最高的DSC（Dice相似系数），甚至优于Combo损失。
• 在BraTS20数据集中，由于梯度不稳定，Dice损失和Tversky损失均失效，而统一焦点损失保持了稳定的训练过程，并实现了更优的分割质量。
• 定性结果表明，使用统一焦点损失生成的分割结果边界更清晰，且假阳性和假阴性更少。
• 该损失函数的训练时间与其它损失函数相比无显著差异，表明其计算高效且具备实际可部署性。
• 尽管对所有类别统一使用单一γ值，统一焦点损失仍优于其他复合损失，表明其对超参数配置具有鲁棒性。