MIA | BayeSeg: Bayesian Modeling for Medical Image Segmentation with Interpretable Generalizability

MIA | BayeSeg: Bayesian Modeling for Medical Image Segmentation with Interpretable Generalizability

论文标题： BayeSeg: Bayesian Modeling for Medical Image Segmentation with Interpretable Generalizability
论文发表期刊： Medical Image Analysis
论文地址： https://arxiv.org/abs/2303.01710
论文代码： https://zmiclab.github.io/projects.html

零、总结

为了探索医学图像域内不变特征的可解释性，提出了可解释泛化性的贝叶斯架构(BayeSeg)。该架构分为图像和标签的统计建模、后验深度推理两部分。首先，对图像进行变量分解，分离出表示形状的变量和表示图像外观的变量，将分割变量限制仅与形状变量相关，并联合图像和其对应的分割标签进行显式统计建模。接着，使用变分贝叶斯框架来近似这些可解释变量的后验分布，并用神经网络实现该框架的训练与推断。因为通过概率图模型显式建模的变量都有实际的物理含义，医生可以通过可视化不同变量的后验分布来判断BayeSeg的决策过程是否正确，模型具有良好的可解释性和泛化能力。该论文获得爱思唯尔出版社、医学影像分析期刊和国际MICCAI学会联合颁发的最佳论文奖，成为第一篇以中国大陆研究机构为第一单位的获奖论文。

一、摘要

由于不同医学成像系统引起的跨域分布转移，导致许多深度学习分割方法在未见过的数据上表现不佳，限制了其在实际应用中的适用性。最近的研究表明，提取域内不变性表示对领域的泛化有好处。然而，域内不变特征(domain-invariant features)的可解释性仍然是一个巨大的挑战。为了解决这个问题，我们提出了一个可解释的贝叶斯框架(BayeSeg)，通过贝叶斯图像建模和标签统计来增强模型在医学图像分割中的泛化性。具体来说，我们首先将图像分解为空间相关变量和空间可变变量，分配分层级的贝叶斯先验来明确地迫使它们分别对域稳定形状和域特定外观信息建模。然后，我们将分割建模为仅与形状相关的局部平滑变量。最后，我们开发了一个变分贝叶斯框架来推断这些可解释变量的后验分布。该框架是用神经网络实现的，因此被称为深度贝叶斯分割。在前列腺分割和心脏分割任务上的定量和定性实验结果表明了该方法的有效性。此外，我们通过解释后验来考察BayeSeg的可解释性，并通过进一步的消融研究分析影响BayeSeg泛化能力的因素。

二、引言

医学图像分割的目的是基于医学成像技术对不同器官的解剖结构进行分类。现代影像技术可以提供多种临床征象，包括全面的病理和结构信息。然而，这对医学图像分割提出了新的挑战。具体而言，由于特定成像系统中的潜在因素，来自不同序列、模式和位置的医学图像的强度可能会有很大差异，如扫描仪、视场、空间分辨率、信噪比和数字处理软件。这些因素可能导致医学图像之间的分布变化，这使得跨序列、跨模态和跨场所分割特别具有挑战性。尽管最近的深度学习方法在计算机视觉和医学图像分析方面取得了很好的表现，但它们容易受到内在分布转移的影响。也就是说，在一个数据集上训练的深度神经网络(DNN)通常不能很好地推广到另一个未见过的数据集上。此外，为DNN收集大量标记的医学图像来防止过拟合的成本很高，因为需要大量的劳动力和有资质的专家。因此，寻求在小数据集上训练可泛化的模型成为医学图像分割的关键问题之一。为了探索可泛化的方法，一个被称为域泛化(domain generalization)的新领域在计算机视觉和医学图像分析中越来越受到关注。

目前的域泛化方法可以分为三类，即基于数据的、基于学习的和基于表示的。基于数据的方法试图通过数据增强或数据泛化来丰富训练数据的多样性。这些方法对于小的分布位移是有效的，但是对于显著的域间隙是脆弱的。基于学习的方法旨在通过特定策略组合来自训练域知识，包括集成学习和元学习。这些方法在特定的目标域上取得了很好的效果，但源域数量的限制会削弱它们的泛化能力。从而开发了基于表示的方法，通过特征对齐提取域不变特征或特征解缠。当前基于表示的方法致力于通过对齐隐藏特征的映射来表示域不变信息。尽管这些方法在计算机视觉中提供了有希望的泛化能力，但很少关注模型的可解释性，直观地说，缺乏以可理解的方式向人类专家提供解释的能力。这可能会引起对临床诊断和治疗中伦理和法律的关注。因此，探索可解释性和可推广的表征对医学图像分割更有吸引力。

