HEMF: an adaptive hierarchical enhanced multi-attention feature fusion framework for cross-scale medical image classification

ML-based methods

To classify medical images, traditional machine learning approaches such as support vector machine (SVM), K-means and K-nearest neighbor (KNN) rely on the extraction of relevant features such as color, texture, and shape, or feature combinations derived therefrom. Rani, Mittal & Ritambhara (2016) use SVM classifiers to detect and classify focal liver lesions. Iwahori et al. (2015) use Hessian filter and K-means++ to detect the polyp region. Manju, Meenakshy & Gopikakumari (2015) extract texture features and use KNN classifier to diagnose prostate disease. These shallow methods typically require fewer computational resources compared to deep learning methods, making them feasible for deployment on low-power clinical hardware. Training converges faster due to simpler architectures, suitable for scenarios with limited time or energy budgets. They often outperform deep learning methods when labeled training data is scarce, as they rely on handcrafted features rather than data-hungry feature learning. Besides, they often provide strong interpretability, aligning with clinical demand for transparency in diagnosis. However, such approaches also possess inherent limitations. Performance heavily depends on manual feature engineering, requiring domain expertise to design optimal descriptors such as texture and shape. Handcrafted features may fail to capture complex spatial hierarchies in high-dimensional images. Linear models struggle with nonlinear feature interactions which are common in medical images. Furthermore, they are susceptible to the curse of dimensionality when dealing with high-dimensional image data.

CNN-based methods

Convolutional neural networks (CNNs) are designed to extract salient features from structured data as multi-dimensional arrays, enabling a highly effective application on image data processing. Convolutional operations inherently capture local receptive fields, ideal for detecting fine-grained pathological features. Weight sharing and pooling operations provide robustness to small spatial shifts, which is critical for anatomical structure alignment. Their inherent ability to automatically learn relevant features has contributed to their effectiveness and widespread use in medical image classification. Shen et al. (2017) propose a multi-crop convolutional neural network, employing a novel multi-crop pooling strategy to automatically extract salient lesion information. Koitka & Friedrich (2016) utilize a single CNN for skin lesion classification, achieving end-to-end training with only images and labels as input. Flayeh & Douik (2024) provide a lightweight model based on autoencoder and CNNs to detect breast cancer. Besides, several studies conduct transfer learning on CNNs, pre-training on a substantial corpus of medical images and then fine-tuning for the downstream classification task. Tajbakhsh et al. (2016) investigate the performance of deep CNNs trained from scratch compared to pre-trained CNNs with layer-wise fine-tuning. Hassan et al. (2020) utilize a pre-trained ResNet50 model for feature extraction and optimization, and then classify with linear discriminator. However, CNNs primarily focus on local spatial feature details, often lack a comprehensive understanding of global semantic information, which is crucial for performance enhancement. Furthermore, transfer learning methods can be time-consuming and necessitate substantial datasets for pre-training. Fine-tuning pre-trained models on small datasets risks overfitting, especially for rare diseases. Besides, Black-box decision-making may undermine clinical trust; while saliency maps such as Gradient-weighted Class Activation Mapping (Grad-CAM) (Selvaraju et al., 2017) and Grad-CAM++ (Chattopadhay et al., 2018) provide post-hoc explanations, they lack causal reasoning aligned with medical knowledge. End-to-end training may ignore established some features, limiting physician acceptance.

Transformer-based methods

In recent years, methods based on self-attention mechanisms, such as vision transformer (ViT) (Dosovitskiy et al., 2020), have demonstrated remarkable performance in medical image classification tasks. ViT splits the input image into patches with position embedding to capture long-range feature dependencies and global semantic information. However, it usually does not pay enough attention to local features and can even disrupt low-level details (Azad et al., 2024). Moreover, it typically requires more training data than CNNs to converge, exacerbating challenges in rare disease classification.

