MST-KDNet: Robust Brain Tumor Segmentation

• The paper introduces MST-KDNet, a teacher–student framework that uses multi-scale transformer distillation to enhance segmentation accuracy under missing-modality conditions.
• Key methodology components include Multi-Scale Transformer Knowledge Distillation, Dual-Mode Logit Distillation, and Global Style Matching to align attention, logits, and feature styles.
• Extensive experiments on BraTS and FeTS benchmarks demonstrate significant improvements in Dice and HD95 metrics compared to previous methods.

MST-KDNet is a teacher–student framework for brain tumor segmentation under missing-modality conditions, introduced in “Bridging the Gap in Missing Modalities: Leveraging Knowledge Distillation and Style Matching for Brain Tumor Segmentation” (Zhu et al., 30 Jul 2025). It is designed to address two problems identified in prior work: tumor boundary segmentation insensitivity and incomplete feature transfer when key imaging modalities are absent. The model couples an identical teacher and student 3D-UNETR backbone with three auxiliary components—Multi-Scale Transformer Knowledge Distillation, Dual-Mode Logit Distillation, and a Global Style Matching Module—so that training exploits all four MRI modalities while inference uses only the student, which may receive fewer modalities (Zhu et al., 30 Jul 2025).

1. Problem setting and overall design

MST-KDNet targets segmentation from multiparametric MRI when one or more modalities are missing. The training and evaluation setting uses four MRI modalities—T1, T1Gd, T2, and FLAIR—and explicitly simulates missing-modality inputs by randomly dropping 1–3 modalities to form 16 input combinations (Zhu et al., 30 Jul 2025). This places the method in the incomplete-modality segmentation regime rather than the fully observed multimodal setting.

The architecture is organized as a teacher–student framework. The teacher network is a 3D-UNETR backbone ingesting all four MRI modalities, while the student network uses an identical backbone but may receive fewer modalities (Zhu et al., 30 Jul 2025). During training both networks run in parallel; only the student is used at inference. This design suggests that MST-KDNet treats complete-modality supervision as a privileged training signal and missing-modality robustness as a distillation target.

Three core modules define the model: Multi-Scale Transformer Knowledge Distillation (MS-TKD), Dual-Mode Logit Distillation (DMLD), and the Global Style Matching Module (GSME) (Zhu et al., 30 Jul 2025). Their division of labor is explicit in the formulation: MS-TKD aligns attention and intermediate embeddings across scales, DMLD aligns output logits using two complementary criteria, and GSME aligns feature style and distribution through adversarial and feature-matching objectives.

2. Multi-scale transformer distillation

The MS-TKD component operates on transformer attention at multiple resolutions. At each Transformer resolution layer $l$, the attention tensor is extracted as $A_l \in \mathbb{R}^{n \times n}$, where $n$ is the spatial sequence length (Zhu et al., 30 Jul 2025). Extreme-value statistics are then computed over heads:

$$A_{\max}^l = \max(A_l, \mathrm{axis}=1), \qquad A_{\min}^l = \min(A_l, \mathrm{axis}=1), \qquad A_{\mathrm{mean}}^l = \frac{1}{n} \sum_{i=1}^{n} A_l^{(i)}.$$

These statistics are used to form three “extreme-value distillation” sequences:

$$EV_1^l = A_{\max}^l \circ A_l, \qquad EV_2^l = A_{\min}^l \circ A_l, \qquad EV_3^l = A_{\mathrm{mean}}^l \circ A_l.$$

The MS-TKD loss aligns both the extreme-value sequences and the intermediate feature embeddings $z^f, z^m$ of teacher and student across $L$ scales:

$$L_{\mathrm{MS\text{-}TKD}} = \alpha \cdot L_{\mathrm{EVD}} + \beta \cdot L_{\mathrm{MST}}$$

