面向部分设备参与的非独立同分布数据的联邦学习算法

邹敏浩,甘中学

(复旦大学 工程与应用技术研究院,上海 200433)

1 引 言

联邦学习由McMahan 等人[1]首先提出,它是一种前景广阔的分布式学习方式,利用边缘设备中的大量分散数据进行分布式训练以获得集中式的机器学习模型,并且不用将数据汇集在一起.通过这种方式,可以解决由边缘设备(如机器人、智能手机和可穿戴设备)生成的大量数据但由于日益增加的隐私保护问题而无法利用的困境[2].在联邦学习过程中,训练开始后服务端首先初始化模型参数,并将初始化的模型参数通过网络分发到客户端(边缘设备)中,客户端的模型加载由服务端发来的参数,并根据本地数据使用随机梯度下降等优化方法对模型进行迭代训练,训练完成之后,将本地的模型参数通过网络传送回服务端,服务端根据收到的模型参数对服务端的全局模型进行更新.等待服务端收到一定的数量的或所有客户端的模型参数或者超过设定好的时间阈值,服务端更新全局模型,并将新更新的全局模型参数再次分发到客户端进行训练,循环往复.

虽然原始数据不再在设备之间传输,但在联邦学习训练期间,模型参数的频繁传输也对网络负载提出了比较高的要求.同时,设备的不同、数据质量的不同和数据分布的差异(通常被概括为客户端异质性)也给联邦学习带来了新的挑战.

本文主要关注以下两个联邦学习的挑战.

1)客户端异质性.数据的异质性和不平衡性是联邦学习中的常见情况[3].客户端异质性包括设备、数据质量、模型和数据分布等方面的差异[4].此外,非独立同分布的数据也有不同的情况,例如标签分布偏差、特征分布偏差和数量偏差等等.这些异质性存在于现实情况中,例如,各个医院拥有患者的不同疾病数据[5]和银行拥有不同客户的信息和数据[6,7]等等,这些异质性给联邦学习算法的收敛性,稳定性和效果等带来了挑战.本文主要关注标签数量不平衡的异质性,其中每个客户端拥有来自所有标签中的固定数量标签的数据样本,这种异质性对算法的收敛性有很大影响.FedAvg[1]首先对客户端异质性进行了研究,在FedAvg中,具有相同标签的数据样本被划分为子集,每个客户端只被分配两个具有不同标签的子集.

2)部分参与.大多数同步联邦学习算法都假设客户端在每一轮培训中都完全参与全局模型更新.然而,由于计算速度或网络速度的差异,客户端的梯度信息总是滞后到达.在实践中,为了提高联邦学习的训练效率,通常只选择部分客户端进行梯度更新,其他掉队客户端由于网络速度慢、客户端之间的本地训练更新速度差或其他原因而不参与本轮训练,并对更新全局模型没有任何贡献.低设备参与率在同质数据分布的联邦学习优化过程的情况下影响的效果有限,但对使用FedAvg进行异质数据分布的训练有很大影响.图1展示了在FedAvg中非独立同分布数据设置下部分设备参与导致的偏差问题的示例.在非独立同分布设置下,朝向全局最优解的方向不同于局部最优解的方向.在全部设备参与下,当前的全局最优解是真正的全局最优解.然而,在部分设备参与设置下,由于局部最优解远离全局最优解,在通信回合中当前的全局最优解可能远离真正的全局最优解,这个偏差会使FedAvg的效果在训练过程中不稳定,严重的可能会影响FedAvg算法的收敛性.因此在非独立同分布数据环境和部分设备参与的情况下,设计一种有效和健壮的联邦学习算法是一个挑战.

图1 非独立同分布设置下部分设备参与导致的偏差示例

在上述分析的基础上,本文提出了一种新颖、高效、健壮的异步联邦优化算法,其中服务端在分发全局模型参数之后继续聚合滞后的客户端的梯度信息,以在下一轮全局模型聚合时修正全局最优值(如图1所示)而不用增加客户端存储或设备之间的通信成本,并可以通过修改联邦学习中的聚合方式从而与现有联邦学习算法的相互兼容.

本文的贡献总结如下:

