Rethinking Architecture Design for Tackling Data Heterogeneity in Federated Learning

Abstract

联邦学习是一种新兴的研究范式，可以在不同组织之间进行机器学习模型的协作训练，同时保持每个机构的数据隐私。尽管最近取得了进展，但仍然存在一些根本性的挑战，例如缺乏收敛性，以及在现实世界的异构设备之间发生灾难性遗忘的可能性。在本文中，我们证明了基于自注意的架构(例如transformer)对分布转移更健壮，从而改善了异构数据的联邦学习。具体来说，我们对一系列联邦算法、真实世界基准测试和异构数据分割的不同神经架构进行了首次严格的实证调查。我们的实验表明，简单地用transformer取代卷积网络可以大大减少之前设备的灾难性遗忘，加速收敛，并达到更好的全局模型，特别是在处理异构数据时。我们发布我们的代码和预训练的模型，以鼓励未来对健壮架构的探索，作为当前优化前沿研究工作的替代方案。

1. Introduction

联邦学习(FL)是一种新兴的研究范式，用于在分布在多个异构设备上的私有数据上训练机器学习模型。FL将每个设备上的数据保密，旨在训练一个全局模型，该模型仅通过通信参数而不是数据本身进行更新。因此，它为跨多个机构的协作机器学习提供了机会，而不会有私人数据泄露的风险[25,36,54]。

事实证明，这在医疗保健[4,7,15,40]、从移动设备学习[17,38]、s-等领域特别有用购物城市[25]，通信网络[49]，在这些地方保护隐私至关重要。尽管FL提供了丰富的机会，但在FL易于应用于现实世界的数据分布之前，仍有一些基本的研究问题需要解决。目前大多数旨在跨非iid设备学习单个全局模型的方法都遇到了挑战，例如并行FL meth-的非保证收敛性和模型权值发散[35,37,68]，以及序列FL方法的严重灾难性遗忘问题[7,16,57]。

虽然大多数研究工作都集中在改进FL中的优化过程上，但我们的论文旨在通过重新思考联邦模型中架构的选择来提供一个新的视角。我们假设Transformer架构[12,61]特别适合于异构数据分布，因为它们对分布偏移[3]具有惊人的健壮性。这一特性导致了变形金刚在自监督学习中的流行，其中异质性通过未标记的预训练数据和标记的测试数据[11]之间的分布转移表现出来，以及在图像和文本等基本异构输入模式上的多模态学习[24,60]。为了研究这一假设，我们在一套联邦算法、真实世界的基准测试和异构数据分割中对几个神经架构进行了第一次大规模的实证基准测试。为了表示Transformer网络，我们在图像任务上使用Vision Transformer[12,41]的标准实现，包括图像分类[31,42]和医学图像分类[27]。

我们的结果表明，VIT -FL(视觉变形联邦学习)在大多数异构设备分裂的设置中表现得特别好，随着异构性的增加，VIT -FL和ResNets[19]的FL之间的差距显著增加。为了理解这些结果，我们发现改进的主要来源在于Transformer模型对异构数据的鲁棒性的增加，这减少了在本质上不同的新设备上训练时对以前设备的灾难性遗忘。变形金刚一起，收敛更快，并达到一个更好的全球模型，适用于大多数设备。通过与专门设计用于对抗异构数据的FL方法进行比较，我们发现VIT -FL在不使用训练启发式、额外的超参数调优或额外的训练的情况下提供了立即的改进。此外，值得注意的是，我们的VIT-FL与现有的基于优化的FL方法是正交的，可以很容易地应用于提高它们的性能。为此，我们得出结论，变形金刚应该被视为未来研究FL问题的一个自然起点。

2. Related Work

联合学习。联邦学习(FL)旨在跨大规模分布式设备[47]在私有数据上训练机器学习模型。为了实现跨异构设备的有效分布式训练，出现了两类方法:(1)并行FL方法涉及以同步或异步的方式并行训练每个本地客户端(如经典的FedA VG[47])，而(2)串行方法以串行和循环的方式训练每个客户端(如循环权重传输(CWT) [7])分裂学习[62])。FedAVG[47]和CWT[7]的示意图如图2所示。FL的核心问题是训练数据在客户端的分布存在数据异质性，导致并行FL方法出现非保证收敛和模型权值发散[21,37,66,68]，串行FL方法出现严重的灾难性遗忘问题[7,16,57]。

