Mamba-Inspired AI Reads Anatomy in Four Directions to Sharpen Medical Image Segmentation

Medical image segmentation has long been a stubborn bottleneck in computer-aided diagnosis. The task sounds simple in principle—trace the outline of an organ or lesion pixel by pixel—but in practice it forces algorithms to contend with complex organ boundary morphologies, structures that vary wildly in scale, and the chronically low contrast that plagues MRI and CT scans. A new open-access study published in Cognitive Computation on 24 October 2026 proposes a way through this thicket. The paper introduces PF-Mamba-UNet, an Anatomical-prior-guided Fusion Mamba-UNet, a segmentation network that fuses the linear-time efficiency of the Mamba state space architecture with attention-based feature fusion and, unusually for deep learning models, explicit anatomical constraints baked directly into the training objective.

The backdrop to this work is a decade-long architectural arms race. Fully convolutional networks pioneered end-to-end pixel-level segmentation, and U-Net—with its symmetric encoder-decoder design and skip connections that marry shallow detail with deep semantics—became the cornerstone of the field. V-Net extended the idea to 3D volumes and introduced the Dice loss to combat class imbalance, while UNet++ and UNet 3+ strengthened feature fusion through nested and full-scale skip connections. When Transformers arrived, models such as Swin-Unet, TransUNet, and UNETR brought self-attention’s ability to capture long-range dependencies into medical imaging. But self-attention carries a heavy price: its computational complexity scales quadratically with sequence length, which becomes punishing when processing large-scale, high-resolution medical images. The field’s urgent need, the authors argue, is an architecture that achieves global dependency modeling and linear computational efficiency at the same time.

That is precisely where Mamba enters. Built on a Selective State Space Model, Mamba achieves sequence modeling with linear time complexity by making its parameters dynamic functions of the input, allowing the model to contextually decide which information to propagate or forget within a sequence. Mamba’s S6 layer overcomes the linear time-invariance constraint of the classical S4 model, and the architecture fuses state space and MLP blocks into a single homogeneous block—linear projection, lightweight convolution, SiLU gating, then selective sequence modeling. VMamba carried this into vision by introducing the Visual State Space module and the SS2D mechanism, which propagates state along multiple directions on image features to capture long-range spatial dependencies without self-attention. A wave of successors—U-Mamba, Mamba-UNet, SegMamba, VM-UNet, and many others—quickly adapted the approach to medical segmentation.

Yet the authors identify three persistent failures in existing models. First, there is an imbalance between local and global features: convolutional networks excel at textures and boundaries but struggle with long-range dependencies, while Mamba-style sequence models enhance global context yet often neglect local details and edge structures, and most employ unidirectional or fixed-direction scanning that limits perception of complex organ geometries. Second, skip connections in U-Net variants suffer from a semantic gap—shallow features carry noise and low-level textures while deep features hold rich semantics but lack spatial detail, and most attention methods fail to synergistically optimize spatial and channel information. Third, mainstream models rely predominantly on data-driven loss optimization, overemphasizing pixel-level alignment while neglecting the boundary smoothness and regional consistency that anatomically plausible segmentation demands, often producing jagged edges, fragmented regions, or anatomically inconsistent pseudo-boundaries.

PF-Mamba-UNet attacks all three problems at once. Its core building block is the Convolution-Reinforced Visual State Space (CR-VSS) module, which adds a Convolutional Refinement Branch operating in parallel alongside the state space feature path. This branch employs lightweight convolutions with inherent spatial local inductive biases to supplement the modeling of edges, textures, and local regional patterns, and its locally enhanced features are fused with the original state-space features. Within the module, a Cross-Space Scan 2D (CSS2D) mechanism systematically supports spatial scanning along horizontal, vertical, diagonal, and anti-diagonal paths. Diagonal scanning, the authors note, particularly excels at capturing organ boundaries and tissue septa with oblique orientations—features often overlooked by conventional horizontal and vertical scans—while the multi-directional mechanism is achieved by transforming and stacking feature maps across different dimensions.

The second innovation is the Coordinated Spatial-Channel Fusion Attention (CSCFA) module, embedded in the skip-connection layers between encoder and decoder. CSCFA first extracts spatial features through a parallel architecture of 3×3 convolution, 5×5 convolution, and dilated convolution with dilation 2, all implemented as depthwise separable group convolutions to reduce computational complexity. Each scale’s feature map undergoes channel-wise average pooling, the resulting response maps are concatenated, and a 7×7 convolution normalized by a sigmoid generates the spatial attention map. Crucially, the module then uses the mean of this spatial attention map as prior information, multiplying it with channel descriptors from global average and max pooling before an MLP performs dimension reduction and expansion to produce channel attention weights. This spatial-first channel recalibration highlights critical anatomical regions while suppressing background noise, and residual connections ensure that features of minute lesions are not entirely masked by attention weights while mitigating gradient vanishing.

