【论文阅读笔记】A Single Stage Knowledge Distillation Network for Brain Tumor Segmentation on Limited MR Ima

Choi Y, Al-Masni M A, Jung K J, et al. A Single Stage Knowledge Distillation Network for Brain Tumor Segmentation on Limited MR Image Modalities[J]. Computer Methods and Programs in Biomedicine, 2023: 107644.

一.核心思想概述

本文提出了一种用于大脑肿瘤分割的单阶段知识蒸馏网络，该网络专门设计用于在有限的磁共振成像(MRI)模态下工作。该方法的核心在于，通过知识蒸馏技术，从完整模态的教师网络（$N_T$）向仅使用有限模态的学生网络（$N_S$）传递信息，在一个训练循环中依次训练两个网络。这种方法利用Barlow Twins损失在潜在空间水平进行知识转移，并采用深度监督的思想在像素级进行知识蒸馏。该研究表明，这种单阶段学习方法在使用FLAIR和T1CE图像模态时，在大脑肿瘤分割任务上表现优异，优于现有的多种状态艺术方法，使其更接近临床实践的需求。

二.模型结构

该框架的关键特点如下：

1. 双网络结构：该框架包括一个教师网络（$N_T$）和一个学生网络（$N_S$）。教师网络接受所有可用的成像模态，而学生网络仅使用有限的模态。这种结构设计是为了从教师网络向学生网络有效地转移知识。

2. 知识蒸馏过程：知识蒸馏是在单阶段框架内进行的，意味着教师和学生网络是同时训练的，但在训练一个时候，固定住另外一个的参数。这与传统的两阶段蒸馏方法不同，在传统方法中，学生网络是在教师网络预训练完成后单独训练的。

3. Barlow Twins损失：为了实现知识蒸馏，作者使用了Barlow Twins损失函数，该函数在潜在空间水平上操作，有助于减少教师和学生网络间的冗余。这是通过比较教师网络和学生网络的潜在特征表示来实现的。

4. 深度监督：此外，为了在像素级进行知识蒸馏，文中采用了深度监督的概念，通过交叉熵损失来训练教师和学生网络的主干网络。这有助于确保网络能够在各个层级上有效地学习和传递知识。

5. 训练方法：在训练过程中，学生网络学习模仿教师网络的输出，而教师网络则利用完整的模态信息来生成高质量的分割结果。通过这种方式，学生网络能够在仅访问有限模态时产生更准确的分割结果。训练学生网络时候用到了三种损失函数：

   • 深度监督Dice损失 ($L_{DSDice}$)：这是一种基于Dice系数的损失函数，用于测量学生网络输出和真实标签之间的相似度。深度监督意味着这种损失不仅在网络的最终输出上计算，还在网络的多个层级上计算，从而促进网络在不同深度上学习更有效的特征表示。

   • 交叉熵损失 ($L_{DSCE}$)：交叉熵损失用于优化学生网络的像素级预测，使其输出与真实标签更加接近。同样，这种损失也采用深度监督方法，以加强网络对细节的学习能力。

   • Barlow Twins损失 ($L_{BT}$)：这是一个用于潜在空间知识蒸馏的损失函数。通过减少教师网络和学生网络之间的冗余信息，它有助于将教师网络中的关键特征有效地转移到学生网络。这种损失函数特别用于处理输入图像的特征表示，以及网络的最终输出。

几个主要组成模块功能概述:

• Distillation on the latent space level:介绍了如何在潜在空间层级上进行知识蒸馏。从教师网络$N_T$中得到两个输出。第一个输出是来自$N_T$瓶颈的特征映射，具有高级属性,另一个是$N_T$最后一个解码器的输出。经过平均池化缩减特征维度并使用简单MLP映射成需要的shape，得到特征$F_T$、Z_T。学生网络也是同样得到$F_S$、$Z_S$，然后是这一步的关键，计算师生网络各自的互相关矩阵$C$。

  $C_F^{i,j}=\frac{{\sum }_b{F}_T^{b,i}{F}_S^{b,j}}{\sqrt{{\sum }_b{\left(F_T^{b,i}\right)}^2}\sqrt{{\sum }_b{\left(F_S^{b,j}\right)}^2}}$ $C_Z^{i,j}=\frac{{\sum }_b{Z}_T^{b,i}{Z}_S^{b,j}}{\sqrt{{\sum }_b{\left(Z_T^{b,i}\right)}^2}\sqrt{{\sum }_b{\left(Z_S^{b,j}\right)}^2}}$

  然后对这两个矩阵应用Barlow Twins loss，该损失是一种用于深度学习中的自监督学习方法，特别适用于减少冗余并提高特征表示的质量。Barlow Twins损失的核心思想是通过使网络的两个不同视角（或变换）的输出尽可能相似，同时保持这些特征表示的独立性，从而提高网络学习到的特征的质量。在知识蒸馏的背景下，Barlow Twins损失被用来优化教师网络（$N_T$）和学生网络（$N_S$）之间的特征表示。具体来说，它促使学生网络的潜在空间特征尽可能接近教师网络的特征，同时确保这些特征不是冗余的或过度相关的。这样做的目的是提高学生网络在有限模态下的性能，使其能够捕获并利用教师网络的复杂特征表示。Barlow Twins损失的数学公式通常涉及到计算两个特征表示之间的相关矩阵，并尝试使这个矩阵接近一个单位矩阵。单位矩阵的特性是对角线上的元素为1（表示每个特征与自身的完美相关），而非对角线元素为0（表示不同特征间的完全不相关）。通过这种方式，Barlow Twins损失鼓励学生网络学习到的特征既与教师网络的特征相似，又能保持特征之间的独立性。

  $L_{BT}=\sum {\left(1-C_{ii}\right)}^2+\lambda {\sum }_{i\ne j}{C}_{ij}^2$

• Distillation on the pixel level：这个很传统，本文是在$N_T$和$N_S$的最后三个解码器的输出之间应用带有深度监督概念的交叉熵损失($L_{DSCE})$。

• Deep supervision:本文中涉及三种情况的深度监督。前两种情况表示使用教师和学生网络的ground-truth标签计算三个解码器级别的骰子损失。在像素级交叉熵损失的知识蒸馏过程中考虑了第三种情况。关于为什么只使用三层解码器，本文中是实验得出，更多的层没有带来显著的性能提升。

• Total Loss：$L_{\text{total }}=\sum \alpha \ast \left(L_{\text{DSDice }}^{N_T}+L_{DSDice}^{N_S}\right)+\beta \ast L_{DSCE}^{N_T\leftrightarrow N_S}+\gamma \ast \left(L_{BT}^{F_T\leftrightarrow F_S}+L_{BT}^{Z_T\leftrightarrow Z_S}\right)$

三.数据集和硬件

数据集用到了公开的BraTS2021和作者自行收集的韩国医院数据

Python 3.8.8 and Pytorch1.8.8 NVIDIA RTX A5000 *4

四.性能效果

与两阶段蒸馏相比：

与SOTA对比：

五.局限性

• 很难准确预测边界粗糙的肿瘤或非常小的肿瘤。
• 在模型错误地将正常组织预测为肿瘤的例子中，运动伪影引起的肿瘤是最常见的，对临床医生的诊断准确性有重要影响。因此，如何开发一个运动工件鲁棒分割模型来克服这些限制是一个值得探索的方向。