Revisiting MAE pre-training for 3D medical image segmentation
本文发表在CVPR-25。
研究目的
自监督训练在3D医学影像里主要受三个限制：
- 预训练数据集规模过小：我们项目中有足够的数据。
- 使用的架构过时或不适合3D医学影像：重点研究以适配项目。
- 评估不充分
本文旨在对这三部分进行标准化改进来使MAE预训练适合3D医学影像病显著提升3D-CNN分割模型的性能。
Datasets
预训练数据集
专有的大脑MRI数据集，该数据集来自超过44个中心，10多种不同MR扫描仪，包含超过9000名患者，约44000张MRI扫描（最后剩余39168张）。
保留T1、T2、T1 FLAIR和T2 FLAIR四个序列的图像。
丢弃任何轴上视野<50 mm的图像，在任何方向上间距> 6.5 mm的图像和文件大小<200 kb的图像（空图像）。
训练数据集
多发性硬化（MS）病变数据集：MS FLAIR（T2W）
脑转移肿瘤数据集：Brain Mets（T1、T2F）
海马体数据集：Hippocampus （T1W）
卒中后病变数据集：Atlas22（T1W）
耳道内和耳道外前庭神经鞘瘤肿瘤数据集：CrossModa（T1W）
训练：验证：测试 = 64：16：20
测试数据集
8个数据集包含颈动脉血管壁、风险器官（OAR）的分割、缺血性卒中病变、前口咽癌和转移性淋巴结、尼日利亚胶质母细胞瘤和高级别胶质瘤成像、脑动脉瘤及其周围脑组织等等。
测试数据集被用来评估当训练微调后分割其他目标结构时，学习的表征的功效。
研究内容
MAE
Masked Autoencoders Are Scalable Vision Learners
MAE采用Transformer架构，实际上就是Mask 掉输入图像的随机的 patches 并进行重建（推荐掩码率75%），但是编码器和解码器是非对称的。
MAE 的具体实现方法是：
首先通过 Linear Projection 和位置编码得到 image tokens。
随机 shuffle 这些 tokens，按照 masking ratio 扔掉最后的一部分。
把 unmasked patches 输出到 Encoder 中，得到这些 tokens 的表征。
把 Encoder 的输出，结合 masked tokens（视为此处需要恢复，后续也是计算mask部分loss），执行 unshuffle操作恢顺序，再一起输入到 Decoder 中。
3DMAE即把VIT改为3D-VIT，每个patch变为
nxnxn大小后平铺输入到transformer中。相较于常用于分割的3D-UNet相关的网络，ViT结构是比较好做MAE的。
3D nnU-Net
nnU-Net实际上使用简单的U-Net架构，主要将复杂的手动方法配置流程系统化为固定参数、基于数据集属性的规则参数以及最少的经验参数进行优化，其处理方法流程如下：
- 从训练数据中提取数据集特征，包括图像大小、体素间距、模态等信息。
- 根据数据集特征,使用一系列经验规则自动配置分割流水线的参数，如预处理、网络拓扑结构、训练策略等。
- 将未涵盖的少量参数(如模型选择、后处理)设置为在训练中经验优化。
- 使用默认的U-Net网络架构模板训练多个模型。
- 通过交叉验证,从训练好的多个模型中经验选择表现最好的单模型或模型组合。
- 如果后处理能提高验证性能,则应用后处理。
MAE 3D nnU-Net
为了在3D-CNN架构使用MAE训练，本文采取了几个措施：
训练前，所有3D MRI图像都会被重采样到统一的物理空间间距
[1x1x1]毫米，并进行Z-score归一化（减均值，除标准差），每个样本固定使用**[160x160x160]大小**的输入。针对CNN的掩码与稀疏化处理：CNN需要规则的网格输入，无法像ViT那样直接丢弃掩码token。为此，论文引入了一套稀疏化处理流程：
稀疏卷积：在编码器部分，使用稀疏卷积进行操作，卷积核在计算时，其覆盖范围内可能同时包含真实像素和被掩码的零值区域。随着网络层数的加深，这些零值会通过感受野“污染”有效特征的计算。 所以需要在每一层卷积操作之后，都重新应用输入时生成的掩码。
掩码令牌：对被掩码的区域填充可学习的掩码令牌，以保持空间结构信息。
致密化卷积：在编码器和解码器之间，应用一个小的**
3x3x3卷积**。（除了最高分辨率外）利用周围有效像素的信息，对掩码令牌所在的区域进行初步的、基于上下文的“填充”或“平滑”，使特征图在空间上更加连续。消融结果
新增的稀疏卷积、掩码令牌和致密化卷积三部进行消融实验，三个部分全部应用后有一定的效果。
实验结果
对比结果
冻结编码器权重是有害的
应该适当减少学习率
Warm-up是必不可少的
和原文推荐的mask ratio差不多，都在60%到75%之间效果最好，本文采用60%-90%的动态mask ratio，和75%相差不大。
S3D超过了development datasets所有的Baseline方法。
最终评估采用Dice相似系数和归一化表面距离来评估，在development和test datasets中的表现（no Dyn.是单独适应每个下游数据集的原始nnU-Net，No Fixed表示使用固定配置进行从零开始训练的nnU-Net）
No Fixed vs No Dyn.：体现了nnU-Net自动化配置流程本身的价值有多大
No Dyn. vs S3D：体现了自监督预训练是否能带来额外的的性能提升
对于动态nnU-Net表现不佳的大多数数据集中，预训练有助于恢复性能。然而，在某些情况下，例如D5，固定配置的nnU-Net仍然是最好的。
ABCD数据集训练结果
使用ABCD dataset（青少年大脑认知发展多模态MRI数据）训练的模型平均DSC低0.6%以及平均NSD低0.6%，说明虽然MAE在一般计算机视觉任务中可以拓展，但由于医学影像领域的独特性，迁移到下游分割任务时可能泛化能力下降
S3D使用40个左右小样本进行微调就和所有样本从头开始训练的效果差距不大（D4在这里差距比较大，因为样本更多）
D138/10/12，D267/17/21，D3166/42/52，D4419/105/131，D5134/34/42
更长的预训练并不会导致性能的提高，在250k steps时反而取得最好效果。
五折验证结果
复制预训练的权重并在解码器预热阶段期间冻结结果最稳定且同样很好。
在不到正常训练时间的15%之后，就可以实现最终性能的很大一部分。