Currently, most large models are pre-trained on non-medical datasets. However, due to the complex structure and features of brain tumor images, these pre-trained models often perform poorly when applied to medical imaging tasks. The introduction of SAM-MED2D, a network pre-trained on large-scale medical image datasets, marks significant progress in this field. However, SAM-MED2D was primarily designed for segmentation tasks and faces limitations when applied to classification, especially in effectively integrating local and global features while managing computational costs. To address these issues, this paper proposes a novel and efficient classification module named Multi-Stage Attention, composed of Spatial Feature Attention and Hierarchical Efficiency Attention. In the Spatial Feature Attention, unlike traditional methods, we first calculate the similarity between <InlineEquation ID="IEq1"> <EquationSource Format="TEX">\( K \)</EquationSource> </InlineEquation> and <InlineEquation ID="IEq2"> <EquationSource Format="TEX">\( V \)</EquationSource> </InlineEquation> instead of the more common approach of calculating the similarity between <InlineEquation ID="IEq3"> <EquationSource Format="TEX">\( Q \)</EquationSource> </InlineEquation> and <InlineEquation ID="IEq4"> <EquationSource Format="TEX">\( K \)</EquationSource> </InlineEquation>. This approach significantly reduces computational complexity. Additionally, to enhance performance, we adjust the direction of query-key pairs by bringing similar pairs closer together and pushing dissimilar pairs further apart, using a novel kernel function to generate the attention map, which improves the diversity of extracted features. In the Hierarchical Efficiency Attention, we partition the features into smaller segments before processing. Traditional methods often input the entire feature set into the multi-head self-attention (MHSA), leading to redundant computation as some heads generate similar outputs. Experimental results demonstrate that our proposed module significantly outperforms existing methods in terms of classification accuracy, recall, and F1-score across two brain tumor datasets. Most notably, theMulti-Stage Attention reduced trainable parameters by 12 million and saved approximately 3G FLOPs, while still achieving a 1.2% improvement in accuracy over vanilla attention module on two datasets.

Multi-stage attention for efficient brain tumor classification with SAM-Med2D

期刊论文
发表日期：

2025年12月15日

查看全文

Hanruo Zhu,
Ziquan Zhu,
Si-Yuan Lu

摘要

Currently, most large models are pre-trained on non-medical datasets. However, due to the complex structure and features of brain tumor images, these pre-trained models often perform poorly when applied to medical imaging tasks. The introduction of SAM-MED2D, a network pre-trained on large-scale medical image datasets, marks significant progress in this field. However, SAM-MED2D was primarily designed for segmentation tasks and faces limitations when applied to classification, especially in effectively integrating local and global features while managing computational costs. To address these issues, this paper proposes a novel and efficient classification module named Multi-Stage Attention, composed of Spatial Feature Attention and Hierarchical Efficiency Attention. In the Spatial Feature Attention, unlike traditional methods, we first calculate the similarity between \( K \) and \( V \) instead of the more common approach of calculating the similarity between \( Q \) and \( K \) . This approach significantly reduces computational complexity. Additionally, to enhance performance, we adjust the direction of query-key pairs by bringing similar pairs closer together and pushing dissimilar pairs further apart, using a novel kernel function to generate the attention map, which improves the diversity of extracted features. In the Hierarchical Efficiency Attention, we partition the features into smaller segments before processing. Traditional methods often input the entire feature set into the multi-head self-attention (MHSA), leading to redundant computation as some heads generate similar outputs. Experimental results demonstrate that our proposed module significantly outperforms existing methods in terms of classification accuracy, recall, and F1-score across two brain tumor datasets. Most notably, theMulti-Stage Attention reduced trainable parameters by 12 million and saved approximately 3G FLOPs, while still achieving a 1.2% improvement in accuracy over vanilla attention module on two datasets.