Self Pre-training with Topology- and Spatiality-aware Masked Autoencoders for 3D Medical Image Segmentation

Masked Autoencoders（MAEs）已被证明在自然和医学图像分析问题中，对于Vision Transformers（ViTs）的预训练是有效的。通过重建可见补丁中的缺失像素/体素信息，ViT编码器可以聚合下游任务的上下文信息。但是，现有的MAE预训练方法，这些方法是专门针对ViT架构开发的，缺乏捕捉几何形状和空间信息的能力，这对于医学图像分割任务至关重要。在本文中，我们提出了一种新的已知MAE的扩展方法，用于自我预训练（即在相同目标数据集上预训练的模型）用于3D医学图像分割。我们提出了一种新的拓扑损失来保留几何形状信息，通过计算输入和重建体积的拓扑特征来学习几何形状信息。我们引入了一个预文本任务，用于预测3D裁剪的中心和八个角的位置，使MAE可以聚合空间信息。我们将MAE预训练策略扩展到混合的最先进的医学图像分割架构，并与ViT一起进行预训练。我们通过将预训练的ViT编码器与我们的预训练SOTA模型相结合，开发了一个用于下游分割任务的精调模型。对五个公共3D分割数据集的广泛实验表明了我们新方法的有效性。

本文旨在提出一种新的MAE预训练方法，用于3D医学图像分割，该方法能够捕获几何形状和空间信息。

本文提出了三个关键思路：1.引入拓扑损失来保留几何形状信息；2.引入预文本任务来预测3D裁剪体的中心和八个角的位置，以便MAE聚合空间信息；3.将MAE预训练策略扩展到混合的SOTA医学图像分割架构中，并与ViT一起进行共同预训练。

本文通过实验验证了新方法的有效性，并在五个公共3D分割数据集上进行了广泛的实验。本文的创新点包括：拓扑损失的引入、预文本任务的使用以及混合SOTA医学图像分割架构的扩展。此外，作者还提供了开源代码。

最近的相关研究包括：《An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》、《Unsupervised Learning of Visual Representations by Solving Jigsaw Puzzles》等。