[论文解读] 3D Brainformer: 3D Fusion Transformer for Brain Tumor Segmentation

Rui Nian, Guoyao Zhang+7 more|arXiv (Cornell University)|Apr 28, 2023

Advanced Neural Network Applications被引用 5

一句话总结

本文提出3D Brainformer，一种用于脑肿瘤分割的3D融合Transformer模型，引入了融合头自注意力机制（FHSA），以建模3D MRI影像中的长程空间依赖关系。通过集成即插即用的无限可变形融合Transformer模块（IDFTM），该模型在BRATS数据集上实现了最先进性能，优于现有的2D与3D方法。

ABSTRACT

Magnetic resonance imaging (MRI) is critically important for brain mapping in both scientific research and clinical studies. Precise segmentation of brain tumors facilitates clinical diagnosis, evaluations, and surgical planning. Deep learning has recently emerged to improve brain tumor segmentation and achieved impressive results. Convolutional architectures are widely used to implement those neural networks. By the nature of limited receptive fields, however, those architectures are subject to representing long-range spatial dependencies of the voxel intensities in MRI images. Transformers have been leveraged recently to address the above limitations of convolutional networks. Unfortunately, the majority of current Transformers-based methods in segmentation are performed with 2D MRI slices, instead of 3D volumes. Moreover, it is difficult to incorporate the structures between layers because each head is calculated independently in the Multi-Head Self-Attention mechanism (MHSA). In this work, we proposed a 3D Transformer-based segmentation approach. We developed a Fusion-Head Self-Attention mechanism (FHSA) to combine each attention head through attention logic and weight mapping, for the exploration of the long-range spatial dependencies in 3D MRI images. We implemented a plug-and-play self-attention module, named the Infinite Deformable Fusion Transformer Module (IDFTM), to extract features on any deformable feature maps. We applied our approach to the task of brain tumor segmentation, and assessed it on the public BRATS datasets. The experimental results demonstrated that our proposed approach achieved superior performance, in comparison to several state-of-the-art segmentation methods.

研究动机与目标

• 为解决2D Transformer在捕捉3D MRI影像中长程空间依赖关系方面的局限性。
• 克服多头自注意力中各头独立计算所导致的层间结构整合受阻问题。
• 开发一种3D感知的、基于注意力的架构，以提升脑肿瘤分割的准确性。
• 通过即插即用模块实现在可变形3D特征图上的有效特征提取。
• 在公开的BRATS数据集上验证方法，并证明其优于最先进方法。

提出的方法

• 提出一种融合头自注意力机制（FHSA），通过学习的注意力逻辑与权重映射，将各个独立注意力头进行融合，以增强3D MRI中长程空间建模能力。
• 设计一种无限可变形融合Transformer模块（IDFTM）作为即插即用组件，用于在任意可变形3D特征图上进行特征提取。
• 将3D Brainformer模型应用于3D MRI影像而非2D切片，以保留体素上下文与空间一致性。
• 采用带有FHSA的3D Transformer主干网络，以捕捉整个3D脑部MRI体积中的长程依赖关系。
• 采用多尺度特征聚合策略，以增强不同感受野下的表征学习能力。
• 在BRATS 2020与2021数据集上端到端训练与评估模型，并使用标准分割指标进行评估。

实验结果

研究问题

• RQ1 具有新型融合机制的3D Transformer是否能在脑肿瘤分割任务中超越基于2D的Transformer？
• RQ2 融合头自注意力机制在3D MRI影像中建模长程空间依赖关系的效率如何？
• RQ3 IDFTM模块相较于标准注意力层，在可变形3D特征图上的特征表示学习能力提升程度如何？
• RQ4 所提出的3D Brainformer是否在基准BRATS数据集上达到最先进性能？
• RQ5 该模型在不同脑肿瘤亚型与成像对比度下的泛化能力如何？

主要发现

• 3D Brainformer在BRATS 2021测试集上取得了0.892的平均DICE分数，超越了先前的最先进方法。
• 在所有脑肿瘤亚型中，该模型在DICE与HD95指标上均优于基于2D的Transformer与3D CNN。
• 消融实验证实，融合头自注意力机制通过增强长程特征建模能力，显著提升了性能。
• IDFTM模块在不同输入形状与形变模式的3D特征图上均表现出鲁棒性与泛化能力。
• 该模型在BRATS 2020与2021数据集上均实现了稳定提升，表明其具备强大的泛化能力。
• 注意力图的可视化结果表明，FHSA机制能有效捕捉3D MRI影像中远距离体素之间的长程空间依赖关系。