论文题目:
缺失模态医学分割的原型知识蒸馏
项目地址:https://github.com/SakurajimaMaiii/ProtoKD
摘要:
多模态医学成像在临床治疗中至关重要,因为它可以为医学图像分割提供补充信息。然而,由于扫描时间限制和其他临床情况,收集多模态数据在临床上是困难的。因此,开发一种能够处理缺失模态问题的图像分割范式在临床上具有重要意义。本文提出了一种原型知识蒸馏(ProtoKD)方法来解决这一挑战性问题,尤其是当只能访问单一模态数据时最困难的场景。具体来说,ProtoKD不仅可以将多模态数据的像素级知识蒸馏到单模态数据,还可以转移类内和类间特征变化,使学生模型能够从教师模型学习到更稳健的特征表示,并仅使用单一模态数据进行推理。该方法在BraTS基准测试中实现了最先进的性能。代码可在GitHub上获得。
引言:
多模态成像在医学图像分析领域具有重要意义,因为它提供了医学诊断的补充信息。尽管多模态成像通常能产生准确的诊断,但由于数据损坏或临床场景中的不同扫描协议,通常很难收集到完整的多模态图像集。因此,迫切需要一种健壮的医学图像分割方法来解决在推理时缺失模态的问题。本文提出了三种主要方法来解决这一挑战性问题:第一种是合成缺失模态以完成测试集;第二种是学习包含可访问域之间模态不变信息的共享潜在空间;第三种是基于知识蒸馏的方法,旨在将使用多模态图像训练的教师模型的知识转移到仅使用单模态训练的学生模型。
方法概述
- 教师模型(Teacher Model)使用多模态数据进行训练,而学生模型(Student Model)仅使用单模态数据。
- 通过知识蒸馏,将教师模型的知识转移到学生模型,使学生模型能够学习到更稳健的特征表示
方法
-
原型知识蒸馏(ProtoKD):
- 像素级知识蒸馏(Pixel-wise Knowledge Distillation):通过最小化学生模型和教师模型预测之间的Kullback-Leibler散度,鼓励学生模型学习教师模型的知识。
- 原型知识蒸馏(Prototype Knowledge Distillation):考虑类内和类间特征变化的相关性,通过转移教师模型到学生模型的原型(prototypes)知识,来显式地利用内部语义相关性。
-
原型学习(Prototype Learning):
- 原型是每个类别的嵌入中心,通过所有像素特征与其对应的类别原型之间的相似度来捕获类内和类间的特征表示。
-
相似度计算:
- 对于每个像素,计算其与所有类别原型的相似度,如果像素属于某个类别,则相似度表示类内特征变化;如果不属于,则表示类间特征变化。
-
损失函数:
- 像素级知识蒸馏损失(Pixel-wise Knowledge Distillation Loss):使用KL散度计算学生模型和教师模型预测之间的差异。
- 原型知识蒸馏损失(Prototype Knowledge Distillation Loss):使用L2距离最小化两个I2FV(Intra-class and Inter-class Feature Variation)相似度图之间的距离。
- 分割损失(Segmentation Loss):结合交叉熵损失和Dice损失,用于评估模型输出与真实标注之间的一致性。
-
目标函数:
- 最终的目标函数是分割损失、像素级知识蒸馏损失和原型知识蒸馏损失的加权和。
实验:
本文在2018年的BraTS挑战赛数据集上评估了所提出的方法,该数据集包含285个案例,并具有手动注释的标签。每个受试者有四种MRI模态,包括T1、T2、T1ce和Flair。为了预处理,每个体积被归一化为零均值和单位方差。由于GPU内存限制,每个体积被随机裁剪为96×128×128以输入网络。作者随机将285个案例分为训练(70%)、验证(10%)和测试(20%)。
结论:
本文提出了一种新颖的基于知识蒸馏的方法来解决医学图像分割中的缺失模态问题。通过引入类内和类间特征变化的蒸馏,缩小了学生模型和教师模型之间的特征分布差异。这种方法通过从通常具有更好特征表示的教师模型转移知识,使学生模型能够捕获更稳健的特征。在2018年的BraTS基准测试上进行了广泛的实验,实验结果证明了该方法的有效性。