A wavelet-enhanced state space model for multi-disease diagnosis across chest X-ray, CT, and MRI

Abstract

Background

Medical image classification faces multiple challenges, including heterogeneous lesion morphology, cross-modal representation discrepancy, and computational efficiency. Traditional convolutional neural networks (CNNs) struggle to model long-range semantic correlations due to their fixed receptive fields, while the quadratic complexity of Transformers hinders the practical processing of high-resolution images. This study proposes the Wavelet Transform-Enhanced Mamba (WTEMamba) model, which integrates wavelet transform with the Selective State Space Model (Mamba) to construct an efficient medical image classification framework.

Methods

The model adopts a three-stage lightweight network to balance performance and efficiency, with core innovations including the design of a Feature Interaction Hybrid Module (FIHM): the convolutional branch extracts local textures, the wavelet-enhanced selective state space branch captures global dependencies and strengthens high-frequency edge features, and channel shuffle is integrated to eliminate information redundancy. Additionally, a multi-scale compressed classifier is built to gradually refine high-level semantic features.

Results

Experiments on 9 medical datasets showed that WTEMamba significantly outperformed baseline models such as CNNs, Transformers, and Mamba-based models. In terms of overall accuracy (OA) and the area under the curve (AUC), WTEMamba improved OA over the best competing model by 0.3% to 2.6%, and achieved AUC improvements of 0.0001 to 0.0125. For example, on the Brain Stroke CT dataset, WTEMamba achieved an OA of 93.45%, outperforming Medmamba by 1.99% (*P<0.05 via bootstrap test), while requiring only 2.1G floating point operations (FLOPs) and reducing the parameter count by over 50% compared to Medmamba. On the challenging breast cancer magnetic resonance imaging (MRI) and Knee X-Ray datasets, WTEMamba exceeded the best competing models by 2.55% and 2.25% in OA, respectively.

Conclusions

The experimental results validate the significant advantages of the WTEMamba model in global dependency modeling and high-frequency detail perception. It effectively balances classification accuracy and computational efficiency, providing a robust and practical solution for medical image analysis. Furthermore, the framework’s outstanding performance on multi-modal data, combined with its exceptional generalization ability, highlights the positive significance of integrating lightweight diagnostic algorithms into medical image analysis. Our code and datasets are available at https://github.com/yuting77777/WTEMamba.

Introduction

In the clinical decision-making chain, the accuracy of medical image classification directly determines the sensitivity of disease screening, the reliability of staging assessment, and the objectivity of treatment efficacy evaluation (1). Current mainstream methods face a dilemma: convolutional neural networks (CNNs), such as the recent lightweight CNN variant CellSage (2) that employs multi-scale convolutions and attention mechanisms, achieve high diagnostic accuracy with an ultra-lightweight architecture of only 3.8 million parameters, yet they are constrained by their local receptive fields, making it difficult to model the complex spatial dependencies between lesions and surrounding tissues. In contrast, vision transformers (3) establish global contextual interactions, but their self-attention mechanism introduces high computational complexity and a rigid demand for large-scale annotated data, which pose bottlenecks in many clinical scenarios.

Recently, structured state-space models like Mamba (4) offer a new paradigm for overcoming the above challenges. This architecture provides three key advantages. At the feature perception level, its selective state mechanism dynamically modulates information flow based on input content, enabling the network to prioritize key diagnostic structures such as lesion boundaries and microcalcifications rather than treating the entire image uniformly. At the computational efficiency level, hardware-optimized parallel scanning enables long-sequence modeling without high computational costs, facilitating real-time analysis of high-resolution three-dimensional (3D) images and multi-phase dynamic contrast-enhanced sequences. At the temporal modeling level, time-varying parameters allow the model to characterize the dynamic evolution of pathological processes, addressing the limitations of traditional static models in analyzing disease progression. Recent studies show that ViM (5) achieves cross-scale dependency modeling with Transformer-like efficiency while retaining CNN-level local feature extraction. Yue et al. (6) first applied Mamba to generalized medical image classification and detection, demonstrating strong performance across various tasks. Liu et al. (7) leveraged cross-level Mamba blocks to achieve robust brain tumor segmentation from incomplete multi-modal magnetic resonance imaging (MRI). For resource-constrained deployment, researchers have developed lightweight networks employing strategies such as skip sampling, directional selection, and pathway simplification to avoid redundancy in full-path scanning. Examples include EfficientVMamba (8), which uses ES2D scanning to reduce State Space Model (SSM) computational burden. DualMamba (9) designs a lightweight Mamba module incorporating dynamic positional encoding, achieving strong performance with a small number of parameters. SCSegamba (10) introduces bottleneck convolutions and gating mechanisms into Mamba to dynamically adjust feature weights, greatly reducing parameters and computation. S³-Mamba (11) removes depthwise convolutional layers from original SSMs and introduces dual residual connections, avoiding over-smoothing and redundancy while reducing information propagation loss. He et al. (12) propose MobileMamba, and experiments verify that its three-stage architecture (the first stage downsample to H/16, with the final output at H/64) outperforms the traditional four-stage architecture: under the same throughput, the accuracy is improved by 0.4%, and under the same accuracy, the throughput is increased by 15%. AttmNet (13) lightweightly integrates Mamba with self-attention via the Att-Mamba module, capturing global context while maintaining computational efficiency. GCNet-Mamba (14) integrates GCNet convolutions with SSM layers and comprehensively outperforms CNN, Transformer, and Mamba models on nine multimodal medical image datasets.

Although state-space models offer unique advantages in processing high-resolution medical images due to their linear complexity, they still encounter three major limitations in handling such data. First, the serialization of two-dimensional (2D) images inevitably disrupts spatial adjacency relationships, reducing the model’s sensitivity to fine morphological features that rely on local neighborhood coordination (e.g., lesion margins and microcalcification clusters), thereby increasing the risk of missing small lesions. Second, existing 2D selective scan mechanisms lack explicit modulation of directional energy in frequency-domain subbands, making it difficult to adaptively distinguish high-frequency pathological contours from background noise. As a result, critical diagnostic signals are easily submerged by imaging artifacts, leading to insufficient robustness in cross-modal scenarios such as computed tomography (CT)/MRI. Third, multi-scale feature fusion faces exponential dimension expansion as the network deepens, and direct aggregation introduces considerable computational redundancy, making it challenging to balance diagnostic accuracy and inference efficiency in lightweight scenarios.

To address the disconnection between frequency-domain perception and spatial-domain modeling, wavelet transform offers a potential solution due to its time-frequency localization properties. It provides localized information in both domains, effectively separating frequency components while preserving spatial structure: low-frequency subbands capture macroscopic anatomical contours, whereas high-frequency subbands contain key diagnostic details such as edges and textures. Inspired by this, researchers have begun integrating wavelet transform into the Mamba framework. For example, MSM-UNet (15) enhances boundary and detail representation by introducing wavelet transform and a dual attention mechanism. Wave-Mamba (16) uses wavelet transform for lossless downsampling and noise separation, providing cleaner inputs. WaveMamba (17) applies wavelet transform for multi-scale decomposition of spatial and spectral features, and EAD-net (18) employs a multi-scale wavelet module for denoising and feature decomposition with an alternating Mamba-CNN encoder. HWA-ResMamba (19) integrates Wavelet Transform into convolutional networks to enhance low-frequency and edge perception and replaces self-attention with Mamba for efficient long-range dependency modeling. This synergy significantly improves the continuity and integrity of coronary arteries.

Despite these improvements, most methods treat wavelet transform merely as a front-end preprocessing tool, serving only to supply smaller or cleaner input sequences to subsequent Mamba modules. Consequently, frequency transformation operates solely at the input stage, lacking parallel interaction and deep synergy with spatial modeling. Frequency features, extracted only shallowly, cannot be progressively reinforced with global semantics as the network deepens, severely undermining the model’s sustained perception of subtle lesion edges and textures. Particularly for cross-modal medical images with distinct physical properties and frequency characteristics, the discriminative information in the frequency domain remains largely underexploited.

To address the above issues, this paper proposes the Wavelet Transform-Enhanced Mamba (WTEMamba), a medical image classification model. To systematically clarify the technical gaps and position our method against existing SSM-based approaches, we summarize the core innovations of representative models in Table 1. Compared with previous works, our core innovation lies in designing a threestage network and constructing a Feature Interaction Hybrid Module (FIHM). In this module, the convolutional branch extracts local texture features, while the waveletenhanced selective statespace branch simultaneously captures global dependencies and strengthens highfrequency edge information. The main innovations are as follows:

Table 1. Comparison of core architectural innovations between WTEMamba and representative SSM-based vision models.

Technical innovations	Medmamba (6)	Wave-Mamba (16)	WaveMamba (17)	GCNet-Mamba (14)	EfficientV Mamba (8)	MobileMamba(12)	WTEMamba (ours)
1. Explicit deep wavelet transform fusion (not just preprocessing)	×	× (only as downsampling and denoising preprocessing)	× (only shallow feature decomposition)	×	×	✔	✔ (deep FIHM fusion)
2. Spatial-frequency primary-auxiliary path interaction (wavelet + channel attention)	×	×	×	×	×	×(only simple additive fusion)	✔ (primaryauxiliary weighting)
3. Three-stage downsampling (8×→16×→32×)	× (four-stage)	× (U-Net structure)	× (single-path)	×	× (four-stage)	✔	✔
4. Multi-scale compression classifier (3-stage feature fusion, progressive dimension reduction)	× (standard linear layer)	×	×	×	×	×	✔
5. Lightweight multimodal medical image classification	✔	× (only low-light enhancement)	× (only hyperspectral)	✔	× (general vision)	× (general vision)	✔ (validated on 9 datasets)
6. Ultra-low params/FLOPs with high accuracy	× (relatively large parameters)	✔	✔	✔	✔	✔	✔ (optimal balance)

FIHM, Feature Interaction Hybrid Module; FLOPs, Floating Point Operations; SSM, State Space Model.

Propose a novel medical image classification method named WTEMamba, which adopts a three-stage network and achieves a good balance between performance and efficiency.

Design an efficient FIHM, integrating a convolutional branch (Conv-Branch), a wavelet-enhanced selective state-space branch (DWTM-Branch), and channel shuffle. This module achieves deep fusion of local feature extraction and global dependency modeling while enhancing the capture of high-frequency edge details.

Design a multi-scale feature vector hybrid classifier for WTEMamba. This classifier is a five-layer network that uses only fully connected layers to reduce computational complexity. Meanwhile, the number of neurons in the hidden layers decreases layer by layer, gradually refining high-level semantic features.

After validation on typical medical image datasets, the results show that all metrics of the WTEMamba model improve compared with the original Medmamba model, which significantly enhances classification performance and further demonstrates the model’s effectiveness and superiority.

Methods

Preliminaries

The structure of the selective state space model is shown in Figure 1. It is the “selective” and “hardware-optimized” version of the state space model, where the input is $x_k$ and the output is $y_k$, both of which are features at the k-th time step.

Figure 1.

Structure of the selective state-space model.

State Space Models (S4)

State Space Models (SSMs) are a mathematical framework for describing dynamic systems. In the continuous domain, SSM maps a 1D input sequence $x(t)\in R^D$ to an output $y(t)\in R^D$ through a hidden state $h(t)\in R^N$, and its continuous system equations are as follows:

$$h^{\text{′}}(t)=Ah(t)+Bx(t)$$ [1]

$$y(t)=Ch(t)$$ [2]

where $A\in R^{N\times N}$ is the state transition matrix, and $B\in R^{N\times D}$ and $C\in R^{D\times N}$ are the input projection matrix and output projection matrix respectively.

Discretization: Introduce the step parameter$\mathrm{\Delta }$, and discretize via the Zero-Order Hold (ZOH) rule:

$$\overline{A}=\exp (\mathrm{\Delta }A),\overline{B}=(\exp (\mathrm{\Delta }A)-I)A^{-1}B,\overline{C}=C$$ [3]

$$\overline{B}\approx (\mathrm{\Delta }A)A^{-1}B=\mathrm{\Delta }B$$ [4]

where$\overline{A}\in R^{N\times N}$, $\overline{B}\in R^{N\times D}$, and $\overline{C}\in R^{D\times N}$. After discretization, the system is converted into a linear recursive form, and the continuous differential equations are transformed into difference equations:

$$h(t)=\overline{A}h(t-1)+\overline{B}x(t)$$ [5]

$$y(t)=\overline{C}h(t)$$ [6]

Selective State Space Models (S6)

Mamba (4) improves the performance of SSMs by introducing the Selective State Space Module (S6), allowing continuous parameters to change with the input. This makes $\overline{B}$, $\overline{C}$, and $\mathrm{\Delta }$no longer fixed, but functions of input x:

$$\overline{B}=s_B(x)$$ [7]

$$\overline{C}=s_C(x)$$ [8]

$$\mathrm{\Delta }={\tau }_A(Parameter+s_A(x))$$ [9]

where$s_B(x)$,$s_C(x)$ and $s_A(x)$ are linear projection functions, and ${\tau }_A$ is a softplus function. This enables the model to dynamically select which information to remember or ignore, making it more suitable for handling spatial structures in visual tasks.

The proposed methods

Overall architecture

This paper designs a three-layer lightweight visual state space model named WTEMamba, whose overall structure is shown in Figure 2. Specifically, WTEMamba consists of a Patch Embedding layer, stacked FIHM modules, a Path Merging layer for downsampling, and a feature classifier.

Figure 2.

The overall architecture of WTEMamba. It consists of two main branches: Conv-Branch is responsible for capturing local spatial features, and DWTM-Branch is responsible for capturing global context and high-frequency detailed features. DWTM-Branch, Wavelet-Enhanced Selective State-Space Branch; FIHM, Feature Interaction Hybrid Module.

The WTEMamba model first used the Patch Embedding layer to divide the input image $x\in R^{H\times W\times 3}$ into 8×8 non-overlapping image patches, and mapped the channel dimension of the image to C. Specifically, it partitioned the original image with a size of 224×224×3 through a convolution operation with both the convolutional kernel and stride set to 8, then linearly projected each 8×8 patch to a 192-dimensional channel count, and finally generated a feature map with a size of 28×28×192.

Subsequently, the model entered a three-stage hierarchical architecture. The core design of this architecture lies in efficiently constructing a multi-scale feature pyramid, gradually fusing local details and global semantic information, and thereby significantly enhancing the model’s capability to capture pathological features of different sizes in medical images.

Stage 1: the model stacked 2 FIHM modules to enhance shallow feature representation at high spatial resolution, effectively preserving details such as edges and textures, and providing rich low-level features for subsequent stages.

Stage 2: the model performed 2× downsampling via Path Merging to compress the feature map size to 14×14, expanded the channel count to 384, and stacked 4 FIHM modules. At medium resolution, this stage balances detail and semantic information, gradually expands the receptive field, enhances the model’s ability to represent medium-scale lesion features, and maintains high computational efficiency.

Stage 3: the model further downsampled to 7×7 resolution, expanded the channel count to 768, and stacked 2 FIHM modules for deep feature extraction. Finally, highly abstract high-level semantic features were formed, enabling the model to understand the complete pathological patterns of the image.

The design of this three-stage architecture differs from the four-stage architectural design adopted by most models (5,20,21). Specifically, in a four-stage network, the input image with dimensions H×W×3 undergoes four consecutive downsampling steps: after the first downsampling, the feature map size becomes $\frac{H}{4}\times \frac{W}{4}\times C_1$; the second downsampling reduces it to $\frac{H}{8}\times \frac{W}{8}\times C_2$; the third downsampling further reduces it to $\frac{H}{16}\times \frac{W}{16}\times C_3$; and finally, the fourth downsampling results in a size of $\frac{H}{32}\times \frac{W}{32}\times C_4$.In contrast, the three-stage network proposed in this study is structurally more compact: it generated a feature map of $\frac{H}{8}\times \frac{W}{8}\times C_1$ during the first downsampling, the second downsampling reduced it to $\frac{H}{16}\times \frac{W}{16}\times C_2$, and the third downsampling further reduced the feature map to $\frac{H}{32}\times \frac{W}{32}\times C_3$. Since the four-stage network retains feature maps with higher resolution in the early stages, it has a higher computational load, leading to relatively slow inference speed. However, the three-stage architecture significantly reduces computational complexity by eliminating one downsampling step, thereby improving inference efficiency. Furthermore, by incorporating the proposed FIHM module, the three-stage network not only ensured faster inference speed but also effectively enhanced feature representation capability and model performance. This was verified in the ablation experiments. For example, on the SARS-COV-2 CT scan dataset, compared with the fourstage network, the three-stage network not only improved the OA metric by 0.22% but also increased the AUC metric by 0.54%, achieving balanced optimization of performance.

At the end of the network, a global average pooling layer was used, which received multi-scale feature maps from the three stages and fused them at a learnable weight ratio, generating a feature vector of size 1×1,344. This cross-stage feature fusion mechanism can adaptively integrate representational information at different levels, strengthen important features and suppress redundant information. The fused features were input into the classifier for disease classification, and the entire process controlled the number of parameters while achieving efficient multi-scale feature interaction and integration, thereby significantly improving the robustness and expressive capability of medical image classification.

Efficient FIHM

As a core component of WTEMamba, FIHM is a lightweight dual-branch block, as illustrated in Figure 2. This method employed channel splitting to map the feature map x into two groups, $x_1$ and $x_2$, and then used Conv-Branch and DWTM-Branch respectively to extract local features and global features from each group. Finally, FIHM adopted channel concatenation to restore the channel dimension size, while channel shuffling performed shuffling on the feature maps in the channel dimension to avoid inter-channel information loss caused by group convolution operations. Meanwhile, the activation functions of Conv-Branch and DWTM-Branch were set to ReLU and SiLU respectively.

Conv-Branch

The First Layer: Batch normalization is applied to the input feature $x_1$ to alleviate internal covariate shift and accelerate training convergence. Then, a 3×3 convolution kernel with stride 1 and padding 1 is used to extract spatial features, which maintains the spatial resolution while leveraging the translation equivariance of convolution to capture local texture structures. Finally, batch normalization and the ReLU activation function are applied again to introduce sparse nonlinearity and enhance the discriminative ability of feature representation.

The Second Layer: Following the design of the first layer, a 3×3 convolution layer with stride 1 and padding 1 is adopted. This part progressively abstracts local patterns through hierarchical convolution, thereby enlarging the effective receptive field and deepening feature representation. Subsequently, batch normalization and ReLU activation are applied to maintain nonlinearity while stabilizing the training process.

The Third Layer: a 1×1 convolution is used for linear recombination and feature fusion across channels. This operation is equivalent to performing a learnable weighted combination along the channel dimension, which enhances cross-channel information interaction while introducing very few parameters, and effectively controls computational complexity.

In the entire convolutional branch, the 3×3 convolution focuses on spatial feature extraction, while the 1×1 convolution is responsible for inter-channel information interaction. Batch normalization between layers ensures training stability, and the ReLU activation function continuously introduces nonlinear transformation capability. This multi-layer convolution stacking structure can progressively extract multi-scale features, from local details to more abstract semantic information, and finally output a feature map with the same size as the original input, laying a foundation for subsequent feature fusion with the DWTM-Branch.

To match convolution operations, permutation was utilized to rearrange the original feature maps. Based on the above, the mathematical form of its Conv-Branch is defined as follows:

$$x_1^p\in R^{\frac{c}{2}\times h\times w}\leftarrow permute(x_1)$$ [10]

$$\widehat{x_1}=BatchNorm(x_1^p)$$ [11]

$$x_1^{\text{′}}=\mathrm{Re}LU(BatchNorm(Con{v}_{3\times 3}(\widehat{x_1})))$$ [12]

$$x_1^{\text{″}}=\mathrm{Re}LU(BatchNorm(Con{v}_{3\times 3}(x_1^{\text{′}})))$$ [13]

$$\overline{x_1}=\mathrm{Re}LU(PWConv(x_1^{\text{″}}))$$ [14]

$$x_1^G\in R^{h\times w\times \frac{c}{2}}\leftarrow permute(\overline{x_1})$$ [15]

DWTM-Branch

This branch is designed to effectively enhance the model’s capability to capture fine structural information such as high-frequency edges in images, while maintaining its advantage in global modeling. Its core lies in processing two complementary information flows in parallel: one focuses on capturing long-range dependencies via the Selective State Space Model (S6), and the other leverages Wavelet Transform to enhance local detail extraction and implicitly expand the receptive field.

The first information flow: In the first information flow, features were first transformed via a linear layer, and then fed into a 2D-Selective-Scan Module (SS2D) (20), as illustrated in Figure 3, to learn global contextual information. Its mathematical expression is given as follows:

Figure 3.

Internal structure of SS2D. SS2D, 2D-Selective-Scan Module.

$$x_2^s=SS2D(Linear(x_2))$$ [16]

$$x_2^o=Linear(SiLU(LayerNorm(x_2^s)))$$ [17]

Among them, the SS2D mechanism is designed to address the directional sensitivity issue when 1D state-space models process 2D images. By introducing the Cross-Scan Module (CSM) (20), this method unfolded 2D spatial features into 1D sequences along four specific directions (top-left to bottom-right, bottom-right to top-left, top-right to bottom-left, and bottom-left to top-right), thereby preserving the spatial structural information of the image. Specifically, for the input feature map $I\in R^{H\times W\times C}$, where $I[h][w]$ represents a token at the h-th row and w-th column in the feature map, the index mapping relationship can be defined as:

$$s_1[i]=I[i\mathrm{mod}W]\left[\lfloor \frac{i}{W}\rfloor \right]$$ [18]

$$s_2[i]=I[(N-i-1)\mathrm{mod}W]\left[\lfloor \frac{N-i-1}{W}\rfloor \right]$$ [19]

$$s_3[i]=I\left[\lfloor \frac{i}{H}\rfloor \right][i\mathrm{mod}H]$$ [20]

$$s_4[i]=I\left[\lfloor \frac{N-i-1}{H}\rfloor \right][(N-i-1)\mathrm{mod}H]$$ [21]

where N=H×W, $0\le i\le N$, $s_1,s_2,s_3,s_4\in R^{N\times C}$denote the expanded and flattened token sequences.