在经典FedA VG算法[47]的最近发展中，包括使用服务器动量(FedA VGM)来缓解每个客户端分布漂移和不平衡[22]，在所有用户之间全局共享小数据子集(FedA VG- share)[68]，使用局部目标的近距离项(FedProx)来减少潜在的权重分歧[37]，或使用其他优化启发方法，如协作重放[52]，无监督对比学习[69]，匹配用户模型[64,65]的特征层，或模型蒸馏[14]来处理异质性。

同时，最近的一些努力旨在缓解连续和串行学习中的灾难性遗忘:限制对先前看到的任务或客户端很重要的权重的更新(弹性权重整合(EWC)[30])，应用深度生成回放来模拟来自先前客户端或任务的数据[52,58]，以及应用循环加权目标来减轻跨标签分布偏误[2]的性能损失等等。然而，所有这些方法都主要集中在改进优化算法上，而没有研究架构设计中提高数据分布变化的鲁棒性的潜力。在我们的工作中，我们展示了架构中的简单选择实际上会产生很大的不同，并且应该成为与当前工作的主要焦点优化方法并行的积极研究领域。

Transformer。Transformer架构首先被提出用于序列到序列的机器翻译[61]，随后在许多NLP任务中建立了最先进的性能，特别是在自监督范式[11]中训练时。最近，变形金刚也被发现广泛适用于涉及图像和视频的任务。例如，Parmar等人[50]将自注意应用于图像的局部邻域，而视觉转换器(ViT)[12]通过直接将具有全局自注意的变压器应用于全尺寸图像，实现了ImageNet分类的最新水平。

相对于语言(即LSTMs[20])和视觉(即cnn[19,34])的经典架构，其令人感兴趣的性能提升最近激发了人们对理解其有效性背后原因的兴趣。在几个特别相关的发现中，ViTs对严重咬合、摄动、区域移位[3,48]以及合成和自然对抗例子[44,51]具有高度鲁棒性。此外，最近的研究表明变压器也适用于异构和多模态数据[24,43,60]。受这些发现的启发，我们为什么假设ViTs也将非常有效地适应FL的数据异质性，并提供详细的实证分析来检验这一假设。

3. Transformers in Federated Learning

在本节中，我们将介绍Transformer体系结构和联邦学习方法的背景。

3.1. Vision Architectures

CNN。对于基于卷积的架构，我们使用ResNet[19]模型族(ResNet-50、ResNet-152和ResNeXt-101 (32x8d))和EfficientNet[59]模型族(EfficientNet- b1、EfficientNet- b5和EfficientNet- b7)，其中包含一系列卷积、ReLU、池化和批处理归一化层。ResNet和EfficientNet是图像分类中最流行的体系结构，也是图像数据FL中使用的标准体系结构[1,39]。

Transformer。作为比较，我们使用视觉变形金刚(ViT(S)， ViT(T)， ViT(B))[12]模型族和Swin变形金刚模型族(Swin(T)， Swin(S)和Swin(B))[41]，它们不使用传统的卷积层。相反，图像特征提取与图像序列化和补丁嵌入策略。每个模型的参数数量见图1。

3.2. Federated Learning Methods

我们应用最流行的并行方法(FedA VG[47])和串行方法(CWT[7])作为训练算法(参见图2中的示意图描述)。
联邦平均算法。FedAVG将每个客户端的局部随机梯度下降(SGD)与平均[47]的迭代模型结合起来。具体来说，在每一轮通信中随机抽取一部分本地客户机，服务器将当前全局模型发送给每个客户机客户。每个选定的客户端然后对其本地训练数据执行本地SGD的E epoch，并将本地梯度发送回中央服务器进行同步聚合。然后，服务器应用平均梯度来更新其全局模型，并重复此过程。

循环权重转移。与FedAVG中每个本地客户端以同步和并行的方式训练不同，CWT中的本地客户端以串行和循环的方式训练。在每一轮训练中，CWT在一个本地客户端上用它的本地数据训练一个全局模型，训练若干个epoch E，然后将这个全局模型转移到下一个客户端进行训练，直到所有本地客户端都在一次[7]上训练完毕。然后，训练过程在客户端中重复循环，直到模型收敛或达到预定义的通信轮数。

4. Experiments

