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

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.

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.

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.

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.