面向数据异构的联邦学习性能优化研究

黄聿辰,赵彦超,郝江山,陈 兵

(南京航空航天大学 计算机科学与技术学院,南京 210000)

0 引 言

联邦学习已经成为现代大规模分布式机器学习中的一个重要方向.联邦学习与传统的集中式学习不同,集中式学习使用存储在中央服务器中的大型数据集对模型进行训练,而在联邦学习中,训练数据可以分布在大量的客户端上[1],比如电话、网络传感器、医院或者各类边缘设备[2],然后通过服务器聚合成完整的全局模型,而不用传输用户数据,从而保护了用户的隐私信息.

联邦学习本地通常采用的神经网络,其处理各种计算机视觉任务非常有效.一般来说,在训练过程中,神经网络都包含所有数据样本.然而,在联邦学习的环境中,各个客户端数据样本量和样本类别都处于变化之中,即数据异构[3].这更类似于生物系统在现实世界中的学习方式.在这种情况下,各个客户端神经网络只能训练各自保有的数据集.因此中心服务器系统面临的主要问题是当神经网络权值适应新的任务时,可以保持对于先前学习到的数据的识别能力.

联邦学习优化的关键挑战是数据分布的异构性.由于参与联邦学习的各个边缘设备所处实际物理环境的不同,能为联邦学习框架提供的数据各不相同,这种不同的数据分布会导致边缘设备在本地训练过程中产生模型偏移.采用这类模型进行聚合,生成的全局模型无法在全类别数据集上有优异表现.最常用的FedAvg[4]通过在与服务器通信之前对可用客户端执行多个本地更新来解决通信瓶颈,即增加本地训练轮数,增大本地训练计算时间在总体任务中的时间占比,从而减少通信开销占比.虽然它在某些应用中取得了成功,但是该方案未充分考虑联邦学习实际应用场景,即各个客户端数据异构,在多轮的本地训练过程中,模型偏移随着迭代而不断加深,导致采用平均算法聚合的全局模型在全类别数据集上性能大幅下降,因此数据异构[5]仍是联邦学习的一个活跃研究领域.

FedAdp[6]考虑了局部模型与全局模型的余弦相似角[7],在聚合过程中按照相似程度赋予权重,该方案一定程度上抑制了局部模型偏移,但是与FedProx相似,仅仅是直接做了全局模型趋同,未能结合各个客户端异构性的数据分布.FEDC[8]提出对运算效率低下的客户端进行低权重聚合,以保证全局模型的性能,但是该方案未考虑各个客户端数据集信息量,即运算效率较低的客户端的数据集对全局模型收敛的重要性未被考虑.Huang[9]提出一种结合本地训练精度与训练频率信息进行加权的聚合方式,可以帮助服务器选择更好的模型进行聚合,但是该权重设置属于启发式方案,并没有对其收敛性进行证明.

因此,由于数据异构造成的局部模型偏移问题极大的引起了全局聚合模型识别性能的下降,针对该问题,本文提出了面向数据异构情况下的联邦学习局部模型偏移的性能优化方案.由于神经网络在对数据集的训练过程中,通常仅有部分参数对数据的特征进行拟合,而另外一部分参数在训练过程中所起到的作用较小,该现象也称为神经网络的过参数化.其中梯度变化量体现了该参数在模型训练过程中起到的作用,积极拟合数据特征的相关参数会在训练中产生更多梯度变化量.利用上述结论,本方案在联邦学习本地训练过程中,在其损失函数上添加待训练模型与剩余各个客户端模型参数之间的差值二范数,即将本地待训练模型与其余客户端模型参数作差并求其二范数,用剩余客户端的模型梯度增益对应作为二范数每个参数位置的权重系数,将该整体形成惩罚项,在优化过程中,对新构建的损失函数进行梯度下降优化,从而保证了所有客户端的模型在本地训练过程中,依然保持着对其他客户端数据分布的拟合能力,有效抑制局部模型偏移.所有参与联邦学习的客户端都能保持对其余所有客户端数据分布的识别能力,将这些客户端的梯度更新进行聚合,生成的全局模型性能有所保证.

