[论文解读] Combo Loss: Handling Input and Output Imbalance in Multi-Organ Segmentation

Saeid Asgari Taghanaki, Yefeng Zheng+6 more|arXiv (Cornell University)|May 7, 2018

Radiomics and Machine Learning in Medical Imaging被引用 4

一句话总结

本文提出 Combo Loss，一种新颖的损失函数，通过在课程学习框架中结合 Dice 损失和加权交叉熵，联合解决多器官医学图像分割中的输入和输出类别不平衡问题。该方法在训练过程中可显式控制假阳性和假阴性的权衡，实现更快的收敛速度，并在 PET、MRI 和超声图像数据集上达到最先进性能，Dice 分数相比 Focal Loss 最高提升 10.8%。

ABSTRACT

Simultaneous segmentation of multiple organs from different medical imaging modalities is a crucial task as it can be utilized for computer-aided diagnosis, computer-assisted surgery, and therapy planning. Thanks to the recent advances in deep learning, several deep neural networks for medical image segmentation have been introduced successfully for this purpose. In this paper, we focus on learning a deep multi-organ segmentation network that labels voxels. In particular, we examine the critical choice of a loss function in order to handle the notorious imbalance problem that plagues both the input and output of a learning model. The input imbalance refers to the class-imbalance in the input training samples (i.e., small foreground objects embedded in an abundance of background voxels, as well as organs of varying sizes). The output imbalance refers to the imbalance between the false positives and false negatives of the inference model. In order to tackle both types of imbalance during training and inference, we introduce a new curriculum learning based loss function. Specifically, we leverage Dice similarity coefficient to deter model parameters from being held at bad local minima and at the same time gradually learn better model parameters by penalizing for false positives/negatives using a cross entropy term. We evaluated the proposed loss function on three datasets: whole body positron emission tomography (PET) scans with 5 target organs, magnetic resonance imaging (MRI) prostate scans, and ultrasound echocardigraphy images with a single target organ i.e., left ventricular. We show that a simple network architecture with the proposed integrative loss function can outperform state-of-the-art methods and results of the competing methods can be improved when our proposed loss is used.

研究动机与目标

• 解决深度多器官分割中输入不平衡（前景体素稀少）和输出不平衡（假阳性与假阴性）的双重挑战。
• 开发一种损失函数，使在训练过程中能显式、可学习地控制假阳性与假阴性之间的权衡。
• 在包括 PET、MRI 和超声在内的多种成像模态中，提升高度不平衡医学影像数据的分割性能。
• 提供一种即插即用的损失函数，无需架构修改即可增强现有模型，提升泛化能力和收敛性。
• 通过可微分的自适应损失直接惩罚假阳性和假阴性，减少对后处理的依赖。

提出的方法

• Combo Loss 将 Dice 相似性系数（DSC）和加权交叉熵（WCE）结合在一个可微分的损失函数中，以稳定训练并避免陷入不良局部极小值。
• 损失函数的公式为：L_combo = α × (1 - DSC) + (1 - α) × WCE，其中 α 控制 DSC 与 WCE 项之间的平衡。
• 在 WCE 项中引入第二个超参数 β，以显式控制假阳性与假阴性惩罚之间的权衡。
• 通过逐步增加 WCE 项的影响，采用课程学习策略，以提升优化稳定性和收敛速度。
• 该损失函数在编码器-解码器网络（如 3D U-Net、V-Net）中端到端应用，无需架构修改。
• 通过交叉验证对超参数 α 和 β 进行调优，以在每个数据集上实现最佳性能。

实验结果

研究问题

• RQ1 统一的损失函数能否有效处理多器官分割中的输入与输出不平衡问题？
• RQ2 结合 Dice 损失与加权交叉熵在不平衡医学图像数据上的模型收敛性和性能方面有何提升？
• RQ3 通过损失函数中的可学习超参数，假阳性与假阴性之间的权衡能在多大程度上实现控制？
• RQ4 所提出的 Combo Loss 是否在多种成像模态中优于现有的损失函数（如交叉熵、Dice 损失和 Focal Loss）？
• RQ5 Combo Loss 是否可作为即插即用的替代方案，无需架构修改即可提升现有分割网络的性能？

主要发现

• 在 PET 数据集中，与标准交叉熵损失相比，Combo Loss 将假阳性率降低了 34.6%，假阴性率降低了 27.8%。
• 在 MRI 数据集中，与 Focal Loss 相比，当应用于 3D V-Net 时，Combo Loss 将 Dice 分数提高了 2.3%，假阴性率降低了 47.4%。
• 在超声图像数据集中，当应用于 3D U-Net 时，Combo Loss 相比 Focal Loss 的 Dice 分数提升了 10.8%。
• 最小化假阳性和假阴性时，最优 β 值在超声图像中为 0.7，在 MRI 中为 0.6，证明了超参数控制的有效性。
• 在 U-Net 和 V-Net 中替换原始损失函数后，Combo Loss 将 MRI 的 Dice 分数从 0.87 提升至 0.91，超声图像中从 0.85 提升至 0.92，显示出一致的性能增益。
• 该方法收敛速度优于标准交叉熵损失，并降低了 Hausdorff 距离中的异常值性能，表明分割更加稳健。