重新思考知识蒸馏中中间层设计用于肾与肝肿瘤分割
计算机视觉与模式识别
2024-05-28 v2 人工智能
机器学习
组织与器官
摘要
知识蒸馏(KD)已在多个领域取得显著成功，但其在医学影像任务(如肾和肝肿瘤分割)中的应用面临挑战。许多现有 KD 方法并非专门针对这些任务设计。此外，主流 KD 方法往往缺乏对学生从教师网络‘蒸馏什么’以及‘从何处蒸馏’的细致考量。这一疏忽可能导致浅层学生层中训练偏差累积，进而削弱 KD 的有效性。为应对这些挑战，我们提出分层选择性反馈蒸馏(HLFD)。HLFD 策略性地将中间层知识蒸馏至更早层，并在特征和像素层面将最终层知识迁移至中间层。该设计使模型能从更早层学习更高质量的表征，从而得到鲁棒且紧凑的学生模型。大量定量评估表明，HLFD 以显著优势超越现有方法。例如，在肾脏分割任务中，HLFD 较无 KD 的学生模型提升超 10%，显著增强其对肿瘤特异性特征的关注。从定性角度看，经 HLFD 训练的学生模型善于抑制无关信息并能精准聚焦于肿瘤特异性细节，这为更高效准确的诊断工具开辟了新路径。代码见 \href{https://github.com/vangorade/RethinkingKD_ISBI24}{此处}。
引用
@article{arxiv.2311.16700,
title = {Rethinking Intermediate Layers design in Knowledge Distillation for Kidney and Liver Tumor Segmentation},
author = {Vandan Gorade and Sparsh Mittal and Debesh Jha and Ulas Bagci},
journal= {arXiv preprint arXiv:2311.16700},
year = {2024}
}
备注
Accepted at ISBI-2024 for Oral Presentation