统计建模能够提高模型的可解释性和泛化能力。受此启发，我们探索了形状信息的显式建模，而不是隐式表示，它代表了器官的结构。如图1所示，两个不同的MR序列之间存在明显的分布偏移，但通过提取形状信息可以很好地减小这种分布偏移。也就是说，提取的形状更有可能是一个域不变的表示，并且从形状中分割具有提高泛化性的潜力。因此，我们将联合描述图像和标签统计，以提取可解释的形状表示，用于泛化分割。

在这项工作中，我们通过联合建模图像和标签统计提出了一个可解释和可推广的贝叶斯分割(BayeSeg)框架。具体来说，我们首先将图像表示为两个变量。一种是表示形状信息的空间相关模型，另一种是表示外观信息的空间变异模型。由于这两个变量都是无约束的，我们进一步分配层次贝叶斯先验来显式描述它们的统计信息。之后，由于难以计算，我们利用变分推理来近似它们的后验。最后，我们建立了一个深度神经网络来实现变分方法，并从给定的图像中推断出形状、外观和分割的后验分布。

这项工作是我们在MICCAI 2022会议论文上的延伸，其中我们提出了一种新的贝叶斯分割框架(BayeSeg)，通过图像和标签统计的联合建模，并通过神经网络实现的变分贝叶斯方法来求解模型。我们的新贡献如下：

1. 我们通过统计描述分割与其真值之间的关系进一步改进了BayeSeg的公式。
2. 我们对心脏和前列腺分割进行了广泛的实验，包括跨模态、跨序列和跨场所设置。此外，我们澄清了分布转移的临床背景，并对域泛化的发展进行了综述。
3. 我们通过可视化和解释从BayeSeg提取的后验来阐述模型的可解释性。此外，我们通过选择统计先验、剪枝概率模型和分析数据分布来研究模型的泛化性。

本文的其余部分组织如下。在第二节中，我们介绍和总结了统计建模和域泛化的相关文献。在第3节中，我们介绍了方法论，包括公式、网络架构和训练策略。第4节提供了实验设置，并验证了BayeSeg的泛化性。在第5节中，我们研究了所提出方法的可解释性和概括性。第6节讨论了BayeSeg的局限性和我们未来的工作。最后，在第7节中对本文进行了总结。

三、方法

这项工作开发了一个可解释的贝叶斯分割(BayeSeg)框架，以增强深度学习模型的可泛化性。该框架主要由两部分组成，即(1)图像和标签的统计建模、(2)后验深度推理。BayeSeg可以有效地增强模型的泛化性，因为它是图像和标签统计的联合建模。虽然传统的变分贝叶斯方法可用于求解由(1)导出的统计模型，但由于数千次在线迭代，它们在推断高维变量时计算成本很高。受深度神经网络学习复杂映射效率的启发，我们提出了通过深度学习来推断后验的方法。为方便起见，本文使用的概念和符号总结于表1。

图2显示了所提出BayeSeg的统计建模和深度推理框架。对于统计建模，我们建立如图2(a)所示的概率图形模型(PGM)。具体而言，我们首先将图像y分解为两个变量。一种是模拟解剖结构信息的空间相关形状x，另一种是模拟图像纹理信息的空间变异外观a。由于器官的结构与域无关，从形状预测标签的模型可以有更好的泛化性。为了分层表征图像和标签的统计特性，形状x以其边界υ为条件，用于检测结构的边缘和分割z，而外观a被设置为由空间变异均值m和空间变异逆方差ρ决定的高斯变量。类似地，分割z被迫局部平滑，并依赖于它的边界w所有类的标记部分π。此外，基于真值标签u被建模为一个依赖于分割的变量，它遵循伯努利分布。最后，将Gamma先验分配给ρ，υ和w，将Beta先验分配给π，将高斯先验分配给m。对于深度推理，我们构建了如图2(b)所示的深度学习框架。首先，我们使用两个CNN分别推断形状和外观的后验分布。其次，我们使用CNN从形状估计分割的后验分布。最后，为了训练这些CNN，将联合使用后验参数来计算最大后验的变分损失，并将其与最大标签似然的交叉熵损失相结合。

