[论文解读] Marginal loss and exclusion loss for partially supervised multi-organ segmentation

Gonglei Shi, Xiao Li+2 more|arXiv (Cornell University)|Jul 8, 2020

Advanced Neural Network Applications参考文献 46被引用 6

一句话总结

该论文提出边缘损失和排除损失，以利用仅部分标注的数据集训练单一的多器官分割网络，其中每幅图像仅标注了部分器官。边缘损失将背景建模为未标注器官与真实背景的合并标签，通过边缘概率实现有效训练；排除损失则强制器官之间的互斥性。该方法在五个基准数据集上实现了最先进性能，且无需额外计算，显著提升了Dice分数并降低了Hausdorff距离。

ABSTRACT

Annotating multiple organs in medical images is both costly and time-consuming; therefore, existing multi-organ datasets with labels are often low in sample size and mostly partially labeled, that is, a dataset has a few organs labeled but not all organs. In this paper, we investigate how to learn a single multi-organ segmentation network from a union of such datasets. To this end, we propose two types of novel loss function, particularly designed for this scenario: (i) marginal loss and (ii) exclusion loss. Because the background label for a partially labeled image is, in fact, a `merged' label of all unlabelled organs and `true' background (in the sense of full labels), the probability of this `merged' background label is a marginal probability, summing the relevant probabilities before merging. This marginal probability can be plugged into any existing loss function (such as cross entropy loss, Dice loss, etc.) to form a marginal loss. Leveraging the fact that the organs are non-overlapping, we propose the exclusion loss to gauge the dissimilarity between labeled organs and the estimated segmentation of unlabelled organs. Experiments on a union of five benchmark datasets in multi-organ segmentation of liver, spleen, left and right kidneys, and pancreas demonstrate that using our newly proposed loss functions brings a conspicuous performance improvement for state-of-the-art methods without introducing any extra computation.

研究动机与目标

• 解决在多个数据集中仅存在部分标注时，训练鲁棒多器官分割模型的挑战。
• 开发能有效利用部分标注医学图像中有限标注信息的损失函数。
• 在不增加计算成本的前提下，提升对胰腺和肾脏等小而复杂的器官的分割准确率与鲁棒性。
• 实现在五个基准数据集联合数据集上的端到端联合训练，其标注完整性各不相同。

提出的方法

• 边缘损失通过将背景标签建模为边缘概率推导得出，即在计算损失前，对所有未标注器官与真实背景的预测概率求和。
• 边缘损失可与标准损失函数（如交叉熵或Dice损失）结合，形成边缘交叉熵损失或边缘Dice损失。
• 引入排除损失，通过测量标注器官与未标注器官预测区域之间的不相似性，强制实现器官间的互斥性。
• 排除损失作为正则化项，惩罚预测结果的重叠，利用了解剖结构天然不重叠的特性。
• 两种损失函数联合使用，并在基于标准U-Net主干网络（如nnUNet）的训练过程中应用，实现在部分标注数据上的端到端学习。
• 该方法无需额外推理计算，且与现有深度学习框架完全兼容。

实验结果

研究问题

• RQ1 边缘损失能否有效处理部分标注医学图像中因合并背景标签而引入的模糊性？
• RQ2 通过排除损失强制实现互斥性，在多器官分割设置中如何提升分割性能？
• RQ3 将边缘损失与排除损失结合，能在多大程度上提升对胰腺和肾脏等小而难分割器官的性能？
• RQ4 能否仅通过这两种新型损失函数，有效训练单一分割网络，以处理五个部分标注基准数据集的联合数据？

主要发现

• 所提方法在所有器官上的平均Dice系数达到0.931，显著优于先前方法（PaNN为0.906，PIPO为0.907）。
• 与PaNN相比，该方法将平均Hausdorff距离降低了24.0%；与PIPO相比，降低了40.0%，表明边界精度更优。
• 仅使用排除损失时，对胰腺和肾脏的性能提升微乎其微，表明其作为辅助损失在与边缘损失结合时效果最佳。
• 随着标注数量减少，性能持续下降，证实了模型性能对标注量的敏感性。
• 可视化结果表明，排除损失能有效减少对邻近结构的错误器官预测，尤其在小器官中效果显著。
• 失败案例揭示了在区分肝脏与相似背景区域，以及脾脏与胰腺的误分类方面仍存在挑战，提示引入形状先验可进一步提升性能。