DAMamba-UNet3D: A Parameter-Efficient Mamba State Space U-Net with Dynamic Adaptive Scan for 3D Medical Image Segmentation

Technical Summary: DAMamba-UNet3D

Problem Statement
3D medical image segmentation requires balancing precise local boundary delineation with coherent global spatial context, all while adhering to strict computational and parameter budgets. Existing architectures face distinct trade-offs:

• Convolutional U-Nets: Offer parameter efficiency and $O(n)$ local mixing but lack explicit global context, requiring depth scaling to capture long-range dependencies.
• Vision Transformers: Provide global reasoning via self-attention but incur prohibitive $O(n^2)$ memory and compute costs for large 3D volumes (e.g., a $128^3$ patch contains $\approx 2.1$ million tokens).
• State-Space Models (SSMs/Mamba): Propagate information globally in $O(n)$ time and memory. However, existing medical SSM segmenters (e.g., SegMamba) rely on fixed scan patterns and large parameter budgets ($\sim 70$M parameters), lacking data-dependent adaptability.

While Dynamic Adaptive Scan (DAS) has been introduced for 2D natural image recognition to learn data-dependent reordering before selective scan, it has not been applied to medical imaging or extended to 3D volumes.

Methodology
The authors propose DAMamba-UNet3D, a hybrid encoder–decoder architecture that integrates Tri-Plane 3D-DAS blocks into a U-Net framework.

1. Architecture Design:

   • Hybrid Structure: The model retains convolutional blocks for the stem (E1), bottleneck, and entire decoder to preserve local details.
   • Encoder Integration: Tri-plane DAMamba blocks are inserted exclusively at encoder stages E2–E4.
   • Compact Variant: The primary model uses a base width of $C=16$, resulting in approximately 5.3 million parameters ($\sim 13\times$ fewer than SegMamba).
   • Large-Scale Variant (DAMamba-L): A wider variant ($\sim 70$M parameters) with stacked DAMamba blocks in the encoder and a convolutional bottleneck, designed to match SegMamba's capacity for direct comparison.

2. Tri-Plane Dynamic Adaptive Scan (DAS):

   • Mechanism: Unlike fixed scans (row-major, zigzag, or multi-orientation flattening), DAS learns input-dependent spatial reordering. An Offset Prediction Network (OPN) predicts 2D offsets from the feature map, which are used to resample features via bilinear interpolation before the selective scan.
   • 3D Extension: The 2D DAS operator is extended to volumes by applying it across three anatomical planes: Axial ($H \times W$), Coronal ($D \times W$), and Sagittal ($D \times H$). The outputs are averaged to form the 3D DAS block:
     $$DAS3D(X) = \frac{1}{3}(DAS_{ax} + DAS_{cor} + DAS_{sag})$$
   • Selective Scan: The reordered feature map is processed by a standard Mamba selective scan block, enabling global context propagation at $O(N)$ cost without quadratic attention maps.

3. Training and Evaluation:

   • Dataset: BraTS 2020 training set (369 subjects, 4 modalities).
   • Loss: Modulated Combo Loss ($L = (L_{CE} + \delta) \cdot (L_{mod}^{Dice} + \delta)$) to focus on hard voxels and boundaries.
   • Protocol: Five-fold subject-level cross-validation with full-volume inference (sliding-window, $128^3$ patches). Metrics include mean Dice and HD95 for Enhancing Tumor (ET), Tumor Core (TC), and Whole Tumor (WT).

Key Contributions

1. First Medical Application of DAS: Transfers Dynamic Adaptive Scan from 2D natural image classification to dense 3D medical segmentation.
2. 2D-to-3D Extension: Lifts the DAS operator to volumetric data via tri-plane adaptive scanning, offering a learnable alternative to hand-designed multi-orientation flattening.
3. Parameter Efficiency: Demonstrates a hybrid U-Net with tri-plane DAS blocks only in the encoder, achieving near-parity with much larger models at $\sim 13\times$ lower parameter cost.
4. Two-Scale Validation: Validates the approach on both a compact scale ($\sim 5.3$M) and a large scale ($\sim 70$M), showing that learned DAS can outperform fixed scanning patterns when architectural placement is optimized.

Results

• Compact Model (DAMamba-UNet3D, $\sim 5.3$M): Achieved a mean Dice of 0.815 $\pm$ 0.013 on BraTS 2020. This is near-parity with the retrained SegMamba (0.824 $\pm$ 0.014) despite having $\sim 13\times$ fewer parameters.
• Large-Scale Model (DAMamba-L, $\sim 70$M): Surpassed the retrained SegMamba with a mean Dice of 0.829 $\pm$ 0.012 (a +0.5 point improvement), with the largest gain observed in the Tumor Core (TC) class (0.842 vs. 0.827).
• Ablation Studies:
  • Placement: Placing DAS blocks only in the encoder (E2–E4) yielded the best results (0.815). Adding DAS to the bottleneck or decoder lowered performance (0.807 and 0.805, respectively), indicating that encoder-only placement is critical.
  • Scan Patterns: Learned tri-plane DAS outperformed fixed scan orders (Raster, Snake, Cross, Tri-orientated) when the encoder placement was matched.

Significance and Claims
The paper claims that learned tri-plane DAS in a hybrid U-Net is a competitive alternative to fixed Tri-orientated Mamba (ToM) scanning for 3D brain tumor segmentation.

• Efficiency: The compact model demonstrates that high performance can be achieved with significantly reduced parameter counts by leveraging adaptive scanning rather than brute-force capacity.
• Performance: At large scales, the learned adaptive mechanism (DAS) can improve upon fixed scanning trajectories (ToM) under identical training and evaluation protocols.
• Modesty: The authors caution that while results on BraTS 2020 are promising, the study is limited to a single dataset. They state that "multi-dataset validation and statistical testing are needed before general claims about DAS versus ToM" can be made. The primary contribution is the extension of DAS to 3D medical volumes and the demonstration of its viability in a parameter-efficient hybrid architecture.