3D Medical Axial Transformer: A Lightweight Transformer Model for 3D Brain Tumor Segmentation

Cheng Liu, Hisanor Kiryu

Keywords: Deep learning, 3D brain tumor segmentation, 3D Transformer, axial attention, self-distillation

TL;DR: A 3D brain tumor image segmentation model based on Axial Attention structure with a novel approach for self-distillation learning

Abstract: In recent years, Transformer-based models have gained attention in the field of medical image segmentation, with research exploring ways to integrate them with established architectures such as Unet. However, the high computational demands of these models have led most current approaches to focus on segmenting 2D slices of MRI or CT images, which can limit the ability of the model to learn semantic information in the depth axis and result in output with uneven edges. Additionally, the small size of medical image datasets, particularly those for brain tumor segmentation, poses a challenge for training transformer models. To address these issues, we propose 3D Medical Axial Transformer (MAT), a lightweight, end-to-end model for 3D brain tumor segmentation that employs an axial attention mechanism to reduce computational demands and self-distillation to improve performance on small datasets. Results indicate that our approach, which has fewer parameters and a simpler structure than other models, achieves superior performance and produces clearer output boundaries, making it more suitable for clinical applications.