AAMamba: Not All Voxels Deserve Equal Tokens in Mamba-Based 3D Medical Image Segmentation

Medical volumes are highly nonuniform. Most voxels belong to homogeneous background or simple tissue, while the evidence needed for accurate segmentation concentrates at lesion margins, organ interfaces, and small irregular structures. State-space models such as Mamba process long 3D sequences efficiently, yet existing Mamba-based segmentation models still build tokens on uniform grids and serialize them in fixed spatial orders. We identify this anatomy-agnostic volume-to-sequence construction as a central bottleneck of 3D state-space segmentation. Uniform tokenization gives background the same representational budget as lesions and boundaries, and naive scan orders can place neighboring voxels far apart once the volume is flattened. We propose AAMamba, an anatomy-aware Mamba framework in which token allocation, sequence serialization, and dense reconstruction are designed jointly. AAMamba estimates a structural density field from intensity deviation, boundary strength, and curvature response, and uses it to assign denser tokens to anatomically complex regions while still covering the whole volume. The adaptive tokens are then ordered by ROI-guided bidirectional Hilbert scanning, which keeps spatial neighbors close in the sequence. ROI guidance is applied only during training to emphasize foreground structures, so inference requires no ROI input. Because adaptive sampling can disturb voxel-level correspondence, a coordinate-aware decoder fuses the adaptive semantic features with grid-aligned spatial references to reduce spatial drift in the dense prediction. AAMamba achieves 84.2% mean Dice on BraTS, 91.85% DSC on ACDC, and 84.26% average Dice on Synapse. The largest gains appear on small or morphologically complex structures, including the enhancing tumor, myocardium, gallbladder, and pancreas, which suggests that anatomy-aware sequence construction is an effective principle for 3D medical state-space models.