林子谦,樊重俊,王 琪

(上海理工大学 管理学院,上海 200093)

1 引 言

深度学习已成功地应用于医学图像的病理检测和诊断中[1],它可以在没有医学专家输入的标准图像处理规则的情况下原始图像中学习并提供高精度的输出.但这需要大量不同的数据才能发挥作用,若数据集差异很小或者疾病数据集很少时,神经网络是没办法有良好的泛化能力的.此外,在训练医学影像模型中也可能依据机构偏见相关的混杂因素,而不是将模型建立在所评估的病理基础上,该模型在检验同一机构的保留数据时,虽然准确度较高,但对外部机构的了解较少,不能很好地扩展到外部机构.理论上,最有效的方法是将有关患者的所有数据集中起来共享到一个中心,并进行训练,能够产生良好的性能.但此方法存在许多挑战,首先,如果患者的数据占用了大量的存储,共享这些数据将是一种麻烦;其次,共享患者数据常常有法律或道德上的障碍,使研究者无法获得所有数据;第三,患者数据很有价值,机构可能不愿意分享数据.

因此联邦学习[2-4]孕育诞生,联邦学习主要是借助客户端-服务器的基本架构,实现多客户端在服务器的协调下,

对训练的模型参数进行聚合与传递,并训练出一个高性能的共享模型.这一概念的正式提出是由谷歌[5]HB Mcmahan等人,他们通过使用边缘服务器架构,将联邦学习用于智能手机上的语言预测模型的更新.随后,发现联邦学习有望打破在医疗领域中的隐私问题,许多学者也开始通过联邦学习算法设计了在脑肿瘤分离[6,7]、乳房密度分类[8]和功能性磁共振成像分析[9]等医学相关研究.

尽管联邦学习在医疗领域有一定的相关研究,但是在实际应用中经常会面临以下问题:1)由于经典医学图像分割模型结构复杂会产生大量的模型参数,若发送如此大量的模型参数,服务器将会产生巨大的通信开销,并且这样的通信开销会随着参与方数量和迭代轮次的增加而增加.此外,如果联邦学校使用同态加密作为安全方案,那么加密后的密文数据在传输效率过程中会产生更多的通信开销.2)现有医学图像分割的训练数据一般都是遵循独立同分布的(Identically Independently Distributions,IID),但是应用到联邦学习当中时大多数都是非独立同分布(non-IID)的,甚至有些医疗机构会传递恶意的模型,这必然会造成全局模型的梯度爆炸与准确度下降的问题.

解决通信开销方面,一般分为通信中的总时耗与通信数据总量这两个问题,主要的解决方法有降低模型的更新频率、模型的压缩、客户端的选择.在降低模型的更新频率方法中,Zhou等[9]提出重叠联邦平均算法,通过集合分层计算策略、数据补偿机制与NAG的算法,实现通信成本的降低;在模型的压缩方法中,Yang等[10]以非线性的方式划分空间,提出了一种基于cosine函数的非线性量化方案,不需要误差反馈等额外梯度恢复信息来调整梯度;在客户端的选择方法中,Nishio T[11]提出了一种联邦客户选择算法,该算法根据累计有效参与值,选择模型迭代效率最高的客户端进行聚合更新,以此优化整个联邦学习算法的收敛效率,进而降低通信代价.

解决数据的非独立同分布主要有3种情况:特征分布偏差、标签分布偏差、数据量的偏差.而在图像分割中经常产生的偏差莫过于不同机构的标记不同,导致的特征分布的偏差,与不同数量的患者数据导致的偏差.因此针对非独立同分布数据的联邦学习问题也引起了研究者的关注,其中一个方向是优化模型聚合过程,Wang等人[12]提出了一种归一化的平均方法办法,该办法消除了目标的不一致性,同时保留了异质性的快收敛.另一个方向是通过知识蒸馏的方法,FedMA[13]完善了通过匹配和平均具有类似特征的隐藏元素来完善聚合过程具有类似特征签名的隐藏元素,将知识提炼整合到联邦学习中的想法也已经得到了验证[14].

