BMDS-Net: A Bayesian Multi-Modal Deep Supervision Network for Robust Brain Tumor Segmentation

TL;DR

BMDS-Net tackles the gap between high segmentation accuracy and clinical reliability in multi-modal brain tumor segmentation by integrating a zero-init multimodal fusion mechanism, a residual-gated deep decoder supervision pathway, and a memory-efficient Bayesian fine-tuning strategy. The two-stage pipeline yields competitive Dice scores while providing voxel-wise uncertainty maps and improved robustness under missing modalities, demonstrated on BraTS 2021 with an exceptionally low calibration error. The approach offers a practical balance between accuracy, stability, and interpretability, achieving robust performance with manageable computational overhead. By unifying robustness and uncertainty calibration within a Transformer-based framework, BMDS-Net advances the feasibility of clinically deployable brain tumor segmentation systems.

Abstract

Accurate brain tumor segmentation from multi-modal magnetic resonance imaging (MRI) is a prerequisite for precise radiotherapy planning and surgical navigation. While recent Transformer-based models such as Swin UNETR have achieved impressive benchmark performance, their clinical utility is often compromised by two critical issues: sensitivity to missing modalities (common in clinical practice) and a lack of confidence calibration. Merely chasing higher Dice scores on idealized data fails to meet the safety requirements of real-world medical deployment. In this work, we propose BMDS-Net, a unified framework that prioritizes clinical robustness and trustworthiness over simple metric maximization. Our contribution is three-fold. First, we construct a robust deterministic backbone by integrating a Zero-Init Multimodal Contextual Fusion (MMCF) module and a Residual-Gated Deep Decoder Supervision (DDS) mechanism, enabling stable feature learning and precise boundary delineation with significantly reduced Hausdorff Distance, even under modality corruption. Second, and most importantly, we introduce a memory-efficient Bayesian fine-tuning strategy that transforms the network into a probabilistic predictor, providing voxel-wise uncertainty maps to highlight potential errors for clinicians. Third, comprehensive experiments on the BraTS 2021 dataset demonstrate that BMDS-Net not only maintains competitive accuracy but, more importantly, exhibits superior stability in missing-modality scenarios where baseline models fail. The source code is publicly available at https://github.com/RyanZhou168/BMDS-Net.

Figures (5)

• Figure 1: The evolving landscape of brain tumor segmentation paradigms. The horizontal axis tracks the architectural shift from local CNN features (Era 1) to global Transformer contexts (Era 2), while the vertical axis measures clinical capability. While Era 2 models achieve SOTA accuracy, they are often fragile. Previous specialized solutions addressed specific issues (e.g., uncertainty) in a fragmented manner. BMDS-Net bridges this gap by unifying high accuracy with the clinical requirements of robustness and confidence calibration.
• Figure 2: Detailed architecture of the proposed BMDS-Net framework. Top Row (Overall Pipeline): The complete workflow processing multi-modal MRI input. The input passes through the Zero-Init MMCF module before entering the Swin UNETR backbone. Note the global attention path (dashed gray line) that transmits $\mathbf{M}_{att}$ to the decoder layers for deep supervision. Bottom Left (MMCF Module): Detailed view of the fusion mechanism. The learnable scalar $\alpha$ is initialized to 0, ensuring the module acts as an identity mapping at the start of training to preserve pre-trained weights. Bottom Right (DDS Module): The Residual-Gated mechanism where the global attention map $\mathbf{M}_{att}$ gates the decoder features $\mathbf{D}_i$. The scalar $\gamma$ (initialized to 0.1) controls the injection of global context into the decoder.
• Figure 3: Qualitative comparison of segmentation results across multiple cases. BMDS-Net shows consistent and accurate segmentation compared to SOTA methods, particularly in preserving the continuity of tumor sub-regions.
• Figure 4: Visual analysis of hyperparameter sensitivity. The plot illustrates the fluctuation in Dice score and stability (standard deviation) as the $\alpha$ value varies.
• Figure 5: Uncertainty analysis. The error map (middle) highlights misclassified voxels (red). The uncertainty map (right) generated by our Bayesian module effectively captures these regions, demonstrating the model's ability to identify its own potential errors.