cs.CVSep 24, 2026

Lightweight Vision Transformer-Based U-Net for Brain Tumor Segmentation from MRI

Authors: Sheekar Banerjee, Md. Srabon Chowdhury, Md. Mahbub Hasan Akash, Ishtiak Al Mamoon

Organizations: Department of CSE, IUBAT, Dhaka, Bangladesh

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Accurate brain tumor segmentation from Magnetic Resonance Imaging is essential for diagnosis, treatment planning, and surgical guidance. Although Convolutional Neural Networks, particularly UNet, have achieved significant success in medical image segmentation, they often struggle to capture the long-range spatial dependencies required to model tumors with irregular shapes and complex boundaries. This paper proposes a lightweight Vision Transformer UNet that combines the hierarchical feature extraction capability of UNet with the global context modeling of Vision Transformers. The proposed architecture incorporates a compact ViT bottleneck within a U-Net encoder-decoder framework, enabling effective learning of both local and global features while maintaining computational efficiency with only 2.6 million trainable parameters. The model was evaluated on the TCGA LGG MRI Segmentation dataset, achieving a mean Intersection over Union of 0.8100 and a Dice score of 0.8446, outperforming the baseline UNet by 3.75% and 3.15%, respectively. Extensive quantitative and qualitative analyses, including confusion matrix evaluation, precision recall curves, per-image performance distribution, and tumor size dependency analysis, demonstrate the effectiveness and robustness of the proposed method for brain tumor segmentation.

Figures & tables

Fig. 1: Example MRI slices with corresponding ground truth tumor masks from the LGG dataset.
Fig. 2: Proposed ViT-UNet architecture. The encoder extracts hierarchical features, the ViT bottleneck captures global context, and the decoder reconstructs the segmentation map using skip connections.
ParameterValue
OptimizerAdamW
Learning Rate1×10−41\times 10^{-4}
Weight Decay1×10−51\times 10^{-5}
Batch Size8
SchedulerCosine Annealing ( Tmax⁡=15T_{\max}=15 )
Loss FunctionBCE + Dice ( α=0.5\alpha=0.5 )
TABLE I: Training Hyperparameters
MetricMeanStd Dev
IoU0.81000.3142
Dice0.84460.2949
Precision0.8051—
Recall0.8260—
Specificity0.9973—
Best Validation IoU0.6527—
TABLE II: Test set performance metrics
Fig. 3: Performance analysis summary showing various metrics and their distributions.
MethodIoUDiceParameters
U-Net [ 3 ]0.78070.81882.4M
TransUNet [ 6 ]0.77700.8210105M
Swin-UNet [ 7 ]0.79000.827027.1M
VT-UNet [ 14 ]0.80100.836025.8M
EF-VPT-Net [ 15 ]0.80700.842018.5M
Proposed ViT-UNet0.81000.84462.6M
TABLE III: Comparison of the proposed method with representative CNN- and Transformer-based segmentation models.
Fig. 4: Pixel-wise confusion matrix heatmap of the proposed model.
Fig. 5: Precision–Recall curve of the proposed model.
Fig. 6: IoU versus tumor area with a linear trend line.
Fig. 7: Distribution of per-image IoU and Dice scores using box and violin plots.
Fig. 8: Qualitative segmentation results: original image, ground truth, prediction, overlay.
Fig. 9: Best and worst segmentation predictions based on IoU score, illustrating the range of performance.
Model VariantViTHybrid LossIoUDice
U-Net Baseline×\times×\times0.78070.8188
U-Net + Hybrid Loss×\times✓\checkmark0.79240.8296
U-Net + ViT Bottleneck✓\checkmark×\times0.80310.8372
Proposed ViT-UNet✓\checkmark✓\checkmark0.81000.8446
TABLE IV: Ablation study of the proposed ViT-UNet architecture on the LGG dataset.
Fig. 10: Training and validation loss, IoU, and Dice over the training epochs.

Explore similar work

CardsList

  1. A Novel Global Context-aware Deep Neural Network for Enhanced Brain Tumor Segmentation using Magnetic Resonance Images

    May 28, 2026Sourjya Mukherjee, Ananya Bhattacharjee, R. MuruganBrain Tumor SegmentationMagnetic Resonance Imaging

    Biomedical Imaging Laboratory, Department of Electronics and Communication Engineering, National Institute of Technology Silchar, Silchar, Assam, India, 788010

  2. Uni-Light: An Ultra-Lightweight Framework via Uncertainty-Aware Knowledge Distillation for Brain Tumour Segmentation

    Date pendingLibing Kuang, Soren Salehi, Ziling Wu +2Brain Tumor SegmentationMagnetic Resonance Imaging

    School of Computer Science, University of Nottingham, Nottingham, UK · Electrical Engineering Department, Sharif University of Technology, Tehran, Iran · University of Pittsburgh, Pittsburgh, PA, USA

  3. From Convolution to Transformer: A Comparative Study of U-Net Variants for Brain Tumor and Retinal Vessel Segmentation

    Jun 20, 2026Khoa Pham, Sindhuja Penchala, Jiacheng Li +2U-NetVessel Segmentation

    Mississippi State University · The University of Alabama