我们的实验旨在回答以下研究问题，这些问题对FL方法的实际部署很重要，同时也有助于我们理解(视觉)Transformer架构。
•与cnn相比，变形金刚能够在FL设置中学习更好的全局模型吗? cnn已经成为FL任务的实际方法(第4.2节)?transformer是否特别能够处理异构数据分区(第4.3.1节)?
•与cnn相比，变压器是否降低了通信成本(第4.3.2节)?
•transformer能否应用于进一步改进现有的基于优化的FL方法(第4.4节)?
•对于从业人员在FL中部署transformer有哪些实用的提示(第4.5节)?

4.1. Experimental Setup
按照[7,21]，我们在研究中评估了Kaggle糖尿病视网膜病变竞赛数据集(记为Retina)[27]、CIFAR-10数据集[31]和现实世界CelebA数据集[42]上的FL。

Retina和CIFAR-10:我们将Retina数据集中的标签二值化为健康(阳性)和患病(阴性)，随机选择6000张平衡图像进行训练，3000张图像作为全局验证数据集，3000张图像作为[7]之后的全局测试数据集。

我们使用CIFAR-10中的原始测试集作为全局测试数据集，从原始训练数据集中留出5000张图像作为全局验证数据集，并使用剩余的45,000张图像作为训练数据集。我们模拟三组数据分区:一个iid数据分区和两个标签分布倾斜的非iid数据分区。视网膜和CIFAR-10中的每个数据分区分别包含4和5个模拟客户端。我们使用每两个客户端之间的平均KolmogorovSmirnov (KS)统计量来衡量标签分布的偏度。KS = 0表示IID数据分区，而KS = 1则是一个非常非IID的数据分区，每个客户端拥有完全不同的标签分布(详细的预处理和数据分区见附录A.1)。

CelebA是一个大规模的人脸属性数据集超过20万张名人照片。我们使用LEAF基准测试[5]提供的联邦版CelebA，它基于身份划分设备。在[5]之后，我们对二元分类任务(存在微笑)进行测试，并丢弃大于8个样本的客户端，以增加任务难度。共纳入227例患者，平均5.34±1.11例，共1213例。
4.2. 使用不同神经架构和(理想的)集中训练的FL的比较:无论应用哪种架构，CWT和FedAVG都可以在IID设置上实现与在中央托管数据(记为Central)上训练的模型相当的结果(图3)。然而，我们观察到在CNNs上的测试准确性显著降低

CWT和FedA VG的异构数据分区，特别是在极端异构的数据分区上(Split 3, KS-1 of CIFAR-10)(图3和图1)。通过简单地用vit替换cnn, CWT和FedA VG即使在高度异构的非iid设置中也成功地保持了模型的准确性。在高度异构的CIFAR-10数据集Split-3, KS-1上，ViT(S)-CWT和ViT(S)-FedA VG的测试精度相对于ResNet(50)-CWT和ResNet(50)-FedA VG分别提高了77.70%和37.34%。因此，VIT特别适合于异构数据。

Comparison with existing FL methods:

我们还比较了VIT -FL与两种最先进的基于优化的FL方法:FedProx[37]和视网膜和CIFAR-10上的FedA VG-Share[68]。我们使用ResNet(50)作为其他比较方法的骨干网，使用ViT(S)作为我们的方法。我们用网格搜索在Split-2数据集上优化最佳参数(FedProx的近项惩罚常数µ)，并将相同的参数应用于所有剩余的数据分区。我们允许每个客户在FedAVGShare中彼此共享5%的数据。如图4所示，在非iid数据分区中，vvi -FL优于所有其他FL方法，特别是在高度异构的非iid设置上。尽管仔细调优了优化参数，FedProx[37]在高度异构的数据分区上仍会出现严重的性能下降。同样，即使在所有客户端之间共享5%的本地数据，FedA VG-Share在高度异构的数据分区Split-3上也会遭受性能下降的影响(CIFAR-10数据集上的Split-3为94.4%，而Split-1为97%)。我们得出结论，简单地使用它优于最近为FL设计的几种方法，这些方法通常需要仔细调整优化参数。请注意，vit的使用与现有的优化方法是正交的，两者的组合可以产生更强的性能(详见第4.4节)。

4.3. Analyzing the Effectiveness of Transformers