本文设计一种针对图像分割领域的联邦蒸馏优化算法(Unet-Federated Distillation Optimize,U-FedDO)算法.具体而言,首先在联邦平均算法(Federated Averaging,FedAvg)的基础上使用教师-学生模型,其中在参与方中有一个经典的Unet模型作为教师模型,还有一个轻量的U-net模型结构作为学生模型,通过教师-学生的框架保证本地参与方性能的稳定;然后本文在服务器上还部署调节器,并设定设置相关的约束条件与权重设置机制.此外,本文还应用在真实数据上设计了IID数据与non-IID数据下的实验,结果充分地表明了该模型在医学图像分割中的有效性与可行性.

本文的组织结构如下,第2节讨论了相关工作,第3节描述了U-FedDO算法的具体细节,第4节给出了实验结果并讨论,在第5节对本文进行总结与讨论.

2 相关工作

U-FedDO是针对图像分割所设计的联邦集成模型,其中包括了医学图像分割模型,联邦优化算法模型,联邦蒸馏算法.下面对3个模型方法简单介绍:

2.1 医学图像分割模型

医学图像分割是一种能够对人体某个组织器官的边界、形状进行测量和标注的技术.传统的图像分割方法通常是基于区域分割法和基于边界分割方法.2015年,全卷积神经网络(Fully Convolutional Networks,FCNs)以一种端到端的形式进行语义分割[15].随后出的U-Net具有灵活、简单并且在样本数据集较少的情况下能获得很好的分割效果,目前已成为医疗图像分割的一类主流方法[16].

U-net结构主要由两部分组成:下采样与上采样.下采样是当图像输入Unet模型后采用一定的卷积用来图像深度的提升与图像特征的提取,接着使用池化操作用于减小图像的面积.在对图像进行4层卷积池化后,进行图像全连接.上采样是因为图像的分割结果需要与输入图像相匹配,需要对最小图像面积进行扩充,将图像通过反卷积操作,并将上方相同层次的图像合并,保留图像原始特征.最后,使用一定的卷积核进行图像卷积,提取图像特征.此外,为了防止传统神经网络在信息传递过程中出现的梯度消失,在Unet结构里的编码器与解码器之间加入了跳跃连接操作.后来在医学图像分割领域中,出现许多改进Unet提升性能与应用场景的图像分割模型,主要包括编、解码器结果的改进、损失函数的改进、瓶颈模块的改进、增强数据流路径、自动搜索结果等.例如在解码器中引入注意力机制的Attention Unet模型,抑制一些不相关分割区域所产生的影响,从而提升性能,或者是net++模型在跳跃连接的方式中加入不同深浅特征的融合,还有像Unet-GAN模型是从内部结构中提出的适应不同提供商图像的适配器.然而,目前基于Unet提升性能的方法主要集中于对Unet结构的本身,却很少有学者通过解决医学图像数据的孤岛的问题,从而提升分割模型的性能.

2.2 联邦优化算法

U-FedDO算法主要是基于在FedAvg算法上进行改进.传统的FedAvg是指参与方在本地数据集上计算梯度后,执行本地训练后并上传模型参数,服务器对各个参与方的模型进行加权平均,从而实现一次系统层面的迭代.但是在FedAvg中,只考虑各数据源的权重与数据源的数据量大小,并未考虑数据源之间数据质量的差异,若出现数据源的数据量很大,但其数据质量很差,这就会造成权重的不合理分配.因此Huang 等[17]在FedAVG基础上提出自适应数据增强的LoAdaBoost FedAvg算法,利用了数据共享策略与全局损失函数的中位数,但是该方法由于损失函数的限制,很难达到较高的性能.此外,Lin使用本地模型的平均对数进行汇总并且聚合[18],来防止聚合过程中梯度遗忘的问题.这些方法虽然有效地解决了一定的异质性给全局模型带来的影响,但是这些模型的前提是假设所有的参与方是诚实,如果存在恶意的参与方或者极端数据时,很容易引起全局模型过程中梯度爆炸的问题,并且使得全局模型低于原本参与方模型的性能.

2.3 联邦蒸馏算法