1)从优化的角度详细的分析了客户端异质性和部分客户端设备参与的情况下对FedAvg联邦学习算法的收敛性和稳定性的影响.

2)提出了一种新的、高效的、异步的、健壮的联邦优化算法--联邦累积学习算法(FedAcc),其中服务端通过聚合延迟的客户端的梯度信息以指导下一轮的客户端梯度更新.

3)基于MNIST、FMNIST和CIFAR-10数据集在不同客户端的参与率的条件上进行了大量的实验.实验结果证明本文的方法能有效的缓解客户端异质性和部分设备参与的问题,并且优于其他对比的方法.

2 相关工作

传统的联邦学习需要频繁的参数交换,对网络吞吐量有严格的要求,并且容易受到不可靠网络和不稳定带宽的影响[8].为了处理联邦学习中的通信瓶颈,一种简单有效的方法是对客户端进行采样以加速收敛.许多方法侧重于如何有效地进行采样,Nguyen等[8]人提出一种快速收敛的联邦学习算法FOLB,它在每一轮模型训练中智能地对客户端进行采样,以优化收敛速度,并且能够通过根据客户端对更新的贡献能力的估计调整聚合来处理客户端的通信和计算的异质性,Xiao等[9]人观察到传统联邦学习的服务器端聚合方法是不合理的,客户端在联邦学习中的贡献可以通过其训练模型的验证精度来区分,基于这一观察,他们提出了一种新的联邦学习算法,基于精度的平均联邦学习算法(ABAVG),该算法改进了传统联邦学习的服务器端聚合方法,并加快了非独立同分布情况下联邦学习的收敛速度.Chen等[10]人提出了一个轻量级的节点选择算法来有效地执行联邦学习任务,该算法在考虑本地计算资源和通信条件的同时,迭代选择异质节点参与全局学习聚合.Wu 和Wang[11]提出了一种优化聚合算法,该算法通过检查局部梯度和全局梯度之间的关系来识别和排除不利的局部更新,从而找出每个全局轮中参与节点的局部更新的最优子集.在精心设计的实验条件及已知先验的知识下,这些基于重要性的采样方式提高了性能,但不一定适用于具有不可预测的实际数据分布环境.

此外一些方法通过异步方式来缓解这个问题.异步方法旨在利用更多的客户端梯度信息,Xie等[12]人提出的FedAsync,一种异步联邦学习算法,通过在新的局部更新和旧的全局模型之间加权平均来更新全局模型,以处理滞后的客户端,从而部署一个高效的联邦学习系统.在FedBuff[13]中,服务端在专用缓冲区中聚合客户端更新,并在异步接收到K个更新时执行服务器端模型更新.Liu等[14]人提出了一种等待策略,动态和自适应地确定每轮的聚合数量,以实现单轮训练时间和期望轮数之间的折衷,从而达到目标精度.Chen等[15]人提出了一种合理限制训练速度不同的客户端传回的参数在全局模型聚合参数中所占比例,并且主动更新延迟客户端使用的聚合参数的异步联邦学习聚合更新方法,从而有效解决客户端训练速度差异对聚合参数造成的负面影响.为了解决滞后设备严重影响FedAvg等常用联邦学习算法收敛的问题,Gu等[16]人研究任意设备不可用情况下的联邦学习算法,并提出一种称为记忆增强联邦平均(MIFA)的算法,它有效地避免了非活跃设备导致的过度延迟,并使用存储在设备的梯度更新纠正梯度偏差.但是,上面这些基于聚合的方法大多是隐式的,没有可解释性.本文的方法虽然也是异步并且具有解释性.

当客户端中的数据是异质的情况,会严重降低联邦学习的性能,收敛速度不稳定且缓慢[17,18].FedGBO由Mills等[19]人提出,它通过在联邦学习的本地训练阶段应用一组全局偏向优化器值来加速联邦学习,这有助于减少客户端对非独立同分布数据的漂移.Karimireddy等[18]人提出SCAFFOLD,它使用控制变量(方差减少)纠正其本地更新中的“客户端漂移”,需要的通信轮数明显减少,并且不受数据异质性或客户端采样的影响,Karimireddy等[20]人提出MIME,它在每个客户端更新步骤中使用控制变量和服务端级的优化器状态(例如动量)的组合,以确保每个本地更新模仿集中式方法在非独立同分布数据上的更新过程,以减轻客户端漂移.然而,这些方法通过向客户端发送额外参数来增加了通信成本.