本文主要贡献有:

1)通过增加本地待训练模型与其余客户端模型参数差值的二范数,并在该二范数上添加模型梯度权重,将整体作为本地训练惩罚项,抑制本地模型对各自数据集的模型偏移,来达到提升全局模型识别精度的效果.

2)在不改变聚合方式的情况下,有效解决了数据异构的问题,使得模型框架更好的适配实际应用环境,实现了优化框架的即插即用.

3)对于提出的优化方案进行了系统的证明,从理论上保证了收敛性和收敛速度.

1 相关工作

大规模机器学习对数据中心服务器的性能需求非常巨大,因此众多分布式学习方案应运而生.随着手机、传感器等边缘计算设备算力的增长和普及,以及将数据移动到数据中心的通信开销较为庞大的问题,在分布式设备组成的网络中进行各自的本地学习来进行统计模型生成的方案变得越来越有吸引力.这个方案被称为联邦学习,它需要解决隐私、数据设备异构以及大规模分布式网络通信开销等新问题.

最近提出的一些优化方法,对上述问题提出了不少有创意的解决方案.如Boyd[10]所提方案以及小批量优化方法[11].在大型网络中,允许不精确的本地更新来平衡通信与计算,和在一个小部分集合的设备组之间进行即时通信.Fedpaq[12]提出了一种周期平均和量化功能的高效通信联邦学习方法,它通过多任务学习框架为每个设备学习独立但是相关的模型并进行周期性更新和发送.尽管该方法在理论上有保证,但这种方法并不能推广到所有问题上.在非凸设置下,一种基于平均原始局部随机梯度下降(SGD)更新的方法[13],实验证明性能良好.SkewScout[14]提出了通过本地训练轮数的限制来控制局部模型偏移的程度,从而提升全局模型的识别精度,但是该方案仅仅考虑了模型性能,并未考虑在联邦学习实际应用场景中,通信开销也是一个重要的问题,压缩本地训练轮数会导致通信开销占比显著增大,从而导致计算与数据传输比例大幅减小.IFCA[15]通过按照模型参数相似程度进行簇划分,在各个簇内进行模型共享,从而使得簇内模型更新方向的一致性,抑制簇内模型偏移.该方案提升了簇间的全局聚合模型性能,但是该方案在实验过程中发现,需要每轮训练过程中保持较高的客户端参与率,否则各个簇内客户端数量不够,无法保持该方案的优化效果,但是联邦学习实际应用场景中,边缘设备的不可用性也是需要考虑的一部分因素.

尽管,最近的一些工作[16,17]探索了数据异构环境下的收敛性保证问题,但他们要求所有设备参与每一轮通信,这在现实联邦学习网络中往往是不可行的.还有一些启发式方法旨在通过共享本地设备数据或服务器端代理数据来解决数据异构[18,19].如Liam[20]该作者通过各个客户端将自己本地数据进行多维度映射,并在该过程中保存数据特征,然后将映射结果传输给服务器,由服务器对所有客户端上报数据进行整合,再下发给各个客户端,客户端进行本地训练时,先在服务器下发的多个客户端映射数据集上进行训练,然后在本地数据集上进行训练,从而减少数据异构带来的全局性能损失.该方案的全局模型性能较高,但是即使对客户端数据进行的映射操作,依然存在暴露本地数据的行为.

因此,这些方法是不现实的.除了增加网络带宽负担之外,根据模型权重推测出数据分布[21]的方式违反了联邦学习的隐私保护要求.将各个客户端的数据进行周期性局部分享的方式[22]则需要详细地生成或收集此类局部辅助数据,计算开销较大,局部辅助数据生成难度较大.

除了数据上的异构性之外,系统上的异构性也是联邦学习中的一个关键问题.由于硬件(CPU、内存)、网络连接(3G、4G、5G、WIFI)和功率(电池水平)的变化,系统属性也有所不同.这些系统级的异构性极大地引发了诸如掉队和容错等问题.在异构网络中,联邦学习优化的一种策略是忽略不受约束的设备,即在训练框架中,放弃聚合无法完成一定训练轮数的设备所提供的梯度更新信息[23].然而,这会对收敛产生负面影响,因为它减少了有助于训练的有效设备的数量,如果被丢弃的设备具有特定的数据特征,则会在全局模型收敛的过程中产生偏差.

