【论文阅读笔记】Prototype knowledge distillation for medical segmentation with missing modality

Wang S, Yan Z, Zhang D, et al. Prototype knowledge distillation for medical segmentation with missing modality[C]//ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2023: 1-5. [开放源码]

一、论文概述

本文介绍了一种名为原型知识蒸馏（Prototype Knowledge Distillation，简称ProtoKD）的方法，旨在解决医学影像分割中的缺失模态问题。在多模态医学成像中，由于扫描时间限制和其他临床情况，往往难以收集完整的多模态数据。ProtoKD方法通过从多模态数据中提取像素级知识以及类内和类间特征变化，将这些知识转移到仅访问单一模态数据的学生模型中。这样，学生模型能够从教师模型学习更鲁棒的特征表示，并仅使用单一模态数据进行推断。该方法在BraTS基准测试中实现了最先进的性能。

二、论文创新点

本文的创新点在于提出了一种新颖的原型知识蒸馏（ProtoKD）方法，用于解决医学影像分割中的缺失模态问题。其主要创新之处包括：

1. 处理单模态数据的挑战： ProtoKD专门针对在只能访问单一模态数据的情况下的医学影像分割，这是一个在临床实践中常见但难以处理的问题。
2. 像素级知识蒸馏： 该方法采用了像素级知识蒸馏，使得学生模型能够从教师模型中学习细粒度的特征表示，这对于医学影像分割任务尤为重要。
3. 类内和类间特征变化的转移： ProtoKD不仅蒸馏像素级知识，还转移类内和类间的特征变化，这有助于学生模型更好地理解和区分不同类别的特征。
4. 原型学习的应用： 该方法利用原型学习来表示每个类别的特征中心，这在少样本学习领域中常见，但在此文中被创新性地应用于知识蒸馏和医学影像分割。
5. 在BraTS基准上的表现： 该方法在BraTS基准测试中展示了其优越性，实现了最先进的性能，这证明了其有效性和实用性。

二、解决缺失模态问题的三类常见方法及其优缺点概述
在本文的引言部分中，作者提到了处理医学影像分割中的缺失模态问题的三种主要方法，以及它们的优缺点：

1. 合成缺失模态： 第一种方法是通过生成模型来合成缺失的模态数据，以完成测试集。这种方法的优点在于它可以直接补充缺失的信息，使得分割模型可以在完整的多模态数据上运行。然而，其缺点是需要额外的训练过程来生成合成数据，而且在只有单一模态可用时，合成多种模态的难度会显著增加。
2. 学习共享潜在空间： 第二种方法旨在学习一个包含模态不变信息的共享潜在空间。这种策略的优点是在多模态数据可用时通常能够取得良好的性能。但是，其缺点是在仅有单一模态可用时性能会大幅下降。
3. **基于知识蒸馏的方法：**通过将从多模态图像训练的教师模型的知识转移到仅使用单一模态训练的学生模型上。这种方法的优点是能够直接从多模态数据中获得的丰富信息中受益，而不需要合成数据。然而，它们的缺点是可能无法充分捕获模态间的复杂关系，特别是在医学图像中，这些关系对于精确的分割至关重要

因此本文要做的事情顺利引出：蒸馏+捕获模态间关系+单一模态可用

三、本文提出的方法

• 像素级知识蒸馏（Pixel-wise Knowledge Distillation）

像素级知识蒸馏是一种常见的知识蒸馏方法，它在医学影像分割中的应用可以被视为像素级的分类问题。在这种方法中，学生模型（单模态输入）被训练以模仿教师模型（多模态输入）的输出。具体来说，通过最小化学生模型和教师模型预测之间的Kullback-Leibler (KL) 散度来实现这一点。这种方法的关键在于使学生模型的预测尽可能接近教师模型的预测，从而学习到教师模型的知识。

• 原型知识蒸馏（Prototype Knowledge Distillation）

  原型知识蒸馏是本文的创新点，它不仅考虑像素级的知识，还关注类内和类间的特征变化。这种方法的核心在于捕捉并传递教师模型中的内部语义关联，这是通过计算每个类别的原型（即类别的特征中心）并分析像素特征与这些原型之间的cos相似度来实现的。原型学习被广泛应用于few-shot learning领域，它代表了每个类的embedding中心，本文引入这种思想。

  • 原型的计算： 对于每个类别 $k$，原型 $c_k$ 是通过对属于该类别的所有像素特征的平均来计算的。公式为其中 $z_i$是像素 $i$的特征嵌入，$y_i$ 像素的真实类别。$c_k=\frac{{\sum }_i{z}_{i}1\left[y_i=k\right]}{{\sum }_{i}1\left[y_i=k\right]}$
  • 类内和类间特征变化（I2FV）： 对于每个像素 $i$，其与类别 $k$的原型 $c_k$ 之间的相似度 $M_k(i)$被计算出来。这个相似度可以代表类内（如果像素属于该类别）或类间（如果像素不属于该类别）的特征变化。$M_k(i)=\frac{z_i^T{c}_k}{\Vert z_i\Vert \Vert c_k\Vert }$,如果像素$i$属于k类，$M_k(i)$表示类内特征变化。如果像素$i$不属于类$k$，则$M_k(i)$可以表示类间特征变化。教师模型和学生模型分别生成$I^{2}FV$映射。
  • 知识转移： 在教师和学生模型中分别计算这些相似度，然后通过最小化两个模型之间的$I^{2}FV$图的$L2$距离来实现知识转移。${\mathcal{L}}_{\text{proto }}=\frac{1}{|\mathcal{N}|K}{\sum }_{i\in \mathcal{N}}{\sum }_{k=1}^K{\Vert M_k^s(i)-M_k^t(i)\Vert }^2$

• 最终损失函数

  几部分的加权：$\mathcal{L}={\mathcal{L}}_{\text{seg }}+\alpha {\mathcal{L}}_{kd}+\beta {\mathcal{L}}_{\text{proto }}$

四、数据集

BraTS 2018

五、实验效果

六、局限性

单一数据集

对比的baseline不够强

类内和类间的知识传递仅采用余弦相似度是否足够？

仅用原型学习，像素均值中心来代表一幅图是否合理？