尽管取得了这些成功,但联邦学习算法仍然面临着部分参与非独立同分布数据所带来的异质性的挑战.与这些方法相比,本文的方法更具鲁棒性和灵活性,并且没有额外的通信成本和存储空间,并且具有高度的可解释性.

3 提出的方法

本节将描述为什么以及如何设计所提出的FedAcc算法,并解释FedAcc如何提高客户端部分参与的收敛性和缓解客户端的带来的异质性.然后,对FedAcc在非独立同分布情况下设置的联邦学习中的效率和性能进行了全面的研究.

3.1 FedAvg

联邦学习中用来解决以下形式的优化问题:

(1)

其中Fi(x)=Ez～Di[f(xi,z)]是第i个客户端的损失函数,z∈Z和Di是第i个客户端的数据分布.在数据异质的情况下,即对于i!=j,Di和Dj可能非常不同,此时函数Fi(x)和f(xi,z)可能是非凸的,这会给联邦学习算法带来挑战.

(2)

(3)

(4)

通过这样的方式就将客户端和服务端的模型参数更新统一到一个双优化器更新的客户端/服务端优化器联邦学习框架中.基于客户端/服务端优化器框架FedAvg算法的其他细节见算法1中的伪代码.

算法1.基于客户端/服务端优化器框架FedAvg算法

输入：本地数据集Di、总客户端数N、通信轮数T、抽样客户端数M、本地训练轮数E、学习率η

输出：最终的模型wT

1.服务端执行:

2.初始化w0

3. fort←0 toT-1 do:

4.St←从N个客户端随机采样M个客户端

5.ns←∑i∈St|Di|

6. fori∈Stin parallel do:

13. 客户端执行:

14. ClientUpdate(i,wt):

15.L(w;b)=∑(x,y)∈bl(w;x;y)

17. for epoch←0 toE-1 do:

18. for each batchb=(x,y) ofDido:

3.2 FedAcc

为了提高联邦学习中的训练效率,通常只选择部分客户端进行梯度更新,而其他由于网络速度慢、客户端之间的本地训练更新速度差异或其他原因而不参与本轮训练的掉队客户端对全局模型的更新没有任何贡献.随着时间的推移,这些掉队者的梯度信息将到达服务器.然而,服务器端已经将聚合模型发送回客户端,因此无法很好地利用到达的梯度信息.因此,一个自然的想法是,到达的梯度信息是否可以在下一轮全局模型更新中重新被使用,因为到达的梯度信息包含了更多的客户端梯度信息,并且有助于纠正由于部分采样而导致的与真实全局最优值的偏差,从而提高训练过程中联邦学习的稳定性和收敛性,同时因为利用了所有客户端的梯度信息,由于部分设备参与导致的异质性将会得到一定的缓解.

(5)

该梯度在t轮中表示真实全局模型和当前全局模型之间的偏差,并当作动量项用来指导下一轮通信中的服务端的梯度更新(见公式(6),算法2第16行).其中α是本轮聚合的采样率同时也是本轮设备的参与率,

(6)

通过这种方式,用一种异步的方式有效的利用了所有客户端的梯度信息来更新全局模型,而不用等到所有客户端模型的参数到达服务端之后再对全局模型进行更新,并且客户端和服务端之间的通信成本没有增加,这意味着该联邦学习算法通信效率是高效的,同时这种异步的方式也提高了联邦学习的训练效率.此外,本文在算法设计和实验过程发现如果交换算法2的第17行和第18行,Δt表示完全参与和部分参与之间的偏差,并且如果Δws依旧根据如算法2第16行计算所得,则该方式的联邦学习算法不会收敛.

算法2.联邦累积学习算法(FedAcc)

输入：本地数据集Di,总客户端数N,通信轮数T,抽样客户端数M,本地训练轮数E,学习率η

输出：最终的模型wT

1.服务端执行:

2.初始化w0

4. fort←0 toT-1 do:

5.St←从N个客户端随机采样M个客户端

6.Ut←其他未被采样的N-M个客户端

7.ns←∑i∈St|Di|,nu←∑i∈Ut|Di|

8. fori∈St∪Utin parallel do:

13. ift=0 then:

14. Δt←0

15.α←1

3.3 FedAdamAcc

虽然FedAcc已经提出,自然而然会考虑到,它是否可以用于其他自适应方法.因此,本文在FedAdam的基础上提出了FedAdamAcc,并在MNIST和FMNIST数据集上对其进行了训练和测试.Adam[21]是一种有效且应用广泛的自适应优化算法.FedAdam[22]是一种自适应联邦学习算法,它将客户端的局部更新和全局模型之间的差异视为伪梯度,并应用自适应梯度下降算法更新全局模型,以提高性能,已被证明是收敛的.

本文将这种累加梯度信息的方式与FedAdam自适应算法相结合,见算法3,其中第19行和第20行的mt和vt分别是梯度和平方梯度的指数移动平均值,并且用参数β1和β2控制这些移动平均的衰减率,β1和β2在这里分别取0.9和0.999,第21行和第22行是用作偏差修正,伪梯度更新由第23行得到.

算法3.FedAdamAcc

输入：本地数据集Di,总客户端数N,通信轮数T,抽样客户端数M,本地训练轮数E,学习率η,全局学习率ηg,衰减参数β1,β2∈[0,1),ε=10-6

输出：最终的模型wT

1.服务端执行:

2.初始化w0

4.fort←0 toT-1 do:

5.St←从N个客户端随机采样M个客户端

6.Ut←其他未被采样的N-M个客户端

7.ns←∑i∈St|Di|,nu←∑i∈Ut|Di|

8. fori∈St∪Utin parallel do:

13. ift=0 then:

14. Δt←0

15.α←1

16.mt-1←0

17.vt-1←0

19.mt←β1·mt-1+(1-β1)·Δws

4 实验设置

本文根据异质联邦学习中对数据集进行划分和任务中最具代表性的情况来评价和衡量本文提出的算法和相关的算法,且希望研究设备的参与率如何影响联邦学习的收敛情况,特别是在非独立同分布数据设置下.为了实现这一目标,本文基于以下设置上进行了大量的实验.

4.1 数据集和模型:

本文在几个公共数据集,MNIST、FMINST和CIFAR-10与其他联邦学习算法进行比较.按照FedAvg的实验设置,将数据集按标签划分为不同的数据子集,每个客户端仅分配两个不同的标签的数据集,标签的数据样本数不固定.

实验发现每个客户端中的数据数量对结果有很大影响,因此本文在CIFAR-10中固定了每个客户端中的标签和数据数量,以减少随机因素的影响.数据集的统计情况汇总在表1中.对于模型,本文使用一个简单的卷积神经网络模型,它有两个5×5卷积层,步幅为1填充为0,然后是2×2池化层(第1个有6个通道,第2个有16个通道)和两个全连接层,并使用ReLU激活函数(第1个有120个神经元,第2个有84个神经元).

表1 数据集的统计情况

表2 用于CIFAR-10数据上训练的简单卷积神经网络模型

4.2 实验环境

本文在基于Pytorch的客户端/服务端优化器的联邦学习框架[4]上实现所有算法.以一定的采样率对客户端进行随机抽样.客户端和服务端都使用回合数来评估训练进度,而不是训练过程中的迭代次数.该框架根据每个客户端的训练样本数对所有客户端的样本数的比值进行加权,而不是根据客户端的数量进行统一加权.由于缺乏计算资源,实验中服务端的聚合和客户端的训练都在一台有1个NVIDIA 1080Ti GPU的机器上完成.

4.3 优化器和超参数

本地训练中使用学习率为0.01、动量为0.9的SGD优化器.默认情况下,批大小设置为64,本地训练回合数设置为5.默认情况下,所有的设备数量总数设置为10.每一轮所有参与的客户端以一定的抽样率(0.1,0.3,0.5,0.7,1.0)被抽样随机参加服务端聚合.为了进行公平比较,研究算法中的通信轮数设置相同,MNIST/FMNIST中为50轮,CIFAR-10中为100轮.

