Source-linked AI summary
Knowledge distillation from multi-modal to mono-modal segmentation networks
Minhao Hu, Matthis Maillard, Ya Zhang, Tommaso Ciceri, Giammarco La Barbera, Isabelle Bloch, Pietro Gori
TL;DR
Clinical constraints frequently make multi-modal imaging unavailable, despite its segmentation benefits. KD-Net transfers knowledge from a multi-modal Teacher to a mono-modal Student using generalized distillation and latent-space matching, and the approach improves mono-modal brain-tumor segmentation across tested architectures.
Problem
Multi-modal segmentation can improve accuracy, but clinical constraints often mean that only one imaging modality is acquired.
Method
KD-Net trains a mono-modal Student to learn from a multi-modal Teacher through generalized knowledge distillation and similar latent representations.
Results
KD-Net achieved better brain-tumor segmentation results than state-of-the-art methods using only T1ce on BraTS 2018.
Takeaways & Limitations
The framework is generic across encoder-decoder architectures and supports mono-modal segmentation when only one modality is available at test time.
Takeaways & Limitations
Synthesizing an additional modality from T1ce produced worse-than-baseline results and was computationally prohibitive.
Abstract
from arXiv · showhide
The joint use of multiple imaging modalities for medical image segmentation has been widely studied in recent years. The fusion of information from different modalities has demonstrated to improve the segmentation accuracy, with respect to mono-modal segmentations, in several applications. However, acquiring multiple modalities is usually not possible in a clinical setting due to a limited number of physicians and scanners, and to limit costs and scan time. Most of the time, only one modality is acquired. In this paper, we propose KD-Net, a framework to transfer knowledge from a trained multi-modal network (teacher) to a mono-modal one (student). The proposed method is an adaptation of the generalized distillation framework where the student network is trained on a subset (1 modality) of the teacher's inputs (n modalities). We illustrate the effectiveness of the proposed framework in brain tumor segmentation with the BraTS 2018 dataset. Using different architectures, we show that the student network effectively learns from the teacher and always outperforms the baseline mono-modal network in terms of segmentation accuracy.
1 Introduction
Multi-modal imaging can improve medical image segmentation, but clinical constraints often leave only one modality available. KD-Net addresses this gap by transferring knowledge from a multi-modal Teacher to a mono-modal Student through generalized distillation and latent-representation matching.
- Multi-modal image fusion can improve segmentation accuracy because different modalities describe different tissues.
- Clinical settings often cannot acquire multiple modalities because of limited physicians and scanners, cost, scan-time constraints, or urgent patient needs.
- Existing approaches synthesize missing modalities or learn modality-invariant feature spaces to support segmentation when modalities are unavailable at inference.
- KD-Net transfers knowledge from a multi-modal network to a mono-modal network using generalized knowledge distillation.
- The Teacher and Student share an encoder-decoder architecture, while the Student learns from the Teacher’s privileged multi-modal information through similar latent representations.
- The framework trains the Teacher with reference-segmentation loss, then trains the Student with KL, KD, and ground-truth losses.
2 KD-Net
KD-Net trains a mono-modal Student from a multi-modal Teacher using soft predictions, latent representations, and reference segmentations. The Student is optimized with complementary losses while the Teacher remains fixed.
- KD-Net setup: The Student receives one modality, while the Teacher receives the concatenation of N modalities for the same sample.The Student modality is indexed by k, a fixed integer between 1 and N.
- Training procedure: The Teacher is trained first, then its frozen weights provide supervision while the Student is trained without back-propagating error into the Teacher.Student training uses knowledge distillation, latent-space dissimilarity, and reference-segmentation losses.
- Knowledge distillation: Soft Teacher predictions transfer additional information, with temperature T controlling target softness and higher T producing softer targets.Soft targets are intended to expose class relations that hard labels may not reveal.
- Knowledge distillation: The distillation loss combines Dice and binary cross-entropy to measure global shape similarity and local pixel-wise agreement with Teacher outputs.Teacher soft labels are binarized for the binary cross-entropy term.
- Latent space: The latent-space loss uses asymmetric KL divergence so the Student bottleneck distribution is encouraged to resemble the Teacher bottleneck distribution.The bottleneck vectors are flattened and normalized before applying the divergence.
- Objective function: The complete Student objective combines distillation, reference-segmentation, and latent-space losses, with λ balancing Teacher versus reference supervision and α scaling KL loss.λ lies in [0, 1], while α is positive.
3 Results and Discussion
On BraTS 2018, KD-Net was evaluated with a four-modality Teacher and a T1ce-only Student, outperforming baseline and comparison models while benefiting from both distillation losses. Qualitative results and additional experiments indicate improved segmentation quality, although modality synthesis performed worse than baseline and was computationally prohibitive.
- Experimental setup: The Teacher used T1, T2, T1ce, and Flair, whereas the Student and baseline received only T1ce.T1ce was selected because it is standard in pre-operative neurosurgery or radiotherapy.
- Model comparison: KD-Net outperformed U-HVED and HeMIS on all three tumor components using T1ce input.The comparison used Dice scores, with U-HVED and HeMIS results taken from prior work without reported standard deviations.
- Ablation study: Both KL and KD losses improved results over the baseline, especially for enhancing tumor and tumor core.The ablation compared the baseline GT loss, KD-Net with only KL, and the complete objective under 0 or 4 skip-connections.
- Qualitative results: Qualitative comparisons showed that KD-Net helped the Student discard outliers and produce higher-quality segmentation labels.The Student appeared to learn more in the enhancing-tumor region highlighted by T1ce.
- Additional observations: Synthesizing Flair from T1ce and using it with T1ce produced worse results than baseline and required prohibitive computational time.Sharing weights between Teacher and Student bottleneck layers produced results almost identical but slightly worse than the proposed framework.
4 Conclusions
KD-Net transfers knowledge from a multi-modal segmentation network to a mono-modal one and improves T1ce-only brain-tumor segmentation on BraTS 2018. The framework is generic across encoder-decoder architectures and is intended for settings where only one modality is available at test time.
- KD-Net combines generalized knowledge distillation with a latent-representation constraint aligning Student and Teacher features.The Teacher uses multiple modalities while the Student uses one.
- On BraTS 2018, KD-Net achieved better T1ce-only segmentation results than state-of-the-art methods.The paper reports gains in segmentation accuracy and robustness to errors.
- The framework is generic and can be applied to any encoder-decoder segmentation network.
- KD-Net is positioned for clinical scenarios where only one imaging modality is available at test time.