FedAPEX：通过联邦加速扰动探索找到更平坦的极小值

《Neurocomputing》：FedAPEX: Find flatter minima via federated accelerated perturbation exploration

　　摘要：在数据异构性条件下，联邦学习面临着一个根本性的权衡：那些旨在加速收敛的方法往往会影响模型的泛化能力，而那些追求平坦、可泛化极小值的方法则通常收敛较慢。基于锐度感知的优化方法虽然能找到更平坦的解，但其标准的单步扰动上升法往往不足以摆脱尖锐的谷地，还可能减缓训练速度。相反，简单应用动量机制来加速局部更新可能会加剧客户端偏差，使全局模型变得不稳定。为了解决这一矛盾，我们提出了FedAPEX这种联邦优化算法，它通过整合双重动量机制，同时实现加速收敛与提升泛化能力。FedAPEX利用全局管理的动量来稳定SAM扰动，引导模型避开尖锐区域，找到真正平坦的极小值。与此同时，另一种由全局动量驱动的下降机制负责协调各客户端的更新，使其与全局优化路径保持一致，从而在加快收敛的同时减少客户端偏差。通过在具有多种异构性的基准数据集上的大量实验，我们证明了FedAPEX的性能显著优于现有的最佳方法。我们的研究结果表明，FedAPEX不仅通过找到明显更平坦的损失函数曲面来实现更高的最终精度，而且通常也需要更少的通信轮次即可达到目标精度。

引言：联邦学习是一种用于大规模机器学习的新兴范式。在中央服务器的协调下，多个拥有私有、分散数据的客户端无需共享原始数据即可共同训练全局模型。联邦学习已在医学分析[12]、[39]、[45]、物联网[IoT][1]、[4]、推荐系统[22]、[23]以及隐私保护型数据安全[2]、[24]、[36]等领域得到广泛研究和应用。在典型的联邦学习场景中，许多客户端（如移动设备或物联网设备）资源有限，这使得算法对通信和计算成本十分敏感。此外，客户端数据往往是非独立同分布的且存在不平衡现象，即数据异构性。这种异构性会导致局部模型过度拟合其本地数据，偏离全局最优解，这种现象被称为“客户端偏差”[17]、[21]、[26]。这种偏差会减缓全局模型的收敛速度，导致更新不稳定，并降低模型的泛化能力，这是联邦学习面临的主要挑战（见图1）。许多有前景的研究通过约束局部目标函数，使局部更新更接近全局模型，以此减轻异构性和偏差的影响[8]、[11]、[21]。近期研究还进一步探讨了损失函数曲面的特性，研究损失表面的锐度与泛化能力之间的关系。研究表明，数据异构性往往会导致模型收敛到尖锐的极小值，而较平坦的极小值则与更强的泛化能力相关[3]、[30]。因此，基于锐度感知优化方法的几种联邦学习方法也被开发出来[3]、[6]、[9]、[30]、[34]。通过将SAM融入局部训练，这些方法能够探索损失值持续保持在较低水平的区域（平坦区域），从而提升模型的泛化能力。然而，单步梯度上升扰动并不能保证一定能找到真正平坦的区域[15]。模型可能在尖锐的“峡谷”中不断振荡，而无法“爬出”到更广阔、更平坦的“盆地”中。因此，联邦学习需要一种更为有效的扰动和探索机制，以便避开尖锐区域并找到可靠的平坦区域。同时，现有的基于SAM的联邦学习方法主要关注解的平坦性和泛化能力，而对加速收敛的重视程度相对较低。另一类研究则致力于加速联邦学习的全局收敛速度。有些方法采用服务器端加速或自适应优化技术[33]、[37]，而另一些方法则引入由全局信息引导的客户端动量机制[16]、[18]、[35]。这些基于加速的技术虽然可以减少通信轮次，但它们并非专门设计用来控制损失函数的锐度或引导模型向更平坦的区域发展。此外，Kang等人[15]指出，SAM扰动中的动量机制能够平滑扰动方向并积累历史梯度，提高模型脱离尖锐区域的可能性。然而，在本地应用动量机制可能会加剧客户端之间的差异，破坏模型的稳定性，进而影响全局收敛速度和性能[42]。因此，一个关键问题是如何在避免模型不稳定的前提下利用动量机制找到平坦区域，同时还能加快全局收敛速度。为了解决这一问题，我们提出了FedAPEX算法，该算法将动量加速与SAM相结合，用于联邦学习优化。在局部更新过程中，我们利用动量来计算SAM扰动，从而弥补SAM在探索能力方面的不足，帮助模型发现更平坦的区域。此外，我们还设计了全局动量机制来加速下降过程。这一机制能够调整并引导每个客户端的更新，减轻直接使用局部动量机制可能带来的不稳定性，同时加快全局收敛速度（见图1）。我们的贡献如下：•我们提出了FedAPEX，这是一种统一的联邦学习锐度感知优化框架，它将全局对齐的扰动动量与基于SAM的下降动量相结合，从而在数据异构性条件下提升对平坦极小值的探索能力并加速收敛速度。•我们对FedAPEX进行了理论分析：我们确定了其O(1/T)的非凸收敛速率，这一速率与目前已知的联邦学习方法的最佳加速速率相当；同时通过局部更新稳定性分析，我们证明FedAPEX能够在不加剧客户端之间差异的情况下实现这种加速效果，确保训练过程的稳定性与高效性。•通过在多种数据集和算法上的大量实验，我们证明FedAPEX的整体性能优于现有的各种基准方法。具体而言，它在多种任务中都能展现出较强的性能或更快的收敛速度，同时还能获得更高的最终测试精度，并引导模型进入更平坦的损失函数区域。