4.4 结果指标

由于客户端的数据分布都是非独立同分布的,部分参与的全局模型的测试精度大部分时间在一定的区间内振荡,且振荡幅度较大.这种不稳定的收敛对于实际场景中的应用是不利的,因为在测试标签未知的情况下很难获得最佳结果和训练回合数.因此,本文使用3个实验结果的最后5轮的平均测试精度作为比较MNIST或FMNIST中研究的算法的衡量标准,以及一个实验结果的最后10轮的平均测试精度(本文已经去除数据集设置的随机因素的影响)作为比较CIFAR-10中所研究算法的指标,并利用箱形图对测试结果进行综合评价.

5 实验结果

5.1 对比实验

本小节将讨论相关联邦学习方法,如FedAvg、FedAsync和FedAcc的各个方面表现.本文比较了这些方法在不同采样率下的收敛性,每个联邦学习算法的测试性能图如图2所示,表3以数字形式总结了测试性能.由于稳定性对于分布式算法至关重要,额外应用的动量项可能会放大算法的不稳定性.本文使用箱图(图2(d))来衡量测试精度的不稳定性.

图2 FedAvg、FedAsync和FedAcc之间在MNIST和FMNIST中不同采样率条件下的平均测试精度直方图

表3 使用简单卷积神经网络的各个方法之间的平均测试精度.对于不同数据集的每个采样率,性能在最佳结果的0.5%以内,以粗体显示.*代表本文提出的方法

本文的方法在MNIST和FMNIST两个简单的数据集在大多数情况下优于其他方法,特别是在低设备采样率下,如图2(a)和图2(b)所示.因为当采样率为1.0时,FedAsync和FedAcc算法实际上等价于FedAvg算法.因此,实验结果的精度的微小差异可被视为随机因素引起的误差.由于设备的部分参与,就可能存在参与的客户端数据分布偏离整体数据分布,极端情况下甚至会使算法不收敛.但采样率的提高会使异步FedAsync算法和本文的FedAcc算法中本轮和上轮算法的梯度信息更接近,梯度更新朝某一方向前进,因此会使整体算法效果偏向某些客户端的数据分布,并可能导致采样率提高,而整体精度略低的情况出现,比如采样率为0.7时的精度略低于采样率为0.5的精度.对于具有简单卷积神经网络的更复杂数据集CIFAR-10,本文得到了相同的结论.箱形图中的方盒从数据的下四分位值延伸到上四分位值,中间的线代表中位数线.上下边缘线表示数据的范围.从图3(b)中可以看出,本文算法的测试精度分布在一个更紧凑的范围内,这意味着本文的算法更稳定.

图3 FedAvg、FedAsync和FedAcc在 CIFAR-10中不同采样率条件下的平均测试精度直方图和箱形图

5.2 基于FedAdam改进性能

本文在MNIST和FMNIST数据集上对FedAdamAcc算法进行训练和测试,训练回合数设置为50轮,并且以平均最后5轮3次实验的结果最为效果对比指标.实验结果表明,基于FedAdam的梯度积累策略算法(FedAdamAcc)也是有效的,并可以取得更好的效果,在低设备参与率的情况效果明显,如图4所示.

图4 FedAvg、FedAdam和FedAdamAcc之间MNIST/FMNIST数据集上中不同采样率的平均测试准确度直方图

然而,由于FedAdam需要在CIFAR-10上进行精细调整才能收敛,并且FedAdam很难获得比FedAvg更好的结果,因此本文没有在CIFAR-10上进行进一步的对比实验.

6 总 结

本文从优化的角度分析了异质情况下部分客户端参与的联邦学习的问题,并提出了一种客户端/服务端优化的联邦学习方法——FedAcc算法.本文的方法基于梯度校正,经过评估可适用于不同的数据集.本文通过实验验证了累加方法对提高部分采样联邦学习算法的收敛性是有效的.通过使用一个简单的客户端/服务器优化器框架,本文可以以一种简单合理的方式将累加方法合并到各种联邦学习算法中.未来的研究方向包括如何在采样率和局部训练速度之间实现更好的平衡,以及针对联邦学习设计合适的模型.