Residual-Gated DDS in BMDS-Net

• The paper introduces DDS, which refines decoder features using residual gating and deep supervision to stabilize training and improve boundary delineation.
• DDS employs a global attention map, 1×1×1 convolution, and a learnable scalar to reweight decoder features at multiple up-sampling stages.
• Empirical results show improved Dice scores and lower HD95 metrics, demonstrating enhanced robustness in scenarios with missing MRI modalities.

Residual-Gated Deep Decoder Supervision (DDS) is a mechanism integrated into Transformer-based encoder–decoder segmentation architectures to stabilize feature learning, refine boundary delineation, and improve robustness under missing-modality scenarios in multi-modal medical imaging. It is a core component of BMDS-Net, a framework designed for robust brain tumor segmentation from multi-modal MRI, and is specifically tailored to address challenges where cross-modal context and precise boundary information must be leveraged without compromising training stability or calibration (Zhou et al., 24 Jan 2026).

1. Architectural Integration and Workflow

BMDS-Net employs a Swin UNETR backbone with the DDS module inserted into each up-sampling (decoder) stage. At each decoder level $i$:

• The decoder feature map $D_i$ with spatial resolution $H_i\times W_i\times D_i\times C_i$ is processed.
• The global attention map $M_\mathrm{att}$, produced by the MMCF encoder, is down-sampled via interpolation to match the resolution of $D_i$.
• The interpolated $M_\mathrm{att}$ undergoes a $1\times1\times1$ convolution followed by a sigmoid activation. The result is scaled by a learnable scalar $\gamma$, incremented by 1, yielding the residual gate $G_i$.
• Decoder features $D_i$ are element-wise multiplied by $G_i$, forming refined features $D_i^\mathrm{ref}$.
• An auxiliary segmentation head processes $D_i^\mathrm{ref}$, yielding an intermediate segmentation output $\hat{Y}^{(i)}$. Each output is included in the training loss, enforcing deep decoder supervision.

The complete block diagram and pseudocode for a decoder stage are as follows:

M_i = Interpolate(M_att, size=D_i.spatial_size) G_i = 1 + γ * sigmoid(Conv1×1×1(M_i)) # residual gate D_i' = D_i ⊙ G_i # gated feature Ŷ^(i)= AuxSegHead(D_i') # deep supervision prediction

Auxiliary segmentation heads are attached immediately after the 32× and 64× up-sampling stages, owing to their need for robust gradient signals and boundary refinement.

2. Mathematical Formulation

Formally, at decoder level $i$:

• Let $$F_\mathrm{in} \equiv D_i \in \mathbb{R}^{C_i\times H_i\times W_i\times D_i}$$ denote the input decoder features.
• $$M_\mathrm{att} \in [0,1]^{C_\mathrm{in}\times H\times W\times D}$$ is the global MMCF attention map.
• $Interp(\cdot)$ is the spatial interpolation operator.
• $\mathcal{P}_\mathrm{proj}(\cdot)$ is a $1\times1\times1$ projection (3D convolution).
• $\sigma(\cdot)$ denotes the sigmoid function.
• $\gamma$ is a learnable scalar (initialized to $0.1$).

Residual-Gated Unit at level $i$: $$G_i = 1 + \gamma \cdot \sigma\bigl(\mathcal{P}_\mathrm{proj}\big(\mathrm{Interp}(M_\mathrm{att})\big)\bigr)\,, \qquad D_i^\mathrm{ref} = D_i \odot G_i$$

Deep Supervision Loss:

Given $L$ levels of supervision, the aggregate DDS loss is

$$\mathcal{L}_\mathrm{DDS} = \sum_{i=1}^L \left[\lambda_i \cdot \mathrm{Dice}\left(Y,\,\mathrm{Softmax}(\hat{Y}^{(i)})\right) + \mu_i \cdot \mathrm{CE}\left(Y,\,\mathrm{Softmax}(\hat{Y}^{(i)})\right)\right]$$

For BMDS-Net: $L=2$ (32× and 64× up-samplings), with weights $\lambda_1=0.4$, $\mu_1=0.4$ (deeper), $\lambda_2=0.2$, $\mu_2=0.2$ (shallower).

3. Implementation Considerations