$$= \alpha \cdot \frac{1}{n} \sum_{i=1}^{3} \sum_{j=1}^{n} (EV_{i,j}^{f} - EV_{i,j}^{m})^2 + \beta \cdot \frac{1}{N} \sum_{i=1}^{3} \sum_{j=1}^{N} (z_{i,j}^{f} - z_{i,j}^{m})^2.$$

Within the architecture, this module is the mechanism for fine-grained alignment of transformer behavior across resolutions. The use of $A_{\max}^l$, $A_l \in \mathbb{R}^{n \times n}$0, and $A_l \in \mathbb{R}^{n \times n}$1 indicates that the method does not distill raw attention alone; it also emphasizes headwise extreme and average responses. A plausible implication is that this formulation is intended to preserve both salient and suppressed attention structure under missing-modality perturbation.

3. Dual-mode logit supervision

DMLD complements feature-level transfer with output-space supervision. Its first mode is Logit Discrepancy, implemented as mean squared error:

$A_l \in \mathbb{R}^{n \times n}$2

Its second mode is a standardized KL divergence. Logits are normalized by Z-score per sample:

$A_l \in \mathbb{R}^{n \times n}$3

The KL loss is then

$A_l \in \mathbb{R}^{n \times n}$4

The combined logit loss is defined as

$A_l \in \mathbb{R}^{n \times n}$5

The defining feature of DMLD is the simultaneous use of MSE and standardized KL rather than either criterion alone. The standardization step makes the KL term depend on normalized rather than raw logits. This suggests that the method is intended to improve transfer when absolute logit scale is unstable across teacher and student, particularly under heterogeneous missing-modality cases. In the ablation study, removing DMLD and replacing it with “SLKD” reduces BraTS 2024 performance from 81.8/59.5/59.8 Dice and 6.6/7.2/6.8 HD95 to 80.0/56.1/55.2 Dice and 8.1/8.7/8.0 HD95 for WT/TC/ET, respectively (Zhu et al., 30 Jul 2025).

4. Global style matching and adversarial alignment

The GSME module aligns deeper feature organization between teacher and student. Encoder feature $A_l \in \mathbb{R}^{n \times n}$6 is max-pooled and fused with transformer output $A_l \in \mathbb{R}^{n \times n}$7; this fused representation is fed to the decoder to produce $A_l \in \mathbb{R}^{n \times n}$8 (Zhu et al., 30 Jul 2025). The module then applies two forms of supervision.

The first is adversarial alignment. A discriminator $A_l \in \mathbb{R}^{n \times n}$9 receives the concatenated feature $n$0, and the adversarial loss is

$n$1

The second is feature matching based on Gram-style relationships. The $n$2 tensors are reshaped into 2D matrices $n$3, and three pairwise matrices are computed:

$n$4

The feature-matching loss is

$n$5

The overall GSME objective is

$n$6

In ablation, removing GSME causes the largest degradation among the three modules on BraTS 2024, reducing WT/TC/ET Dice to 78.3/55.1/53.4 and increasing WT/TC/ET HD95 to 9.6/9.7/9.5 (Zhu et al., 30 Jul 2025). This indicates that the style-matching and adversarial terms are central to the model’s missing-modality robustness. A plausible interpretation is that GSME compensates for modality-dependent feature shifts that are not fully addressed by attention and logit distillation alone.

5. Objective function and training protocol

The full training objective is a weighted sum of distillation, style alignment, and segmentation terms:

$n$7

where $n$8 is the standard soft-Dice segmentation loss (Zhu et al., 30 Jul 2025).

The reported training configuration uses BraTS 2024 with 1,350 cases and FeTS 2024 with 1,250 cases, each split into 80% train and 20% test (Zhu et al., 30 Jul 2025). Preprocessing consists of random flips, rotations, cropping, and resizing volumes to $n$9. Missing-modality simulation is performed by randomly dropping 1–3 modalities to form 16 input combinations.