Subsequently, each sequence was processed through the S6 block to generate token sequences $s_1^{\text{′}},s_2^{\text{′}},s_3^{\text{′}},s_4^{\text{′}}\in R^{N\times C}$. This S6 block is a selective state-space model based on Mamba, as shown in Figure 3, which dynamically adjusts the state transition process through input-dependent parameters to distinguish and retain key information. The processed sequences $s_1^{\text{′}}$ to $s_4^{\text{′}}$ were then reconstructed into 2D feature maps $I_1^{\text{′}}$ to $I_4^{\text{′}}$ through inverse operations, where their pixel positions are determined via inverse mapping:

$$I_1^{\text{′}}[i][j]=s_1^{\text{′}}[j\times W+i]$$ [22]

$$I_2^{\text{′}}[i][j]=s_2^{\text{′}}[N-1-j\times W-i]$$ [23]

$$I_3^{\text{′}}[i][j]=s_3^{\text{′}}[i\times H+j]$$ [24]

$$I_4^{\text{′}}[i][j]=s_4^{\text{′}}[N-1-i\times H-j]$$ [25]

Finally, the feature maps from the four directions were fused via element-wise addition:

$$I^{\text{′}}=I_1^{\text{′}}+I_2^{\text{′}}+I_3^{\text{′}}+I_4^{\text{′}}$$ [26]

An output with the same size as the input was obtained. This design ensures that each pixel can fuse global multi-directional contextual information, overcomes directional bias while maintaining linear computational complexity, and significantly enhances the model’s ability to model long-range dependencies in medical images.

The second information flow: To capture high-frequency details more precisely, the same feature map $x_2$ was linearly transformed into $x_2^{\text{′}}$, followed by feature fusion of frequency information (WFM), as shown in Figure 4. The main path of this module utilized wavelet transform to capture the frequency-domain features of images; the auxiliary path, on the other hand, provided supplementary global information statistics through a lightweight attention mechanism. These two paths were fused with weights that clearly distinguished between primary and secondary roles, jointly enhancing the model’s ability to perceive and utilize frequency information. The ultimate goal is to improve the quality of visual feature representation via such frequency-domain enhancement. Considering the need for high efficiency with minimal overhead here, this paper adopted the Daubechies 4 (db4) wavelet for the transformation. The db4 wavelet is one of the most widely used members of the Daubechies wavelet family; it possesses orthogonal and non-redundant properties, enabling accurate signal reconstruction without occupying additional storage. Its compact support structure provides excellent locality in both time and spatial domains, which can effectively capture signal and image details with minimal edge effects. Furthermore, its smoothness is superior to that of low-order Haar wavelets, avoiding block distortion.

Figure 4.

The structure of WFM module. WFM, Wavelet Frequency Module.

Specifically, given the input feature map$x_2^{\text{′}}\in R^{h\times w\times \frac{c}{2}}$, the wavelet transform decomposes it into four sub-bands: one low-frequency sub-band (LL) and three high-frequency sub-bands (LH, HL, HH). Each sub-band has a spatial resolution of$\frac{h}{2}\times \frac{w}{2}$, preserving the same channel dimension$\frac{c}{2}$. By concatenating these four sub-bands along the channel dimension, the resulting wavelet-domain feature map $x_2^w\in R^{\frac{h}{2}\times \frac{w}{2}\times 2c}$ is obtained. This transformation yields a 4-fold increase in channel count (from$\frac{c}{2}$to 2c) while reducing the spatial resolution by half in each dimension, as formalized below:

$$x_2^w=WT(x_2^{\text{′}},[f_{LL},f_{LH},f_{HL},f_{HH}])$$ [27]

where $f_{LL}$ is the low-pass filter and $f_{LH},f_{HL},f_{HH}$ are high-pass filters corresponding to horizontal, vertical, and diagonal edge directions, respectively.

In the wavelet domain, local 1×1 convolutional operations were employed to further process multi-band features. While maintaining computational efficiency, this achieves feature compression and information refinement, effectively enhancing the feature representation capability. The number of channels of the processed sub-band features was reduced to c/2. For reconstruction, before the Inverse Wavelet Transformation (IWT), the dimension was restored to 2c through channel expansion to meet the processing requirements of IWT. Finally, the original spatial resolution was reconstructed via IWT, and the output feature $x_2^{IW}\in R^{h\times w\times \frac{c}{2}}$ was obtained:

$$x_2^{IW}=IWT(\varphi (Conv(x_2^w)),[f_{LL},f_{LH},f_{HL},f_{HH}])$$ [28]

where $\varphi$denotes channel expansion.

Meanwhile, the processing workflow of the auxiliary path was as follows: first, it extracted channel-wise statistical information from the original feature map via global average pooling; subsequently, a lightweight network composed of two 1×1 convolutional layers learned the nonlinear relationships between channels—this network was designed to first perform dimension reduction and then restore the original number of channels, aiming to reduce parameters and introduce nonlinearity; finally, it generated channel attention weights via the Sigmoid activation function, and these weights were multiplied with the original feature map by means of a broadcasting mechanism to achieve adaptive feature enhancement.

$$x_2^{CA}=CA(x_2^{\text{′}})$$ [29]

