Neural Memory Self-Supervised State Space Models With Learnable Gates
Zhihua Wang, Yuxin He, Yi Zhang, Tao He, Jiajun Bu · IEEE Signal Processing Letters · 2025
Discrete Ordinary Differential Equations (ODEs) have been employed to develop lightweight deep neural networks in recent years. In this letter, we introduce a novel lightweight UNet variant called the neural memory Self-Supervised State Space Model (nmS4M-UNet) for medical image segmentation, where discrete ODEs serve as the decoder. The proposed nmS4M block has learnable gates and performs multi-head computation to enhance memory updates. Additionally, the nmS4M-UNet incorporates a self-supervised learning branch to improve feature extraction capabilities. The intermediate features are reused as partial input to the decoder, helping to mitigate network overfitting. The nmS4M-UNet reduces the number of parameters by 29.70% compared to the standard UNet. Experimental results on the PH2, ISIC2018, and BU-COCO datasets demonstrate that the proposed nmS4M-UNet achieves performance comparable to state-of-the-art models.