【论文阅读笔记】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$、ZSZ_S