CNN-transformer methods integrate the advantages of CNNs and ViTs, enabling the extraction and fusion of both local and global features. Dai, Gao & Liu (2021) use CNNs and ViTs to classify parotid gland tumors and knee injuries. Zhang, Liu & Hu (2021) fuse transformers and CNNs for medical image segmentation. Dai et al. (2021) combine the strengths from transformers and CNNs for image classification. Several studies investigate integrating transformers with other architectures or mechanisms for specific tasks. Chen et al. (2021) fuse transformers and U-net architecture for medical image segmentation. Wang et al. (2024) combine transformer with multi-granularity patch embedding and self attention for medical time-series classification. However, the aforementioned works perform feature fusion at a single scale of the samples, failing to integrate multi-scale representations, which hinders their ability to effectively capture the hierarchical and complex structures present in medical images. Recently, several studies investigate hierarchically extracting and fusing local and global features to suppress the incorporation of noise. HiFuse (Huo et al., 2024) designs a hierarchical feature fusion module named Hierarchical Feature Fusion (HFF) to fuse local and global features from the multi-layer encoders. However, its feature characterization capabilities can be further enhanced, and the model also exhibits a relatively large number of parameters and FLOPs. Based on HiFuse, CAME (Li, 2023) introduces an external attention mechanism to learn the relation of different samples. However, its attention mechanism is designed as the fourth branch besides the local, global, and feature fusion branches, resulting in a complex model and disruption on local and global features.

To mitigate the limitations of model complexity and insufficient feature characterization capabilities, we propose an adaptive hierarchical enhanced multi-attention feature fusion framework named HEMF. We use enhanced local feature (ELF) and enhanced global feature (EGF) extraction modules to hierarchically and parallelly extract multi-scale features. We also design a hierarchical enhanced feature (HEF) fusion module, which includes Mixed Attention mechanism and Squeezed Inverted Residual Multi-Layer Perceptron (SIRMLP). The proposed HEF is capable of fusing local and global representations from various scales, learning the relation of different samples and high-dimensional representations of the fused features.

Overview

In this section, we describe the details of HEMF. As depicted in Fig. 1, HEMF is built upon three distinct branches: a local branch (marked with blue), dedicated to extracting enhanced local features (ELF); a global branch (marked with green), designed to extract enhanced global features (EGF); and a feature fusion branch (marked with red), to fuse the hierarchical features obtained from local and global branches.

The model takes as input an image of size 224 $\times$ 224 $\times$ 3. Initially, the image is convolved using a convolutional kernel with a stride of 4 and a kernel size of 4, and increasing to 96 channels. The architecture subsequently employs parallel four-layer local and global feature extraction branches, each dedicated to capturing the distinct scale-specific representations. Within each layer of the local and global branches, the feature maps first undergo a layer normalization (LN) before being fed into four stacked ELF extraction blocks and three stacked EGF extraction blocks, respectively. We use four ELF extraction blocks and three EGF extraction blocks in each layer to achieve a characteristic balance of different scales. In the feature fusion branch, the hierarchical enhanced feature (HEF) fusion module of each layer receives the outputs from the corresponding ELF and EGF extraction blocks, as well as the output from the previous HEF layer, to perform enhanced feature fusion. The fusion output serves as the input for the next HEF block. After each layer, the resolution of the feature map is reduced to 1/4, while the number of channels is doubled. Following the four-layer feature fusion, the resulting 7 $\times$ 7 $\times$ 768 feature map sequentially undergoes a global average pooling, a layer normalization, and finally a linear layer to obtain the classification result.

Enhanced local feature extraction module

The enhanced local feature (ELF) extraction module is designed to extract local detailed features from the image, which enables the model to understand the fine-grained details of the samples and plays a crucial role in image classification. As illustrated on the right of Fig. 2, the ELF extraction module comprises a 3 $\times$ 3 standard convolution, a 3 $\times$ 3 depthwise convolution, and a 1 $\times$ 1 pointwise convolution. The input feature map undergoes a 3 $\times$ 3 standard convolution, followed by a layer normalization. It sequentially undergoes the 3 $\times$ 3 depthwise convolution and the 1 $\times$ 1 pointwise convolution. After applying the Gaussian Error Linear Unit (GELU) activation function, the resulting feature map is added to the original input feature map via a skip connection. Compared with the local feature block in HiFuse (Huo et al., 2024), a 3 $\times$ 3 standard convolution is introduced, as seen in the left of Fig. 2. The ELF extraction module employs the depthwise and pointwise convolutions to reduce computational cost while utilizing standard convolution to enhance local feature characterization capability. The skip connection retains detailed spatial and semantic information from the input feature map. The ELF extraction module is formulated as follows:

(1) $$l_i=\mathrm{L}\mathrm{N}(f^{3\times 3}(L_{i-1}))$$

(2) $$L_i=\mathrm{G}\mathrm{E}\mathrm{L}\mathrm{U}(f^{1\times 1}(f^{\mathrm{d}\mathrm{e}\mathrm{p}\mathrm{t}\mathrm{h}3\times 3}(l_i)))+L_{i-1}$$where $f^{n\times n}$ is the convolution operation with a kernel size of $n\times n$, LN( $\cdot$) indicates layer normalization operation, $L_{i-1}$ and $L_i$ indicate the input and output of the ELF block, respectively.

Enhanced global feature extraction module

The enhanced global feature (EGF) extraction module is designed to extract global semantic features from the image, which facilitates the model to comprehend the global semantic information of the samples. Due to the significant intra-class variability and inter-class similarity often observed in medical images, global features are crucial for medical image classification.

As a powerful computer vision architecture, ViT captures global semantic information by employing sophisticated spatial transformations and learning long-range feature dependencies. However, as the core module of ViT, the self-attention mechanism lacks explicit spatial prior information and exhibits high quadratic complexity. To improve the performance of ViT, Retentive Networks Meet Vision Transformers (RMT) (Fan et al., 2024) proposes a novel self-attention mechanism called Manhattan Self-Attention (MaSA). MaSA transforms the unidirectional and one-dimensional temporal decay into bidirectional and two-dimensional spatial decay, introducing rich and explicit spatial prior information into ViT and reducing the computational complexity to linear. The Manhattan Self-Attention can be formulated as follows:

(3) $$\mathrm{M}\mathrm{a}\mathrm{S}\mathrm{A}(X)=(\mathrm{S}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}(Q{K}^{\mathrm{T}})\odot D^{2d})V$$

(4) $$D_{nm}^{2d}={\gamma }^{|x_n-x_m|+|y_n-y_m|}$$where Q, K and V indicate the query matrix, key matrix and value matrix, respectively. $\odot$ indicates the hadamard product, $\gamma$ indicates the decay value, and ( $x_n,y_n$) indicates the unique two-dimensional position coordinate within the plane of the $n$-th token.

Inspired by RMT, we propose a transformer-like enhanced global feature (EGF) extraction module. As illustrated on the right of Fig. 2, the EGF extraction module consists of a position embedding block, an RMT block, as well as a linear layer. The input feature map undergoes once position embedding (only once even with multiple stacked EGF blocks) and is then fed into the RMT block. After a skip connection and layer normalization, the output is passed to the linear layer, followed by another skip connection and layer normalization. Compared with the global feature block in HiFuse (Huo et al., 2024), we add skip connections with layer normalization and use RMT block instead of the Window-based Multi-head Self-Attention (W-MSA) (Liu et al., 2021) to enhance global feature characterization capability, as seen in the left of Fig. 2. The EGF extraction module leverages the RMT block to capture global semantic information and long-range feature dependencies, while the linear layer enhances global feature characterization. The two skip connections combined with layer normalization preserve detailed spatial and semantic information from the input data, mitigating overfitting to some extent. The EGF extraction module is formulated as follows:

(5) $$g_i=\mathrm{L}\mathrm{N}(\mathrm{R}\mathrm{M}\mathrm{T}(\mathrm{P}\mathrm{o}\mathrm{s}\mathrm{E}\mathrm{m}\mathrm{b}(G_{i-1}))+\mathrm{P}\mathrm{o}\mathrm{s}\mathrm{E}\mathrm{m}\mathrm{b}(G_{i-1}))$$

