Rethinking Intermediate Layers design in Knowledge Distillation for Kidney and Liver Tumor Segmentation

TL;DR

The paper tackles the challenge of applying knowledge distillation to kidney and liver tumor segmentation by introducing Hierarchical Layer-selective Feedback Distillation (HLFD), which blends feature- and pixel-level guidance across multiple teacher-to-student layer mappings. HLFD decomposes distillation into FLFD and PLFD, incorporating unified and individual transfers with multi-task losses to encourage early-layer quality representations and accurate pixel-level predictions. Across KiTS and LiTS datasets, HLFD outperforms non-KD students and existing KD methods, achieving notable Dice score gains and improved volume accuracy, while qualitative analysis confirms sharper tumor focus and reduced irrelevant information. This approach yields a robust, compact student model with practical potential for efficient, accurate clinical diagnostics, and code is publicly available for replication.

Abstract

Knowledge distillation (KD) has demonstrated remarkable success across various domains, but its application to medical imaging tasks, such as kidney and liver tumor segmentation, has encountered challenges. Many existing KD methods are not specifically tailored for these tasks. Moreover, prevalent KD methods often lack a careful consideration of `what' and `from where' to distill knowledge from the teacher to the student. This oversight may lead to issues like the accumulation of training bias within shallower student layers, potentially compromising the effectiveness of KD. To address these challenges, we propose Hierarchical Layer-selective Feedback Distillation (HLFD). HLFD strategically distills knowledge from a combination of middle layers to earlier layers and transfers final layer knowledge to intermediate layers at both the feature and pixel levels. This design allows the model to learn higher-quality representations from earlier layers, resulting in a robust and compact student model. Extensive quantitative evaluations reveal that HLFD outperforms existing methods by a significant margin. For example, in the kidney segmentation task, HLFD surpasses the student model (without KD) by over 10\%, significantly improving its focus on tumor-specific features. From a qualitative standpoint, the student model trained using HLFD excels at suppressing irrelevant information and can focus sharply on tumor-specific details, which opens a new pathway for more efficient and accurate diagnostic tools. Code is available \href{https://github.com/vangorade/RethinkingKD_ISBI24}{here}.

Figures (3)

• Figure 1: The input X and augmentation $X'$, undergo encoding by both a pre-trained teacher encoder and a randomly initialized student encoder, resulting in representations $z^{t}_{early}$, $z^{t}_{\text{mid}j}$, $z^{t}_{ter}$, and $z^{s}_{early}$, $z^{s}_{\text{mid}j}$, $z^{s}_{ter}$, respectively. These representations contribute to feature-level loss functions, $\mathcal{L}_{UFD}$ and $\mathcal{L}_{IFD}$. Additionally, the teacher decoder decodes $z^{t}_{ter}$, producing representations $p^t{early}$, $p^t_{\text{mid}j}$, $p^t_{ter}$, which are utilized in pixel-level loss functions, $\mathcal{L}_{UPD}$ and $\mathcal{L}_{IPD}$. The training process is further enhanced with the inclusion of a supervised focal dice loss ($\mathcal{L}_{seg}$).
• Figure 2: The green color highlights the regions of interest (ROI) representing tumors. The segmentation maps are presented for both KiTS (first two rows) and LiTS (last two rows) datasets, with 'G.T.' denoting the ground truth.
• Figure 3: Gradient-activated class maps for KiTS19, featuring CAM results for both the final and first layers. Remarkably, HLFD focuses on the tumor region, maintaining effectiveness in suppressing irrelevant information as the process advances to the final layer.