[论文解读] MS-KD: Multi-Organ Segmentation with Multiple Binary-Labeled Datasets

Shixiang Feng, Yuhang Zhou+3 more|arXiv (Cornell University)|Aug 5, 2021

Advanced Neural Network Applications参考文献 13被引用 7

一句话总结

该论文提出 MS-KD，一种多教师单学生知识蒸馏框架，通过使用多个预训练的单器官分割网络（教师）和一个单个多器官学生网络，训练一个多器官分割模型。通过应用基于区域的 logits 与特征级监督——其中每个教师在其对应器官区域以及所有教师一致的共享背景区域上指导学生——该方法实现了优越性能，在基准数据集上的平均 DSC 提升至 0.43%，超越了单个模型和当前最先进方法。

ABSTRACT

Annotating multiple organs in 3D medical images is time-consuming and costly. Meanwhile, there exist many single-organ datasets with one specific organ annotated. This paper investigates how to learn a multi-organ segmentation model leveraging a set of binary-labeled datasets. A novel Multi-teacher Single-student Knowledge Distillation (MS-KD) framework is proposed, where the teacher models are pre-trained single-organ segmentation networks, and the student model is a multi-organ segmentation network. Considering that each teacher focuses on different organs, a region-based supervision method, consisting of logits-wise supervision and feature-wise supervision, is proposed. Each teacher supervises the student in two regions, the organ region where the teacher is considered as an expert and the background region where all teachers agree. Extensive experiments on three public single-organ datasets and a multi-organ dataset have demonstrated the effectiveness of the proposed MS-KD framework.

研究动机与目标

• 为解决 3D 多器官分割的高标注成本问题，利用现有具有二值标签的单器官数据集。
• 开发一种知识蒸馏框架，使单个学生模型能够从多个针对不同任务的教师模型中学习，每个教师模型均在不同的单器官数据集上训练。
• 通过多个教师提供的软监督，克服现有方法的局限性，如基于条件的方法计算效率低下或依赖低质量伪标签。
• 通过基于区域的监督，整合来自多样化单器官专家的知识，提升多器官分割的泛化能力和性能。

提出的方法

• 该框架使用 K 个预训练的单器官分割网络作为教师，一个多器官分割网络作为学生。
• 应用基于区域的监督：每个教师在自身器官区域以及所有教师一致的背景区域上监督学生。
• logits 与特征级监督：通过维度转换将教师的二分类 logits 映射至学生 K+1 类输出空间，随后应用蒸馏损失。
• 特征级监督：使用分布匹配损失，将学生的中间特征分布与教师的对齐。
• 学生通过组合 logits 与特征级损失端到端训练，实现来自多个任务特定教师的知识迁移。
• 该方法避免迭代推理或动态头生成，确保训练和推理阶段计算成本均较低。

实验结果

研究问题

• RQ1 仅使用多个具有二值标注的单器官数据集，能否有效训练单个多器官分割模型？
• RQ2 如何在无需任务特定条件或迭代推理的情况下，有效将多个任务特定教师模型的知识迁移至单个学生模型？
• RQ3 基于区域的监督（结合 logits 与特征）是否优于标准伪标签或一致性蒸馏方法？
• RQ4 所提出的框架在仅使用单器官数据训练的情况下，能否在未见的多器官数据集上实现良好泛化？

主要发现

• 当在三个单器官数据集上训练时，MS-KD 在平均 DSC 上相比单个单器官模型绝对提升 0.80%（91.10% vs. 90.30%）。
• 引入特征级监督后，平均 DSC 进一部提升 0.22% 至 91.32%，证明中间特征蒸馏的显著优势。
• 在 BTCV 多器官基准上，MS-KD 的平均 DSC 比单个模型高出 1.32%，展现出强大的泛化能力。
• 尽管无需额外推理时间或模型参数，MS-KD 仍优于当前最先进方法，包括 Co-training（平均 DSC 提升 0.43%）和 DoDNet。
• 定性结果表明，由于多个教师提供的平滑监督，MS-KD 生成的边界更精确、更平滑，尤其在胰腺等困难器官上表现更优。
• 不确定性分析显示，教师在器官边界处最不确定，而 MS-KD 有效利用这种软监督提升了边缘检测性能。