(6) $$G_i=\mathrm{L}\mathrm{N}(\mathrm{L}\mathrm{i}\mathrm{n}\mathrm{e}\mathrm{a}\mathrm{r}(g_i)+g_i)$$where $G_{i-1}$ and $G_i$ indicate the input and output of the EGF block, respectively.

Hierarchical enhanced feature fusion module

Feature fusion contributes to enhance model performance and robustness by integrating features from diverse sources, hierarchical levels, or modalities (Ngiam et al., 2011). To fuse the enhanced local and global features, as well as utilize the high-dimensional fused features at each layer, we propose a hierarchical enhanced feature (HEF) fusion module with the similar three-branch structure of HiFuse (Huo et al., 2024). The core components of the module are a novel attention mechanism and an inverted residual block named Mixed Attention and Squeezed Inverted Residual Multi-Layer Perceptron (SIRMLP), respectively. The HEF fusion module can adaptively fuse the enhanced local and global features from the same layer and the fused features from the preceding layer. As shown in Fig. 3, the module consists of a local branch, a global branch, and a main branch. In the local and global branches, the input feature maps undergo the Mixed Attention, which can capture the relation of different samples and focus on critical features and spatial regional information. In the main branch, the input feature map undergoes operations such as convolution, average pooling, and concatenation. The resultant feature map undergoes the SIRMLP, which can efficiently learn the high-dimensional representation of the fused features with a low cost. The HEF fusion module is formulated as follows:

(7) $$h_x=\mathrm{A}\mathrm{v}\mathrm{g}\mathrm{P}\mathrm{o}\mathrm{o}\mathrm{l}(f_{\mathrm{d}\mathrm{i}\mathrm{m}\times 2}^{1\times 1}(HE{F}_{i-1}))$$

(8) $$h_y=\mathrm{G}\mathrm{E}\mathrm{L}\mathrm{U}(f^{1\times 1}(\mathrm{L}\mathrm{N}(\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{c}\mathrm{a}\mathrm{t}[h_x,f^{1\times 1}(L_i),f^{1\times 1}(G_i)])))$$

(9) $$HE{F}_i=\mathrm{S}\mathrm{I}\mathrm{R}\mathrm{M}\mathrm{L}\mathrm{P}(\mathrm{L}\mathrm{N}(\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{c}\mathrm{a}\mathrm{t}[\mathrm{M}\mathrm{i}\mathrm{x}\mathrm{A}\mathrm{t}\mathrm{t}\mathrm{n}(G_i),\mathrm{M}\mathrm{i}\mathrm{x}\mathrm{A}\mathrm{t}\mathrm{t}\mathrm{n}(L_i),h_y]))+h_x$$where $L_i$ and $G_i$ indicate the output of multiple stacked ELF and EGF blocks, respectively. $HE{F}_{i-1}$ indicates the output of the previous HEF layer and $HE{F}_i$ indicates the output of the fusion. Additionally, as a general module, the proposed HEF fusion module can be further adopted in other downstream tasks for efficient feature fusion.

Datasets

ISIC2018 dataset

The ISIC2018 dataset (Codella et al., 2019) consists of $10,015$ images from seven categories of skin lesions, including melanocytic nevi (NV), dermatofibroma (DF), melanoma (MEL), actinic keratosis (AKIEC), benign keratosis (BKL), basal cell carcinoma (BCC) and vascular lesions (VASC). Due to its inter-class similarity and data imbalance, this dataset serves as an important benchmark for medical image classification tasks. The specific class distribution is detailed in Table 4. Representative image samples from the ISIC2018 dataset are displayed in Fig. 4A. The original size of these images is $650\times 450$ pixels. For our experiments, we resize all the images to $224\times 224$ pixels. Following the data partition strategy of HiFuse (Huo et al., 2024) and ResGANet (Cheng et al., 2022), we allocate 70% of the samples for training and validation, while the remaining 30% for testing.

