Brain-SAM: A SAM-based Model Tailored for Brain MRI Lesion Segmentation

Brain-SAM 技术摘要

问题陈述
磁共振成像（MRI）是神经影像学的基石，然而对脑结构和病变的精确分割仍然具有挑战性。虽然像 Segment Anything Model (SAM) 及其继任者 SAM2 这样的基础模型在自然图像领域展示了卓越的泛化能力，但其在脑部 MRI 中的应用面临着显著障碍。现有的医学适配模型（如 MedSAM、SAM-Med3D）通常是在包含心脏、肝脏和肺部等结构的混合器官数据集上进行训练的。这些模型难以适应大脑独特的解剖学特征，例如复杂的组织结构、模糊的灰白质边界以及高度多样化的病变形态。此外，大多数基于 SAM 的适配方案将 3D 体积扫描转换为 2D 切片，忽略了切片间的空间上下文信息。此外，对手动交互式提示（prompt）的依赖限制了其可扩展性，因为临床医生可能难以为癫痫病灶等细微病理提供准确的提示。最后，原始 SAM 的掩码解码器（mask decoder）是为具有清晰边界的 2D 自然图像设计的，在处理具有模糊边缘和组织异质性的脑部 MRI 数据时往往表现不佳。

方法论
作者提出了 Brain-SAM，这是一个专门为 3D 脑部 MRI 病变分割量身定制的模型。该架构通过三个主要的修改扩展了 SAM2 框架：

1. 图像编码器的 3D 适配： 将原始的 2D Hiera 编码器（一种分层视觉 Transformer）进行了适配，使其能够直接处理 3D 体积数据。这涉及扩展四个关键模块：

   • 块嵌入（Patch Embedding）： 用 3D 卷积取代 2D 卷积，将体积输入划分为体素块（voxel patches）。
   • 位置嵌入（Position Embedding）： 从 2D 扩展到 3D 表示，以捕捉深度、高度和宽度方向的空间关系。
   • Token 序列（Token Sequence）： 沿所有三个维度展平 3D 体积，以创建适合体积注意力机制的更长 Token 序列。
   • Q-Pooling： 用 MaxPool3d 取代 MaxPool2d，以实现跨深度和空间维度的特征下采样。

2. 受 UNETR 启发的解码器： 为了解决原始 SAM 解码器在处理精细局部纹理和边界细节方面的局限性，作者引入了一个受 UNETR 启发的解码器。该组件通过多尺度特征融合，将 Transformer 的全局建模能力与 CNN 的局部精细表示相结合。与需要分辨率转换的标准 UNETR 解码器不同，该设计利用 Hiera 编码器固有的分层特征进行直接融合和渐进式上采样。该架构支持双头输出：一个用于全自动分割，另一个用于生成基于提示的分割所需的中间特征。

3. 训练与微调策略：

   • 预训练： 模型使用 SAM2 权重进行初始化，并在 BraTS2021 数据集（1,000 例，4 种模态）上进行预训练阶段，专门针对自动分割任务以建立稳健的医学表示。
   • 基于适配器（Adapter）的微调： 为了在适应多种脑部病变数据集（肿瘤、中风、癫痫）的同时防止灾难性遗忘，作者在每个 Hiera 块之前插入了轻量级的适配器模块（带有 3 维卷积的瓶颈架构）。在微调期间，骨干网络保持冻结，仅更新共享的适配器参数。
   • 提示机制： 使用点提示（point prompts）作为主要的交互方式，并将其扩展到 3D。掩码生成策略经过修改：不再使用点积运算中的更新后的 Key，而是将处理后的 Query 与来自类 UNETR 解码器的中间特征表示进行融合。
   • 精炼： 后处理精炼机制利用用户点击，通过 3D 高斯核来局部调整掩码概率。

核心贡献

1. SAM2 的 3D 扩展： 本研究成功地将 SAM2 框架从 2D 自然图像扩展到了 3D 医学图像分析，实现了直接的体积处理，而无需进行逐切片近似。
2. 混合解码器架构： 引入的类 UNETR 解码器使模型能够在统一的架构内支持全自动和交互式分割，有效地捕捉了局部精细细节和全局上下文依赖。
3. 全面的数据集策划： 作者策划了一个涵盖脑肿瘤、中风和癫痫的大规模数据集，使模型能够支持广泛的临床应用。

结果
研究人员在涵盖六种 MRI 模态的七个公共数据集（BraTS-PED, BraTS-SSA, BraTS-MET, BraTS-MEN, ISLES2022, ATLAS v2.0, OPMRI）上进行了广泛实验。Brain-SAM 与特定任务模型（nnU-Net, UNETR, Swin UNETR）以及基础模型（SAM-Med3D, VISTA3D）进行了对比。

• 整体性能： Brain-SAM 取得了优于所有基准模型的 Dice 分数和更低的 95% 豪斯多夫距离（HD95）。具体而言，与 nnU-Net 相比，它在癫痫病灶上的 Dice 提升了 22%，在脑转移瘤上提升了 9%，在脑膜瘤上提升了 6%。
• 小病灶分割： 该模型在分割小病灶（如癫痫）方面展现出显著优势，Dice 分数比强基准模型高出 15%–18%。与其它方法相比，它在处理小目标（<5,000 体素）时，在召回率和精确度之间保持了更好的平衡。
• 消融实验：
  • 将 3D-ViT 骨干网络替换为 3D-Hiera 后，小病灶的 Dice 分数提升了高达 14%，并显著降低了 HD95。
  • 集成类 UNETR 解码器后，与不含该组件的版本相比，平均 Dice 分数提升了 8%。
  • 基于适配器的微调策略被证明是非常高效的参数利用方式（39.4M 可调参数，而全量微调为 72.02M），同时达到了与全量微调相当的性能。

意义与主张
论文声称 Brain-SAM 为下游脑部 MRI 分析任务提供了一个稳健的预训练模型。通过解决脑组织特有的解剖学和结构性挑战（如模糊边界和小尺寸病灶），它比现有的通用基础模型提供了更可靠、更准确的分割。作者认为，这种能力可以支持更好的临床决策，特别是在癫痫和脑肿瘤等的早期诊断和治疗规划方面。该模型兼具自动模式和提示引导模式的功能，旨在减少人工标注工作量，同时保持高精度。作者指出，目前的研究局限于单模态数据，并建议未来的工作可以探索多模态集成，以进一步增强病变表征能力。