本文中,解决通信开销问题的主要工具是知识蒸馏模型.它是Hinton提出的一种师生结构,用于在深度学习间传递知识.尽管框架最初只是为了在同一数据集上的边缘转移,但是之后,通过在学生网络的训练过程中,引入了一种正则化,挖掘两个模型传递过程中的潜力.它主要思想是将大型神经网络学习到的知识提炼并传授给小模型中,从而实现从大型网络压缩成小型网络的目的.Jeong等学者[19]就提出了联邦蒸馏(Federated Disillation,FedDisill)相对较实用且效果较好的一种方法,它能够大大减少通信代价的同时适用于缺少标签的异质数据、异构模型的场景,基本概念就是只交换本地局部模型的输出而非交换传统联邦学习采用的全局模型参数,这使得输出的模型参数尺寸比传统的小得多,从而可以减少通信开销.随后,就有Gihun Lee[20]学者认为全局模型的收敛性受到数据异质性的影响,全局模型会对有偏差的客户端的产生梯度偏移,于是提出了联邦本地自我蒸馏(Federated Local Self-Distillation,FedLSD)模型用来缓解梯度偏移的问题,并提高其性能.但是联邦蒸馏的不足之处在于只能缓解异质性的数据产生梯度遗忘的问题,并不能解决出现恶意或极端模型下的梯度爆炸的问题,并且当两个网络模型大小相差太大时,知识蒸馏的效果会失效,因此联邦蒸馏的收敛性和应用性研究需要进一步研究.

3 U-FedDO方法

基于FedAvg算法在non-IID的数据集导致性能下降的问题,本文提出了一种在服务器筛选的机制,并应用在教师学生模型中.本节分别讨论性能下降对应的数学模型与调节器上的约束条件,并且介绍调节器的筛选方式、参与方中Unet教师学生结构的变换、U-FedDO算法的工作流程.

3.1 数学模型

FedAvg是由随机梯度下降算法(stochastic gradient descent,SGD)推演而来的,假设有K客户参加联邦学习,则每个参与方在本地数据上使用SGD计算梯度,将模型参数上传给服务器,然后服务器将收到模型参数加权平均并分配给参与方,经过一定次数的迭代后停止的.各参与方传递的梯度下降算法的优化目标为:

(1)

(2)

本文中,联邦蒸馏的核心思想是客户端在本地的自我蒸馏,若全局模型聚合过程中对有偏见的局部分布进行拟合,对应的权重将发生偏移,这很容易导致性能的降低.因此,可以通过蒸馏的方法保留关于本地数据的分布,当全局模型聚合导致偏移时,本地的教师模型可以纠正全局模型的分布.以下损失函数LFed DO是交叉熵LCE和蒸馏损失LLSD的线性组合:

LFed DO=(1-β)LCE+βLLSD

(3)

其中,q(x,w)是每个本地客户模型和全局模型在最后一轮的softmax概率与单次真实标签py(x).LLSD是教师模型预测与学生模型预测之间的KL-Divergence蒸馏损失,具体公式为(4):

(4)

其中c∈[C]类的预测概率,C是分类的数量.蒸馏的损失是将对数概率softmax函数z除以τ来软化其预测概率.具体如公式:

(5)

联邦学习服务器的权重根据FedAvg算法计算.因此结合公式(2)、公式(3)可以得出,联邦蒸馏模型的更新参数公式(6):

(6)

(7)

(8)

D[f(x,wt+1),y]>D[f(x;w(k)),y]-d

(9)

其中约束条件(8)保证每次筛选中聚合的客户端参数参与聚合并且所有的权重和为1,约束条件(9)中的D为评价指标,d为评价指标的训练过程中的相对误差值.

3.2 算法模型

本次设计的U-FedDO算法主要有服务器、调节器、参与方的学生模型与参与方的教师模型.服务器负责全局权重的聚合与输出给本地模型,调节器是本地传递的权重进行筛选,参与方中有教师模型保留本地模型的权重,学生模型主要负责训练服务器传递的权重与结合本地教师模型的损失.