最接近本文的工作的是Prox[24],该文章的作者提出的FedProx算法,就像本文的算法一样,也使用参数差值的二范数作为惩罚项.然而,与本文的算法不同的是,在FedProx中,惩罚项是各向同性的,其直接将本地模型与全局模型更新方向进行同步,并未考虑全局模型本身已经是各个局部模型在本地异构数据上的模型偏移之后的聚合结果,所以此时的全局模型更新方向并不能代表最优解方向.Ohad[25]提高了处理二次目标的收敛速度,证明了其可以通过提升数据容量大小,从而达到加速收敛.最近的一项工作[26]证明了FedAvg变体的对数据异构的收敛速度,它还为实践中已知的一种现象提供了一个理论解释,即当局部迭代次数过高时全局性能会下降,因为在数据异构情况下,过多的本地迭代会导致局部模型偏移,所以之后的全局聚合模型无法在全类别数据集上有优异性能.

2 系统模型

本节将详细介绍联邦学习中存在的数据异构问题带来的影响,分为两个小节对系统建模以及问题定义进行介绍.

2.1 系统建模

在概率论与统计学中,独立同分布(Independent and identically distributed,缩写为IID)是指一组随机变量中每个变量的概率分布都相同,且这些随机变量互相独立.

在机器学习任务中,训练数据集的独立同分布是常规假设,常用的标准数据集如MNIST,共包含6万张图片,分为10类,每一类6千张图片,即均匀划分,在集中式训练过程中,每一个mini-batch都是均分采样,符合独立同分布的假设.

但是在联邦学习实际应用场景中,由于边缘设备本地所包含数据集各不相同,各客户端数据集Dk,其中k为客户端编号.Di∪Dj=0或者Di∩Dj→0,即客户端数据的样本交集为0或者交集很少.在这种数据分布下,本地客户端的训练会产生模型更新的偏移,如图1所示.

图1 模型更新偏移Fig.1 Model update offset

因此,参与训练的各个客户端本地模型更新的聚合方向与全局理论最优解方向不一致,就会导致最终生成全局模型识别精度的下降.

2.2 问题定义

为了修正本地模型对各自数据集在训练过程中产生的模型偏移,即保证模型在本地训练过程中,能减少模型在其余客户端模型关键参数位置的变化.该问题可以表述为:

θi为客户端i的本地待训练模型.θk为其余客户端模型,其中k为除i以外的客户端.Gk为其余客户端本地更新梯度增益,用来标识其余客户端模型的关键参数,使得关键参数在客户端i的本地优化过程中得到更高的权重,从而保留该参数的更新方向.

数学表述为:

argminθiL(θi)+λ∑k∈S(θi-θk)Gk(θi-θk)

(1)

将该参数作为惩罚项加入到本地训练过程中的损失函数中进行梯度下降优化,在优化过程中就可以使得每一个客户端在本地进行模型的训练与更新的同时,保持其余客户端的模型更新方向,使得每一个客户端都不会产生对自己本地数据分布的严重模型偏移,从而在后续的全局聚合中,能够有效提高生成的全局模型对于全类数据集的识别精度.

3 局部模型偏移抑制

该章节介绍了优化方案在本地的实现过程,算法的整体流程,以及对优化方案的理论证明.

3.1 局部优化

神经网络对于数据样本的训练普遍存在过参数化现象,即对于某些特定数据分布的样本集合,神经网络中的神经元参数只有部分对于该数据分布的拟合起到作用,可将其称为关键参数.因此在各个客户端本地训练过程中,应当减少对于其余客户端关键参数的更新,从而保证各个客户端本地训练完毕之后对于其余客户端数据依然可以有效识别.

本文通过限制局部更新使其接近其余客户端模型,在局部子问题中添加一个约束项.客户端不仅仅最小化常规局部损失函数Ls,而是使用以下新构建损失函数Lt, s进行约束更新:

Lt, s(θ)=Ls(θ)+λ∑j∈S(θ-θt-1,j)Gt-1,j(θ-θt-1,j)

(2)

其中,Ls(θ)为客户端集合,S为客户端集合,t为全局训练轮数,j为其余客户端标识,θt-1,j为上一轮全局训练的各个客户端模型,其中j∈S,Gt-1,j为上一轮全局训练的各个客户端梯度,其中j∈S.用本地模型与、其余客户端模型参数作差并取其二范数,在该值各个参数位置用模型梯度作为权重,将其加入损失函数进行SGD,可以减少训练过程中其余客户端关键参数位置的变化幅度,从而保留其余客户端信息量.

本文将该优化方案简称为FedOC.维护FedOC所需的历史数据的存储和传输,存在数据量大以及数据传输成本高的问题.然而,通过对损失函数的计算整合,可以避免这些潜在的问题.损失函数Lt,s可以被重新整理为:

(3)

即应用多项式乘法规则,将两个多项式逐个相乘,其中常数C为θt-1,j和Gt-1,j两个项的相乘结果,因为这两个项为上一轮运算结果,因此其结果为常数,故用常数C表示.

服务器只需要维护并向边缘设备传输除θ之外的2个与θ维度相同的元素:

ut-1=∑j∈S(Gt-1,j),vt-1=∑j∈S(Gt-1,j)θt-1,j

(4)

按照公式(3),每个客户端本地训练过程,仅需要公式(4)的两个模型系数即可.

需要注意的是,该方案只需要从边缘设备发送与本地梯度相关的聚合信息到服务器.在隐私性方面,它与经典的FedAvg算法没有区别.图2给出了框架结构图.

图2 局部模型偏移抑制框架Fig.2 Local model offset suppression framework

3.2 算法流程与描述

3.2.1 算法流程

算法.局部模型偏移抑制(FedOC)

输入：K为客户端总数;T为全局通信轮数;θ为训练模型;N为所有客户端训练样本总数;nk为客户端k的样本数量;

1.ForkinKdo:

2. 服务器发送训练模型θ给客户端k

4.服务器侧运算参数u0,v0

5.End for;

6.Fort=0…T-1 do:

7. 服务器随机选择的K子集kt

8.服务器向kt发送全局模型θt,参数ut,vt

9.本地客户端优化新构建损失函数

10.θt+1,j=argminθFk(θ)

11.=L(θ)+λθTutθ-2λθTvt+C

12.客户端上传本地梯度Gt,j

13.服务器接收各个客户端上传梯度,生成本地模型

14.运算新一轮的ut+1,vt+1

16.End for;

3.2.2 算法描述

正式训练之前,进行预训练的数据收集,由参与联邦学习训练的客户端首先完成1轮初步训练,并上传各自训练梯度,由服务器进行对应模型更新.在正式的联邦学习训练过程中,服务器选择客户端集合的任意子集参与本轮联邦学习训练,并向子集中的客户端传递预训练过程中所有客户端上报的参数θ0,u0,v0,各个参与本轮训练的客户端用指定参数构建损失函数,完成SGD优化,上传本轮更新梯度,由服务器更新对应本地模型和全局模型,并运算出下一轮需要提供给客户端进行损失函数构建的参数进行下发,不断迭代,直到训练达到全局轮数为止.

最后,由于FedOC只对FedAvg进行轻量级修改,仅仅在本地训练过程中采用了新构建的损失函数进行优化,对于聚合方面并未有所改动.因此可以使用在FedAvg成立的推理过程,并且使FedOC轻松地集成到现有的系统中.特别是,本文提出,FedAvg是FedOC的一个特殊情况,即λ=0.

3.3 收敛性及收敛速度证明

3.3.1 定义,假设以及定理介绍

接下来对框架的收敛性以及收敛速度进行证明,首先,将介绍一些定义,假设和定理.

定义1.(光滑性)函数f具有Lipschitz连续梯度,其中常数L>0(即f是L-平滑的):

(5)

定义2.(强凸性)函数f是μ强凸的,其中μ>0

(6)