$$CA(x)=Sigmoid(Con{v}_2(\mathrm{Re}LU(Con{v}_1(AvgPool(x))))\otimes x$$ [30]

The main path weight was set to 0.7 and the auxiliary path weight to 0.3, clearly defining the primary-secondary functional relationship. This configuration ensures that the core frequency-domain analysis capability of wavelet transform occupies a dominant position, allowing the model to mainly rely on its precise mathematical transformation properties for feature extraction. Meanwhile, the global frequency perception capability of the auxiliary path is retained with an appropriate weight, enabling it to provide necessary supplementary information without significantly interfering with the main path. This not only prevents the auxiliary path from accounting for an excessive proportion but also compensates for the global statistical properties that wavelet transform may overlook through a lightweight enhancement mechanism. On the basis of a clear distinction between primary and secondary roles, balanced complementarity of capabilities is achieved, and its mathematical expression is as follows:

$$x_2^{WFM}=0.7\times x_2^{IW}+0.3\times x_2^{CA}$$ [31]

$$x_2^f=Linear(LayerNorm(SiLU(x_2^{WFM})))$$ [32]

Finally, the output $x_2^G\in R^{h\times w\times \frac{c}{2}}$ of the DWTM-Branch was formed by fusing the global feature map $x_2^o$ extracted by SS2D and the local detail feature map $x_2^f$ enhanced by WFM via simple element-wise addition. Subsequently, it was further optimized through SiLU, layer normalization, and a linear layer, resulting in the final $x_2^G$:

$$x_2^G=Linear(LayerNorm(SiLU(x_2^o+x_2^f)))$$ [33]

This fusion mechanism effectively integrates global dependencies and high-frequency edge information enhanced by wavelet transform, significantly enhancing the model’s ability to perceive fine-grained structures in images.

In summary, the output of FIHM can be expressed as:

$$y=x\oplus ︎g(Concat(x_1^G,x_2^G))$$ [34]

where g denotes channel shuffling.

Classifier

To adapt to the multi-classification task of medical images, this model employed a fully connected classification network to process the weighted fused multi-scale features $y_c\in R^{1\times 7C}$. This network comprised three compressed hidden layers (with the number of neurons being 5C, 3C, and C respectively), which gradually refined high-level semantic information and enhanced the discriminative ability for pathological features. The output layer was configured with m neurons (where m denotes the number of disease categories) and adopted the Softmax activation function to generate a normalized probability distribution. This distribution intuitively reflected the confidence level of the input image belonging to each disease category, providing an interpretable basis for clinical decision-making.

During the training process, network parameters were optimized via the backpropagation algorithm, with the cross-entropy loss function minimized {Eqs. [35,36]}. This loss function forces the model to focus on learning the discriminative boundaries between classes, making it particularly suitable for class imbalance scenarios commonly encountered in medical images (e.g., insufficient samples of rare lesions). The index of the maximum value in the finally output probability vector corresponds to the predicted disease label, and this method can be directly integrated into the clinical auxiliary diagnosis process.

$$Y^{\text{′}}=classifier(y_c)$$ [35]

$$L=-\frac{1}{N}\sum _{i=1}^N{Y}_i^{\text{′}}\cdot \log (Y_i)$$ [36]

In Eq. , $Y^{\text{′}}$ is a vector of length m, and the index of its maximum value corresponds to the predicted label. In Eq. , L denotes the loss function, N represents the number of training images, and Y is a one-hot vector that indicates the ground-truth label.

Datasets

In this study, 9 publicly available medical image datasets were employed to evaluate the effectiveness and potential of WTEMamba in medical image classification, as presented in Table 2. These datasets cover 2 X-ray, 2 CT, 2 MRI, and 3 extended datasets of other modalities. The study was conducted in accordance with the Declaration of Helsinki and its subsequent amendments.

Table 2. The detailed descriptions of 9 datasets used in this work.

Name	Imaging modality	Classes	Dataset size	Training/validation/test
CPN X-Ray	Chest X-ray	3	5,228	3,659/627/942
Knee X-Ray	Knee X-ray	3	1,947	1,362/233/352
SARS-COV-2 Ct-scan	Pulmonary CT	2	2,482	1,736/297/448
Brain Stroke CT	Brain CT	3	6,650	4,654/798/1,198
Breast Cancer Patients MRI	Breast MRI	2	1,480	1,036/176/268
Brain Tumor MRI	Brain MRI	4	7,023	4,914/843/1,266
Kvasir	Gastrointestinal endoscope	8	4,000	2,800/480/720
Fetal-Planes-DB	Maternal-fetal ultrasound	6	12,400	8,679/1,861/1,860
HAM10000	Human skin smartphone	7	10,015	7,007/1,203/1,805

CT, computed tomography; MRI, magnetic resonance imaging.

X-ray modality datasets

X-ray images are grayscale images that create contrast through X-ray absorption differences among different tissues, thereby visualizing the internal bone structures of the human body and certain physiological lesions, including:

CPN X-Ray (22,23): the CPN X-Ray is a chest X-ray image database, which is categorized into three classes: COVID-19, normal, and viral pneumonia, and contains a total of 5,228 images. Specifically, the training set contains 3,659 images, the validation set 627 images, and the test set 942 images.

Knee X-Ray (24): the Knee X-Ray is a knee joint X-ray image dataset, which includes three classes: normal, osteopenia, and osteoporosis. Among them, the training set contains 1,362 images, the validation set 233 images, and the test set 352 images.

CT modality datasets

CT images utilize X-ray beams for tomographic scanning and are reconstructed by computers into 3D density images, which can clearly display the fine structures of bones and organs, including:

SARS-COV-2 CT-scan (25,26): the SARS-COV-2 CT-scan is a CT scan dataset of patients, which is categorized into two classes: whether a patient has COVID-19 or not. Among them, the training set contains 1,736 samples, the validation set 297 samples, and the test set 448 samples.

Brain Stroke CT (27): the Brain Stroke CT is a head CT scan dataset, which is categorized into three classes: ischemic stroke, hemorrhagic stroke, and non-stroke. Among them, the training set contains 4,654 samples, the validation set 798 samples, and the test set 1,198 samples.

MRI modality datasets

MRI images utilize strong magnetic fields and radiofrequency pulses to detect the signals of hydrogen atoms in water molecules within the human body, thereby generating high-contrast soft tissue images and having no ionizing radiation, including:

Breast cancer patients MRI: the breast cancer patients MRI is an MRI scan dataset of patients’ breasts, which is categorized into two classes: healthy (benign) and diseased (malignant). Among them, the training set contains 1,036 samples, the validation set 176 samples, and the test set 268 samples.

Brain Tumor MRI: the Brain Tumor MRI dataset comprises three datasets figshare, SARTAJ, and Br35H, totaling 7,023 human brain MRI images across four classes: glioma, meningioma, non-tumor, and pituitary tumor. The training set contains 4,914 samples, the validation set 843, and the test set 1,266.

Other modality datasets

Besides the aforementioned modalities, other medical images such as endoscopic, ultrasound, and dermatoscopic images utilize optical or acoustic technologies to directly visualize internal body cavities, fetal development, and skin surface conditions, thereby assisting in diagnosis, including:

Kvasir (28): Kvasir contains multiple classes of gastrointestinal anatomical landmarks, pathological findings, or endoscopic examination processes, with hundreds of images per class and 4,000 samples in total. The training set includes 2,800 samples, the validation set 480, and the test set 720.

Fetal-Planes-DB (29): the Fetal-Planes-DB is a large-scale dataset of maternal-fetal screening ultrasound images, comprising six categories: four common fetal anatomical planes, the maternal cervix, and a general category for other planes. It contains 12,400 samples in total, with 8,679 for training, 1,861 for validation, and 1,860 for testing.

HAM10000 (30,31): HAM10000 is a dataset of mobile-captured skin disease images covering seven types, totaling 10,015 samples. The training set contains 7,007 samples, the validation set 1,203, and the test set 1,805.

Results

This section provides a detailed description of the experimental setup and evaluates the proposed model on 9 public datasets.

Experimental details

To verify the effectiveness and competitiveness of WTEMamba in medical image classification, several visual networks with demonstrated strong performance in medical imaging tasks, along with baseline models, were selected as references. The dataset was split into training, validation, and test sets using four fixed random seeds (42, 123, 2024, and 2026) to ensure reproducibility. During data preprocessing, random cropping to 224×224 was applied to the training set, while fixed resizing was used for the validation and test sets, all sharing the same normalization. Datasets were loaded via ImageFolder, and a category-to-index mapping dictionary was constructed and saved as a JSON file, with output statistics including class counts and distribution. The batch size was set to 64. The AdamW optimizer was employed with an initial learning rate of 0.0001, β1=0.9, β2=0.999, and weight decay of 1e−4, along with the cross-entropy loss function. Training was conducted for a maximum of 150 epochs, with early stopping triggered when validation accuracy showed no improvement over a consecutive number of epochs.

Hardware configuration: The operating system is Ubuntu 20.04.3, and the GPUs include four NVIDIA A800 80GB PCIe cards. Software environment: We adopt the PyTorch network framework, including Python 3.10.7, PyTorch 1.13.0, and CUDA 12.1.

Evaluation metrics

To comprehensively quantify the model performance, this study adopted 6 standard evaluation metrics:

Overall accuracy (OA): it describes the overall prediction accuracy of the model, which is calculated as the proportion of correctly predicted samples to the total number of samples. The higher the OA value, the better the overall prediction performance of the model. The formula is as follows:

$$OA=\frac{TP+TN}{TP+TN+FP+FN}$$ [37]

Precision: it measures the reliability of positive predictions and is defined as the proportion of true positives to all positive predictions. The higher the Precision value, the more reliable the positive prediction results are. The formula is as follows:

$$Precision=\frac{TP}{TP+FP}$$ [38]

Sensitivity: it reflects the model’s ability to identify positive samples, which is calculated as the proportion of true positives to the actual number of positive samples. The higher the Sensitivity value, the stronger the model’s ability to identify positive samples. The formula is as follows:

$$Sensitivity=\frac{TP}{TP+FN}$$ [39]

Specificity: it evaluates the model’s ability to exclude negative samples and is defined as the proportion of true negatives to the actual number of negative samples. The higher the Specificity value, the stronger the model’s ability to exclude negative samples. The formula is as follows:

$$Specificity=\frac{TN}{TN+FP}$$ [40]

F1-score: it is the harmonic mean that integrates Precision and Sensitivity, used to balance the contributions of these two metrics. The higher the F1-score value, the better the model’s comprehensive performance in Precision and Sensitivity. The formula is as follows:

$$F1-score=\frac{2\times Precision\times Sensitivity}{Precision+Sensitivity}$$ [41]

AUC: it characterizes the overall discriminative ability of a model under different classification thresholds. For binary classification tasks, AUC is computed by plotting the true positive rate (TPR) against the false positive rate (FPR) and measuring the area under the curve. For multi-class datasets, we adopt the One-vs-Rest (OvR) strategy, where AUC is calculated for each class against all others, and the final reported value is the macro-average across all classes.

Furthermore, this study simultaneously analyzed the model’s floating-point operations (FLOPs) and parameters (Params) to assess its feasibility under resource-constrained conditions.

Experimental results

The performance of WTEMamba

To comprehensively and objectively evaluate the performance of the WTEMamba model, we conducted comparative experiments on nine datasets, with the compared models roughly categorized into three architectures: CNN, Transformer, and Mamba. All results are reported as mean±standard deviation over multiple runs, where the standard deviation reflects the performance fluctuations of the model under different data splits. The 95% confidence intervals for F1, OA, and AUC were estimated using the bootstrap percentile method. When WTEMamba achieved the best performance, statistical significance against the runner-up was assessed via a paired bootstrap test: in each resampling iteration, a bootstrap sample of the same size as the original fixed test set was drawn with replacement, the metric difference between WTEMamba and the compared model was computed on this sample, and the null hypothesis was tested using the two-sided quantiles of the empirical distribution of these differences. Significance levels are denoted as *P<0.05, **P<0.01, and ***P<0.001.

Table 3 presents the evaluation results of WTEMamba and the reference models on the CPN X-Ray and Knee X-Ray datasets. Specifically, on the CPN X-Ray dataset, WTEMamba achieved an OA of 97.07%, and its precision and F1-score, among other metrics, were significantly better than those of all compared models. It is worth noting that despite having only 11.0M parameters and 2.1G FLOPs, its performance still surpassed that of Medmamba-S (6) and EfficientVMamba-S (8), which have larger parameter counts, demonstrating a better balance between computational efficiency and performance. On the Knee X-Ray dataset, the model performed well in terms of F1-score and AUC value, indicating that it has good generalization ability across different types of X-ray image classification tasks.

Table 3. Performance comparison between WTEMamba and reference models on X-ray datasets.

Dateset	Model	FLOPs (G)	Params	P (%)	Se (%)	Sp (%)	F1 (%)	AUC	OA (%)
CPN X-Ray	DacNet (32)	2.9	8.1	96.93±0.28	96.82±0.30	98.07±0.15	96.85±0.30 (96.32–97.48)	0.9956±0.0004 (0.9949–0.9961)	96.79±0.31 (96.22–97.44)
	WTConvNeXt-S (33)	8.8	54.0	95.10±0.28	94.98±0.33	97.43±0.16	95.02±0.31 (94.44–95.61)	0.9933±0.0014 (0.9906–0.9958)	94.92±0.32 (94.35–95.55)
	GazeGNN (34)	2.1	13.5	96.13±0.62	96.07±0.58	97.97±0.30	96.09±0.60 (95.04–97.19)	0.9954±0.0010 (0.9933–0.9972)	95.99±0.61 (94.84–97.17)
	Swin Transformer-S (35)	8.7	48.8	96.01±0.69	95.89±0.74	97.89±0.38	95.93±0.74 (94.57–97.26)	0.9949±0.0014 (0.9926–0.9971)	95.83±0.76 (94.55–97.28)
	EfficientViT-M4 (36)	0.3	8.8	95.04±1.01	94.78±0.63	97.41±0.52	95.22±0.57 (94.28–96.21)	0.9927±0.0024 (0.9887–0.9967)	95.15±0.77 (93.86–96.41)
	PVT v2-B1 (37)	2.1	13.2	94.57±1.00	94.43±0.91	97.11±0.51	94.54±0.79 (93.18–95.96)	0.9915±0.0030 (0.9859–0.9967)	94.54±0.60 (93.44–95.65)
	MobileMamba-S6 (12)	1.0	15.0	96.09±0.29	95.95±0.38	97.92±0.20	95.97±0.36 (95.35–96.61)	0.9945±0.0012 (0.9925–0.9965)	95.89±0.38 (95.19–96.62)
	EfficientVMamba-S (8)	1.4	13.2	95.01±0.73	94.94±0.16	97.43±0.43	94.74±0.44 (93.95–95.54)	0.9935±0.0030 (0.9885–0.9982)	94.66±0.44 (93.85–95.46)
	Medmamba-S (6)	3.5	22.8	96.40±0.73	96.42±0.66	98.06±0.34	96.40±0.72 (95.15–97.59)	0.9950±0.0006 (0.9925–0.9964)	96.33±0.71 (95.12–97.49)
	WTEMamba	2.1	11.0	97.38±0.69†*	97.37±0.74†*	98.12±0.91†*	97.50±0.78 (96.16–98.78)†*	0.9958±0.0020 (0.9927–0.9990)†*	97.07±0.81 (95.69–98.38)†*
Knee X-ray	WTConvNeXt-T (33)	4.5	30.0	75.56±0.62	70.61±0.36	85.24±0.40	71.93±0.19 (71.61–72.25)	0.8744±0.0005 (0.8735–0.8752)	73.22±0.44 (72.45–73.93)
	GazeGNN (34)	2.1	13.5	74.27±0.50	66.27±0.45	83.79±0.37	67.49±0.39 (66.88–68.16)	0.8606±0.0009 (0.8589–0.8623)	70.33±0.52 (69.30–71.40)
	Swin Transformer-S (35)	8.7	48.8	75.97±0.32	66.16±0.56	83.26±0.53	67.65±0.37 (67.02–68.30)	0.8654±0.0009 (0.8637–0.8673)	69.54±0.44 (68.75–70.28)
	EfficientViT-M3 (36)	0.3	6.9	74.41±0.44	69.56±0.44	84.71±0.52	70.41±0.56 (69.45–71.39)	0.8719±0.0013 (0.8695–0.8743)	71.50±0.47 (70.58–72.08)
	PVT v2-B1 (37)	2.1	13.2	79.32±0.47	72.57±0.37	85.60±0.53	73.77±0.40 (73.08–74.44)	0.8658±0.0012 (0.8635–0.8678)	73.57±0.44 (72.75–74.38)
	MobileMamba-T4 (12)	0.7	14.2	75.42±0.52	73.96±0.48	86.51±0.48	74.01±0.47 (73.09–74.78)	0.8817±0.0012 (0.8796–0.8840)	74.51±0.71 (73.24–75.62)
	EfficientVMamba-T (8)	0.9	6.2	75.15±0.32	73.46±0.52	87.20±0.41	73.88±0.31 (73.32–74.37)	0.8824±0.0007 (0.8811–0.8837)	75.55±0.40 (74.78–76.26)
	Medmamba-T (6)	2.0	14.5	76.35±0.67	72.14±0.34	86.25±0.67	73.60±0.38 (72.94–74.23)	0.8781±0.0025 (0.8738–0.8824)	74.41±0.52 (73.57–75.22)
	WTEMamba	2.1	11.0	80.26±0.41†*	75.93±0.45†*	87.86±0.22†*	77.15±0.56 (76.06–78.13)†**	0.8853±0.0031 (0.8794–0.8912)†*	77.80±0.47 (76.84–79.43)†**

^†, best value in each column. Statistical significance is assessed using a bootstrap test compared with our model, and asterisks indicate significance levels relative to our model (*, P<0.05; **, P<0.01). Params, P, Se, Sp, and F1 respectively denote parameter size, precision, sensitivity, specificity, and F1-score. AUC, area under the curve; FLOPs, Floating Point Operations; OA, overall accuracy.

The ROC curves of the CPN X-Ray and Knee X-Ray datasets under different classification models are shown in Figure 5. Analysis of the ROC curves reveals that on the CPN X-Ray dataset, our model’s curve is the closest to the top-left corner, with its AUC value reaching 0.9958. On the Knee X-Ray dataset, when the FPR is low, the curve of the model in this study fluctuates slightly, while in the later stage, it becomes flat and remains at a high level. This indicates that WTEMamba has stronger discriminative ability in distinguishing positive and negative samples.

Figure 5.

ROC curves of different models on the X-ray datasets. (A) CPN X-Ray, (B) Knee X-Ray. AUC, area under the curve; ROC, receiver operating characteristic.

Table 4 presents the evaluation results of WTEMamba and reference models on the SARS-COV-2 CT-scan and Brain Stroke CT datasets. It can be seen that WTEMamba achieves an OA of 94.43% on the SARS-COV-2 CT-scan dataset, which is higher than that of Lung-EffNet (39) and Medmamba-T (6). Its specificity was particularly outstanding, demonstrating an efficient ability to capture subtle lesion features in COVID-19 CT images. However, its sensitivity is relatively lower than that of some comparative models, which may be attributed to the model’s emphasis on high-frequency edge enhancement and the inherent class imbalance in the dataset, leading to a trade-off that slightly compromises the recall of positive samples. On the Brain Stroke CT dataset, WTEMamba achieved an OA of 93.45%, and it performed particularly well in terms of sensitivity and F1-score, indicating that the model has strong discriminative ability for multi-classification tasks of stroke types.

Table 4. Performance comparison between WTEMamba and reference models on CT datasets.

Dateset	Model	FLOPs (G)	Params	P (%)	Se (%)	Sp (%)	F1 (%)	AUC	OA (%)
SARS-COV-2 CT-scan	ResNet50 (38)	4.1	25.6	92.99±0.43	93.29±0.29†	92.80±0.56	93.20±0.54 (92.32–93.97)	0.9772±0.0005 (0.9764–0.9779)	93.24±0.38 (92.78–93.75)
	Lung-EffNet (39)	1.4	6.5	96.42±0.53	90.45±0.72	96.47±0.50	93.31±0.80 (92.35–94.30)	0.9809±0.0017 (0.9787–0.9832)	93.44±0.95 (92.23–94.72)
	WTConvNeXt-T (33)	4.5	30.0	94.35±0.50	89.54±0.81	94.29±0.42	91.85±0.57 (91.09–92.64)	0.9634±0.0020 (0.9604–0.9653)	91.77±0.53 (91.08–92.36)
	PVT v2-B1 (37)	2.1	13.2	96.33±0.61	88.64±0.72	96.44±0.74	92.64±0.69 (91.54–93.53)	0.9579±0.0013 (0.9561–0.9591)	92.84±0.62 (91.74–93.72)
	MobileMamba-T4 (12)	0.7	14.2	92.69±0.56	90.93±0.43	92.79±0.45	91.66±0.69 (90.97–92.44)	0.9573±0.0034 (0.9529–0.9608)	91.66±0.58 (91.01–92.38)
	EfficientVMamba-S (8)	1.4	13.2	94.99±0.47	92.89±0.52	95.26±0.65	93.85±0.60 (92.82–94.91)	0.9843±0.0010 (0.9828–0.9853)	94.04±0.57 (93.05–94.95)
	Medmamba-T (6)	2.0	14.5	93.53±0.58	88.68±0.35	93.94±0.45	91.10±0.30 (90.54–91.66)	0.9682±0.0032 (0.9630–0.9731)	91.66±0.38 (91.09–92.24)
	WTEMamba	2.1	11.0	96.90±0.47†*	91.19±0.53	96.85±0.51†*	94.12±0.75 (92.90–95.32)†*	0.9867±0.0007 (0.9852–0.9881)†*	94.43±0.42 (93.74–95.03)†*
Brain Stroke CT	ResNet50 (38)	4.1	25.6	88.79±0.51	88.68±0.57	94.76±0.71	88.17±0.72 (87.30–88.96)	0.9768±0.0039 (0.9669–0.9827)	91.30±0.65 (90.53–92.03)
	EfficientViT-M4 (36)	0.3	8.8	89.82±0.53	86.44±0.84	93.83±0.62	87.74±0.67 (86.96–88.49)	0.9684±0.0045 (0.9637–0.9749)	91.14±0.66 (90.29–92.01)
	WTConvNeXt-S (33)	8.8	54.0	86.66±0.54	82.44±0.53	92.46±0.49	84.23±0.46 (83.74–84.87)	0.9548±0.0056 (0.9485–0.9627)	89.40±0.69 (88.40–90.31)
	PVT v2-B2 (37)	4.0	24.8	88.77±0.98	88.74±0.61	95.13±0.71	88.67±0.82 (87.45–89.51)	0.9756±0.0056 (0.9698–0.9813)	92.17±0.99 (91.21–93.68)
	MobileMamba-S6 (12)	1.0	15.0	88.72±0.83	83.39±0.56	92.79±0.23	85.37±0.42 (84.81–85.94)	0.9590±0.0055 (0.9535–0.9668)	89.80±0.38 (89.19–90.34)
	EfficientVMamba-S (8)	1.4	13.2	90.57±0.52	87.60±1.02	94.35±0.65	88.68±0.63 (88.00–89.37)	0.9760±0.0048 (0.9699–0.9799)	92.47±1.11 (91.52–93.86)
	Medmamba-S (6)	3.5	22.8	90.06±0.40	87.40±0.87	93.82±0.40	88.36±0.76 (87.49–89.25)	0.9687±0.0027 (0.9657–0.9717)	91.46±0.43 (90.94–91.99)
	WTEMamba	2.1	11.0	91.74±0.54†*	89.47±1.01†*	95.15±0.61†*	90.30±0.96 (88.56–92.04)†*	0.9770±0.0058 (0.9678–0.9862)†*	93.45±0.68 (92.36–94.53)†*

^†, best value in each column. Statistical significance is assessed using a bootstrap test compared with our model, and asterisks indicate significance levels relative to our model (*, P<0.05; **, P<0.01). Params, P, Se, Sp, and F1 respectively denote parameter size, precision, sensitivity, specificity, and F1-score. AUC, area under the curve; CT, computed tomography; FLOPs, Floating Point Operations; OA, overall accuracy.

The ROC curves of the SARS-COV-2 CT-scan and Brain Stroke CT datasets under different classification models are shown in Figure 6. Analysis of Figure 6 shows that the curve of the SARS-COV-2 CT-scan dataset rises rapidly in the low FPR region, with an AUC of 0.9867; the curve of the Brain Stroke CT dataset remains at a high level overall, with an AUC of 0.9770, significantly outperforming other comparative models. This indicates that WTEMamba has a stronger ability to differentiate between ischemic and hemorrhagic strokes.

Figure 6.

ROC curves of different models on the CT datasets. (A) SARS-COV-2 CT-scan, (B) Brain Stroke CT. AUC, area under the curve; CT, computed tomography; ROC, receiver operating characteristic.

Table 5 presents the evaluation results of WTEMamba and the reference models on the Breast Cancer Patients MRI and Brain Tumor MRI datasets. It can be observed that the proposed model achieved an OA of 87.86% on the Breast Cancer Patients MRI dataset, with outstanding specificity, which effectively enhanced the reliability of distinguishing benign and malignant lesions. However, the sensitivity is relatively low, primarily because the wavelet-enhanced branch amplifies ambiguous high-frequency components in ill-defined soft-tissue lesions, while the heterogeneous malignant cases in this MRI dataset further hinder the recall of positive samples. On the Brain Tumor MRI dataset, WTEMamba attained an OA of 96.69% and an F1-score of 96.57%, comprehensively outperforming LCDEiT (43) and Medmamba-S (6), and demonstrating robust multi-class tumor discrimination performance.

Table 5. Performance comparison between WTEMamba and reference models on MRI datasets.

Dateset	Model	FLOPs (G)	Params	P (%)	Se (%)	Sp (%)	F1 (%)	AUC	OA(%)
Breast Cancer Patients MRI	WTConvNeXt-T (33)	4.5	30.0	65.21±0.90	89.80±0.73	51.53±1.20	75.49±0.91(74.43–76.56)	0.7661±0.0062 (0.7564–0.7723)	70.54±1.13 (69.55–71.69)
	BTS-ST (40)	9.6	55.7	85.54±0.87	78.52±0.86	86.46±0.73	81.74±0.52 (81.05–82.24)	0.9220±0.0021 (0.9194–0.9246)	82.40±0.75 (81.28–83.24)
	Swin Transformer-S (35)	8.7	48.8	75.42±0.88	80.35±0.87	73.64±1.03	77.84±0.60 (77.02–78.44)	0.8550±0.0057 (0.8468–0.8595)	77.33±0.90 (76.68–78.63)
	EfficientViT-M2 (36)	0.2	4.2	82.29±0.84	90.17±0.87†	80.55±1.30	87.40±0.57 (86.72–88.11)	0.9407±0.0060 (0.9335–0.9485)	85.31±0.78 (84.10–86.20)
	PVT v2-B1 (37)	2.1	13.2	85.71±1.16	83.54±0.76	86.34±0.56	84.85±0.62 (84.13–85.39)	0.9379±0.0078 (0.9284–0.9482)	84.88±1.17 (83.96–85.54)
	MobileMamba-T4 (12)	0.7	14.2	90.03±1.26	76.55±1.06	91.79±0.65	82.95±0.71 (82.42–83.87)	0.9075±0.0089 (0.8955–0.9174)	84.36±1.19 (83.10–85.68)
	EfficientVMamba-T (8)	0.9	6.2	75.30±0.69	88.90±0.87	70.88±1.14	81.44±1.26 (80.07–82.79)	0.8935±0.0108 (0.8815–0.9069)	79.81±0.99 (78.48–81.05)
	Medmamba-T (6)	2.0	14.5	78.83±1.33	89.25±1.23	76.50±0.68	83.48±0.75 (82.33–84.14)	0.9290±0.0039 (0.9250–0.9344)	82.88±0.42 (82.37–83.46)
	WTEMamba	2.1	11.0	91.89±0.40†*	84.38±0.93	93.24±0.59†**	88.01±0.70 (87.09–88.84) †*	0.9532±0.0015 (0.9512–0.9550) †*	87.86±0.44 (87.34–89.36) †**
Brain Tumor MRI	MBTFCN (41)	4.1	8.2	94.85±0.93	94.31±0.97	98.23±0.32	94.23±0.73 (93.44–95.08)	0.9970±0.0012 (0.9952–0.9978)	94.56±1.06 (93.75–96.03)
	WTConvNeXt-S (33)	8.8	54.0	94.85±0.76	94.81±0.91	98.31±0.27	94.80±0.91 (93.54–95.72)	0.9948±0.0018 (0.9921–0.9967)	94.99±0.89 (93.87–95.89)
	Swin Transformer-S (35)	8.7	48.8	94.37±0.60	94.40±0.62	98.24±0.21	94.33±0.59 (93.54–94.98)	0.9934±0.0009 (0.9922–0.9943)	94.51±0.58 (93.78–95.20)
	MultiFeNet (42)	18.8	17.1	95.13±0.73	94.84±0.74	98.31±0.35	94.78±0.84 (93.96–95.90)	0.9965±0.0012 (0.9948–0.9978)	95.10±0.68 (94.60–95.98)
	LCDEiT (43)	6.1	3.4	96.00±0.77	95.82±0.91	98.65±0.23	95.77±0.98 (94.63–96.94)	0.9967±0.0006 (0.9959–0.9971)	96.01±0.93 (95.02–96.98)
	EfficientViT-M4 (36)	0.3	8.8	94.83±0.86	94.74±0.89	98.32±0.27	94.74±0.89 (93.54–95.70)	0.9945±0.0018 (0.9922–0.9966)	95.02±0.90 (93.84–95.80)
	PVT v2-B1 (37)	2.1	13.2	95.12±0.56	94.95±0.56	98.43±0.17	94.98±0.55 (94.38–95.70)	0.9961±0.0010 (0.9946–0.9968)	95.16±0.52 (94.56–95.82)
	MobileMamba-S6 (12)	1.0	15.0	93.58±0.83	93.50±0.91	97.72±0.53	93.51±0.92 (92.58–94.77)	0.9922±0.0012 (0.9912–0.9939	93.71±0.88 (92.70–94.90)
	EfficientVMamba-S (8)	1.4	13.2	95.29±0.94	95.24±0.93	98.46±0.35	95.21±1.03 (93.75–96.00)	0.9949±0.0016 (0.9926–0.9963)	95.35±1.02 (93.96–96.40)
	Medmamba-S (6)	3.5	22.8	95.42±0.98	95.51±1.10	98.58±0.37	95.40±1.08 (93.86–96.33)	0.9964±0.0019 (0.9935–0.9980)	95.60±1.00 (94.03–96.62)
	WTEMamba	2.1	11.0	96.33±0.69†*	96.36±1.05†*	98.64±0.82†*	96.57±0.90 (95.76–97.60) †**	0.9971±0.0011 (0.9955–0.9979) †*	96.69±0.64 (95.69–97.40) †*

^†, best value in each column. Statistical significance is assessed using a bootstrap test compared with our model, and asterisks indicate significance levels relative to our model (*, P<0.05; **, P<0.01). Params, P, Se, Sp, and F1 respectively denote parameter size, precision, sensitivity, specificity, and F1-score. AUC, area under the curve; FLOPs, Floating Point Operations; MRI, magnetic resonance imaging; OA, overall accuracy.

The ROC curves of the breast cancer patients MRI and brain tumor MRI datasets under different classification models are shown in Figure 7. For the Breast Cancer Patients MRI dataset, when the FPR is 0.07, the TPR of its curve reaches 0.85, with an AUC of 0.9532; the curve of the Brain Tumor MRI dataset remains smooth and at a high level overall, with an AUC of 0.9971, which is significantly better than that of MultiFeNet (42) and Swin Transformer-S (35). This indicates that WTEMamba has a relatively strong ability to differentiate between benign and malignant lesions.

Figure 7.

ROC curves of different models on the MRI datasets. (A) Breast Cancer Patients MRI, (B) Brain Tumor MRI. AUC, area under the curve; MRI, magnetic resonance imaging; ROC, receiver operating characteristic.

Table 6 presents the evaluation results of WTEMamba and reference models on the Kvasir, Fetal-Planes-DB, and HAM10000 datasets. Specifically, on the Kvasir, Fetal-Planes-DB, and HAM10000 datasets, even for images of diverse modalities, WTEMamba still outperformed all comparative models in metrics such as OA and AUC. Meanwhile, it also achieved low values in terms of FLOPs and parameter counts. This fully demonstrates the model’s strong adaptability to multiple types of medical image tasks.

Table 6. Performance comparison between WTEMamba and reference models on other datasets.

Dateset	Model	FLOPs (G)	Params	P (%)	Se (%)	Sp (%)	F1 (%)	AUC	OA(%)
Kvasir	ResNet50 (38)	4.1	25.6	79.24±1.22	79.16±1.15	97.19±0.33	78.83±1.14 (77.88–80.43)	0.9785±0.0035 (0.9744–0.9830)	79.13±0.88 (78.26–80.34)
	WTConvNeXt-S (33)	8.8	54.0	76.88±1.12	76.85±0.95	96.68±0.28	76.75±0.88 (76.09–78.03)	0.9715±0.0042 (0.9673–0.9770)	76.84±0.96 (75.97–78.19)
	EfficientViT-M4 (36)	0.3	8.8	79.71±0.42	80.19±1.05	97.32±0.36	80.44±0.86 (79.36–81.49)	0.9806±0.0042 (0.9766–0.9865)	80.46±0.98 (79.53–81.81)
	PVT v2-B1 (37)	2.1	13.2	79.88±0.85	80.81± 0.71	97.43± 0.35	80.51±0.57 (79.60–81.42)	0.9809±0.0027 (0.9767–0.9851)	81.04±1.03 (79.40–82.68)
	MobileMamba-S6 (12)	1.0	15.0	79.96±0.75	79.52±0.99	97.22±0.21	79.28±0.59 (78.16–80.32)	0.9783±0.0038 (0.9732–0.9819)	79.50±0.44 (78.70–80.24)
	EfficientVMamba-S (8)	1.4	13.2	77.05±0.86	77.12±0.97	96.61±0.65	77.00±0.93 (75.67–77.91)	0.9778±0.0034 (0.9731–0.9808)	77.04±0.89 (75.86–77.92)
	Medmamba-T (6)	2.0	14.5	80.86±1.08	80.09±0.92	97.28±0.23	79.94±0.93 (78.54–80.85)	0.9785±0.0020 (0.9758–0.9809)	80.29±0.92 (78.91–80.95)
	WTEMamba	2.1	11.0	82.08±0.85†*	82.01±1.04†*	97.50±0.40†*	81.87±1.34 (80.56–82.80)†*	0.9833±0.0035 (0.9780–0.9865) †*	81.81±1.33 (80.34–82.79) †*
Fetal-Planes-DB	ResNet50 (38)	4.1	25.6	93.13±0.49	90.37±1.57	98.42±0.37	91.34±0.93 (90.68–92.64)	0.9941±0.0010 (0.9924–0.9951)	93.03±0.37 (92.27–93.58)
	WTConvNeXt-B (33)	15.5	93.0	91.79±0.74	89.00±0.87	98.23±0.38	90.26±0.80 (89.55–91.37)	0.9927±0.0027 (0.9888–0.9949)	92.33±0.99 (91.22–93.69)
	Swin Transformer-B (35)	15.4	86.7	91.38±0.54	88.55±0.52	98.14±0.34	89.73±0.71 (88.62–90.32)	0.9919±0.0025 (0.9890–0.9952)	91.73±0.58 (90.86–92.31)
	EfficientViT-M5 (36)	0.5	12.4	90.94±0.74	86.10±1.25	97.71±0.49	88.21±0.48 (87.55–88.68)	0.9893±0.0014 (0.9876–0.9910)	90.07±0.38 (89.63–90.90)
	PVT v2-B3 (37)	6.7	45.2	92.05±0.60	91.68±0.88	98.48±0.30	91.85±0.50 (91.15–92.27)	0.9932±0.0012 (0.9915–0.9942)	93.02±0.53 (92.45–93.74)
	MobileMamba-B1 (12)	1.5	17.1	90.88±0.59	88.86±1.01	98.08±0.30	89.79±0.52 (89.14–90.44)	0.9904±0.0008 (0.9894–0.9914)	91.23±0.49 (90.60–91.88)
	EfficientVMamba-B (8)	4.1	34.3	91.85±0.69	93.56±0.75	98.64±0.32	92.62±0.99 (91.21–93.53)	0.9933±0.0016 (0.9915–0.9954)	93.48±0.39 (92.81–93.82)
	Medmamba-B (6)	7.4	47.1	92.24±0.74	92.81±1.46	98.70±0.26	93.14±0.89 (91.83–94.10)	0.9936±0.0023 (0.9903–0.9954)	93.83±0.62 (93.04–94.55)
	WTEMamba	2.1	11.0	93.39±0.60†*	93.81±0.97†*	98.71±0.32†*	93.51±0.80 (92.20–93.97) †*	0.9944±0.0036 (0.9893–0.9980) †*	94.39±0.66 (93.75–95.23) †*
HAM10000	VGG16 (44)	15.7	134.3	58.36±1.08	53.36±1.20	94.17±0.47	55.32±0.56 (54.43–56.17)	0.9354±0.0040 (0.9290–0.9418)	77.63±0.87 (76.23–78.99)
	WTConvNeXt-B (33)	15.5	93.0	58.20±1.08	54.99±0.80	94.51±0.35	56.10±0.72 (55.35–57.06)	0.9306±0.0035 (0.9273–0.9357)	78.55±0.57 (77.93–79.34)
	Swin transformer-B (35)	15.4	86.7	62.73±1.23	56.32±1.36	93.64±0.89	58.95±0.99 (57.96–79.33)	0.9430±0.0046 (0.9364–0.9472)	78.95±0.99 (77.96–80.33)
	PVT v2-B3 (37)	6.7	45.2	66.31±1.28	59.06± 1.27	94.75±0.33	61.95±1.01 (60.96–63.20)	0.9474±0.0012 (0.9460–0.9492)	79.55±0.87 (78.35–80.28)
	MobileMamba-B1 (12)	1.5	17.1	56.14±1.11	53.10±1.08	94.38±0.36)	51.84±0.91 (50.97–53.08)	0.9416±0.0036 (0.9378–0.9465)	78.56±0.71 (77.74–79.41)
	EfficientVMamba-B (8)	4.1	34.3	64.24±0.91	58.84±1.44	94.65±0.61	60.84±0.84 (60.16–62.09)	0.9462±0.0031 (0.9423–0.9505)	80.72±0.92 (79.81–81.95)
	Medmamba-B (6)	7.4	47.1	64.42±0.87	62.87±0.72	95.30±0.93	63.58±0.75 (62.55–64.32)	0.9403±0.0023 (0.9381–0.9434)	81.89±0.97 (80.66–82.91)
	WTEMamba	2.1	11.0	68.14±0.80†**	65.36±0.72†*	95.53±0.43†*	66.36±1.04 (65.11–67.78)†**	0.9486±0.0015 (0.9462–0.9498)†*	83.35±0.82 (82.15–84.16)†**

^†, best value in each column. Statistical significance is assessed using a bootstrap test compared with our model, and asterisks indicate significance levels relative to our model (*, P<0.05; **, P<0.01). Params, P, Se, Sp, and F1 respectively denote parameter size, precision, sensitivity, specificity, and F1-score. AUC, area under the curve; FLOPs, Floating Point Operations; OA, overall accuracy.

The ROC curves of the Kvasir, Fetal-Planes-DB, and HAM10000 datasets under different classification models are shown in Figure 8. It can be observed that the ROC curve of WTEMamba demonstrates strong performance across all three tasks. For the Kvasir dataset, when FPR <0.05, TPR exceeds 0.8, with an AUC of 0.9833; the ROC curve of the Fetal-Planes-DB dataset remains stable and at a high level overall, achieving an AUC of 0.9944; although the ROC curve of the HAM10000 dataset shows slight fluctuations due to class imbalance, its AUC still reaches 0.9486. These results confirm the model’s generalization ability on medical images of different modalities.

Figure 8.

ROC curves of different models on other datasets. (A) Kvasir, (B) Fetal-Planes-DB, (C) HAM10000. AUC, area under the curve; ROC, receiver operating characteristic.

Lastly, we conduct a comprehensive analysis of the performance of the WTEMamba model across various modalities. Figures 9-12 present the OA and AUC values of different models under different modalities.

Figure 9.

Performance of different models on the X-ray datasets. (A) CPN X-Ray, (B) Knee X-Ray. AUC, area under the curve; OA, overall accuracy.

Figure 10.

Performance of different models on the CT datasets. (A) SARS-COV-2 CT-scan, (B) Brain Stroke CT. AUC, area under the curve; CT, computed tomography; OA, overall accuracy.

Figure 11.

Performance of different models on the MRI datasets. (A) Breast Cancer Patients MRI, (B) Brain Tumor MRI. AUC, area under the curve; MRI, magnetic resonance imaging; OA, overall accuracy.

Figure 12.

Performance of different models on other datasets. (A) Kvasir, (B) Fetal-Planes-DB, (C) HAM10000. AUC, area under the curve; OA, overall accuracy.

Comprehensive result analysis of the X-ray modality

In the X-ray modality, WTEMamba achieved an OA of 97.07% on the CPN X-Ray dataset, outperforming the second-ranked DacNet by 0.28%. On the Knee X-Ray dataset, its OA exceeded that of EfficientVMamba-T by nearly 2 percentage points, with an AUC increase of 0.0029. This advantage stems from the model’s efficient fusion of local details and global semantics. Moreover, although WTEMamba outperformed other models on the Knee dataset in terms of OA and AUC, its overall performance was lower than on the CPN dataset, likely due to more subtle and morphologically variable lesion features, as well as potential class imbalance or image quality differences.

Comprehensive result analysis of the CT modality

In the CT modality, WTEMamba outperformed the best comparative model, EfficientVMamba, on both the SARS-COV-2 CT-Scan and Brain Stroke CT datasets, with OA improvements of 0.39% and 0.98%, and AUC improvements of 0.0024 and 0.0010, respectively. This advantage stems from the high contrast and the clear anatomical structures of CT images. Due to the varying X-ray absorption coefficients of different tissues, distinct grayscale differences and boundaries exist between them. For example, COVID-19-related lesions such as ground-glass opacities exhibit density differences from normal lung tissue, while hemorrhagic and ischemic regions in brain CT show prominent contrast. The core FIHM module in WTEMamba establishes long-range dependencies through the SS2D branch to capture the overall pathological distribution, and leverages the WTM branch to enhance high-frequency edge and detail features, which aligns well with the physical characteristics of CT images. This global-local synergy enables the model to fully exploit the discriminative information in CT data, thereby achieving efficient and accurate classification.

Comprehensive result analysis of the MRI modality

In MRI modality tasks, WTEMamba performed prominently on the Breast Cancer Patients MRI and Brain Tumor MRI datasets. On the Breast Cancer Patients MRI dataset, WTEMamba outperformed EfficientViT by 2.55% in OA and 0.0125 in AUC. On the Brain Tumor MRI dataset, WTEMamba outperformed Medmamba by 1.09% in OA and 0.0007 in AUC. MR images feature relatively high soft tissue contrast, blurred lesion boundaries, and more subtle texture variations. While enhancing discriminative features, the wavelet transform may also strengthen semantically ambiguous local details. Although these improvements alleviate semantic modeling issues, the complex morphological features and pixel-level semantic variations in MRI still pose challenges for representational learning. For example, the intricate tissue structures and relatively large intra-class variation in breast MRI may explain why its performance is slightly lower than that of brain tumor tasks. Nevertheless, this still demonstrates WTEMamba’s capability and adaptability in processing complex MR images.

Comprehensive result analysis of other modalities

In other medical modalities, WTEMamba also exhibits relatively good cross-modal adaptability. Leveraging the three-stage feature pyramid and the FIHM module, the model simultaneously captures subtle mucosal lesions in the digestive tract, easily confused fetal ultrasound planes, and fine dermatological differences. By retaining local discriminative features while integrating global context, it achieves stable and accurate classification across complex medical images with numerous categories and small inter-class differences.

Notably, it achieves comparatively high performance with only 11.0M parameters and 2.1G FLOPs, demonstrating a clear efficiency advantage. In CT tasks, the model’s capabilities complement the image characteristics, showing relatively good application potential; in Xray, MRI, and other modalities, it also exhibits reliable generalization and practical value, offering an important reference for lightweight multimodal medical image analysis.

Visual interpretability of WTEMamba

In this study, the GradCAM method was used to perform visualization analysis on typical images from the 9 datasets, and the resulting heatmaps are shown in Figure 13. With the help of GradCAM, WTEMamba successfully generated locally meaningful heatmaps (with red, yellow, and blue indicating high, medium, and lowcorrelation regions, respectively), and the predicted confidence scores are annotated below each subfigure. To highlight the comparison with relevant models, we select Medmamba as the baseline and MobileMamba, which also adopts a threestage lightweight architecture, as representatives. From the confidence scores, it can be seen that WTEMamba achieves relatively high confidence on datasets such as CPN X-Ray, but some comparison models (e.g., Medmamba) achieve higher scores on Kvasir, indicating that our model still has room for optimization. Meanwhile, the confidence on the Breast Cancer Patients MRI dataset is relatively low, and the corresponding attention regions in the heatmaps are somewhat dispersed, which can be regarded as a challenging case for the current architecture. Nevertheless, the model can still focus on the lesion regions relatively accurately on datasets such as chest Xray, head CT, and dermoscopy, with little interference from background.

Figure 13.

The visual heatmaps of WTEMamba using Grad-CAM on 9 datasets. CT, computed tomography; MRI, magnetic resonance imaging.

Ablation study

To verify the effectiveness of each module in the model proposed in this paper, ablation experiments were conducted on five datasets. Table 7 details the performance changes of the WTEMamba model when using different modules, with a focus on the channel shuffle (Shuffle), the DWTM-Branch (DWTM), the simplified three-stage network (TSN), and the enhanced classifier (EC). If the DWTM-Branch is not included, the model contains only the Conv-Branch.

Table 7. Ablation experiments on the WTEMamba model across different datasets.

Dataset	Model	Shuffle	DWTM	TSN	EC	Evaluation metrics
						P (%)	Se (%)	Sp (%)	F1 (%)	AUC	OA (%)
CPN X-Ray	Experiment 1	×	×	×	×	94.72	94.64	96.74	94.65	0.9946	94.54
	Experiment 2	✔	×	×	×	95.16	94.95	97.41	95.01	0.9946	94.90
	Experiment 3	✔	✔	×	×	96.34	96.27	97.76	96.60	0.9949	96.26
	Experiment 4	✔	✔	✔	×	97.07	97.12	98.03	97.14	0.9951	97.00
	Experiment 5	✔	✔	✔	✔	97.38†	97.37†	98.12†	97.50†	0.9958†	97.07†
SARS-COV-2 CT-scan	Experiment 1	×	×	×	×	93.93	88.94	94.14	91.36	0.9681	91.52
	Experiment 2	✔	×	×	×	94.50	91.15	94.59	92.79	0.9743	92.86
	Experiment 3	✔	✔	×	×	96.24	90.71	96.40	93.39	0.9781	93.53
	Experiment 4	✔	✔	✔	×	93.42	94.25†	93.24	93.83	0.9835	93.75
	Experiment 5	✔	✔	✔	✔	96.90†	91.19	96.85†	94.12†	0.9867†	94.43†
Brain Tumor MRI	Experiment 1	×	×	×	×	93.97	93.85	98.07	93.82	0.9937	94.15
	Experiment 2	✔	×	×	×	95.35	95.27	98.51	95.30	0.9960	95.50
	Experiment 3	✔	✔	×	×	95.50	95.22	98.47	95.28	0.9955	95.42
	Experiment 4	✔	✔	✔	×	96.31	96.24	98.63	96.27	0.9967	96.45
	Experiment 5	✔	✔	✔	✔	96.33†	96.36†	98.64†	96.57†	0.9971†	96.69†
Kvasir	Experiment 1	×	×	×	×	79.88	78.91	96.99	78.57	0.9765	78.91
	Experiment 2	✔	×	×	×	80.21	79.84	97.12	79.79	0.9790	79.84
	Experiment 3	✔	✔	×	×	81.36	80.83	97.26	80.52	0.9806	80.83
	Experiment 4	✔	✔	✔	×	81.70	81.39	97.34	81.41	0.9807	81.39
	Experiment 5	✔	✔	✔	✔	82.08†	82.01†	97.50†	81.87†	0.9833†	81.81†
Fetal-Planes-DB	Experiment 1	×	×	×	×	93.09	90.60	98.41	91.72	0.9933	92.96
	Experiment 2	✔	×	×	×	93.16	92.72	98.45	92.40	0.9940	93.38
	Experiment 3	✔	✔	×	×	93.53†	90.65	98.48	91.96	0.9938	93.33
	Experiment 4	✔	✔	✔	×	92.29	92.90	98.59	92.42	0.9940	93.49
	Experiment 5	✔	✔	✔	✔	93.39	93.81†	98.71†	93.51†	0.9944†	94.39†

^†, best value in each column. P, Se, Sp, and F1 respectively denote precision, sensitivity, specificity, and F1-score. AUC, area under the curve; CT, computed tomography; DWTM, wavelet-enhanced selective state-space; EC, enhanced classifier; MRI, magnetic resonance imaging; OA, overall accuracy; TSN, three-stage network.

Shuffle (Experiment 2): brought stable gains, with OA improvements ranging from 0.36% to 1.35%, and an average improvement of about 0.88% across the five datasets. The benefits were more pronounced on the CT and MRI datasets, indicating that channel shuffling is effective in alleviating the information isolation caused by grouped convolutions.

DWTM (Experiment 3): the contribution varied across datasets. It yielded significant improvements on CPN X-Ray and Kvasir, but caused slight decreases on Brain Tumor MRI and Fetal-Planes-DB. This suggests that wavelet enhancement is particularly advantageous for images with clear structures and sharp edges (e.g., Xray, endoscopy), but may introduce redundant highfrequency components in softtissue or ultrasound images, requiring subsequent modules for synergistic correction.

TSN (Experiment 4): Brain Tumor MRI showed the largest improvement, followed by CPN X-Ray and Kvasir, while the improvements on SARS-COV-2 CT-scan and Fetal-Planes-DB were relatively modest. Overall, this validates the enhancing effect of the three-stage structure on complex lesion representation.

EC (Experiment 5): the enhanced classifier was most prominent on Fetal-Planes-DB, with SARS-COV-2 CT-scan and Kvasir also showing clear improvements, whereas the gain on CPN was minimal. This indicates that the progressive dimensionality-reduction classifier has a stronger ability to optimize decision boundaries for samples with semantic ambiguity and small interclass differences (e.g., multi-class ultrasound).

The complete model achieved superior performance on all datasets, verifying the synergistic effectiveness of all modules. In particular, although DWTM introduced slight fluctuations when added alone on a few datasets, it ultimately yielded positive gains when combined with TSN and EC, demonstrating the necessity of joint multimodule optimization.

Figure 14 presents the radar charts of ablation experiments across datasets. Overall, on the CPN X-Ray, Brain Tumor MRI, Kvasir, and Fetal-Planes-DB datasets, the radar charts all approach regular pentagons, indicating that the five evaluation dimensions achieve relatively balanced optimization, and the globallocal feature fusion effectively enhances diagnostic robustness. Among them, the F1-score axes on CPN X-Ray and Brain Tumor MRI exhibit particularly notable expansion, reflecting the enhanced representation of complex lesions by WTEMamba. It is worth noting that on the SARS-COV-2 CT-Scan dataset, although the complete model’s radar chart is the fullest across all dimensions, the Sensitivity axis shows a relative retraction while the Specificity and Precision axes expand markedly, suggesting that the model makes a certain trade-off in recall of positive samples while improving negative sample discrimination, yet the overall balance remains superior to other configurations. This phenomenon does not appear on the Kvasir and Fetal-Planes-DB datasets, verifying the adaptive adjustment capability of the module combination across different modalities.

Figure 14.

Radar chart of ablation experiments on different datasets. (A) CPN X-Ray, (B) SARS-COV-2 CT-scan, (C) Brain Tumor MRI, (D) Kvasir, (E) Fetal-Planes-DB. P, Se, Sp, and F1 respectively denote precision, sensitivity, specificity, and F1-score. CT, computed tomography; MRI, magnetic resonance imaging.

Hyperparameter analysis

Table 8 summarizes the performance of different main-to-auxiliary path weight ratios on the three medical image datasets. The experimental results show that the weight ratio of 0.7 for the main path and 0.3 for the auxiliary path achieves relatively high overall metrics across multiple datasets. On the CPN X Ray dataset, this configuration yields higher OA, F1, and AUC than the other four groups, although its Specificity (98.12%) is slightly lower than the 98.44% obtained with the 0.5/0.5 setting, indicating that while this weight ratio improves overall discriminative ability, it does not achieve the optimal performance in terms of negative sample discrimination. On the SARS COV 2 CT and Brain Tumor MRI datasets, the 0.7/0.3 combination also shows leading OA and AUC, but its Specificity differs little from that of the 0.6/0.4 or 0.5/0.5 configurations. These experiments verify the effectiveness of the primary auxiliary fusion strategy in cross modal tasks, and further confirm that the 0.7/0.3 allocation is more adaptable to the pathological feature variations across different diseases.

Table 8. Performance comparison of different main-to-auxiliary path weight ratios on three medical image classification datasets.

Dataset	Main path	Auxiliary path	Evaluation metrics
			P (%)	Se (%)	Sp (%)	F1 (%)	AUC	OA (%)
CPN X-Ray	0.5	0.5	97.00	96.94	98.44†	96.97	0.9956	96.92
	0.6	0.4	96.25	95.99	98.26	96.04	0.9951	95.96
	0.7	0.3	97.38†	97.37†	98.12	97.50†	0.9958†	97.07†
	0.8	0.2	96.12	95.94	97.93	95.97	0.9942	95.91
	0.9	0.1	96.29	96.26	98.10	96.28	0.9940	96.22
SARS-COV-2 CT-scan	0.5	0.5	96.22	88.56	96.45	92.23	0.9720	92.46
	0.6	0.4	96.28	90.05	96.45	93.06	0.9697	93.22
	0.7	0.3	96.90†	91.19†	96.85†	94.12†	0.9867†	94.43†
	0.8	0.2	96.84	90.05	96.77	93.78	0.9785	93.97
	0.9	0.1	94.74	89.55	94.92	92.07	0.9767	92.21
Brain Tumor MRI	0.5	0.5	96.31	96.28	98.46	96.38	0.9961	96.52
	0.6	0.4	93.96	93.88	98.07	93.89	0.9932	94.15
	0.7	0.3	96.33†	96.36†	98.64†	96.57†	0.9971†	96.69†
	0.8	0.2	93.08	92.92	97.73	92.90	0.9893	93.13
	0.9	0.1	94.94	94.89	98.38	94.90	0.9952	95.10

^†, best value in each column. P, Se, Sp, and F1 respectively denote precision, sensitivity, specificity, and F1-score. AUC, area under the curve; CT, computed tomography; MRI, magnetic resonance imaging; OA, overall accuracy.

Performance analysis of Mamba

To validate the effectiveness of SS2D in the FIHM module, we replace it with a standard Transformer (w/ Transformer) and compare the two on five datasets. Table 9 shows that the Mamba-based WTEMamba consistently outperforms the Transformer counterpart across all comparisons. For example, OA on CPN X-Ray rises from 95.59% to 97.07%, and on Fetal-Planes-DB from 90.16% to 94.39%. Figure 15 visualises the percentage improvements for each metric. On CPN X-Ray, OA, F1 and AUC improve by 1.5%, 1.9% and 0.2%, respectively. On SARS-COV-2 CT-scan, AUC gains 0.6%. On Fetal-Planes-DB, OA improves by 4.7%. This superiority stems from the input-dependent dynamic parameterisation of Mamba’s selective state-space model, which adaptively decides which pathological features to retain or ignore. This capability is crucial for medical images with blurred lesion boundaries and subtle inter-class differences. Meanwhile, its linear computational complexity avoids the quadratic overhead of Transformers, enabling efficient processing of high-resolution feature maps.

Table 9. Mamba vs. Transformer: quantitative comparison under the proposed framework.

Dataset	Model	Evaluation metrics
		P (%)	Se (%)	Sp (%)	F1 (%)	AUC	OA (%)
CPN X-Ray	w/ Transformer	95.71	95.64	97.78	95.65	0.9934	95.59
	w/ Mamba (ours)	97.38†	97.37†	98.12†	97.50†	0.9958†	97.07†
SARS-COV-2 CT-scan	w/ Transformer	96.24	90.71	96.40	93.39	0.9813	93.53
	w/ Mamba (ours)	96.90†	91.19†	96.85†	94.12†	0.9867†	94.43†
Brain Tumor MRI	w/ Transformer	95.21	94.96	98.40	95.02	0.9956	95.18
	w/ Mamba (ours)	96.33†	96.36†	98.64†	96.57†	0.9971†	96.69†
Kvasir	w/ Transformer	81.58	80.78	97.25	80.53	0.9785	80.78
	w/ Mamba (ours)	82.08†	82.01†	97.50†	81.87†	0.9833†	81.81†
Fetal-Planes-DB	w/ Transformer	90.43	87.58	97.75	88.56	0.9896	90.16
	w/ Mamba (ours)	93.39†	92.81†	98.71†	93.51†	0.9944†	94.39†

^†, best value in each column. P, Se, Sp, and F1 respectively denote Precision, Sensitivity, Specificity, and F1-score. AUC, area under the curve; CT, computed tomography; DWTM, wavelet-enhanced selective state-space; EC, enhanced classifier; MRI, magnetic resonance imaging; OA, overall accuracy; W/, with.

Figure 15.

Percentage improvement of Mamba over Transformer on the five datasets. (A) CPN X-Ray, (B) SARS-COV-2 CT-scan, (C) Brain Tumor MRI, (D) Kvasir, (E) Fetal-Planes-DB. AUC, area under the curve; CT, computed tomography; MRI, magnetic resonance imaging; OA, overall accuracy.

Discussion

WTEMamba achieves better classification performance on nine public medical image datasets covering Xray, CT, MRI, endoscopy, ultrasound, and dermoscopy. On CPN X-Ray, it reaches an OA of 97.07% and an AUC of 0.9958, significantly outperforming competitors such as DacNet, Swin Transformer, and Medmamba. On SARS-COV-2 Ct-Scan and Brain Tumor MRI, it achieves OAs of 94.43% and 96.69%, respectively, also outperforming existing methods. For datasets with many classes and subtle inter-class differences, such as Breast Cancer Patients MRI and Fetal Planes DB, WTEMamba maintains a stable advantage, fully demonstrating its strong adaptability to multi-modal medical images. Ablation studies validate the effectiveness of channel shuffling, the waveletenhanced branch, the three-stage network, and the enhanced classifier, with the wavelet branch contributing most significantly to improvements in specificity and precision. Replacing SS2D in FIHM with a standard Transformer leads to clear drops in OA, AUC, and other metrics across five datasets, confirming that the selective statespace mechanism of Mamba, with its input-dependent dynamic parameterisation and linear complexity, is more suitable for modelling long-range semantic dependencies in medical images. Current limitations include relatively lower sensitivity on the SARS-CoV-2 CT-scan and Breast Cancer MRI datasets, and all experiments rely solely on public benchmarks. In clinical practice, performance may fluctuate due to variations in equipment, protocols, noise, and artifacts. To address this, we plan to collect multi-centre clinical data to systematically assess and improve model robustness under heterogeneous conditions, thereby improving clinical reliability.

Conclusions

This study proposes the WTEMamba framework, which aims to address the trade-off challenge between global dependency modeling and computational efficiency in medical image classification. By adopting a three-stage lightweight network architecture, this framework not only significantly improves the model’s inference speed but also more effectively meets the requirements of clinical real-time diagnosis. The core contribution lies in the design of the FIHM, which innovatively integrates the local feature extraction capability of the convolutional branch, the global dependency modeling advantage of the selective state-space model, and the enhancement mechanism of wavelet transform for high-frequency details. While significantly enhancing classification performance, it overcomes the quadratic computational complexity bottleneck of traditional Transformer-based models. Experimental results demonstrate that WTEMamba exhibits outstanding performance across nine medical image datasets covering multiple modalities including X-ray, CT, and MRI. Furthermore, while maintaining high accuracy, the model achieves a substantial improvement in inference speed, providing a new technical pathway and solution for efficient and reliable computer-aided diagnostic systems in practical medical scenarios. Future work will focus on three aspects: (I) further improving the model’s generalization ability to high-resolution pathological images and exploring deployment optimization for edge computing devices; (II) deeply analyzing the model’s decision-making mechanism by integrating interpretability analysis tools; and (III) further investigating the transfer potential of WTEMamba in downstream tasks such as medical image segmentation, detection, and reconstruction, so as to promote the widespread adoption of selective state-space models in the field of intelligent healthcare.

Supplementary

The article’s supplementary files as

Ethical Statement: The authors are accountable for all aspects of the work in ensuring that questions related to the accuracy or integrity of any part of the work are appropriately investigated and resolved. The study was conducted in accordance with the Declaration of Helsinki and its subsequent amendments.