- The paper introduces a novel framework combining Hybrid Mamba, attention-guided mask refinement, and uncertainty fusion to enhance segmentation with minimal annotations.
- It leverages selective state-space modeling and query isolation to address noisy support masks and support-query domain shift in clinical MRI data.
- Experimental results on BraTS 2020 show significant improvements in Dice similarity and boundary accuracy compared to traditional and attention-based methods.
RUFNet for Few-Shot Brain Tumor Segmentation: Query-Guided Mask Refinement and Uncertainty Fusion via Hybrid Mamba
Motivation and Challenges
Few-shot segmentation in neuro-oncology addresses the annotation bottleneck in clinical MRI by leveraging minimal labeled data to delineate tumor regions. The primary challenges are (i) contamination of prototype construction due to noisy support masks, (ii) inter-patient variation leading to support-query domain shift, and (iii) lack of explicit pixel-wise confidence in predictions. Existing methods either suffer from data hunger (e.g., attention-based architectures), fail to robustly model support-query correspondence, or lack uncertainty estimation for reliable inference.
RUFNet Architecture and Methodological Contributions
RUFNet introduces an end-to-end framework composed of three principal components:
- Hybrid Mamba Backbone: Employs selective state-space modelling for long-range support-query dependencies with linear complexity. The backbone incorporates a Support Reset Module (SRM) to prevent forgetting class priors, and a Query Isolation Module (QIM) to restrict intra-query variance, preserving spatially extended medical image priors.
- Attention-Guided Mask Refinement Module (AGMR): Refines noisy support masks by query-driven cross-attention, recalibrating the mask based on query features. Mask embeddings are concatenated with aligned support/query features, generating soft masks that encode "partial foreground" semantics. This reduces prototype degradation and enables more consistent and accurate region delineation, particularly in boundary and ambiguous areas.
- Uncertainty-Aware Posterior Fusion Module (UAPF): Estimates pixel-wise log-variance to provide spatially adaptive fusion. Uncertainty-based weights balance meta-predictions and query-aligned priors, pushing predictions towards reliable regions and modulating ambiguous areas. A spatial refinement block enhances coherence in the final output.
The combination allows robust few-shot segmentation by explicitly handling noisy labeling, domain shift, and prediction uncertainty within the medical imaging context.
Experimental Results and Quantitative Analysis
The BraTS 2020 dataset was used for evaluation, with preprocessing to normalize slice selection and reduce annotation noise. RUFNet was compared in both 1-way 1-shot and 1-way 5-shot configurations.
- Ablation Study: AGMR and UAPF individually improved Dice similarity coefficient (DSC) and Hausdorff distance (HD) over Mamba-only baselines, but their combination yielded the strongest gains (1-way 1-shot: DSC 84.3%, HD 10.55 mm; 1-way 5-shot: DSC 86.1%, HD 7.67 mm).
- Comparative Performance: Relative to PANet, SENet, RPNet, and more recent registration- or alignment-based methods, RUFNet delivers significantly higher DSC (by ≥9 points) and more stable boundary localization (lower HD and reduced variability). Early prototype-based approaches suffer from severe mask degradation and unstable segmentation (DSC ≤36%, HD ≥60 mm), while supervised alignment methods improve consistency (DSC up to 75-77%) but lack spatial uncertainty modulation.
- Visualizations: RUFNet masks are notably closer to ground-truth delineations in both tumor localization and boundary shapes, even under support set variation, confirming robustness in scenarios with scarce support annotations.
Theoretical and Practical Implications
The RUFNet approach demonstrates that query-guided support mask recalibration and pixel-wise uncertainty fusion are not only viable but essential for robust few-shot segmentation in the presence of noisy medical annotation and inter-domain variation. Selective state-space interaction via Hybrid Mamba addresses computational scalability and preserves long-range dependencies, which are often neglected by traditional attention mechanisms.
Practically, RUFNet enables more trustworthy tumor segmentation in rare or under-annotated settings, with explicit handling of boundary uncertainties and a reduction in annotation errors. The method can be directly integrated into clinical pipelines where voxel-level annotation is expensive and time-consuming.
Limitations and Future Directions
The evaluation is restricted to binary foreground segmentation on 2D BraTS slices without multi-center external validation. Further research should extend RUFNet to 3D multi-class segmentation, incorporate multi-scanner data, and explore calibration of uncertainty estimates across heterogeneous clinical cohorts.
There is scope for integration with active learning frameworks, hybrid registration approaches, and fine-grained anatomical priors, potentially improving generalization and supporting clinical deployment in diverse settings.
Conclusion
RUFNet represents a technically rigorous advancement in few-shot brain tumor segmentation, combining Hybrid Mamba state-space modeling with query-guided mask refinement and uncertainty-aware posterior fusion. The demonstrated improvement in accuracy and stability over existing methods underscores the necessity of both semantic alignment and uncertainty modeling for clinical-grade MRI tumor segmentation under limited annotation. The architecture offers a scalable template for future research in robust and trustworthy few-shot medical imaging.