相关工作：联邦学习允许多个客户端在不共享私有数据的情况下共同训练全局模型。McMahan等人[25]首次提出了联邦学习以及FedAvg算法，该算法通过多次局部更新和全局平均操作实现了高效的通信式训练。然而，由于客户端之间存在数据异构性（非独立同分布数据），会导致“客户端偏差”，进而减缓全局收敛速度并降低模型的泛化能力。为了解决这一问题，FedProx[21]引入了近似项，用以限制局部更新之间的差异。

问题设定与符号说明：我们考虑一个标准的联邦学习场景，其中包含n个客户端（用k标识），这些客户端共同训练一个全局模型。每个客户端k都拥有一个私有的本地数据集Dk。联邦优化的目标是最小化全局损失函数minwf(w)=1n∑k=1nfk(w)，其中fk(wk)=E(x,y)～Dk[?k((x,y);wk)]表示与客户端k的本地数据集相关的经验损失函数，wk则是该客户端的本地模型参数。由于受到计算和通信条件的限制……

理论分析：我们从两个互补的角度对FedAPEX进行了分析。首先，我们研究了在部分客户端参与情况下它的非凸收敛速率，结果表明双重动量机制能够保持联邦学习方法已知的最佳加速速率。接着，我们分析了在数据异构性条件下的局部更新稳定性，结果表明额外的动量机制并不会比FedSAM造成更大的客户端间差异。综合这些结果可以看出，FedAPEX同时满足了两种已知的理论保障条件。

实验设置：为了验证我们所提出的FedAPEX算法的有效性，我们在四个基准数据集上进行了实验，并在数据分布复杂且客户端参与度有限的情况下，将其与多种联邦学习基准方法进行了比较。所有实验都在一台配备英特尔至强E5 - 2698 v4处理器、英伟达GeForce RTX 3090显卡以及256GB内存的工作站上完成，实验所用的软件为PyTorch 2.4.1，配套的CUDA版本为12.1，cuDNN版本为9.1.0。除了文中始终使用的主要实验设置外，我们还有……

结论：在本文中，我们提出了FedAPEX这种联邦学习算法，它不仅能够提升模型的泛化能力，通常还能实现更快或与现有方法相当的收敛速度。通过整合双重动量机制——一种用于引导基于锐度感知的扰动，另一种用于推动全局对齐的优化步骤——FedAPEX能够有效地在数据异构性较强的联邦学习环境中探索复杂的损失函数曲面。我们的大量实验表明，FedAPEX的整体性能优于现有的各种方法。

CRediT作者贡献说明：徐远：撰写——初稿、可视化、验证、软件、方法论、概念构思。陈学松：撰写——审阅与编辑、可视化、软件、方法论、概念构思。马丽燕：撰写——审阅与编辑、方法论、概念构思。李芳：撰写——审阅与编辑、监督、资源协调、项目管理、方法论、资金获取、概念构思。

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

致谢：本研究得到了上海市自然科学基金（项目编号22ZR1419500）以及上海市科学技术委员会（项目编号22DZ2229014）的支持。徐远于2020年获得四川大学学士学位，目前就读于华东师范大学数学科学学院，攻读硕士学位。他的研究兴趣包括联邦学习和差分隐私。

徐远|陈学松|马丽燕|李芳
华东师范大学数学科学学院，中国上海200241