Astraea论文阅读笔记

联邦学习——Astraea

在开始个性化联邦学习之前，我也看了一些基于全局模型的联邦学习论文，这里挑两篇写个总结，一篇是《Self-Balancing Federated Learning With Global Imbalanced Data in Mobile Systems》，该论文提出Astraea框架来缓解设备之间数据标签不平衡和局部不平衡问题；第二篇是《Node Selection Toward Faster Convergence for Federated Learning on Non-IID Data》，该论文设计了一个选择参与节点的概率模型。通过排除不利于全局收敛的节点参与模型聚合，提高模型收敛速度和预测精度。
这两篇论文都是基于全局模型的联邦学习，第一篇把所有节点的数据用于模型聚合，第二篇则是选择性的让部分节点参与模型聚合
这里先总结一下第一篇，第二篇的内容放在下一篇文章中。

《Self-Balancing Federated Learning With Global Imbalanced Data in Mobile Systems》

现有研究存在的问题

现有研究假设，尽管设备上的数据量可能不平衡，但全局数据分布可以被认为是平衡的。然而，实际在大多数分布式移动设备的真实场景中，全局数据分布是不平衡的（数据标签不平衡）。

本文研究内容

设计了一个新的自平衡联邦学习框架，名为Astraea。Astraea框架通过两个策略来缓解不平衡数据集的问题。
第一，在训练模型之前进行数据扩展以缓解全局不平衡。
第二，使用一些中介，根据中介之间的KL散度（KLD）分配每个中介所负责的客户端，并重新安排客户端的模型训练。通过将具有倾斜数据的客户端结合起来，可能能够实现新的局部平衡。
这样说可能有些抽象，没关系，接着往下看！

理论与方法介绍

联邦学习中不平衡数据的训练：
要考虑数据是分布式，私有的，带有噪声和异常值等，可以将数据不平衡分为以下三类：
1)大小不平衡，即每个设备(或客户端)上的数据大小不均匀；
2)局部不平衡，即非独立同分布(Non-IID)，每个设备不遵循同一数据分布；
3)全局不平衡，指所有设备的数据都是标签不平衡的。
文章中通过具体的实验证明，相比于其他两种不平衡的情况，全局不平衡对联邦学习中模型的精度损失影响更大。考虑到用户的数据隐私，上传或共享用户的本地数据都是不可行的。为了应对这一挑战，本文提出了一个名为Astraea的自平衡联邦学习框架，该框架通过基于z分数的数据扩展和基于中介的多客户端重调度来提高训练精度。（P3-4）
Astraea框架的设计目标是缓解训练数据的全局不平衡和局部不平衡，恢复精度。

数学推导：

接下来用数学推导验证
1)不平衡数据不可能通过模型训练得到最优参数，这个证明有点长就不在此赘述，有兴趣可以直接看论文。
2)用数学归纳法证明了当满足客户端数据分布和测试集数据分布一致这一初始条件时，联邦学习可以恢复模型的精度

根据上述推到，训练集和测试集的数据分布差异是造成联邦学习准确性下降的原因之一。因此，为了实现一个新的局部均衡，本文提出了Astraea框架来对少数类进行数据扩展，并创建中介来合并多个数据倾斜的客户端。

Astraea框架

为了解决精度下降的问题，应该重新平衡每个客户端的训练数据。一种方法是重新分配客户的本地数据，直到实现均匀分布。然而，共享数据会引发隐私问题，并导致较高的通信开销。另一种方法是依次更新全局模型。每个客户端根据最新的全局模型进行训练，并按顺序将其更新参数应用于全局模型。这意味着该方法的通信开销和时间开销是联邦学习的K倍。结合以上两种思路，本文提出在联邦学习的服务器和客户端之间增加中介，重新平衡训练。

Astraea框架如上图所示，该框架由服务器、中介和客户端三部分组成。
服务器负责维护一个全局模型w0，将该模型部署到中介，并使用FedAvg同步聚合来自中介的参数更新Dw4、Dw7、Dw8。
客户端可以是移动电话或物联网设备。图中的四个形状代表四类数据。根据数据分布的特点，客户端可以分为三类:
数据均衡的客户端：有足够的均衡数据，可以准备模型训练 (例如图3中的客户端E和F)。
数据量较少的客户端：数据量相对较小，很难参与模型训练。
倾斜的客户端：有足够的训练数据，但偏向持有某些类别的数据，这导致本地数据不平衡(即客户A-D，G，H)。
中介有两个职能：
1） 重新安排三类客户的训练。例如，客户端G具有标签为0和标签为1的数据，客户端H具有标签为2和标签为3的数据。中介可以结合G和H的训练，实现一个局部数据均衡。
2） 中介需要使所负责的客户端的整体数据分布接近统一数据分布，并根据KL散度来衡量局部均衡情况。
这里的中介是虚拟组件，它可以直接部署在服务器或移动边缘计算(MEC)服务器上，以减少通信开销。

Astraea的训练过程

Algorithm 1 是训练过程的伪代码。首先，FL服务器初始化全局模型，然后启动新的通信轮次r，并将全局模型发送给中介。接下来，每个中介m协调自己所负责的客户端进行训练，同时并行计算参数更新并发送给服务器。FL服务器收集所有中介的参数更新，通过FedAvg来聚合这些参数。最后，服务器将全局模型进行更新，并结束这一轮通信。更新的模型是下一轮(即n+1)通信的初始模型。看着有点懵的话可以结合下面的运行流程一起理解。

Astraea的运行流程

