[Paper Review] A Novel Vision Transformer with Residual in Self-attention for Biomedical Image Classification

Arun Sharma, Nishchal K. Verma|arXiv (Cornell University)|Jun 2, 2023

Brain Tumor Detection and Classification4 citations

TL;DR

This paper proposes a novel Vision Transformer (ViT) architecture with residual connections in multi-head self-attention, where the best attention head—selected via L1 norm of attention scores—is added as a residual to the final output. The method improves feature representation in biomedical image classification, achieving 93.38% accuracy on blood cell images and 96.15% on brain MRI tumor detection, outperforming standard ViT and convolutional SOTA models.

ABSTRACT

Biomedical image classification requires capturing of bio-informatics based on specific feature distribution. In most of such applications, there are mainly challenges due to limited availability of samples for diseased cases and imbalanced nature of dataset. This article presents the novel framework of multi-head self-attention for vision transformer (ViT) which makes capable of capturing the specific image features for classification and analysis. The proposed method uses the concept of residual connection for accumulating the best attention output in each block of multi-head attention. The proposed framework has been evaluated on two small datasets: (i) blood cell classification dataset and (ii) brain tumor detection using brain MRI images. The results show the significant improvement over traditional ViT and other convolution based state-of-the-art classification models.

Motivation & Objective

• To address the challenge of limited and imbalanced biomedical image datasets in deep learning.
• To improve feature representation and generalization in vision transformers for medical image classification.
• To introduce a residual mechanism in multi-head self-attention that selects the most informative attention head based on L1 norms of attention scores.
• To enhance model robustness and reduce overfitting in small biomedical datasets.
• To demonstrate superior performance over standard ViT and state-of-the-art CNN-based models on two biomedical image classification benchmarks.

Proposed method

• The proposed method computes the L1 norm (Manhattan norm) of attention scores from each multi-head attention head.
• The attention head with the highest L1 norm is selected as the 'best' head, representing the most informative attention pattern.
• The output of this best head is added as a residual connection to the final projected attention output after the linear projection layer.
• This residual connection integrates the most discriminative attention features into the final representation, enhancing feature learning.
• The modified attention block is integrated into a pre-trained ViT-base model, followed by fine-tuning on biomedical datasets.
• The method maintains computational efficiency, with complexity nearly identical to standard ViT due to low-overhead sorting and norm operations.

Experimental results

Research questions

• RQ1 Can selecting the most informative attention head via L1 norm of attention scores improve ViT performance on small biomedical datasets?
• RQ2 Does integrating the best attention head as a residual connection enhance feature representation and classification accuracy?
• RQ3 How does the proposed ViT variant compare to standard ViT and state-of-the-art CNN models in terms of accuracy and generalization on imbalanced biomedical image data?
• RQ4 Does the residual mechanism reduce overfitting and improve generalization on limited training data?
• RQ5 Can the proposed attention mechanism be effectively applied to diverse biomedical image tasks such as blood cell classification and brain tumor detection?

Key findings

• The proposed ViT achieved 93.38% accuracy on the blood cell classification dataset, significantly outperforming the next best model (ViT at 88.38%).
• On the brain MRI tumor detection dataset, the proposed method reached 96.15% accuracy, surpassing the standard ViT (92.15%) and ResNet18 (92.16%).
• The F1-score on the blood cell dataset reached 0.94, indicating strong balance between precision and recall, compared to 0.89 for standard ViT.
• The model showed reduced gap between training and validation accuracy curves, indicating better generalization and lower overfitting.
• The confusion matrices confirmed that the proposed model had fewer misclassifications than the pre-trained ViT baseline on both datasets.
• The complexity analysis confirmed that the method adds minimal computational overhead, maintaining O(n²d) complexity similar to standard ViT.