• Parameter Initialization: $\gamma$ is initialized at $0.1$ so that $G_i \approx 1$ at the outset, rendering decoder behavior initially equivalent to vanilla Swin UNETR. The projection and auxiliary heads are zero-initialized (bias $= 0$, small weights), supporting stable early optimization.
• Gradient Propagation: Auxiliary losses from all decoder levels back-propagate through $D_i^\mathrm{ref}$, modulating both $D_i$ and $G_i$ for joint refinement.
• Placement of Supervision: Supervision heads are added after coarsest up-sampling stages (32×, 64×) to provide early and coarse boundary cues, which are critical for spatial detail recovery.

4. Empirical Performance and Observed Benefits

• Stable Feature Learning and Delineation: Contextual gating with $M_\mathrm{att}$ injects multi-modal reliability signals into decoder, accentuating regions where boundary information is difficult to recover. The residual gate ($G_i \approx 1+$) preserves initial network behavior, mitigating vanishing gradients.
• Quantitative Gains: DDS alone (in baseline+DDS configuration) improves full-modality segmentation metrics:

Metric	Baseline	DDS Only
WT Dice	0.9279	0.9312
TC Dice	0.9111	0.9144
ET Dice	0.8629	0.8718
HD95 (WT, mm)	2.30	2.10
HD95 (TC, mm)	2.39	1.93
HD95 (ET, mm)	3.84	2.83

• Resilience to Missing Modalities: Under scenarios with missing MRI inputs (e.g., Missing-T1ce: 0.848 baseline vs. 0.865 DDS Only Dice; Missing-T2: 0.364 baseline vs. 0.369 DDS Only Dice), DDS improves segmentation robustness by leveraging the global attention map to reweight features and compensate for partial data.
• Ablation Evidence: DDS is isolated as the prime contributor to boundary-sensitive performance metrics. When combined with MMCF, DDS provides stability under missing modalities, with only a minor reduction in peak Dice scores compared to DDS alone.

5. Relationship to Training and Loss Functions

In Stage 1 (deterministic pre-training), the total BMDS-Net loss combines main and deep supervisions:

$$\mathcal{L}_\mathrm{total} = \mathcal{L}_\mathrm{seg} + 0.2 \cdot \mathcal{L}_\mathrm{distill}$$

where

$$\mathcal{L}_\mathrm{seg} = \mathrm{DiceCE}(Y_\mathrm{final}, Y_\mathrm{gt}) + \sum_{i=1}^L \left[\lambda_i \cdot \mathrm{DiceCE}(\hat{Y}^{(i)}, Y_\mathrm{gt})\right]$$

and

$$\mathrm{DiceCE}(\hat{Y}, Y) = \mathrm{Dice}(\mathrm{Softmax}(\hat{Y}), Y) + \mathrm{CE}(\mathrm{Softmax}(\hat{Y}), Y)$$

$\mathcal{L}_\mathrm{distill}$ is a self-distillation term aligning the $\ell_2$ norm of refined decoder outputs with the interpolated attention map.

In Stage 2 (Bayesian fine-tuning), the final segmentation layer is replaced by a BayesianConv, and only the ELBO is minimized: $$\mathcal{L}_\mathrm{ELBO} = \mathrm{DiceCE}(Y_\mathrm{final}, Y_\mathrm{gt}) + \beta_\mathrm{KL} \cdot D_\mathrm{KL}\left[q(\mathcal{W}) || p(\mathcal{W}) \right]$$ Deep supervision is omitted in this stage, leaving the encoder and decoder (with DDS-refined representations) frozen.

6. Significance and Context within Multi-Modal Medical Segmentation

Residual-Gated Deep Decoder Supervision is motivated by the clinical need for robustness and reliability in the presence of missing and corrupted imaging modalities. By utilizing a multi-modal global attention map for decoder feature gating and by enforcing coarse-to-fine auxiliary supervision, DDS addresses both vanishing gradient challenges and the brittleness of prior Transformer-based models. Notably, in the robust segmentation of brain tumors (as benchmarked on BraTS 2021), DDS demonstrates empirical improvements in both hard boundary detection and resilience to input sparsity. A plausible implication is that similar residual-gated deep supervision may generalize to other multi-modal segmentation domains where cross-modal reliability and boundary accuracy are critical (Zhou et al., 24 Jan 2026).