The third component is the most conceptually distinctive: Anatomical Prior Regularization, a parameter-free paradigm that integrates two differentiable constraints into training. The Boundary Smoothness Anatomical Constraint rests on the observation that anatomical edges typically exhibit strong gradient changes while tissue interiors show relatively uniform intensity. The constraint computes discrete gradients of the predicted segmentation mask in horizontal and vertical directions, averages them into a segmentation gradient strength, computes the image’s own intensity gradient after grayscale conversion, and penalizes the absolute deviation between the two—so boundaries appearing in non-edge regions incur a penalty, as do failures to respond at true edges. The Region Consistency Anatomical Constraint, grounded in the spatial continuity of biological tissues, applies 3×3 average pooling to each class probability map and requires each pixel’s predicted probability to stay close to its neighborhood average, reducing holes and islands in segmentation results. The overall loss combines a base loss of equal parts binary cross-entropy and Dice loss with these constraints at weights of 0.8, 0.1, and 0.1 respectively, a configuration validated by grid search on the ACDC dataset.

The experimental evidence spans three public datasets covering distinct challenges: ACDC, a cardiac MRI dataset of 150 patients from the University Hospital of Dijon with annotations of the left ventricle, right ventricle, and myocardium; Synapse, an abdominal contrast-enhanced CT benchmark of 30 cases with 3,779 axial images annotated across eight organs including the notoriously blurred pancreas and gallbladder; and PROMISE12, the MICCAI 2012 prostate MRI challenge with multi-center T2-weighted volumes. On ACDC, PF-Mamba-UNet achieves a Dice coefficient of 0.9082, mIoU of 0.8356, and boundary metrics of 1.0873 for HD95 and 0.2825 for ASD, ranking first in accuracy, precision, and specificity. On Synapse it attains a Dice of 0.7632 and mIoU of 0.6601 with HD95 of 29.1827 and ASD of 9.3111, and on PROMISE12 a Dice of 0.8559 with HD95 of 3.2145—outperforming U-Net, SwinUNet, Mamba-UNet, Mamba-Sea, KM-Unet, and CSWin-Unet across the board.

Ablation studies dissect where the gains come from. A direction-wise analysis on ACDC shows that horizontal-only scanning yields a Dice of 0.9011 and HD95 of 1.6749 at 8.8 GMac FLOPs; adding vertical, diagonal, and finally anti-diagonal scanning progressively lifts performance to the full configuration’s Dice of 0.9082 and HD95 of 1.0873 at 11.3 GMac, though marginal gains shrink as directions accumulate—motivating future adaptive direction selection. Component-wise removal confirms that CSS2D most significantly improves boundary quality, CSCFA primarily enhances cross-layer semantic fusion, and the anatomical prior constraints reduce boundary artifacts. The authors are candid about failures: on the Synapse dataset, the right kidney shows worse HD95 and ASD than Mamba-UNet, and the gallbladder a worse HD95, likely because generic anatomical prior constraints may over-smooth local structures in organs with high shape variability, small extent, or ambiguous contrast.

Efficiency and statistical rigor round out the picture. PF-Mamba-UNet requires 19.14 million parameters, a 73.03 MB model size, 11.3 GMac FLOPs, 18.6 ms per image latency, and 53.8 frames per second throughput on a single RTX 3090—compared with Mamba-UNet, a 46.6% reduction in parameters and model size, 20.4% fewer FLOPs, and 16.6% lower latency. Case-level paired statistical validation on ACDC shows a Dice of 0.9071 ± 0.0322 versus 0.8957 ± 0.0446 for Mamba-UNet, supported by paired t-test (p = 0.0262) and Wilcoxon signed-rank test (p = 0.0018), while the CSWin-UNet comparison is borderline on the Wilcoxon test (p = 0.0566) and should be interpreted with caution. The model is not extremely lightweight for mobile devices, the authors concede, and it currently models 2D slices only, with generic rather than organ-specific priors. Future work will explore structured pruning, INT8 quantization, learnable spatial gating for scan directions, organ-specific shape priors developed with clinical radiologists, and extension to 3D volumetric and multimodal segmentation—steps that could determine whether anatomically constrained Mamba networks become a practical fixture in real clinical pipelines.

Subject of Research: Artificial Intelligence in Medical Imaging

Article Title: PF-Mamba-UNet: An Anatomical-Prior-Guided Medical Image Segmentation Network Based on Mamba and Spatial-Channel Collaborative Attention Fusion

Article References: Xiao, L., Liao, K., Zhou, B., & Fan, C. (2026). PF-Mamba-UNet: An Anatomical-Prior-Guided Medical Image Segmentation Network Based on Mamba and Spatial-Channel Collaborative Attention Fusion. Cognitive Computation, 18(1), 121. https://doi.org/10.1007/s12559-026-10667-7

Image Credits: AI Generated

DOI: 10.1007/s12559-026-10667-7

Keywords: PF-Mamba-UNet, Anatomical-Prior-Guided, Medical, Image, Segmentation, Network, Based, Mamba, Spatial-Channel, Collaborative, Attention, Fusion