客户端中需要将Unet的模型压缩并参与全局训练,但传统的蒸馏法并不足以将Unet模型中的U型结构蒸馏成一个非常小的U型网.因此在学生模型的设计过程中,先是遵循了经典Unet的编码器-解码器的结构、编码器为4个下采样包括3×3的卷积和2×2的池化层、单通道图像作为输入,但不同的是在将传统Unet的卷积层数从64改成16,同时保持收缩路径翻倍的趋势,具体如表1 所示.上采样为解码器也是通过4个2×2的池化层,同时保持传统Unet在每个池化层后面采用Relu激活函数的结构,来防止梯度消失的问题,与降低计算量的问题.最后一层采用softmax激活函数对像素进行分类.在医学图像分割任务中,输入图像与输出掩码在外观表现上是有差异的,但在底层结构和特征上是一致的.在连续的卷积过程中会导致某些潜在特征被忽略,因此在特征映射掩码过程中,在下采样和上采样之间加入跳跃连接可以保持高级特征并映射到分割掩码.在医学图像分割中,通常只有病灶与背景两种分类,损失函数则为二元交叉熵(Binary Cross Entropy,BCE),其公式为(10):

表1 Unet卷积层数对比表Table 1 Unet convolutional layer comparison table

(10)

其中,N为输入图像像素总数,gi是第i个像素真实的标签,0表示背景,1为肺部.pi是对应像素预测为前景的概率.当前景像素远远小于背景像素数量时,即gi=0的数量远大于gi=1的数量,在优化过程中gi=0的部分便会占据主导,使得模型严重偏向背景.因此,在客户端中主要有两个模型进行训练,第一是传统的64-Unet模型作为教师模型在本地训练,第二是改良后的学生16-Unet模型参与全局训练,通过公式(4)的蒸馏损失函数,学生模型能够继承教师模型训练概率分布,保证性能不会因卷积层的降低而下降.

服务器主要功能是将被选择的参与方上传的本地模型进行聚合,并将重复发放给各个参与方.服务器中就需要三个阶段分别为初始化、聚合、评估.在初始化阶段,服务器首先将医疗机构加入联邦学习模型训练任务中,同时服务器需要分配信息参数发送给参与方如学习率、迭代次数、优化器等.在聚合阶段,服务器需要接受客户端上传的模型参数,并使用将收到的模型参数进行聚合,聚合公式为(11):

(11)

调节器主要功能是筛选与重新平衡.首先,通过服务器设置的约束机制判断对各个参与方传递的模型参数,若为有利方则参加聚合,若为有害方则剔除.同时调节器还根据参与方的数量重新分配模型的权重,若参与方多次被筛选出去时,调节器将不再参加后续模型的迭代,从而减少通信开销.此外,在公式(9)的筛选给定相对误差值d,缓解模型聚合偏移带来的局部收敛的情况.

3.3 工作流程

如图1,显示U-FedDO框架的训练、传递、蒸馏、筛选、聚合的过程.首先,服务器进行初始化设定模型参数,学习率迭代次数等,接着客户端训练一个教师64-Unet模型,同时接受服务器传来的初始化全局16-Unet模型的参数,开始进行训练.然后,客户端将训练好的参数发送给服务端,服务端中通过调解器去筛选能够提高准确度的参与方,并更新每个参与方在服务器聚合的权重.将聚合后的模型参数再依次分配给参与方进行下轮训练,直到模型收敛,保存模型并结束本次的联邦学习.此外,调解器是虚拟组件,它可以直接部署在联邦服务器上以减少通信开销.具体如算法1所示.

图1 U-FedDO结构图Fig.1 U-FedDO structure diagram

算法1.基于Unet的联邦蒸馏学习

输入:全局权重,本地迭代次数,本地数据,学习率

1.在服务器执行

2.初始化原模型参数W0,将原始的模型参数W0,广播给所有的参与方

3.For 每一全局模型更新轮次t=1,2,…,do

4. 调节器通过筛选确定C个参与方的集合

5. For 所有参与方C并行地 do

8. End for

10.If 模型收敛

12.End if

13.在调节器执行执行

14.从参与方获得最新的本地模型参数与服务器上的测试集

15.For 所有参与方C

16. ifD[f(x;wt+1),y]>D[f(x;w(k)),y]-d

17.将可以用模型参数存入参数Ct