该定义旨在允许本地客户在灵活地执行训练,以便每个本地目标都可以收敛.局部计算与通信的数量可以通过调整局部迭代的数量来调整,即更多的局部迭代表明更精确的局部解.

此外,本文介绍以下假设:

假设1.目标函数f(x)是有界的,即最优解的有界性minf(θ)=f(θ*)>-∞.

假设1显然满足,因为目标函数f(x)存在一个最小值.

假设2中关于随机梯度平方范数的条件是常规的.与随机梯度的期望的假设相比,这是一个更加合理的假设.

(7)

假设3确保了局部梯度的可估计性.

3.3.2 定理以及其证明

结合上述定义和假设,本文将继续介绍并在后续证明了一些有用的定理.

定理1.根据定义3,新构建本地损失函数具有γ-不确定解,因此可得:

(8)

其中c为客户端数量.

证明：

对于定理1,即全局梯度的有界性.

根据定义3可得:

全局梯度可表示为:

将级数展开可得:

根据假设2,得证:

证明完毕.

定理2.根据定义2,如果f(x)是μ-强凸,则:

(9)

证明：

对于定理2,即本地损失相比于最优解的有界性.

由于根据强凸函数线性结合的性质可得,新构建损失函数同样为μ-强凸.

定义：

令Γ′(θ′)=0

可得:

因此可得:

带入公式可得:

整理结果,得证:

证明完毕.

定理3.根据定义1,定义2,定义3,假设2,以及假设3,可得,在T轮全局迭代之后,FedOC收敛于全局最优解模型θ*.

(10)

证明：

对于定理3,即FedOC收敛于全局最优解模型θ*,全局损失的稳定下降.

根据定义1,可得:

根据假设3,可得:

不等式两边同时减去:

得证:

证明完毕.

定理4.根据定理3结果可得,全局轮数可由相关常数提供有界性保证,即:

(11)

其中ξ为初始损失与最优损失差值的线性相关值.

证明：

根据定理3的结果,即:

令t+1=T,1-2μησ=q

存在常数ξ,使得:

不等号两边取对数,并令:

可得:

等式两边同时除以logq,

得证:

证明完毕.

4 实验与分析

本节介绍了优化方案所采用的实验设置,各项对比实验、消融实验和实验结果的分析.

4.1 实验基础设置

本文的算法框架基于开源pytorch框架,采用RTX2080作为模型运算单元,并且实现了FedProx,sharedata两种框架进行对比实验.在联邦学习实际应用场景中,各个边缘设备算力参差不齐,因此选择对算力要求相对较低的全连接网络作为训练模型.

采用MNIST,FMNIST,CIFAR10这3种常用数据集进行实验.训练样本为50000个,测试样本为10000个.实际训练过程中由于MNIST和FMNIST的样本维度为(28,28,1)而CIFAR10为(32,32,3),因此为了方便模型进行数据输入,将CIFAR10所有样本数据进行灰度化.这样仅需在模型输入阶段将输入维度由784修改为1024即可.

首先将客户端设置为100,在MNIST,FMNIST,CIFAR10这3种数据集上进行与FedProx,sharedata的对比实验.各个客户端在训练开始之前,随机分配总类别的其中3种,每一种类别随机选取3000个样本,形成总计9000个,3大类的数据异构分布场景.

实验过程中,对数据分布不断做随机划分,进行多次实验,选取平均值作为实验结果.

4.2 实验结果分析

客户端数量和数据集的不同都会对实验结果产生影响,因此设计了相关对比实验.消融实验为未采用优化方案的常规FedAvg和采用了优化方案的FedOC进行对比.

4.2.1 客户端数量,数据集对比实验

客户端数量为100的3种数据集下的各个优化方案对比实验结果见图3.可以发现,各个优化算法在CIFAR10数据集上的表现相对MNIST,FMNIST要略差一些,这是由于CIFAR10的数据样本相对复杂,简单的全连接网络对于其特征的拟合没有对于MNIST以及FMNIST那么完善.

图3 客户端数量为100且在3种数据集下对比实验Fig.3 Comparative experiments on three datasets with 100 clients

