Hybrid Multihead Attentive Unet-3D for Brain Tumor Segmentation

TL;DR

Brain tumor segmentation remains challenging due to heterogeneity and multimodal MRI variability. The authors propose Hybrid Multihead Attentive U-Net 3D, integrating multihead attention into a 3D U-Net to capture long-range spatial dependencies across modalities. On BraTS 2020, the approach outperforms baselines such as SegNet, FCN-8s, and Dense121 U-Net, achieving high Dice coefficients and robust metrics, with class-specific DSC reaching near-perfect values in some tests. The work demonstrates strong potential for clinical impact in diagnosis, treatment planning, and monitoring, while acknowledging the need for validation on larger datasets and real-world deployment.

Abstract

Brain tumor segmentation is a critical task in medical image analysis, aiding in the diagnosis and treatment planning of brain tumor patients. The importance of automated and accurate brain tumor segmentation cannot be overstated. It enables medical professionals to precisely delineate tumor regions, assess tumor growth or regression, and plan targeted treatments. Various deep learning-based techniques proposed in the literature have made significant progress in this field, however, they still face limitations in terms of accuracy due to the complex and variable nature of brain tumor morphology. In this research paper, we propose a novel Hybrid Multihead Attentive U-Net architecture, to address the challenges in accurate brain tumor segmentation, and to capture complex spatial relationships and subtle tumor boundaries. The U-Net architecture has proven effective in capturing contextual information and feature representations, while attention mechanisms enhance the model's ability to focus on informative regions and refine the segmentation boundaries. By integrating these two components, our proposed architecture improves accuracy in brain tumor segmentation. We test our proposed model on the BraTS 2020 benchmark dataset and compare its performance with the state-of-the-art well-known SegNet, FCN-8s, and Dense121 U-Net architectures. The results show that our proposed model outperforms the others in terms of the evaluated performance metrics.

Figures (10)

• Figure 1: Visualization of Brain Tumor Segmentation using NeuroLearn Library.The figure showcases different visualizations of the input medical images and corresponding segmentation masks from the BraTS 202020 dataset. These multiple rows in the figure highlight the representation of a FLAIR image in different masks i.e. epi, anat and roi.
• Figure 2: Data preprocessing workflow, step by step approach from data loading to saving the new dataset as numpy volumes.
• Figure 3: U-Net 3D + MHA architecture for brain tumor segmentation. It features an encoding and decoding path with multihead attention modules to define features and a softmax layer to produce a segmentation map highlighting different tumor regions. Implemented using TensorFlow and Keras, the model accurately identifies and analyzes brain tumors from 128x128x128 pixel input images with three channels.
• Figure 4: Comparison of the training and validation accuracies with respective losses at LR=0.001 with a batch size of 16 at 50 epochs.
• Figure 5: Comparison of the training and validation accuracies with respective losses at LR=0.001 with a batch size of 8 at 50 epochs.