Astraea的运行流程包括初始化、重新平衡、训练和聚合

初始化：在初始化阶段，FL服务器首先等待移动设备加入FL模型训练任务。设备通过将其本地数据分布情况发送到服务器①，加入训练任务。在确定参与训练的设备(客户端)之后，FL服务器初始化神经网络模型的参数和优化器，并收集参与者的本地数据分布情况。
重新平衡：分两步重新平衡FL的训练，即数据扩展和下采样、重调度。②-④
首先执行基于z分数的数据扩展和下采样，以缓解训练数据的全局不平衡。然而，数据扩展和下采样有可能加剧局部不平衡，并将导致精度损失。因此，本文提出由中介来异步接收并应用来自客户端的更新来平均本地不平衡数据。通过重新调度客户端的训练，中介可以得到一个更平衡的模型。通过实验证明，与仅使用其中一种策略相比，扩展和下采样策略和重调度策略的结合可以获得更高的精度，两种策略的伪代码如Algorithm 2 和Algorithm 3所示，具体看论文，这里就不赘述了。


训练：在每一轮通信开始时，每个中介将模型发送给所负责的客户端⑤。每一个客户用小批量SGD作E次迭代来训练模型，并将更新后的模型返回给相应的中介。本地迭代次数仅影响每个客户的训练时间，不会增加额外的通信开销
中介收到更新后的模型⑥并将其发送给下一个等待训练的客户端。每个客户端都完成一轮训练称为一轮中介迭代。Astraea多次重复这个过程。然后，所有的中介将模型的参数更新发送到FL服务器(⑦)。在中介迭代期间，通信开销和模型精度之间有一个权衡，例如可以通过增加中介迭代次数来减少通信开销，以加快模型收敛。（中介一次性迭代多轮，也就是客户端训练多次后再向中介提交，可以减少中介和客户端之间的通信次数）
聚合：首先，FL服务器根据FedAvg聚合所有更新。然后，FL服务器将更新后的模型发送给中介，并开始下一轮同步。Astraea和标准FL算法在模型聚合阶段的主要区别是Astraea可以实现局部平衡。因此，Astraea的聚合模型比标准的联邦学习算法的效果更加平衡。

验证分析与实验效果

使用top-1作为评估CNN模型的指标,用该方法与FedSGD在不平衡的数据集EMIST上的训练效果作比较，结果显示模型的预测精度由74.85%提升至79.24，提升了4.39%。

使用top-1作为评估CNN模型的指标,用该方法与FedSGD在不平衡的数据集CINIC-10上的训练效果作比较，结果显示模型的预测精度提升了6.51%。

参数$T_d$的值是可调的，我们可以根据具体任务调整它的值。然而，如图5和6所示，过小的值是不合适的，3.0至3.5的范围更为推荐，这也是基于z分数的异常值检测算法的推荐设置。
随着每轮参与训练的客户端数量的增加，模型的训练收敛速度加快，但随着迭代次数的增加会出现过拟合的问题，这可以通过正则化来加以解决。
提高每个中介负责的客户端数量的阈值无助于提高模型的精度。
当更多的客户参与培训或更多的客户被分配给中介时，中介可以实现更好的部分平衡。

为研究中介和扩充策略的参数对平衡度的影响，本文对使用FedAvg，只使用扩充策略，只使用重调度策略，使用Astraea这四种方法后的所有客户端或中介的KL散度进行对比。结果表明，每轮参与训练的客户端数量的增加或提高每个中介负责的客户端数量的阈值可以减小KLD范围。这表明，当更多的客户参与训练或更多的客户被分配给中介时，中介可以实现更好的局部平衡。总之，随着训练规模的扩大，Astraea的准确性有所提高。图8b中的结果证明了类似的结论，与FedAvg或单一策略情况相比，Astraea使KLD显著降低。

Astraea开销问题

最后，关于Astraea的开销问题主要可以分为三种，即时间，存储空间和通信。
时间开销：有三项主要任务需要额外的时间:数据扩展和下采样、重调度、中介的额外训练时间。由于在初始化阶段只进行一次数据扩展和下采样，因此其在整个训练过程的时间消耗可以忽略不计。重调度的时间复杂度为O(C2)。时间开销主要在于模型训练，FL中，每一轮通信花费的时间是E*T，E是迭代次数，T是一次客户端训练所需的时间。在Astraea中，每次同步所花费的时间是$E_{my}E\ast T$
存储开销：Astraea要求客户端提供额外的存储空间来存储扩展的样本，存储空间和模型精度的权衡图10所示

通信开销：表3中的实验结果表明，Astraea实际上具备比FL更有效的通信。这是因为Astraea可以提高模型收敛速度。例如，在EMNIST中训练CNN模型，以top-1准确率为标准，联邦学习达到74%的准确率需要861MB的通信开销，而Astraea只需要473MB的通信开销。

总结

联邦学习（FL）是一种很有前途的分布式机器学习框架，具有隐私保护的优点。然而，FL不能很好地处理不平衡的数据集。本文研究了不平衡的训练数据对FL的影响，以及由全局不平衡引起的8.44%的准确度损失。作为一个解决方案，本文基于全局模型，提出了一个自平衡的FL框架Astraea，该框架重新平衡了训练数据。实验结果表明，Astraea的top-1在不平衡的EMNIST上提高了+4.39%(恢复了52.0%的精度损失)，在不平衡的CINIC-10上提高了+6.51%(恢复了46.9%的精度损失)。最后测量了Astraea的开销，表明其通信是高效的。