3.1 图像和标签的统计建模

图像和标签的统计建模旨在利用概率方法探索其内部统计信息，这在医学图像分割中很有用。许多先前的研究表明，图像统计在图像恢复、图像分类和图像分割中很重要，但纹理图像的建模仍然具有挑战性。医学图像分割不仅需要对图像统计进行描述，还需要对标签统计进行表征。因此，我们提出将图像建模和标签统计相结合用于医学图像分割。

表2显示了m，ρ，υ，w和π的选择先验。具体来说，我们分配高斯先验于m，分配Gamma先验于ρ，υ和w，分配Beta先验于π。这里，1表示一个所有元素都为1的向量，I表示一个单位矩阵。注意，对于π，我们选择Beta分布而不是Dirichlet分布，因为前者是(3)的共轭先验，这可以大大简化我们下面的公式。通过上述统计建模，BayeSeg可以提供可解释的表示，提高医学图像分割的泛化能力。

3.2 图像和标签的变分推理

3.3 变分损失的展开

3.4 神经网络和训练策略

四、实验

在本节中，我们通过两个任务评估BayeSeg框架在跨序列、跨场所和跨模态场景下的领域泛化性能，即第4.2节中的前列腺分割和第4.3节中的心脏分割。此外，我们在第4.4节中进行了消融研究。

4.1 数据集

对于前列腺分割，我们使用了来自NCI-ISBI 2013、I2CVB和PROMISE12三个公开数据集的T2前列腺MRI图像，如表4所示。在预处理过程中，我们将所有图像重新采样到0.36458×0.36458 mm的固定间距，并裁剪出每个图像直方图的前0.5%。然后，将三维图像重新格式化为二维切片，过滤掉不含前列腺区域的切片。切片被中间裁剪为384×384并调整大小为192×192。采用Z-score归一化和随机数据增强，包括仿射变换、弹性变换和加性高斯噪声。

在心脏分割方面，我们的目标是利用分割模型来解析左心室(LV)、右心室(RV)和心肌(Myo)三种结构。如表4所示，我们使用了5个数据集中的6个域，即MSCMR，EMIDEC，ACDC，MMWHS和CASDC 2013。特别是在专家的监督下，我们自己对EMIDEC RV和CASDC的真值图进行了注释。我们已经筛选和过滤了切片，以便在数据集之间对齐z轴分辨率。然后，将所有剩余的切片重新采样到1.36719 × 1.36719 mm的固定像素间距，然后裁剪或填充到212×212，以关注心脏的ROI。我们使用旋转、缩放和仿射变换来增强数据，以减少训练过程中的过拟合。通过z-score归一化，所有切片归一化到平均值为零，标准差为1。注意，CASDC的CT图像在归一化前的Housefield值通过[-200,300]窗口进行了对比度增强。

4.2 前列腺分割

为了研究跨场所场景下模型的泛化性，我们使用由6个域组成的3个前列腺分割数据集进行训练和测试，如表4所示。在训练阶段，我们将来自RUNMC的数据作为源域，将其30个案例随机分成3个组，21个案例用于训练，3个案例用于验证，6个案例用于测试。我们在前列腺分割中配置了广泛使用的EfficientNet-b2作为我们的分割主干。Adam对网络进行了1200 epoch的优化，初始学习率为3e−4，在第1000 epoch时学习率降至3e−5。我们在验证集上选择最佳模型进行测试。在测试过程中，我们将BMC、BIDMC、HK、UCL和I2CVB的数据作为目标域。为了进行比较，我们选择了五种方法，包括标准经验风险最小化和其他四种域泛化方法，即Cutout、IBN-Net、RandConv和DSU。值得注意的是，为了进行公平的比较，这些方法与我们的方法使用相同的分割主干和训练策略来实现。之后，我们计算每个受试者的Dice得分(0-100)，以评估这些方法的性能。

