[论文解读] Improving Calibration and Out-of-Distribution Detection in Medical Image Segmentation with Convolutional Neural Networks

Davood Karimi, Ali Gholipour|arXiv (Cornell University)|Apr 12, 2020

Radiomics and Machine Learning in Medical Imaging被引用 5

一句话总结

本文提出使用异质医学影像数据进行多任务学习，以改进基于CNN的分割模型在校准和分布外（OOD）检测方面的表现。通过在多种器官和影像模态上联合训练单一模型，实现了更优的校准预测，并提出一种基于CNN特征图的谱分析方法，其在OOD检测性能上优于基于不确定性的方法，在关键实验中达到100%准确率。

ABSTRACT

Convolutional Neural Networks (CNNs) have shown to be powerful medical image segmentation models. In this study, we address some of the main unresolved issues regarding these models. Specifically, training of these models on small medical image datasets is still challenging, with many studies promoting techniques such as transfer learning. Moreover, these models are infamous for producing over-confident predictions and for failing silently when presented with out-of-distribution (OOD) data at test time. In this paper, we advocate for multi-task learning, i.e., training a single model on several different datasets, spanning several different organs of interest and different imaging modalities. We show that not only a single CNN learns to automatically recognize the context and accurately segment the organ of interest in each context, but also that such a joint model often has more accurate and better-calibrated predictions than dedicated models trained separately on each dataset. Our experiments show that multi-task learning can outperform transfer learning in medical image segmentation tasks. For detecting OOD data, we propose a method based on spectral analysis of CNN feature maps. We show that different datasets, representing different imaging modalities and/or different organs of interest, have distinct spectral signatures, which can be used to identify whether or not a test image is similar to the images used to train a model. We show that this approach is far more accurate than OOD detection based on prediction uncertainty. The methods proposed in this paper contribute significantly to improving the accuracy and reliability of CNN-based medical image segmentation models.

研究动机与目标

• 解决在小样本数据上训练的基于CNN的医学图像分割模型存在的校准性差和预测不可靠的问题。
• 通过在推理阶段检测分布外（OOD）数据来提升模型可靠性，因为模型在此类情况下常无声失效。
• 探究在异质数据集（不同器官、不同模态）上进行联合训练是否能相比专用模型提升泛化能力与校准性。
• 开发一种比现有基于不确定性的方法更准确的 OOD 检测新方法，用于医学图像分割。

提出的方法

• 在涵盖多种器官和影像模态的异质医学影像数据集混合体上训练单一CNN模型，不引入模态或器官的额外输入信息。
• 基于最后卷积层特征图的功率谱，对特征图进行谱分析，以计算 OOD 分数（OODM）。
• 在 OODM 上设定阈值 τ，利用分布内数据进行校准，将测试图像分类为分布内或 OOD。
• 使用 AUC、准确率、灵敏度和特异性等指标，将所提出的 OODM 方法与基于预测不确定性的检测方法进行比较。
• 在多个数据集上联合训练模型（例如，MRI 中的肝脏、胎儿 MRI 中的皮层板），以评估泛化能力和校准性能。
• 使用 Dice 相似性系数（DSC）和校准度量（如 ECE）来评估分割准确率和置信度校准性。

实验结果

研究问题

• RQ1 在异质医学影像数据（不同器官、不同模态）上训练的单一CNN模型，其分割性能是否可与在单一数据集上训练的专用模型相当或更优？
• RQ2 在异质数据上进行多任务学习是否相比在单一数据集上训练能带来更好的预测校准性？
• RQ3 对CNN特征图进行谱分析是否能有效检测医学图像分割中的分布外（OOD）图像？
• RQ4 所提出的基于谱分析的 OOD 检测方法在准确率上相比基于不确定性的 OOD 检测方法表现如何？
• RQ5 该模型是否能在无显式模态或器官信息的情况下，泛化到不同数据分布（如不同年龄组、不同成像协议）？

主要发现

• 在异质数据集（如 MRI 中的肝脏、胎儿 MRI 中的皮层板）上训练的单一CNN模型，其分割性能与在单一数据集上训练的专用模型相当或更优。
• 在异质数据上训练的模型表现出显著更优的校准预测能力，表现为更低的期望校准误差（ECE）。
• 基于CNN特征图谱分析的 OOD 检测方法在肝脏分割任务中，对 Liver-MRI-SPIR 数据的 OOD 检测实现了 100% 准确率、100% 灵敏度、100% 特异性，AUC 达到 1.00。
• 相比之下，基于不确定性的 OOD 检测方法在相同实验中仅达到 64% 准确率、61% 灵敏度、60% 特异性，AUC 为 0.65。
• 分布内数据（Liver-MRI-DUAL-In 和 Liver-MRI-DUAL-Out）的 OODM 值始终低于阈值 τ=0.0057，而 OOD 数据（Liver-MRI-SPIR）的 OODM 值显著更高，实现了完美分离。
• 该模型在测试时成功分割了 Liver-MRI-DUAL-Out 数据集（DSC=0.886），且其 OODM 值分布与训练时的分布内数据相似，证实了对分布偏移的鲁棒性。