AnatoMask: Enhancing Medical Image Segmentation with Reconstruction-guided Self-masking 论文阅读

摘要:由于标记数据的稀缺性，自监督学习(self-supervised learning, SSL)通过从未标记数据中提取语义表示，在三维医学图像分割中受到了广泛关注。在SSL策略中，掩码图像建模(MIM)通过重建随机掩码图像来学习细节表示，显示了其有效性。然而，传统的MIM方法需要大量的训练数据才能达到良好的性能，这对医学成像仍然是一个挑战。由于随机掩模对医学图像中的所有区域进行均匀采样，可能会忽略关键的解剖区域，从而降低预训练的效率。我们提出了一种新的MIM方法AnatoMask，它利用重建损失来动态识别和掩盖解剖上重要的区域，以提高预训练效果。
AnatoMask采用自蒸馏方法，模型学习如何找到更重要的区域来屏蔽，以及如何重建这些被屏蔽的区域。为了避免次优学习，Anatomask使用掩蔽动力学函数逐步调整预训练难度。我们在4个具有多种成像模式(CT、MRI和PET)的公共数据集上评估了我们的方法。与现有的SSL方法相比，AnatoMask展示了卓越的性能和可扩展性。
代码可在https://github.com/ricklisz/AnatoMask上获得。
关键词:医学图像分割·自监督学习·蒙面图像建模

1介绍

虽然监督学习在3D医学图像分割中显示出巨大的前景，但由于缺乏标记数据，其在临床环境中的潜力受到阻碍。

作为替代方案，自监督学习通过从未标记数据中提取语义表示在医学图像分析中显示出巨大的前景[23,35,48,50]。在各种SSL方法中，掩膜图像建模在自然图像中取得了很好的效果[16]，并在医学图像中得到了应用

分析[8,9,49]。MIM利用一个简单的任务，要求模型重建被屏蔽图像的部分，从而学习图像的详细表示[16,45]。

然而，MIM在自然图像中的成功在很大程度上依赖于大规模的训练数据[46]，这对医学成像提出了重大挑战。

虽然许多MIM方法利用随机掩蔽策略并在自然图像中取得了成功[3,16]，但我们认为这种策略没有利用医学图像中解剖区域的固有属性，导致预训练效果较差。在医学图像中，解剖区域，如器官和肿瘤，被认为比不太相关的充满空气和充满液体的区域更重要。由于解剖区域的特征是不同的形状和对比度，这些区域在被掩盖时直观地更难以重建，需要在预训练期间进行更多的指导。因此，以往的方法需要更多的训练数据才能达到预期的效果，这极大地限制了它们在医学相关任务中的应用。因此，为了设计一个有效的医学图像MIM框架，识别这些复杂的解剖区域是至关重要的，并创建一个更具挑战性的任务。

为了解决这些挑战，我们提出了AnatoMask，这是一个重建引导的MIM框架，模型从更多的信息掩模中学习，以提高预训练效果。观察到解剖区域通常具有较高的重建损失(图1 a和b)，如果我们可以增加这些区域的屏蔽概率，也会鼓励网络更有效地从这些区域提取信息。因此，我们建议通过重建损失来识别解剖学上重要的区域。我们不是使用预训练的模型来识别这些区域，而是提出一种自蒸馏方法，其中模型可以首先充当教师来生成具有解剖学意义的面具，然后充当学生来从这些面具中学习。为了确保教师能够适应学生的表现，我们使用学生的指数移动平均(EMA)权重来更新教师模型。这样，我们不仅鼓励模型学习如何重建这些区域，而且还鼓励模型学习在哪里找到这些重要的区域。这种双重学习方法对于发展更深层次的解剖学理解至关重要，特别是在数据受限的情况下。最后，我们提出了一种易-难掩蔽策略，即随着学生的进步，教师不断增加MIM任务的难度，避免模型一开始过度关注困难区域。我们对TotalSegmentator[40]进行了自监督学习，TotalSegmentator是一个用于体积医学图像分割的综合CT数据集。我们采用以前的SOTA架构STUNet[20]，并通过对TotalSegmentator[40]、FLARE22[26]、AMOS22[22]和AutoPETII[11]进行微调来评估预训练效果。我们首先比较了随机掩蔽和我们的AnatoMask之间的差异(图1c和d)，然后比较了AnatoMask与随机掩蔽(SparK[37])的训练效率(图1e)。我们的方法不断提高训练效率，优于随机掩蔽。我们还将我们的方法与以前最先进的(SOTA) SSL方法进行了比较。我们发现我们的方法表现出卓越的性能和扩展能力。