表5报告了Dice得分的平均值和标准差。总的来说，BayeSeg的平均Dice得分为77.5，比第二好的方法高出7.7。对于带有野偏置场的BIDMC数据，六种方法中有四种无法分割前列腺，因为它们的Dice得分的平均值和标准差具有接近的数值。在先进的方法中，我们的方法超过了20种第二好的方法，性能下降最小。对于轻微偏置场的HK数据，在5个先进的方法中，BayeSeg的性能下降幅度是最小的。在其他场所中，所有方法都获得了可比较的结果，性能差距很小，这意味着目标和源之间的域转移较少。我们可视化了图3 (a)所示的三个场所的中位数和最差情况。可以观察到，只有两种方法(DSU和BayeSeg)在所有情况下都成功。但是，在I2CVB的最坏情况下，DSU会产生一个断开的分割，这违反了前列腺的解剖结构。相反，BayeSeg保留所有情况下的连接分割结果，受益于空间相关建模。

4.3 心脏分割

除了跨场所场景，我们进一步探索了模型在跨序列和跨模态设置下的通用性，使用五个心脏分割数据集，包括六个域，如表4所示。在训练阶段，我们将MSCMR的LGE作为源域，将其45个案例随机分成3个组，25个案例用于训练，5个案例用于验证，15个案例用于测试。之后，我们遵循会议论文中的设置，选择U-Net作为分割主干。最后，使用Adam优化器对BayeSeg进行了2000次的训练，学习速率为1e−4，并在验证集中选择了最优模型进行测试。在测试中，我们将MSCMR的T2、EMIDEC的LGE、ACDC的bSSFP、MMWHS的bSSFP和CASDC的CT作为靶结构域。为了进行比较，我们采用了与4.2节中前列腺分割相同的比较方法和评价方案。

为了定量地评估模型的通用性，我们在表6中报告了平均值、标准差和Dice得分的下降。从表中可以看出，BayeSeg在目标域上的性能与源域最接近，Dice只下降了5.5%。特别是对于交叉序列场景(MSCMR的T2)，所有对应方法得到的Dice分数都小于36.3，其中性能下降最小的高于44.5%。相比之下，BayeSeg在MSCMR的T2域中表现最佳，超越了第二好的方法，其性能下降了9.9%。对于跨场所(EMIDEC的LGE)和跨模态(CASDC的CT)场景，BayeSeg分别比次优方法高出10.2%和3.6%。为了进行定性评估，我们在图3 (b)中分别可视化了代表跨序列、跨场所和跨模态的三个领域的典型案例。从图中可以看出，在6个域中的其中5个，所提出的BayeSeg提供了相当或更好的性能。然而，在EMIDEC LGE的最坏情况下，所有方法都失败了，这表明我们的方法仍有改进的空间。

4.4 消融实验

为了探讨BayeSeg随机映射和变分损失的有效性，我们在本小节中对心脏分割任务进行了消融研究。特别地，我们评估了不同BayeSeg设置的泛化性。为了减少随机性的影响，当训练过程超过1000个epoch时，我们每100个epoch进行一次测试，通过平均10次评估得到最终结果。这种严格的方法使我们能够彻底分析单个组件对BayeSeg整体泛化性的影响。

如表7所示，变分损失是影响BayeSeg泛化性的关键因素。该形状借助变分损失对感兴趣器官的域不变结构信息进行建模，增强了在不可见域上分割的泛化性。此外，可以看到单独引入随机映射并不能带来显著的性能改进。然而，值得注意的是，使用随机映射是应用变分损失的先决条件。在变分损失的指导下，随机映射使BayeSeg能够看到给定图像的形状方面的无限样本，从而也有助于提高在不可见域的性能。

五、结论

在这项工作中，我们提出了一个贝叶斯分割框架(BayeSeg)，通过图像和标签统计的联合建模来提高医学图像分割的可解释性和泛化能力。具体来说，我们将图像分解为外观和形状，其中分配分层级的贝叶斯先验，迫使外观和形状分别建模特定于领域的外观和领域稳定的形状信息。然后，通过标签统计建模，从形状中生成分割预测。此外，我们还提出了一种由神经网络实现的变分方法来推断后验。在前列腺分割和心脏分割任务上的定量和定性实验结果表明了该方法的有效性。此外，我们进一步解释了BayeSeg提取的整体后验，并通过消融研究验证了显式先验建模和图像与标签统计联合建模对提高泛化能力的好处。