The optimization protocol uses batch size $$A_{\max}^l = \max(A_l, \mathrm{axis}=1), \qquad A_{\min}^l = \min(A_l, \mathrm{axis}=1), \qquad A_{\mathrm{mean}}^l = \frac{1}{n} \sum_{i=1}^{n} A_l^{(i)}.$$0 for $$A_{\max}^l = \max(A_l, \mathrm{axis}=1), \qquad A_{\min}^l = \min(A_l, \mathrm{axis}=1), \qquad A_{\mathrm{mean}}^l = \frac{1}{n} \sum_{i=1}^{n} A_l^{(i)}.$$1 epochs with Adam, learning rate $$A_{\max}^l = \max(A_l, \mathrm{axis}=1), \qquad A_{\min}^l = \min(A_l, \mathrm{axis}=1), \qquad A_{\mathrm{mean}}^l = \frac{1}{n} \sum_{i=1}^{n} A_l^{(i)}.$$2, $$A_{\max}^l = \max(A_l, \mathrm{axis}=1), \qquad A_{\min}^l = \min(A_l, \mathrm{axis}=1), \qquad A_{\mathrm{mean}}^l = \frac{1}{n} \sum_{i=1}^{n} A_l^{(i)}.$$3, and $$A_{\max}^l = \max(A_l, \mathrm{axis}=1), \qquad A_{\min}^l = \min(A_l, \mathrm{axis}=1), \qquad A_{\mathrm{mean}}^l = \frac{1}{n} \sum_{i=1}^{n} A_l^{(i)}.$$4 (Zhu et al., 30 Jul 2025). The loss weights are $$A_{\max}^l = \max(A_l, \mathrm{axis}=1), \qquad A_{\min}^l = \min(A_l, \mathrm{axis}=1), \qquad A_{\mathrm{mean}}^l = \frac{1}{n} \sum_{i=1}^{n} A_l^{(i)}.$$5, $$A_{\max}^l = \max(A_l, \mathrm{axis}=1), \qquad A_{\min}^l = \min(A_l, \mathrm{axis}=1), \qquad A_{\mathrm{mean}}^l = \frac{1}{n} \sum_{i=1}^{n} A_l^{(i)}.$$6, $$A_{\max}^l = \max(A_l, \mathrm{axis}=1), \qquad A_{\min}^l = \min(A_l, \mathrm{axis}=1), \qquad A_{\mathrm{mean}}^l = \frac{1}{n} \sum_{i=1}^{n} A_l^{(i)}.$$7, and $$A_{\max}^l = \max(A_l, \mathrm{axis}=1), \qquad A_{\min}^l = \min(A_l, \mathrm{axis}=1), \qquad A_{\mathrm{mean}}^l = \frac{1}{n} \sum_{i=1}^{n} A_l^{(i)}.$$8, with temperature $$A_{\max}^l = \max(A_l, \mathrm{axis}=1), \qquad A_{\min}^l = \min(A_l, \mathrm{axis}=1), \qquad A_{\mathrm{mean}}^l = \frac{1}{n} \sum_{i=1}^{n} A_l^{(i)}.$$9; $$EV_1^l = A_{\max}^l \circ A_l, \qquad EV_2^l = A_{\min}^l \circ A_l, \qquad EV_3^l = A_{\mathrm{mean}}^l \circ A_l.$$0, $$EV_1^l = A_{\max}^l \circ A_l, \qquad EV_2^l = A_{\min}^l \circ A_l, \qquad EV_3^l = A_{\mathrm{mean}}^l \circ A_l.$$1, $$EV_1^l = A_{\max}^l \circ A_l, \qquad EV_2^l = A_{\min}^l \circ A_l, \qquad EV_3^l = A_{\mathrm{mean}}^l \circ A_l.$$2, and $$EV_1^l = A_{\max}^l \circ A_l, \qquad EV_2^l = A_{\min}^l \circ A_l, \qquad EV_3^l = A_{\mathrm{mean}}^l \circ A_l.$$3 are tuned on a held-out set. The explicit weighting scheme shows that the auxiliary style term is downweighted relative to MS-TKD, DMLD, and Dice, even though its ablation effect is substantial.

