[论文解读] Trustworthy Medical Segmentation with Uncertainty Estimation.

Giuseppina Carannante, Dimah Dera+3 more|arXiv (Cornell University)|Nov 10, 2021

Adversarial Robustness in Machine Learning参考文献 44被引用 7

一句话总结

该论文提出了一种用于医学图像分割的贝叶斯深度学习框架，通过一阶泰勒展开近似传播模型参数的均值与协方差，以最大化证据下界（evidence lower bound, ELBO），实现不确定性估计。该方法生成分割图与不确定性图，展现出在MRI和CT数据集上对噪声和对抗性攻击的优越鲁棒性，高不确定性值表明模型在输入受损或分布外时信心下降。

ABSTRACT

Deep Learning (DL) holds great promise in reshaping the healthcare systems given its precision, efficiency, and objectivity. However, the brittleness of DL models to noisy and out-of-distribution inputs is ailing their deployment in the clinic. Most systems produce point estimates without further information about model uncertainty or confidence. This paper introduces a new Bayesian deep learning framework for uncertainty quantification in segmentation neural networks, specifically encoder-decoder architectures. The proposed framework uses the first-order Taylor series approximation to propagate and learn the first two moments (mean and covariance) of the distribution of the model parameters given the training data by maximizing the evidence lower bound. The output consists of two maps: the segmented image and the uncertainty map of the segmentation. The uncertainty in the segmentation decisions is captured by the covariance matrix of the predictive distribution. We evaluate the proposed framework on medical image segmentation data from Magnetic Resonances Imaging and Computed Tomography scans. Our experiments on multiple benchmark datasets demonstrate that the proposed framework is more robust to noise and adversarial attacks as compared to state-of-the-art segmentation models. Moreover, the uncertainty map of the proposed framework associates low confidence (or equivalently high uncertainty) to patches in the test input images that are corrupted with noise, artifacts or adversarial attacks. Thus, the model can self-assess its segmentation decisions when it makes an erroneous prediction or misses part of the segmentation structures, e.g., tumor, by presenting higher values in the uncertainty map.

研究动机与目标

• 解决深度学习模型在临床环境中对噪声或分布外医学图像敏感导致的脆弱性问题。
• 克服现有分割模型仅提供点估计而缺乏不确定性量化的问题。
• 开发一种方法，通过参数分布矩估计预测不确定性，实现对分割决策的自我评估。
• 通过将高不确定性与错误预测或缺失结构（如肿瘤）相关联，提升深度学习模型在医学影像中的可信度。

提出的方法

• 采用编码器-解码器架构进行医学图像分割，并结合贝叶斯神经网络原理。
• 应用一阶泰勒级数展开，传播模型参数分布的一阶与二阶矩（均值与协方差）。
• 通过最大化证据下界（ELBO）学习给定训练数据的模型参数后验分布。
• 生成两个输出：预测的分割图与基于预测分布协方差矩阵生成的对应不确定性图。
• 利用预测分布的协方差矩阵量化不同空间位置上分割决策的不确定性。
• 在MRI和CT数据集上训练与评估该框架，以评估其在噪声和对抗性扰动下的鲁棒性。

实验结果

研究问题

• RQ1 贝叶斯深度学习框架是否能提升医学图像分割模型在噪声或受损输入下的鲁棒性？
• RQ2 通过参数分布矩进行的不确定性估计是否能有效识别分割预测中低置信度区域？
• RQ3 与标准模型相比，所提出的不确定性图在检测对抗性攻击或图像伪影方面表现如何？
• RQ4 当模型未能分割出如肿瘤等结构时，该框架是否能通过高不确定性区域突出显示自身预测的低置信度？
• RQ5 不确定性量化在多大程度上增强了深度学习模型在临床部署中的可信度？

主要发现

• 所提出的框架在基准MRI和CT数据集上，相较于最先进分割模型，展现出更强的噪声与对抗性攻击鲁棒性。
• 不确定性图在受噪声、伪影或对抗性扰动影响的图像块上始终分配更高的不确定性值。
• 当模型未能分割出如肿瘤等结构时，不确定性图在这些区域表现出升高的值，表明其具备对低置信度的自我评估能力。
• 预测分布的协方差矩阵能有效捕捉不同空间位置上分割决策的不确定性。
• 该框架的不确定性估计可实现对分布外输入的可靠检测，支持更安全的临床部署。
• 该方法在不增加推理时间或架构复杂度的前提下，仅通过贝叶斯参数估计即实现了分割可靠性的提升。