an interpretable vision transformer framework for automated brain tumor classification

1. Problem Statement

Brain tumors are among the most critical neurological conditions, where early and accurate diagnosis directly impacts patient survival. While Magnetic Resonance Imaging (MRI) is the gold standard for detection, manual interpretation by radiologists is:

• Time-intensive and cognitively demanding.
• Subject to inter-observer variability, especially for small or atypically located tumors.
• Limited by resource scarcity, as specialist neuroradiology expertise is often unavailable in resource-limited settings.

Existing deep learning solutions, primarily Convolutional Neural Networks (CNNs), rely on local receptive fields, which limits their ability to model the long-range structural dependencies often diagnostically relevant in brain MRI. Furthermore, many existing models lack clinical interpretability, making it difficult for clinicians to trust or audit AI-driven diagnoses.

2. Methodology

The authors propose a comprehensive deep learning framework based on the Vision Transformer (ViT-B/16) architecture, enhanced with specific preprocessing, augmentation, and training strategies tailored for medical imaging.

A. Dataset and Preprocessing

• Dataset: A publicly available dataset of 7,023 MRI scans divided into four classes: Glioma (23.1%), Healthy (28.5%), Meningioma (23.4%), and Pituitary Tumor (25.0%).
• Split: Stratified 80/10/10 split (Train/Val/Test) preserving class proportions.
• Preprocessing (CLAHE): The authors apply Contrast Limited Adaptive Histogram Equalization (CLAHE) to all images before training. Operating on local 8×8 tiles in the CIE LAB color space (luminance channel only), CLAHE enhances local contrast and accentuates tumor boundaries that are often invisible in standard normalization, addressing the low-contrast nature of MRI acquisitions.

B. Model Architecture

• Backbone: ViT-B/16 (Vision Transformer Base), pretrained on ImageNet-21k.
• Mechanism: Input images (224×224) are split into 196 non-overlapping 16×16 patches. These are projected into 768-dimensional tokens with a learnable [CLS] token and positional embeddings.
• Encoder: 12 Transformer Encoder blocks utilizing Multi-Head Self-Attention (12 heads) to capture global pairwise relationships between all image patches simultaneously.
• Head: A classification head consisting of a Dense layer (256 units, GELU), Dropout (0.3), and a final Dense layer (4 classes).

C. Training Strategy

The framework employs a sophisticated two-stage fine-tuning and regularization pipeline:

1. Two-Stage Fine-Tuning:
   • Stage 1 (Warm-up): Backbone frozen; only the classification head is trained for 5 epochs to prevent catastrophic forgetting.
   • Stage 2 (Full Fine-tuning): All parameters unfrozen with discriminative learning rates (1e-5 for backbone, 1e-4 for head) and Cosine Annealing decay.
2. Augmentation:
   • Pixel-level: Random flips, rotations (±15°), affine translations, zoom, and contrast jitter.
   • Sample-level: MixUp and CutMix are applied per batch to improve generalization and force class-discriminative feature learning.
3. Regularization:
   • Exponential Moving Average (EMA): Shadow weights are maintained (decay=0.999) to stabilize optimization.
   • Test-Time Augmentation (TTA): At inference, predictions are averaged over five views (original, flips, rotations, contrast-enhanced) to boost performance.

D. Interpretability

• Attention Rollout: Unlike CNNs that use Grad-CAM, this framework uses Attention Rollout to visualize the model's decision process. It recursively multiplies attention matrices across layers (accounting for residual connections) to trace information flow from input patches to the [CLS] token, generating clinically interpretable heatmaps.

3. Key Contributions

1. High-Performance ViT Pipeline: Achieved 99.29% test accuracy on a 4-class brain tumor dataset, outperforming all surveyed CNN-based baselines.
2. MRI-Specific Preprocessing: Demonstrated the efficacy of CLAHE in enhancing tumor boundaries for transformer-based models.
3. Advanced Augmentation: Validated the combined use of MixUp and CutMix within a ViT-based medical imaging pipeline.
4. Robust Training Protocol: Introduced a two-stage fine-tuning strategy with discriminative learning rates and EMA regularization specifically for medical transfer learning.
5. Clinical Interpretability: Provided Attention Rollout visualizations that align with clinical reality (e.g., focusing on diffuse infiltrative regions for gliomas and symmetric structures for healthy brains), addressing the "black box" issue in medical AI.

4. Results

• Overall Performance:
  • Accuracy: 99.29%
  • Macro F1-Score: 99.25%
  • Precision/Recall: Macro averages of 99.26% and 99.25% respectively.
• Per-Class Performance:
  • Healthy: Perfect classification (100% Precision, Recall, F1).
  • Meningioma: 100% Recall (only 2 misclassifications as pituitary).
  • Glioma: 98.15% Recall with 100% Precision.
  • Pituitary: 98.86% Recall.
• Comparison: The proposed method outperformed the strongest prior CNN baseline (ResNet-50 based) by ~1.8% and the strongest prior ViT baseline by 1.19%, translating to approximately 120 fewer misclassifications per 7,000 patients.
• Visual Analysis: Attention heatmaps correctly localized:
  • Glioma: Diffuse infiltrative regions in frontal/temporal lobes.
  • Meningioma: Well-circumscribed peripheral masses with dural tails.
  • Healthy: Symmetric deep brain structures (corpus callosum, basal ganglia).

5. Significance and Future Work

This paper establishes that Vision Transformers, when combined with domain-specific preprocessing (CLAHE) and rigorous regularization, can surpass traditional CNNs in brain tumor classification. The integration of Attention Rollout is a critical contribution, providing the clinical interpretability necessary for deployment in real-world medical settings where clinicians must validate AI decisions.

Future Directions:

• External validation on multi-center datasets to test generalizability across different MRI acquisition protocols.
• Integration of multi-sequence MRI modalities (T1, T2, FLAIR, DWI).
• Extension to tumor grading (WHO Grade I–IV).
• Prospective clinical validation with neuroradiologists to assess utility as a "second-reader" diagnostic tool.