18. End if

19.确定参与存入的参数与数量Ct,并发生给服务器

20.End for

21.在参与方更新

24.for 从1到迭代次数S的每一个本地迭代I do

25. for 从1到批量数量B=的批量序号b do

28. End for

29.End for

4 实验及结果分析

4.1 实验数据与配置参数

本次分别使用COVID-19[21]扫描影像与HAM10000[22]数据集.其中,COVID-19扫描影像中有20张已标记的图像,共切片为3530张图片,左右肺与感染由两名放射科医生标记,并由经验丰富的放射科医生进行验证;HAM10000数据集包含10000张七种互斥类的皮损区域皮肤镜图像,且至少在3位皮肤科专家共识的基础上对其标记,并通过金标准组织病理学的证实.

本文采用留出法策略,即分别在K-1源域训练,并在一个未见过的目标域上进行测试.因此,将COVID-19扫描影像与HAM10000数据划分为本地训练集,本地验证集、本地测试集、全局验证集、全局测试集.在IID数据集的实验设计中,将两个数据集分别以10%的全局验证、10%的全局测试,并将剩下的80%的数据随机均分至3个医疗机构,并且机构也将分配的数据以8∶1∶1的比例分为本地训练、本地验证、本地测试.由于Unet的图像分割模型本身可适用于小数据集,且U-FedDO就是解决医院数据不足与数据不平衡的问题,则在non-IID数据的实验设计中,将两个数据集中随机抽取以4∶2∶2∶1的比例的小数据存放在4个不同的机构,其余数据与IID数据集一致.此外,所有的训练集与测试集之间是互斥的关系.

实验硬件配置为:CPU8核,内存28G,显卡 RTX 3080,软件系统包括Ubuntu21.0,CUDA8.0,本次以单机的模式进行评估.在联邦学习过程中,所有客户使用相同的超参数设置,本地模型使用GDB优化器进行训练,批量大小为1,动量为0.9.步骤大小和学习速率都被设置为0.01,图片压缩为0.5尺寸,知识蒸馏温度参数τ根据经验设定为10,共训练了50个全局轮次,每个全局轮次中的局部历时迭代设置为1.

4.2 评价指数

为了对分割性能进行评测,本文采用的评价指标包括Dice相似系数来衡量分割模型的性能,这是一个用于医学图像分割的标准评价指标.具体的计算公式如下:

(12)

其中Sseg为预测值,Ssef为真实值.

4.3 消融研究

“消融研究”在机器学习中经常被用一些来描述去除网络的某些部分的过程,以便更好地理解网络的行为.U-FedDO模型中的3个关键问题:1)教师-学生的16Unet模型对性能与通信开销的影响情况;2)U-FedDO在数据平衡与非平衡下的性能情况;3)参与的医疗机构数量增加情况下,U-FedDO性能变化.

1)在模型参数方面,如表2所示传统的深度神经网络模型通常有数百万,数千万甚至上亿的参数量,因此发送如此多的浮点数值将给协调方产生巨大的通信开销,而本文设计的轻量16-Unet模型明显不管是模型大小还是参数数量都远小于其他模型,同时也比传统的Unet模型减少了17倍的参数.训练时间方面,在512×512像素的数量为1000的切片图片中,传统的64-Unet模型单次训练时间为84s,而16-Unet模型的单次训练时间为63s,训练时间降低约为25%.

表2 经典CNN模型与Unet的模型大小与参数数量Table 2 Model size and number of parameters of classical CNN model and Unet

在性能方面,如图2所示,在64-Unet为经典的Unet模型,而16-Unet为经过知识蒸馏后的模型,可以判断的出模型在局部验证普遍比在全局验证中的准确率更高,而轻量的16-Unet在教师-学生模型下,性能不会大量的下降.

图2 16-Unet与64-Unet在测试集中的性能对比图Fig.2 Performance of 16-Unet vs.64-Unet in the test set

