MoFedAGR：在联邦学习中利用自适应梯度正则化与全局动量来缓解客户端漂移问题

《Neural Networks》：MoFedAGR: Mitigating Client Drift with Adaptive Gradient Regularization and Global Momentum in Federated Learning

　　摘要：联邦学习是一种具有隐私保护功能的新型分布式机器学习框架，但它容易受到异构数据的影响。异构数据会导致客户端模型在本地训练时过度拟合本地数据集，从而偏离全局最优解，这种现象被称为客户端漂移。为了解决客户端漂移的问题，我们从优化和泛化角度出发进行研究。我们全面考虑了训练过程中客户端漂移的影响，并将其量化为聚合误差。首先，我们提出了自适应梯度正则化方法，该方法基于梯度正则化，根据本地模型与全局模型之间参数方差的大小为不同参数施加不同的正则化强度，从而减轻由聚合误差引起的性能下降，帮助模型收敛到更平坦的最小值。为了获取本地模型与全局模型之间的方差以计算自适应梯度正则化项，我们从服务器端引入全局动量作为全局梯度的近似值，并进一步将其用作梯度校正项。接着，我们提出了MoFedAGR方法，该方法结合了梯度校正项与自适应梯度正则化项，有助于客户端模型收敛到一致的平坦最小值。我们还给出了所提算法的理论收敛性界。此外，在多个图像分类数据集上的实验表明，我们的算法显著提升了模型性能，并具备较强的泛化能力。

引言：联邦学习由McMahan等人于2017年提出，它允许多个客户端在不暴露其私有数据集的情况下协同训练模型。与传统分布式计算框架不同，联邦学习采用在客户端进行多次本地训练的方式，以提高训练速度并减少通信开销。本地训练完成后，客户端将各自的模型上传至服务器进行平均聚合，从而获得用于后续训练周期的新全局模型。联邦学习被广泛应用于医疗、金融和自动驾驶等多个领域。然而，由于客户端数据的异构性，客户端模型往往倾向于拟合本地数据集，从而偏离全局最优解，这种现象被称为客户端漂移。在异构数据集上进行联邦学习会导致模型性能不佳且收敛速度缓慢，这成为了限制联邦学习应用的一个重大瓶颈。在物联网等现实世界的通信带宽受限情况下，客户端只能有限度地参与联邦学习，这就导致了采样不平衡问题，进一步加剧了数据异构性的影响。为防止本地模型过度拟合局部最优解，许多研究致力于保持各本地模型之间的一致性。例如，可以在参数空间中引入正则化项来维持客户端模型之间的一致性，或者通过模型蒸馏将客户端模型的特征空间与全局模型对齐。这些算法试图减少本地模型之间的差异，但忽视了尖锐的局部损失函数曲线会放大客户端漂移带来的误差这一事实。仅通过最小化经验风险最小化损失得到的客户端模型损失函数曲线相对较为尖锐。在情况多变的情况下，本地模型之间的差异愈发明显，从而导致聚合后的全局模型性能显著下降。为使本地模型的损失函数曲线更加平坦，我们在本地训练过程中采用了基于梯度正则化的方法。梯度正则化将梯度范数作为损失函数中的惩罚项，通过惩罚训练过程中的梯度范数，模型能够收敛到损失函数曲线较为平坦的区域，从而具备良好的泛化性能和抗干扰能力。由于计算梯度正则化项需要二阶信息，直接计算其海森矩阵在计算上成本较高。实际上，梯度正则化项的导数通常是通过有限差分法来近似的。虽然传统的梯度正则化能使客户端模型收敛到平坦的最小值，但由于客户端漂移的存在，聚合后的全局模型可能会大幅偏离这些平坦的局部最小值，从而影响全局模型的泛化能力。为进一步降低客户端漂移带来的性能下降，我们提出了自适应梯度正则化方法。首先，我们分析了聚合过程对本地模型的影响，随后提出以聚合误差作为衡量聚合后模型性能下降的指标。该误差项与客户端模型与全局模型之间的差异以及局部梯度的大小呈正相关。将此误差项作为正则化项，我们就得到了自适应梯度正则化项。具体而言，自适应梯度正则化项定义为梯度与方差项乘积的范数。与传统梯度正则化不同，我们对那些出现显著偏差的参数施加更强的正则化强度，以防止在客户端模型聚合后性能出现大幅下降。我们还加入了来自服务器端的动量项，作为全局梯度的近似值。通过比较本地梯度与全局梯度之间的差异，可以计算出方差项。通过采用自适应梯度正则化，我们既能够保持全局模型的泛化能力，又能减少客户端漂移带来的性能恶化。为促使客户端模型向全局最优解靠拢，我们还同时利用服务器端的动量来修正客户端模型的梯度。我们的主要贡献如下：•我们提出了自适应梯度正则化方法，它是梯度正则化的一种改进形式。该方法通过限制客户端漂移带来的误差，提升了聚合后全局模型的性能，促进了全局模型向平坦最小值的收敛。•我们提出了MoFedAGR方法，该方法结合了服务器端动量校正与自适应梯度正则化，有助于客户端模型收敛到一致的平坦最小值。我们还对MoFedAGR在平滑非凸场景下的收敛性进行了分析。•我们在CIFAR-10/100和TinyImagenet数据集上进行了大量实验，验证了我们所提算法的性能。与几种经典基准方法相比，我们的方法将准确率提升了高达10.14%。

