U-MLLA: A Cognitive-Inspired Enhancement of Linear Attention for Medical Image Segmentation

Medical image segmentation is fundamental to computer-assisted diagnosis but faces challenges across diverse imaging modalities. Linear attention mechanisms succeed in natural images but are limited in medical segmentation due to insufficient spatial dependency and tissue heterogeneity modeling. Research indicates that successful dense prediction requires balanced permutation variance, strong inductive capabilities, and precise absolute position information. Current linear attention approaches satisfy the first two requirements but critically lack the third, significantly impacting medical segmentation where spatial localization is essential. To address these limitations, we propose U-MLLA, which integrates U-Net with mamba-like linear attention (MLLA) for multiscale feature and context capture. We further introduce complementary conditional and absolute positional encoding (APE) to compensate for position information deficits in linear attention. Experiments show U-MLLA provides robust features, and the complementary strategies significantly improve multi-organ and tumor segmentation. APE particularly excels with complex structures requiring precise boundary delineation. This cognitively inspired architecture adapts 93% of ImageNet-1k weights and increases effectiveness. Comprehensive evaluations across six challenging datasets (e.g., FLARE22, AMOS22CT/MR, ACDC) and 24 tasks, U-MLLA achieves state-of-the-art performance with an average DSC of 88.32%, outperforming nnUNetV2-2D and SwinUNetR by 4.37% and 1.98%. These results highlight U-MLLA’s potential for clinical applications that require precise anatomical delineation, where APE is essential for maintaining spatial context and differentiating similar structures. The code is available at https://github.com/csyfjiang/U-MLLA.