Table 4:

Data distribution of ISIC2018, Kvasir and COVID-19 CT datasets.

ISIC2018		Kvasir		COVID-19 CT
NV	6,075	Z-line	500	COVID	349
MEL	1,113	Pylorus	500	NonCOVID	397
BKL	1,099	Cecum	500
BCC	514	Esophagitis	500
AKIEC	327	Polyps	500
VASC	142	Ulcerative colitis	500
DF	115	Dyed and lifted polyps	500
		Dyed resection margins	500
Seven classes	10,015	Eight classes	4,000	Two classes	746

DOI: 10.7717/peerj-cs.3181/table-4

Metrics

We use accuracy (Acc), precision (Prec), recall, F1-score, Matthews Correlation Coefficient (MCC), kappa, Receiver Operating Characteristic (ROC) curve, and the area under the ROC curve (AUC) as evaluation metrics. Here, accuracy represents the proportion of correctly classified instances out of the total number of instances. Precision measures the correct predicted proportion of all instances predicted as positive. Recall quantifies the proportion of actual positive instances that are correctly identified. F1-score is the harmonic mean of precision and recall, with higher values indicating better classification performance. The Matthews Correlation Coefficient (MCC) is a metric for evaluating classification performance in binary classification tasks. Its value ranges from −1 to +1, with higher values indicating better classification performance. MCC incorporates all four confusion matrix categories to provide a balanced performance measure, making it particularly suitable for imbalanced datasets. Kappa coefficient is used to measure the consistency of classifiers. The value of Kappa coefficient ranges from 0 to 1, where 1 indicates perfect consistency and 0 indicates random consistency. The ROC curve is a plot with the false positive rate (FPR) on the x-axis and the true positive rate (TPR) on the y-axis, which is less affected by imbalance class distributions. A larger area under the ROC curve (AUC) indicates superior classification performance. These metrics are formulated as follows:

(20) $$Acc=\frac{TP+TN+FP+FN}{TP+TN}$$

(21) $$Prec=\frac{TP}{TP+FP}$$

(22) $$Recall=TPR=\frac{TP}{TP+FN}$$

(23) $$F1=\frac{2\times (Prec+Recall)}{Prec\times Recall}$$

(24) $$MCC=\frac{TP\times TN-FP\times FN}{\sqrt{(TP+FP)(TP+FN)(TN+FP)(TN+FN)}}$$

(25) $$kappa=\frac{P_o-P_e}{1-P_e}$$

(26) $$FPR=\frac{FP}{FP+TN}$$

(27) $$AUC=\sum _{i=1}^{n-1}\frac{TP{R}_i+TP{R}_{i+1}}{2}\times (FP{R}_i+FP{R}_{i+1})$$where TP, TN, FP and FN indicate true positive, true negative, false positive, and false negative, respectively. $P_o$ represents the observed agreement rate (the empirical classification accuracy), and $P_e$ denotes the expected chance agreement rate (the accuracy achievable through random predictions).

Implementation details

The experiments are conducted on the Ubuntu 20.04.6 LTS operating system. GPU acceleration is provided by an NVIDIA Tesla T4 with 16 GB memory. The software environment includes Python 3.9, PyTorch 1.13.1, and CUDA 11.7.

The experimental parameters are set as follows. The initial and final learning rates are 1e−4 and 1e−6, respectively. AdamW optimizer with the cosine annealing learning rate scheduling strategy is adopted. We train 150 epochs with a batch size of 32. We use the categorical cross-entropy (CCE) loss function defined by the following formula:

(28) $$L=-\frac{1}{N}\sum _{n=1}^N\sum _{i=1}^C{y}_i\log \widehat{y_i}$$where N represents the number of samples, C represents the number of classes, $y_i$ is the true label for sample $i$, and $\widehat{y_i}$ is the predicted label for sample $i$. To accelerate training, mixed-precision training is utilized. Our code can be available from https://github.com/Esgjgd/HEMF. More parameter settings are detailed in Table 5.

