[论文解读] Stochastic Segmentation Networks: Modelling Spatially Correlated Aleatoric Uncertainty

Miguel Monteiro, Loïc Le Folgoc+6 more|arXiv (Cornell University)|Jun 10, 2020

Explainable Artificial Intelligence (XAI)参考文献 48被引用 79

一句话总结

随机分段网络（SSNs）在对数空间中使用低秩多变量正态分布建模分割图的联合分布，允许多种空间相关的假设并在二维和三维医学图像中实现高效推断。

ABSTRACT

In image segmentation, there is often more than one plausible solution for a given input. In medical imaging, for example, experts will often disagree about the exact location of object boundaries. Estimating this inherent uncertainty and predicting multiple plausible hypotheses is of great interest in many applications, yet this ability is lacking in most current deep learning methods. In this paper, we introduce stochastic segmentation networks (SSNs), an efficient probabilistic method for modelling aleatoric uncertainty with any image segmentation network architecture. In contrast to approaches that produce pixel-wise estimates, SSNs model joint distributions over entire label maps and thus can generate multiple spatially coherent hypotheses for a single image. By using a low-rank multivariate normal distribution over the logit space to model the probability of the label map given the image, we obtain a spatially consistent probability distribution that can be efficiently computed by a neural network without any changes to the underlying architecture. We tested our method on the segmentation of real-world medical data, including lung nodules in 2D CT and brain tumours in 3D multimodal MRI scans. SSNs outperform state-of-the-art for modelling correlated uncertainty in ambiguous images while being much simpler, more flexible, and more efficient.

研究动机与目标

• 动机：在图像分割中捕捉固有不确定性的需求，即存在多种可能的边界（例如医学成像）
• 提出一个轻量级的概率框架，在不改变网络架构的情况下建模联合标签图分布。
• 实现从单次前向传播生成多个时空上相干的分割假设。
• 展示在二维CT肺结节和三维多模态MRI脑肿瘤分割任务中的适用性。

提出的方法

• 将对数几率图建模为低秩加对角的多变量正态分布：η|x ~ N(μ(x), Σ(x))，其中 Σ = P P^T + D。
• 将 μ(x)、P 和 D 参数化为神经网络的输出，从而实现高效计算和可扩展性。
• 在对数空间中用蒙特卡洛采样近似对 η 的不可积积分。
• 使用重参数化技巧通过蒙特卡洛样本反向传播，并采用类似交叉熵的目标函数。
• 推断在一次前向传播中生成对数几率图的多个样本；平均对数几率提供最可能的分割。

实验结果

研究问题

• RQ1 通过低秩对数空间分布对空间相关的本质不确定性进行建模，是否能够产生连贯的多种分割假设？
• RQ2 相对于确定性和现有概率分割方法，SSN 在医学成像任务上是否提高了预测性能和不确定性校准？
• RQ3 该方法如何扩展到高维的3D数据并与常用分割骨干网络结合？

主要发现

• SSNs 在建模模糊医学图像的相关不确定性方面优于最先进方法。
• 在肺结节分割任务中，低秩协方差模型在预测性能和不确定性校准方面优于对角协方差变体与确定性基线。
• SSNs 能在 2D 和 3D 中实现多种合理分割的采样，几乎不需要改变架构，且计算开销适中。
• 在 BraTS 2017 三维脑肿瘤分割中，随机模型在 Dice 分数上与确定性模型相当，同时提供一个分割的分布。
• 该框架允许推断后对样本进行操作（例如调整类别存在性）以及基于温度的多样性控制。