鉴于这些有希望的实证结果，我们现在进行了仔细的实证分析，以揭示究竟是什么导致了Transformers成绩的提高。

4.3.1 Transformers generalize better in non-IID settings

FL的分布式特性意味着跨客户端的数据分布可能存在很大的异质性。

先前的研究表明，使用FedA VG或CWT训练FL模型分别会引起权重发散和灾难性遗忘等问题[30,57]。我们认为，cnn中使用的局部卷积已被证明更多地依赖于局部高频模式[13,26,63]，可能对异构设备特别敏感。这个问题在医疗保健数据的FL中尤其普遍，因为不同机构捕获的输入图像由于不同的医学成像协议可能在局部模式(强度、对比度等)上存在显著差异[16,55]，以及由于用户说话[33]、键入[17]和书写[28]的习惯而在自然数据分割中存在显著差异。
另一方面，vit使用自我注意来学习全局交互[53]，与cnn相比，它已被证明对局部模式的偏见较小。这一特性可能有助于它们对分布的惊人健壮性位移和对抗性扰动[3,48]。为了进一步分析变形金刚在异构数据间的泛化能力，我们设计了以下实验:跨异构设备的灾难性遗忘:cnn通常在分布外数据上工作得更差。这种现象在连续FL法CWT中尤为严重。由于其顺序和串行的训练策略，在CWT范式中训练cnn通常会导致在非iid数据分区上的灾难性遗忘:在具有不同数据分布的新客户端上更新几次后，模型在先前客户端上的性能突然下降[3,48]。这将导致较差和较慢的收敛，这在FL中是不受欢迎的。在迁移学习文献中也发现了类似的遗忘问题[8,9,56]。

我们在CIFAR-10数据集的Split-3上评估CWT，以说明这种灾难性遗忘现象。在图5中，我们绘制了随着更多的客户端参与CWT学习，Client-3验证数据集(与其训练数据集共享相同的数据分布)上预测精度的演变。当在Client-3上将训练良好的模型转移到Client-4上时，之前Client-3验证数据集上的预测精度会突然急剧下降(从> 98%下降到< 1%精度)。然而，以ViT为骨干训练的模型(ViT(S)CWT)能够将知识从Client-3转移到Client4，同时只丢失少量Client3上的信息(保持98%的准确性)。因此，vit可以更好地泛化到新的数据分布，而不会忘记旧的数据分布。

我们进一步比较了ViT(S)-CWT与专门用于减轻灾难性遗忘的优化方法EWC30。

CWT在CIFAR-10的Split-3上的串行训练可以看作是一个增量类学习任务，其中每个客户端包含数据集中的一个类的独占子集。

每个客户端模型将相同的分类器共享到标准化的联合标签空间[23]。然而，从图5中可以看出，EWC勉强解决了高度异构的数据分区上的灾难性遗忘问题，这也与[23]中报告的结果相匹配。该实验进一步证明了ViT的有效性，超越了为FL设计的优化方法。

2. Generalization of VIT-FL on real-world federated datasets:

训练有素的联邦模型应该在其他未见客户端的非分布测试数据集上表现良好。为了测试变形金刚的可泛化性，我们将其应用于现实世界联邦CelebA数据集[42]和将其与ResNet对应的FedProx[37]和FedA VG-Share进行比较[68]。我们在表1中报告了使用不同FL方法训练的模型对来自所有本地客户端的测试数据的并集的测试精度。我们的VIT-FL方法优于最先进的FL方法，也减少了方差。这表明变形金刚学习了比cnn更好的全局模型

3.VIT -FL在极端大规模环境下的泛化:为了验证VIT-FL在更大规模的现实世界分布式学习环境下的有效性，其中涉及数千个客户端，我们进一步将不同的FL方法应用于Retina和CIFAR-10数据集上的极端边缘情况。这里的边缘情况定义为一个客户机只持有一个数据样本，这在医疗保健领域非常常见，因为患者只持有属于自己的一个数据样本。这导致了大量的异构客户端:Retina有6000个，CIFAR-10有45000个。从表2中可以看出，vit在这种极端异构的边缘情况下仍然学习了一个有前途的全局模型，显著优于ResNet模型(Retina上从50%到80%，CIFAR-10上从30%到90%)。

4.3.2 Transformers converge faster to better optimum