分析整体实验结果得出结论,共享数据的优化方案在各个数据集上的表现,都要优于其他两个优化方案.这是因为共享数据方案使得联邦学习系统拥有趋于独立同分布的数据集分布,因此模型可以在训练中充分拟合数据特征而不会产生局部模型偏移,更加接近集中式学习方案,从而提升识别精度.通常,数据共享型优化方案可以作为其他优化方案的目标值.但是由于联邦学习对于隐私保护的需求,这类对客户端样本信息进行公开或者部分公开的做法并不严谨.

再将客户端数量调整为200进行实验,观察各优化算法在客户端数量变化之后的性能,实验结果如图4所示.

图4 客户端数量为200且在3种数据集下对比实验Fig.4 Comparative experiments on three datasets with 200 clients

根据实验结果得出结论,客户端数量增加之后,各个优化手段在3种数据集上的识别精度都有所提升.在实验室环境中,公共数据集样本种类有限,随着客户端数量增加,随机分发数据时,拥有相同种类甚至同一样本的客户端数量会成比例上升,各个客户端的数据异构性产生弱化,因此随着客户端数量上升,全局模型拟合能力会随之提升.

将本方案与现有方案FedProx进行比较,可以看出,在识别精度方面,要优于FedProx 5个百分点.FedProx 在本地损失函数上仅添加了本地待训练模型参数与当前全局轮数下发的全局模型的差值二范数,此时的全局模型是由上一轮的各个本地客户端模型按照平均算法聚合生成的,然而各个本地模型因为其数据异构性质,会产生较大的局部模型偏移,采用该本地模型进行聚合生成的全局模型和理论全局模型最优解存在差异,笼统地将本地待优化模型参数与全局模型参数做同步的方案并不精确,只是在各向同性上抑制了本地客户端的偏移,而FedOC充分考虑各个客户端在本地训练过程中其余客户端的更新方向,因此在最终的识别精度上FedOC也要更高.

4.2.2 消融实验

该小节实验中加入了未采用任何针对非独立同分布数据集进行优化的常规FedAvg训练框架作为消融实验进行对比,用户数为100,3种数据集对比实验结果如图5所示.

图5 客户端数量为100且在3种数据集下消融实验Fig.5 Ablation experiments on three datasets with 100 clients

用户数为200,3种数据集对比实验结果如图6所示.

图6 客户端数量为200且在3种数据集下消融实验Fig.6 Ablation experiments on three dataset with 200 clients

可以看出,未采用优化手段的FedAvg识别精度上升缓慢,这是由于局部模型偏移引起的全局模型更新方向偏移,最终远离理论全局最优解方向,因此识别性能提升缓慢.

4.2.3 结果分析

1)本方案不涉及客户端数据的共享,在隐私保护方面与常规FedAvg相同,可采用所有适用于FedAvg的隐私保护方案.

2)本方案仅仅对FedAvg进行轻量化改造,并未对聚合过程进行变化,仅在本地训练过程中对损失函数进行优化,因此适用于FedAvg以及其变体场景,实现即插即用.

3)本方案充分考虑系统中所有客户端的数据分布,并在训练中保持其更新方向,因此在全局模型性能上优于如FedProx等现有算法.

4)常规FedAvg算法在非独立同分布数据集上性能较差,模型偏移导致其收敛缓慢.实验中其余优化算法均可在规定全局训练轮数中收敛.

5)在训练总样本类别固定的情况下,增加系统内客户端数量,可以提升数据分布相似性浓度,从而提高全局模型性能.

表1为实验对比数据.

表1 实验对比数据Table 1 Experimental comparision data

5 结 语

本文主要提升了联邦学习在数据异构条件下的全局模型性能.基于局部梯度信息量来保存各个客户端上对拟合起重要作用的关键模型参数,从而抑制本地客户端模型训练时的模型偏移.虽然本地训练过程中加入了额外的惩罚项,并且服务器端需要进行梯度的矩阵运算,一定程度上延长了训练时间,但是全局模型性能的提升是显著的.下一步可以考虑在矩阵运算部分进行轻量化,以提升训练效率.