点击下方卡片，关注「集智书童」公众号
想要了解更多：
前沿AI视觉感知全栈知识👉「分类、检测、分割、关键点、车道线检测、3D视觉（分割、检测）、多模态、目标跟踪、NerF」
行业技术方案 👉「AI安防、AI医疗、AI自动驾驶」
AI模型部署落地实战 👉「CUDA、TensorRT、NCNN、OpenVINO、MNN、ONNXRuntime以及地平线框架」
欢迎扫描上方二维码，加入「 集智书童-知识星球 」，日常分享论文、学习笔记、问题解决方案、部署方案以及全栈式答疑，期待交流！
免责声明
凡本公众号注明“来源：XXX（非集智书童）”的作品，均转载自其它媒体，版权归原作者所有，如有侵权请联系我们删除，谢谢。
这篇论文的主要贡献是提出了一种名为FedBrain-Distill的联邦学习方法，该方法在保持用户隐私的同时，确保了联邦学习客户端在模型架构上的独立性。FedBrain-Distill使用了一个教师集，将他们的知识蒸馏为一个简单的学生模型。
论文对FedBrain-Distill进行了评估，结果表明，在独立数据（IID）和相同分布（IID）的数据以及非独立分布（非IID）的数据上，它具有高精确度和低通信成本。需要注意的是，作者使用了Dirichlet分布将数据划分为IID和非IID数据。。
联邦学习（FL）：在这种方法中，多个客户端在不共享原始数据的情况下共享模型更新。
知识蒸馏（KD）：这是一种机器学习技术，将一个大型模型（教师）的复杂知识传递给一个较小的模型（学生）。
脑癌分类：这是一个关键的医学应用场景，需要解决隐私问题和数据可访问性。
非独立同分布（非IID）数据：在这种数据分布下，不同客户端的数据集是相互独立的，每个客户端看到的数据集不同。
1 Introduction
最近，由于其对患者生存率产生重大影响，脑肿瘤的分类已成为当务之急。如果诊断不准确，这种现象可能会导致不可逆转的神经损伤，在一些情况下甚至可能导致死亡。脑肿瘤是医学患者中最具有挑战性的问题之一，其发病率在医疗患者中迅速增长。许多因素使一种肿瘤与其他肿瘤区分开来，其中一些因素包括肿瘤的大小、其生长模式以及其恶性程度。
随着医疗成像和机器学习的技术进步，磁共振成像（MRI）可以轻易地被机器学习算法（如卷积神经网络（CNNs））利用。这些机器学习模型已经表明了在脑肿瘤分类中的重要性。然而，在部署这些模型时，作者面临着意外的后果。首先，由于在多个机器学习模型之间分享用户敏感数据，患者的隐私被侵犯。换句话说，数据被分配给需要训练的尽可能多的模型。其次，作者忽略了CNN模型的高准确性所涉及的计算复杂性。第三，即使作者不考虑这两个障碍，某些医疗机构在利用CNN模型时可能需要不同的架构。
知识蒸馏（KD）最近作为一种与架构无关的解决方案而出现，因为它能够从复杂的教师模型向简单的学生模型传递知识[1]。这种技术也可以在联邦学习（FL）设置中使用，称为联邦知识蒸馏（FKD）[2]。因此，FKD将能够解决上述挑战。首先，CKD通过构建基于FL的架构来保护用户的敏感数据。此外，与FL中的默认设置不同，FKD技术允许客户端在模型架构上相互独立。因此，知识从单一或多个复杂教师模型中蒸馏到一个较简单的学生模型，确保患者敏感数据得到保护，计算负担由有能力和资源处理大量数据的复杂教师承担。此外，由于FL设置中的主要标准之一是通信成本，FKD确保客户端间的通信尽可能地优化，因为教师和学生模型只传递软标签，而不是模型参数[3]。
在这篇论文中，作者提出了一种名为FedBrain-Distill 的新的、通信高效的方法，用于脑肿瘤分类任务，该方法将来自复杂教师模型的知识蒸汽到一个更简单的学生模型。作者利用多个预训练的VGGNet16模型进行教师训练。此外，作者利用Dirichlet分布在这些教师之间创建非独立且同分布（non-IID）数据。对于FedBrain-Distill，作者利用两种不同的场景，其中2位和5位教师的知识被蒸馏到一个简单的学生模型。作者在实验阶段尝试这两种设置，分别对应教师之间的IID和非IID数据分布。这种设置让作者在数据高度偏斜于特定类数据时观察到FedBrain-Distill的效率。在作者的实验中，作者使用了Figshare脑肿瘤数据集，该数据集包括三类肿瘤：脑膜瘤、胶质瘤和垂体瘤。
本文的结构如下：在II节中，作者讨论了提出相似方法的相关工作。在III节中，作者介绍了FedBrain-Distill的架构，并讨论了如何从教师和学生模型中提取知识。在IV节中，作者提供了实验结果。
II Related Work
许多研究已经致力于利用预训练模型进行脑肿瘤分类。尽管这项工作不是基于知识蒸馏（KD）技术，Deepak和Ameer [4]的工作利用GoogLeNet模型从脑部 MRI 图像中提取特征。事实上，他们的研究整合了多个深度 CNN 模型，并与 GoogLeNet 相结合，以在Figshare MRI 数据集上实现高精度。Wu等人[5]提出了FedKD，该方法解决传统FL方法中固有的通信开销问题。他们提出了一种自适应互惠式知识蒸馏与动态梯度压缩相结合的方法来降低通信开销，同时不牺牲模型性能。Rahimpour等人 [6] 提出了一个跨模态蒸馏方法，该方法涉及使用多序列 MRI 数据训练教师模型和使用单序列数据训练学生模型。他们使用了2018年BraTS数据集和自定义U-net模型的内部数据集。Shakib Khan等人 [7] 和Adepu等人 [8] 利用DL模型来实现教师模型的部署，以检测和分类黑色素瘤。
Qin等人 [9] 提出了一种新型的结构，用于从教师到学生的整体结构有效地传播医学图像分割知识。Wang等人[10] 利用条件概率表示从教师模型中提取知识，实现腹部内脏器官的多任务分割。Wang等人 [11] 提出了一种新方法，通过使用生成模型来实现教师模型，以减少来自不相关局部模型的错误预测的影响，从而增强聚合模型的整体性能。
Zhang等人 [12]的工作解决数据异构挑战，通过为教师使用生成模型来创建合成数据并传递模型知识和蒸馏，以实现聚合。Chen等人 [13] 提出MetaFed，为不同联邦之间的FL创建了一个框架，利用循环知识蒸馏来积累通用知识并实现每个联邦的个性化模型。最后，Viet等人 [14] 利用多个预训练模型在FL设置中使用Figshare数据集对脑肿瘤进行分类。他们使用VGGNet16在IID和非IID数据上实现了惊人的准确性。
III Proposed Method
FL（Flattened Language）范式的出现，为研究界开启了大门，使得数据从服务器传输到代码的方式发生了改变，变成了代码从服务器传输到数据。换句话说，一个名为聚合器的中央服务器协调从一组客户端收集的单一全连接模型的学习过程[15]。图1展示了这个过程。
FedBrain-Distill 是基于 FL 架构构建的。FedBrain-Distill 将一组教师的知识提取为一个简单的学生模型。这个简单的模型可以在缺乏训练所需硬件的任何客户端上部署。这一观点与 FL 设置中的聚合器-客户端关系相同。因此，在 FedBrain-Distill 中的教师是客户端，而学生模型位于聚合器的最上层。图2说明了 FedBrain-Distill 整体架构。以下作者将详细讨论每个部分。
Preprocessing
教师和学生模型的成功都依赖于数据预处理，因为它可以解锁脑图像的丰富特征，从而实现有效的训练。因此，所有肿瘤图像都采用了 Contrast-Limited Adaptive Histogram Equalization (CLAHE) 技术 [15] 进行了归一化、重排和增强。这样，每个图像中的肿瘤区域（肿瘤 Mask ）都得到了增强，使模型能够更有效地提取特定的肿瘤区域。图3展示了未经过增强的归一化图像和 Reshape 前的增强图像之间的差异。所有肿瘤图像的形状都为512 x 512。FedBrain-Distill将所有图像 Reshape 为224 x 224 x 3，因为所有教师模型所使用的VGGNet16模型仅接受后两种形状的输入层。
Data
在FedBrain-Distill中，作者对IID和非IID数据进行了评估，以确保学生结果和相关性异构数据的可靠性。FedBrain-Distill通过从Dirichlet分布中抽样类先验来实现IID和非IID数据分布[17]。当接近0时，分割变得更 heterogeneous，而当接近无穷大时，分割变得更 uniform。FedBrain-Distill用表示IID设置，而对于非IID设置，则使用。
Public and Private Dataset
在联邦学习（FL）环境中，隐私是非常重要的，因此所有客户端都倾向于依赖他们自己的本地数据集。FKD设置中的这个数据集被称为私有数据集。另一方面，FKD的最终目的是基于公共数据集，从教师提炼出知识传递给学生，这需要一个可以在整个联邦内所有参与者之间共享的第二数据集。在提炼阶段，该数据集被称为公共数据集，它将在教师和学生的模型之间共享[3]。
Teacher Models
每一教师模型在其自己的私有数据上进行训练，并具有相应的标签。每个教师模型的最终目标是优化其局部交叉熵损失函数：
其中表示教师模型对输入的输出对数似然值。最优的教师模型参数是：
。
一旦每个教师模型得到训练，软标签将在公共数据集上生成：
这个softmax函数在神经网络（NN）中使用，其结果通常是类的概率，通过使用softmax层。这一层的任务是将logits 转换为概率，作者称之为软标签。温度控制概率分布的软硬程度。越高，概率分布越软。对于学生的模型来说，设定为1。此外，对于IID和非IID数据分割方案，分别设置为10和20。这样做的原因是，处理非IID数据时，使用更高的温度可以使来自不同教师的噪声软标签变得平滑。平滑的分布有助于学生模型从教师模型的集合中学习更通用的表示。一旦所有教师模型生成的所有软标签，它们将在聚合器上进行聚合，通过平均所有教师的软标签：
其中是教师模型的总数量。
Student Model
学生模型 通过混合蒸馏损失（带有聚合软标签）和学生损失（带有公共数据集的真正标签）。学生模型的总训练损失为：
其中 是一个超参数，允许学生模型确定学生损失和蒸馏损失的比例。FedBrain-Distill 使用 ，以便其受到蒸馏损失的90%影响。作者需要注意到这个参数与Dirichlet分布 参数的区别。学生损失是公众数据集的真实标签的交叉熵损失：
其中 是公众数据集的样本数量， 是对应的真正标签。蒸馏损失是聚合软标签和学生模型在公众数据集上的预测之间的Kullback-Leibler（KL）散度：
学生模型旨在最小化总损失：
最终，通过最小化 ，学生模型有效地从聚合软标签提供的丰富信息中学习，捕捉到从多个教师模型中提取的知识。
Student Model Architecture
作者使用VGGNet16作为所有教师的模型。VGGNet16具有1.38亿参数[19]。表1显示了作者的学生模型的结构概述，从表中可以看出，该模型只有95,434个参数。与VGGNet16相比，作者的模型要简单得多。
IV Experimental Results
图4：在非IID数据上，使用2位和5位教师进行10轮通信后，学生模型在测试集上的准确率增加。
图5：在IID数据上，使用2位和5位教师进行10轮通信后，学生模型在测试集上的准确率增加。
图6：在非IID数据上，使用2位和5位教师进行10轮通信后，学生模型的总损失逐渐下降。
图7：在IID数据上，使用2位和5位教师进行10轮通信后，学生模型的总损失逐渐下降。
作者的实验均在Google Colab Pro环境中使用Python 3进行实现，使用了TensorFlow和Keras库。使用FigsShare的脑癌数据集来训练学生模型和教师模型。该数据集包含3064个癌肿瘤图像，分为三种不同类型的肿瘤：脑膜瘤、胶质瘤和垂体瘤。通过将参数设置为10000和0.5来分别实现IID和非IID数据。
接着，作者将教师的数据集划分为两部分，其中80%的原数据将用于训练集。对于IID和非IID数据，该比例保持一致。对于公共数据集，作者使用原始数据的50%创建了一个公共训练集。最后，对于测试集，作者使用原数据的一半创建了一个独立的测试集，以便在学生模型上评估性能。作者分别使用2位和5位教师在两种不同的设置下进行了10轮通信。此外，作者对FedBrain-Distill进行了与Viet等人[14]研究的比较，他们利用FL设置，在Figshare脑癌数据集上使用VGGNet16进行肿瘤分类。由于原始模型结构和训练数据集相同，这使得它能与基准方法进行公平的比较。
从图4和图5可以看出，学生在使用2位和5位教师进行10轮通信后，在非IID数据的测试集上的准确率逐渐提高。然而，在IID数据中，尤其是当教师数量为2时，收敛速度似乎快于教师数量为非IID数据的情况。图6和图7显示了学生在10轮通信后，总损失逐渐减少。和准确率一样，学生在IID数据上的总损失也呈平滑下降趋势。
表2展示了使用2位教师和FL设置的FedBrain-Distill在与Viet等人采取相同数据集、方法和训练次数的情况下，之间的比较。在非IID数据中，FedBrain-Distill能够超过Viet等人的方法。即使在经过100轮训练后，FedBrain-Distill也仍然落后于Viet等人，但是在通信上传和下载方面，FedBrain-Distill显著低于Viet等人的方法。
表3展示了使用5位教师和FL设置的FedBrain-Distill在与Viet等人采取相同数据集、方法和训练次数的情况下，之间的比较。可以看出，在非IID数据的前10轮中，FedBrain-Distill的性能更好。然而，在将近2%的速度上，FedBrain-Distill相对于Viet等人使用100轮训练的全局模型有所失误。与FedBrain-Distill相比，他们的方法在非IID数据上的通信成本和收敛时间的投入更大。Viet等人的工作已经消耗了403 Mb，仅用于上传模型参数，在FL设置中，当客户端数量增加时，这会成为一个挑战。此外，他们通过100轮进一步训练后，在非IID数据上的精度仅提高了2%。由于作者方法中的潜力，作者可以通过考虑数据增强和其他统计方法，以及教师软标签的聚合，进一步优化FedBrain-Distill以实现更准确的学生模型。
总的来说，作者认为通过考虑提高教师模型的性能、学生模型的不同架构、软最大函数中的温度改变等因素，可以实现FedBrain-Distill的改进，以在IID和非IID数据上达到更好的准确率。作者还发现，利用具有不同预训练模型的多个教师，如U-Net或ResNet，也可以提高FedBrain-Distill的性能。
V Conclusion
肿瘤分类一直是医疗机构关注的重心。其中最有效的方法之一是利用深度学习方法在MRI图像中进行肿瘤类型分类。然而，深度学习方法在跨不同医疗机构之间共享数据时，可能侵犯用户的隐私。为了更好地进行脑肿瘤分类，作者提出了FedBrain-Distill。
FedBrain-Distill不仅保护了用户的隐私，同时还解决了一些联邦学习的局限性。这些问题包括通信效率、依赖模型架构和收敛时间。与非IID老师进行比较时，FedBrain-Distill展示了令人鼓舞的结果。
此外，FedBrain-Distill仍可以通过实施不同的策略进行优化，其中一些策略包括考虑不同的学生模型架构、更改softmax函数中的温度T、更改总损失公式中的α参数，以及最后使用U-Net或ResNet等预训练模型训练教师。
扫码加入👉「集智书童」交流群
（备注： 方向+学校/公司+昵称 ）
点击下方“ 阅读原文 ”，
了解更多AI学习路上的 「武功秘籍」