相关工作：联邦学习：FedAvg由McMahan等人于2017年首次提出，该框架采用多次本地训练和平均聚合的方式。由于数据存在异构性，客户端模型之间的一致性不足会降低全局模型的性能。许多研究致力于保持客户端模型之间的一致性。Li等人于2020年提出的FedProx方法通过引入近似项，防止本地模型偏离全局模型。Karimireddy等人于2020年提出的SCAFFOLD方法则利用方差来处理这一问题。

在数据异构性条件下的联邦学习：考虑到联邦学习的架构，其目标是在训练数据集上最小化全局模型的损失：minω∈Rdf(ω)=1N∑i=1Nfi(ω),fi(ω)=Eξi～Di(fi(ω,ξi))，其中f(ω)和fi(ω)分别为全局损失函数和本地损失函数，Rd是参数空间，N为客户端总数，ξi表示来自遵循本地数据分布Di的本地训练数据集的随机样本。在数据存在异构性的情况下，每个客户端模型都是独立优化的。

自适应梯度正则化：方程2表明，r轮后的聚合误差项可大致视为本地模型梯度与本地模型与全局模型之间方差项的乘积。本地训练完成后，可通过全局模型与本地模型更新后梯度之间的差异来量化方差项：?r≈∑iS?fi(ωir)(ωr?ωir)=∑iS?fi(ωir)(gr?gir)，其中?fi(ωir)是本地模型的梯度，gr和gir分别是全局模型和客户端模型的更新值。

实验部分：在本节中，我们首先介绍了实验设计，包括数据集、基准方法、模型架构以及超参数的选择。随后，我们结合实验结果对各项指标进行了进一步分析和评估。

结论：在本文中，我们提出了一种自适应梯度正则化方法，用于减轻聚合误差并提升模型的泛化能力；同时，我们引入了全局动量校正项，以保持客户端模型之间的一致性，从而缓解客户端漂移问题。此外，我们还为部分客户端参与和非凸场景下的方法提供了收敛性证明。在真实世界数据集上的全面测试表明，我们的方法有效降低了数据异构性带来的负面影响。

关于写作过程中生成式AI和AI辅助技术的声明：在撰写本文的过程中，作者使用了ChatGPT和DeepSeek将中文翻译成英文。在使用这些工具之后，作者对内容进行了必要的审阅和编辑，并对最终发表的内容负全责。

作者贡献声明：王翔：撰写——初稿、验证、方法论；田磊：撰写——审阅与编辑、验证；甘嘉豪：验证；杨晨：方法论、形式化分析；林飞龙：撰写——审阅与编辑、资金筹集；李明路：撰写——审阅与编辑、监督。

利益冲突声明：作者声明不存在任何可能影响本文研究成果的已知财务利益或个人关系。

致谢：本研究得到了中国国家自然科学基金项目62273310和62320106006的支持，以及浙江省自然科学基金项目LZ26F020004的支持。

王翔|田磊|甘嘉豪|杨晨|林飞龙|李明路