7.Harmonious parameters and performance: Lightweight convolutional stage and local feature weighted fusion MLP for medical image segmentation
Published in Biomedical Signal Processing and Control, 2024
Transformer-based models is widely applied in the field of medical image processing in recent years, thanks to its advantage in capturing global representations. However, while reaping high performance, the significant computational complexity, high training cost, and redundant dependencies of the model cannot be ignored. Therefore, our attention is drawn to the Multilayer Perceptron (MLP) as the complex self-attention alternative. We propose a lightweight image segmentation model based on depth-wise separable convolution and MLP, named UConvNeXt. It comprises both initial convolutional stages and MLP stages in the latent stage. Specifically, we employ large-scale kernel depth-wise separable convolution to replace the convolution blocks in UNet. The parameter is significantly reduced while maintaining the performance of the original model. In contrast to static feature fusion, we propose a novel local feature weighted fusion MLP (LFWF-MLP) module. In this module, contextual information is captured by shifting all tokens along the spatial directions. Feature correlations between different positions are weighted fused, and key regions receive increased attention, ultimately enhancing the segmentation performance. The experimental results demonstrate that, compared to models with similar parameter levels, more powerful segmentation performance is exhibited by UConvNeXt. Moreover, when compared to models with higher parameters than ours, comparable or even slightly superior segmentation results are achieved by UConvNeXt. In terms of model scale, compared to the UNet, parameters is reduced by 17 times, and computational complexity is lowered by 14 times, and inference speed is improved by 3 times in UConvNeXts.
Recommended citation:
Harmonious parameters and performance: Lightweight convolutional stage and local feature weighted fusion MLP for medical image segmentation, Y.-X. Chen, Y.-J. Xiong*, X.-H. Qiu and C.-M. Xia*, Biomedical Signal Processing and Control, 2024, 98: 106726
Download Paper