[论文解读] MIS-FM: 3D Medical Image Segmentation using Foundation Models Pretrained on a Large-Scale Unannotated Dataset

Guotai Wang, Jianghao Wu+4 more|arXiv (Cornell University)|Jun 29, 2023

Radiomics and Machine Learning in Medical Imaging被引用 23

一句话总结

MIS-FM 引入 Volume Fusion，一种用于3D医学影像分割的自监督预训练策略，利用未注释数据来对整个分割模型进行预训练，以及 Parallel Convolution and Transformer Network (PCT-Net) 用于在下游任务的有效迁移。

ABSTRACT

Pretraining with large-scale 3D volumes has a potential for improving the segmentation performance on a target medical image dataset where the training images and annotations are limited. Due to the high cost of acquiring pixel-level segmentation annotations on the large-scale pretraining dataset, pretraining with unannotated images is highly desirable. In this work, we propose a novel self-supervised learning strategy named Volume Fusion (VF) for pretraining 3D segmentation models. It fuses several random patches from a foreground sub-volume to a background sub-volume based on a predefined set of discrete fusion coefficients, and forces the model to predict the fusion coefficient of each voxel, which is formulated as a self-supervised segmentation task without manual annotations. Additionally, we propose a novel network architecture based on parallel convolution and transformer blocks that is suitable to be transferred to different downstream segmentation tasks with various scales of organs and lesions. The proposed model was pretrained with 110k unannotated 3D CT volumes, and experiments with different downstream segmentation targets including head and neck organs, thoracic/abdominal organs showed that our pretrained model largely outperformed training from scratch and several state-of-the-art self-supervised training methods and segmentation models. The code and pretrained model are available at https://github.com/openmedlab/MIS-FM.

研究动机与目标

• 激励在大规模未标注的3D医学图像上对分割模型进行预训练，以在注释数据有限时提升性能。
• 提出 Volume Fusion 作为自监督伪分割前置任务，在无需人工注释的情况下产生成对的输入和分割标签。
• 介绍 PCT-Net，一种结合并行卷积与 Transformer 区块的网络，用于高效迁移到多尺度分割任务。
• 展示在头颈、胸部和腹部器官分割的效果，结合大规模预训练。
• 发布对 110k CT 扫描的预训练模型，以实现更广泛的临床迁移。

提出的方法

• Volume Fusion (VF): 将一个扫描中的前景子体积与另一个扫描中的背景子体积进行融合，使用离散融合系数图，将体素级融合类别转化为伪分割前光任务。
• 融合系数 alpha_i 来自离散集合 {0, 1/K, ..., 1}，为体素创建 C = K+1 类；模型预测对应于 alpha_i 的体素类别 Y。
• 预训练目标 L_sup 将 Dice 损失和交叉熵损失结合起来，在伪分割任务上端到端训练完整分割模型。
• PCT-Net: 三层金字塔结构，利用双分支 PCT 块（局部卷积分支和全局自注意力分支）将局部 CNN 特征与长程 Transformer 上下文融合。
• Embedding 模块在高分辨率使用二维卷积，在较低分辨率使用三维卷积，以处理各向异性的3D CT数据；在多尺度产生带深度监督的预测。
• 训练设置使用 110k 未标注的 CT 体积（PData-110k）进行预训练，随后在下游分割数据集上进行微调，使用标准 Dice 和 ASSD 指标。

Figure 1 : Overview of our proposed self-supervised pretraining method based on volume fusion.

实验结果

研究问题

• RQ1 VF 预训练是否在下游3D CT任务上提升分割性能，相比从头开始训练和其他 SSL 方法？
• RQ2 融合参数 K 如何影响迁移性能和在 VF 预训练中的上下文学习？
• RQ3 预训练数据规模（1k、10k、110k 体积）对下游分割性能有何影响？
• RQ4 PCT-Net 是否能在不同尺度的器官（头颈、胸部、腹部）上高效利用 VF 预训练的特征？

主要发现

• VF 预训练在 SegTHOR 的胸部器官分割上始终优于从头训练和其他 SSL 方法，Dice 达到 88.30%，ASSD 1.78 mm。
• 在 MICCAI 2015 Head-Neck 数据集上，使用 VF 的 PCT-Net 达到更高的平均 Dice(82.74) 和更低的 ASSD(0.77 mm)，优于 nnU-Net、TransUNet、UNETR++ 等基线。
• 当 K=4 时，VF 取得最佳权衡，达到最低 ASSD(1.78 mm) 且 Dice 表现对比测试数据集具有竞争力；较大 K 值未必总是提升性能。
• 提升预训练数据规模（1k → 10k → 110k）通常提升下游 Dice，对 3D U-Net 和 PCT-Net 都有收益；使用 PCT-Net 架构时收益更大。
• 与 Patch Swapping、Model Genesis 和 MIM 相比，VF 在 SegTHOR 上获得最高的平均 Dice(88.30%) 和最低的 ASSD(1.78 mm)。
• VF+PCT-Net 相较于单独 VF 或基线架构，展现出下游结果的改进，体现了更好的迁移与上下文学习。

Figure 2 : 2D illustration of fused images with different $K$ values. The colorful images are fusion coefficient maps, and they are converted into label maps for the pretext segmentation task.