6. Quantitative performance, ablations, and computational profile

The paper reports average Dice (%) and HD95 (mm) across all 16 missing-modality scenarios for BraTS 2024 and FeTS 2024 (Zhu et al., 30 Jul 2025).

Dataset	Dice (WT / TC / ET)	HD95 (WT / TC / ET)
BraTS 2024	81.8 / 59.5 / 59.8	6.6 / 7.2 / 6.8
FeTS 2024	88.4 / 84.3 / 73.4	5.9 / 5.7 / 5.4

These results are presented as evidence that MST-KDNet surpasses current leading methods in both Dice and HD95 scores, particularly in conditions with substantial modality loss (Zhu et al., 30 Jul 2025). The summary further states that extensive experiments on two large public benchmarks show significant gains of 2–6 pp Dice and 1–3 mm HD95 over prior art, even in extreme one-modality cases (Zhu et al., 30 Jul 2025).

The ablation study on BraTS 2024 quantifies the contribution of each module.

Variant	Dice (WT / TC / ET)	HD95 (WT / TC / ET)
Full MST-KDNet	81.8 / 59.5 / 59.8	6.6 / 7.2 / 6.8
− MS-TKD	79.8 / 54.4 / 54.2	7.5 / 8.3 / 7.8
− GSME	78.3 / 55.1 / 53.4	9.6 / 9.7 / 9.5
− DMLD (SLKD)	80.0 / 56.1 / 55.2	8.1 / 8.7 / 8.0

The complexity profile is also specified. MST-KDNet has approximately 45 M parameters, an inference time of approximately 0.7 s per volume of size $$EV_1^l = A_{\max}^l \circ A_l, \qquad EV_2^l = A_{\min}^l \circ A_l, \qquad EV_3^l = A_{\mathrm{mean}}^l \circ A_l.$$4 on a Tesla V100, a training overhead of +15% runtime versus vanilla UNETR due to student/teacher forward-passes and discriminators, and a memory peak of approximately 12 GB GPU for batch size 1 (Zhu et al., 30 Jul 2025). This places the method in a regime of moderate deployment cost relative to its 3D transformer-based backbone and dual-network training procedure.

7. Interpretation and position within missing-modality segmentation

MST-KDNet can be read as a composite strategy for missing-modality segmentation in which three alignment levels are optimized simultaneously: multi-resolution attention and feature alignment through MS-TKD, output-space alignment through DMLD, and distributional/style alignment through GSME (Zhu et al., 30 Jul 2025). The model’s reported robustness and generalization potential are tied to this decomposition rather than to a single distillation loss.

A common misconception in this area is that missing-modality robustness is reducible to logit distillation or to naive teacher–student transfer. The ablation results argue against that interpretation: removing any one of MS-TKD, GSME, or DMLD degrades both Dice and HD95, and removing GSME produces particularly large HD95 increases (Zhu et al., 30 Jul 2025). This suggests that boundary quality, represented here by HD95, depends on feature-level and style-level alignment in addition to class-score matching.

Another potential misconception is that the teacher is required at deployment. In MST-KDNet, both networks run in parallel only during training; only the student is used at inference (Zhu et al., 30 Jul 2025). The method therefore separates privileged complete-modality supervision from the operational setting, which is the missing-modality student.

Within the limits of the reported evidence, MST-KDNet is characterized by a specific synthesis: fine-grained, multi-resolution attention alignment, both MSE and flexible KL-based logit alignment, and a style-matching adversarial-plus-Gram loss to handle any missing-modality setting (Zhu et al., 30 Jul 2025). The reported BraTS 2024 and FeTS 2024 results, together with the complexity profile and released source code, position it as a concrete implementation of transformer-based distillation for clinically relevant incomplete-modality segmentation.