一个强大的FL方法不仅应该在IID和非IID数据分区上健壮地执行，而且还应该具有较低的通信成本，以便在通信有限的带宽上进行部署。通信成本由收敛前的轮数和收敛前的轮数决定模型参数。我们计算了达到一个预先定义的目标测试集准确率为一个集中训练的ResNet(50)预测准确率的95%所需的通信轮数。具体来说，我们将Retina和CIFAR-10数据集的目标精度分别设置为77.5%和91.5%。我们将串行CWT方法上的一个通信回合定义为跨所有联邦本地客户端的一个完整训练周期。

从图4和表3中可以看出，所有评估的FL方法在同构数据分区上快速收敛到目标测试性能。然而，ResNet(50)-FedA VG和ResNet(50)-CWT的收敛速度随着异构性的增加而下降，甚至在高度异构的数据分区上达到一个平台(永远达不到目标精度)。相比之下，VIT-FL在异构数据上仍然快速收敛。例如，在CIFAR-10上的异构数据分区Split-2和Split-3上，由于严重的灾难性遗忘，ResNet(50)-CWT完全发散，而ViT(S)-CWT在34轮和85轮通信后达到了目标性能。

4.4. In Conjunction with Existing Methods

由于我们对架构选择的调查在很大程度上与现有的基于FL方法的优化正交，因此我们的发现可以很容易地与后者结合使用。我们将Vision transformer与基于优化的方法(FedProx[37]和FedAVG-Share[68])结合起来，并将其应用于Retina和CIFAR-10数据集。从表3和图6，当应用到现有的FL优化

4.5. Take-aways for Practical Usage

局部训练周期:标准使用E表示局部模型通过其局部数据集的轮数。已知E强烈影响FedA VG[47]和CWT[7]的性能。我们对局部训练周期E对VITFL的影响进行了实验研究。对于ViT(B)-FedA VG，我们考虑E∈{1,5,10}，对于ViT(B)-CWT，我们考虑E∈{1,5}。从图7中，我们发现ViT表现出与CNN相似的现象，即较大的E加速了ViT(B)-FedA VG在同质数据分区上的收敛，但可能导致异构数据分区上的最终性能下降。

类似地，ViT(B)-CWT也有利于每个客户端之间的频繁传输速率，就像ResNet(50)-CWT[7]在非iid数据分区上一样。因此，我们建议用户在同质数据上应用大E以减少通信，而在高度异构的情况下应用小E (vv - feda VG E≤5,vv - cwt E = 1)。

预训练对VIT- fl的影响:有证据表明，从零开始训练[12]时，VIT通常需要更大量的训练数据才能比cnn表现得更好。我们通过实验研究了预训练对vvi - fl的影响。我们使用FedA VG作为训练算法，使用Swin(T)[41]作为骨干网，并在CIFAR-10上进行测试。我们在训练期间应用与[41]相同的增强和正则化策略，并将最大通信轮数设置为300。如表4所示，在理想的中央托管和FL设置下，从头训练Swin(T)的性能都会下降。尽管如此，它的性能在高度异构的数据分区上，从头训练时的Split-3(64.50%)比ResNet(50)-FedA VG(图3上的59.68%)要好得多。在实际应用中，建议用户将VIT作为首选，因为在应用预训练模型时，VIT- fl的性能始终优于cnn(图1和图3)。如果没有大规模的预训练数据集，可以选择自监督预训练[6,18]。

其他训练提示:FL中VIT的训练策略可以直接继承VIT训练，如使用线性热身和学习率衰减，梯度剪切。相对较小的学习率和梯度范数剪辑是CWT中稳定VIT训练的必要条件，特别是在高度异构的数据分区中。梯度范数剪辑也有助于在异构数据上训练具有cnn的FL，因为它已被证明可以减少局部更新与当前全局模型[37]之间的权重差异。更多的一般提示和实验分析请参见附录B.1。

5. Conclusion

尽管FL最近取得了进展，但在处理异构数据时，在收敛和遗忘方面仍然存在挑战。与以前改进优化的方法不同，我们通过重新思考FL中的架构设计提供了一个新的视角。使用transformer对异构数据和分布偏移的鲁棒性，我们进行了广泛的分析，并证明了transformer在减轻灾难性遗忘、加速收敛以及达到并行和串行FL方法的更好最优方面的优势。我们发布代码和模型是为了鼓励在优化的同时开发健壮的体系结构。