BMDS-Net: Bayesian Multi-Modal Deep Supervision

• BMDS-Net is a deep learning framework that integrates deterministic feature learning with Bayesian probabilistic inference for robust brain tumor segmentation from multi-modal MRI.
• The architecture enhances a Swin UNETR backbone with novel MMCF and DDS modules, improving modality fusion, boundary localization, and resilience to incomplete inputs.
• Empirical validation on BraTS 2021 shows competitive Dice scores, reduced Hausdorff distance, and effective uncertainty maps that inform clinical decision-making.

BMDS-Net (Bayesian Multi-Modal Deep Supervision Network) is a robust deep learning framework for brain tumor segmentation from multi-modal MRI, explicitly addressing clinical challenges of missing modalities and the need for calibrated uncertainty estimation. BMDS-Net enhances a Swin UNETR encoder–decoder backbone with two novel architectural modules—Zero-Init Multimodal Contextual Fusion (MMCF) and Residual-Gated Deep Decoder Supervision (DDS)—and introduces a memory-efficient Bayesian fine-tuning stage for probabilistic inference. Empirical validation on BraTS 2021 demonstrates BMDS-Net's stability in the face of corrupted or missing imaging modalities and its utility for uncertainty-aware clinical deployment (Zhou et al., 24 Jan 2026).

1. Architectural Overview

BMDS-Net is architected on a Swin UNETR backbone, leveraging Transformer-based global context extraction. The sequence of computational steps is as follows: multi-modal MRI inputs pass through the MMCF module, modulating each modality’s contribution; the fused data feeds into a Swin Transformer encoder for hierarchical feature extraction. Decoder blocks with skip connections reconstruct pixel-level spatial detail. DDS modules attach auxiliary segmentation heads with residual-gated feature modulation at deep decoder layers (32× and 64× resolutions). Initial training is deterministic; afterwards, the final segmentation head is replaced with BayesianConv3d for stochastic, uncertainty-aware prediction.

The following table summarizes the high-level component roles:

Component	Purpose	Notes
MMCF	Dynamic modality weighting, modality robustness	Zero-initialization for stable training
Swin Transformer Encoder	Hierarchical, global context extraction	Self-attention
Decoder + Skip Connections	Spatial detail reconstruction	Cascaded upsampling layers
DDS	Boundary sharpening, training stabilization	Deep layer auxiliary losses
BayesianConv3d Head	Uncertainty estimation	Variational inference, MC sampling

BMDS-Net’s modular design enables seamless integration of deterministic feature learning and probabilistic inference in a two-stage training regime.

2. Zero-Init Multimodal Contextual Fusion (MMCF)

The MMCF module addresses clinical variability by learning to reweight and fuse multi-modal MRI channels. Given four MRI modalities as input $X$, MMCF applies a feature-encoder $\mathcal{F}_{enc}$ yielding intermediate representation $F_{feat}$. Two convolutional “heads” compute:

• Multimodal spatial attention: $$M_{att} = \sigma(\mathcal{C}_{att}(F_{feat})) \in [0,1]^{4 \times H \times W \times D}$$
• Auxiliary uncertainty map: $U_{map} = \sigma(\mathcal{C}_{unc}(F_{feat}))$

Fusion is realized via zero-initialized scalar residual gating:

$$X_{fused} = X + \alpha \cdot (X \odot M_{att}),\qquad \alpha(t=0) = 0$$

where $\odot$ denotes channelwise multiplication. Zero-initialization ensures the initial fused output matches the original input, avoiding large weight perturbations and enabling transfer learning without destabilizing gradients. This mechanism allows the model to adaptively suppress or enhance modality contributions, improving robustness to missing or corrupted inputs (Zhou et al., 24 Jan 2026).

3. Residual-Gated Deep Decoder Supervision (DDS)

The DDS mechanism strengthens segmentation accuracy and boundary sharpness by residually modulating deep decoder features with global attention. For each decoder stage $i$, spatially-resized $M_{att}$ ($M_{att}^{(i)}$) controls feature refinement:

$$G_i = 1 + \gamma \cdot \sigma(\mathcal{P}_{proj}(M_{att}^{(i)})),\qquad \gamma(t=0) = 0.1$$

$D_i^{ref} = D_i \odot G_i$

Auxiliary segmentation heads placed at decoder depths 32× and 64× further provide deep supervision. The composite loss during deterministic pre-training is:

$$\mathcal{L}_{seg} = \mathcal{L}_{DiceCE}(Y_{final},\,Y_{gt}) + \sum_{i\in\{32\times,64\times\}} \lambda_i\,\mathcal{L}_{DiceCE}(Y_{aux}^{(i)},\,Y_{gt})$$

with $\lambda_{32\times}=0.4$, $\lambda_{64\times}=0.2$. A bidirectional feature distillation loss aligns encoder attention with decoder activations:

$$\mathcal{L}_{distill} = \sum_i \|\mathcal{N}(\|D_i^{ref}\|_2) - \mathcal{N}(\mathrm{Interp}(M_{att}))\|_2^2$$

Total pre-training loss: $$\mathcal{L}_{total} = \mathcal{L}_{seg} + 0.2\mathcal{L}_{distill}$$.

These strategies jointly enhance boundary localization and training stability, particularly under partial input corruption.

4. Bayesian Fine-Tuning Strategy

BMDS-Net employs a lightweight Bayesian fine-tuning stage to augment the deterministic backbone with voxelwise uncertainty calibration. The final deterministic conv3d layer is replaced with BayesianConv3d, whose weights $W$ follow a variational posterior $q_\theta(W)=\mathcal{N}(\mu,\sigma^2)$ with $\sigma = \log(1+\exp(\rho))$. Weight sampling uses the reparameterization trick: $W = \mu + \sigma \odot \epsilon$, $\epsilon \sim \mathcal{N}(0, I)$. Training optimizes the evidence lower bound (ELBO):

$$\mathcal{L}_{ELBO} = \mathcal{L}_{DiceCE}(X,Y;W) + \beta_{KL} D_{KL}(q_\theta(W) \parallel p(W)),\; p(W)=\mathcal{N}(0,I)$$

A single Monte Carlo sample is used per training pass. During inference, $T=20$ MC samples produce predictive mean and per-voxel variance:

$$\sigma^2_{pred}(v) = \frac{1}{T} \sum_{t=1}^T [p_t(v) - \bar{p}(v)]^2$$

The resulting uncertainty maps correlate strongly with error regions, directly supporting cautious clinical review and deployment safety. The entire fine-tuning process is memory-efficient and incurs minimal runtime overhead (Zhou et al., 24 Jan 2026).

5. Quantitative Evaluation and Ablation

Empirical results on the BraTS 2021 validation set compare BMDS-Net to Swin UNETR and ablations thereof. BMDS-Net achieves competitive Dice scores and consistently reduced Hausdorff Distance (HD95), especially in clinically sensitive tumor regions.

Model	WT Dice	WT HD95 (mm)	TC Dice	TC HD95	ET Dice	ET HD95
Swin UNETR (baseline)	0.9279	2.30	0.9111	2.39	0.8629	3.84
BMDS-Net (full)	0.9293	2.27	0.9098	2.22	0.8675	3.27

In missing-modality scenarios (Dice mean ± std):

• Missing T1ce: Swin UNETR 0.848±0.152; BMDS-Net 0.868±0.137
• Missing T2: Swin UNETR 0.364±0.100; BMDS-Net 0.388±0.115

Ablation studies indicate DDS contributes maximal peak Dice and boundary refinement, while the combination of MMCF and DDS yields the best robustness to missing modalities. Inference efficiency is high: BMDS-Net processes $128^3$ inputs at 4.89 FPS (baseline: 5.34 FPS; MMCF adds ~15 ms, DDS negligible).

6. Practical Implications

BMDS-Net’s uncertainty maps (ECE=0.0037) strongly associate with actual segmentation errors, allowing radiologists to efficiently identify regions requiring manual assessment. The network’s resilience to missing sequences directly addresses operational realities in clinical radiology, where incomplete or corrupted MRI scans are frequent. The two-stage (deterministic + Bayesian) training pipeline balances accuracy, robustness, and computational feasibility, facilitating real-world deployment without prohibitive resource demands.

7. Implementation Recipe and Code Access

The canonical two-stage training procedure is as follows:

for epoch in range(E1): for x, y in D: F_feat = F_enc(x) M_att = sigma(C_att(F_feat)) x_fused = x + alpha * (x * M_att) z = Encoder(x_fused) D = Decoder(z) for each decoder stage i: G_i = 1 + gamma * sigma(P_proj(Interp(M_att))) D_i_ref = D_i * G_i logits_main = FinalHead(D_refined) logits_aux = AuxHeads(D_refined) L_seg = L_DiceCE(logits_main, y) + sum(lambda_i * L_DiceCE(logits_aux[i], y)) L_distill = sum(norm(L2(D_i_ref)) - norm(Interp(M_att)))**2 Backpropagate L_seg + 0.2 * L_distill Replace final conv with BayesianConv3d(q_theta(W)=N(mu,sigma²)) for epoch in range(E2): for x, y in D: W ~ q_theta(W) logits = M_Bayes(x; W) L_ELBO = L_DiceCE(logits, y) + beta_KL * D_KL(q(W) || p(W)) Backpropagate L_ELBO

The official source code is available at https://github.com/RyanZhou168/BMDS-Net for reproducibility (Zhou et al., 24 Jan 2026).