[论文解读] Multiclass MRI Brain Tumor Segmentation using 3D Attention-based U-Net

Maryann M. Gitonga|arXiv (Cornell University)|May 10, 2023

Advanced Neural Network Applications被引用 5

一句话总结

本文提出一种基于3D注意力机制的U-Net架构，将三种MRI模态（T2-FLAIR、T1CE、T2）融合为单一的4D输入，以提升多类别脑肿瘤分割性能。通过在跳跃连接中集成3D通道与空间注意力模块，模型增强了对肿瘤区域的关注，同时抑制了背景组织，最终在BraTS 2021测试集上达到0.9864的Dice系数，超越了先前的最先进方法。

ABSTRACT

This paper proposes a 3D attention-based U-Net architecture for multi-region segmentation of brain tumors using a single stacked multi-modal volume created by combining three non-native MRI volumes. The attention mechanism added to the decoder side of the U-Net helps to improve segmentation accuracy by de-emphasizing healthy tissues and accentuating malignant tissues, resulting in better generalization power and reduced computational resources. The method is trained and evaluated on the BraTS 2021 Task 1 dataset, and demonstrates improvement of accuracy over other approaches. My findings suggest that the proposed approach has potential to enhance brain tumor segmentation using multi-modal MRI data, contributing to better understanding and diagnosis of brain diseases. This work highlights the importance of combining multiple imaging modalities and incorporating attention mechanisms for improved accuracy in brain tumor segmentation.

研究动机与目标

• 通过在3D U-Net架构中集成注意力机制，提升多模态MRI中多区域脑肿瘤分割的准确性。
• 通过在训练过程中抑制无关背景特征，降低计算开销并提升模型泛化能力。
• 通过将T2-FLAIR、T1CE和T2图像堆叠为单个输入张量，实现对多模态MRI扫描的一次性分割。
• 通过使用α=0.7、β=0.3的Tversky损失函数，解决肿瘤分割中的类别不平衡问题。
• 证明3D注意力门在增强小而复杂肿瘤亚区域特征表示方面的有效性。

提出的方法

• 该模型采用3D U-Net结构，通过编码路径提取特征，通过解码路径实现精确分割，两者通过跳跃连接相连。
• 在跳跃连接中插入3D注意力门，基于通道与空间重要性动态加权特征图，使用查询（g）和键（x）向量进行计算。
• 注意力门通过3D卷积、ReLU激活函数和Sigmoid函数计算可学习的注意力图，以调制特征响应。
• 输入为形状为(3M × L × W × S)的堆叠4D张量，将三种非本征MRI模态（T2-FLAIR、T1CE、T2）融合为一个输入体积，以获得更丰富的特征表示。
• 网络使用Tversky损失（α=0.7，β=0.3）进行训练，以平衡精确率与召回率，并在BraTS 2021上通过Dice系数进行评估。
• 进行了两次训练试验：一次使用ReLU，一次使用LeakyReLU，以缓解死亡ReLU问题，结果显示LeakyReLU性能更优。

Figure 1: 3D Attention U-Net Architecture

实验结果

研究问题

• RQ1 与标准3D U-Net及其他最先进模型相比，基于3D注意力机制的U-Net是否能提升多类别脑肿瘤分割的准确性？
• RQ2 将三种MRI模态（T2-FLAIR、T1CE、T2）融合为单一输入体积，是否能增强特征表示并提升分割性能？
• RQ3 3D注意力机制在训练和推理过程中通过抑制无关背景特征，在多大程度上降低了计算负担？
• RQ4 在脑肿瘤分割任务中，α=0.7、β=0.3的Tversky损失函数在处理类别不平衡问题方面有多有效？
• RQ5 在该3D医学图像分割设置中，用LeakyReLU替代ReLU是否能改善模型收敛性与最终分割性能？

主要发现

• 所提出的3D注意力U-Net在BraTS 2021验证集第127个周期时达到0.9562的Dice系数，优于基线ReLU模型。
• 在测试集上，模型取得0.9864的Dice系数，显著超过标准3D U-Net在相同数据集上报告的0.92。
• 采用LeakyReLU激活的模型性能优于ReLU，测试集上Tversky损失为0.0136，Dice系数为0.9864。
• 注意力机制有效抑制了背景组织，增强了对肿瘤亚区域（包括坏死、水肿和增强肿瘤）的关注。
• 堆叠的多模态输入（T2-FLAIR、T1CE、T2）提供了更丰富的空间与强度特征，实现在单次前向传播中更精确的分割。
• 3D注意力门的集成提升了模型泛化能力，并通过在上采样过程中优先关注相关特征，降低了计算负载。

Figure 2: Visual representation of 3D Attention mechanism