2)U-FedDO的分割效果主要在两个方面分别是IID数据与non-IID数据.一是在IID数据集下,如图3所示,两个数据集在FedAvg、U-FedDO、FedDistill中的差距不大,但是由于FedDistill中的本地教师模型会保留本地的性能较差的模型参数与轻量Unet的模型的影响下,性能会有所下降,但是U-FedDO在有效地筛选情况下,性能与FedAvg几乎一致.同时对比图2中单个机构的性能可以得到以下结论,当用单源数据训练的模型在部署到未知的领域时不能获得良好的结果,但是在更多的客户参与的联邦训练时,泛化性能会增加,这是由于从多个来源聚集的数据可以覆盖一个更全面的数据分布.

图3 3种方法在IID数据下的性能对比图Fig.3 Performance of 16-Unet vs.64-Unet in the test set

二是在non-IID数据下,如图4所示,4个医疗机构的准确率也是各不相同,本文设计的U-FedDO算法在两个数据集中的non-IID数据性能是最高的,最高机构的性能还高约10%,且远比FedAvg的准确度高,这取决于调节器在聚合过程中,有效地筛选出其他恶意的参与方模型,保证在服务器在聚合过程始终对全局模型的优化的效果.

图4 4个医疗机构与3个联邦模型的分割箱型图Fig.4 Split box plot of four medical institutions with three federal models

为了展现显示效果,图5显示了COVID-19扫描影像与HAM1000的部分测试集的分割效果,可以看到在本次实验的数据集中,U-FedDO分割的效果在整体的形状和位置上都已经十分接近手工标注,对于局部的图像信息也展现得较为清晰与详细,线条相对逼真与准确,但是在边界上手工标注的会比U-FedDO更圆滑一点.而相比与FedAvg与FedDistill两个算法由于极端数据的影响,在COVID-19扫描影像虽然整体的肺部结构大致能分割出来,但是在肺部外无相关的也进行了误判分割,显得分割效果有杂乱且没有目标;在HAM1000的数据集中,FedAvg与FedDistill分割效果也低于U-FedDO,特别是当图像中还有毛发干扰时,分割效果的差距尤为明显.

图5 3种方法不在平衡数据下的分割对比图Fig.5 Comparison of the segmentation of the three methods not under the balanced data

3)当不同医院数量且医院间的数据分布不平衡的情况下,U-FedDO和FedAvg的性能会有怎么样的变化.图6展示了图像分割的结果,参与联邦训练的医疗机构依次从1逐渐增加到16,并且图片数量不同且每个参与方的准确率也各不相同.U-FedDO模型在1至16中的准确率始终高于传统的FedAvg,并且随着参与方数量的变化趋于平稳上升.FedAvg由于存在大量的数据不平衡的参与方所导致的准确率不高的情况,特别是当数量到9与10时,参与方存在极度不平衡状态且准确率不高的情况,由此可以判断使用FedAvg会明显产生梯度爆炸的问题.本次实验再次证明了U-FedDO在数据不平衡情况下,能够有效地防止梯度爆炸的情况,并且泛化能力方面具有稳定的功效.

图6 医疗机构不同数量性能对比图Fig.6 Comparison of different numbers of medical institutions

5 结束语

本文中,基于在医疗图像分割在联邦模型中两大问题:通信开销和数据不平衡,设计一个U-FedDO算法,主要是应用了模型的知识蒸馏与服务器上部署调解器的两个方法,使用知识蒸馏压缩方法在经典图像分割模型Unet的结构特性减少其卷积层数,设计一个轻量的Unet模型.在调解器方面,通过约束条件与重新设计聚合权重的方法保持全局模型不会因为极端数据而梯度爆炸.实验结果表明,在使用改进后的Unet蒸馏方法,模型以极小的性能下降的代价,有效将单次训练的模型参数降低了16倍并且训练时间减少约25%,大量的降低了联邦训练过程的总时耗与通信数据总量.在IID数据集的实验中,FedAvg、U-FedDO、FedDistill性能并没有很大差异.在non-IID的数据中,U-FedDO模型明显比FedAvg的准确更高,且随着参与方的数量增加,U-FedDO也比FedAvg得性能表现的更稳定.

研究发现,本文只针对横向联邦学习的不均衡数据做探讨,在下一阶段的工作中,将深入研究纵向联邦学习在图像分割的应用.