Results

To evaluate the classification performance, we compare HEMF with advanced models in recent years, including Conformer (Peng et al., 2021), ConvNext (Liu et al., 2022), PerViT (Min et al., 2022), Focal (Yang et al., 2022), UniFormer (Li et al., 2023), BiFormer (Zhu et al., 2023), HiFuse (Huo et al., 2024) and CAME (Li, 2023). We conduct experiments on the ISIC2018, Kvasir and COVID-19 CT datasets using identical experimental settings.

Results on COVID-19 CT dataset

Following the data partitioning strategy in HiFuse Huo et al. (2024) and He et al. (2020), we train our models from scratch, allocating 60% of the data for training, 15% for validation, and 25% for testing. As shown in Table 8, HEMF achieves the highest accuracy (82.26%). Compared to the previous best model CAME, HEMF achieves improvements of 4.80% in accuracy, 4.70% in precision, 4.59% in recall, and 4.75% in F1-score, demonstrating a significant improvement.

ROC curves and confusion matrices

Figure 5 illustrates the ROC curves and confusion matrices for HEMF on the ISIC2018, Kvasir and COVID-19 CT datasets. As shown, HEMF exhibits remarkable classification performance across all three datasets. On the ISIC2018 dataset, despite the similarity between MEL and NV, HEMF can effectively distinguish them. On the Kvasir dataset, HEMF can validly differentiates two similar types of dyed and lifted polyps and dyed resection margins.

To further validate consistency, we evaluate the MCC and kappa coefficient on ISIC2018 and Kvasir datasets. Table 9 shows that HEMF outperforms other models in both metrics, demonstrating robust classification consistency.

Ablation study

To evaluate the impact of each key component of HEMF and gain a deeper comprehension, we conduct an ablation study on the ISIC2018 dataset, examining the contributions of the enhanced local feature (ELF) block, enhanced global feature (EGF) block, Multi-Head External Attention (MHEA), Channel and Spatial Attention (CA&SA), and Squeezed Inverted Residual Multi-Layer Perceptron (SIRMLP). As shown in Table 10, metrics are relatively low when only the ELF block or EGF block is used. The recall improves significantly when both ELF and EGF blocks are employed. Adding MHEA or CA&SA results in a moderate increase in accuracy. Further improvements in F1-score, precision, and recall are observed when all three attention mechanisms are used together (Mixed Attention). Finally, the addition of SIRMLP (i.e., the complete HEMF model) leads to a substantial improvement. We believe that SIRMLP can effectively learn and leverage the high-dimensional representation of the fused features, significantly enhancing the characterization capability of HEMF.

Heatmap visualization

To further validate the capability of HEMF in capturing discriminative features, we employed Grad-CAM++ (Chattopadhay et al., 2018) to generate heatmaps by visualizing the output of the final layer (excluding the linear layer). Figure 6 presents representative visualization results across three medical image datasets: (A) ISIC2018, (B) Kvasir, and (C) COVID-19 CT. The generated heatmaps demonstrate that HEMF consistently localizes pathological regions with high precision, indicating its effectiveness in hierarchically integrating both global contextual information and local discriminative features. Multi-scale fusion mechanism enables the model to reliably identify and characterize lesion patterns across diverse medical image modalities.

Limitations

The proposed HEMF model achieves state-of-the-art performance in medical image classification, underscoring the promise of hierarchical multi-scale feature extraction and fusion. However, several potential limitations should be noted:

• While HEMF demonstrates superior model parameter efficiency and computational economy relative to existing approaches, further architectural refinements could potentially enhance its compactness.

• The current HEMF model employs the conventional categorical cross-entropy loss function. Exploring alternative or adaptive loss functions may lead to further performance improvements, particularly in response to varying dataset distributions.

• The current HEMF implementation does not incorporate any data augmentation techniques such as CutMix and MixUp. Integrating advanced augmentation strategies could enhance model generalization, improve robustness to perturbations, and mitigate class imbalance.