[论文解读] Reducing the Communication Cost of Federated Learning through Multistage Optimization.

Charlie Hou, Kiran Koshy Thekumparampil+2 more|arXiv (Cornell University)|Aug 16, 2021

Privacy-Preserving Technologies in Data参考文献 28被引用 6

一句话总结

本文提出了一种用于联邦学习的多阶段优化框架，通过先应用本地优化方法直至达到由数据异质性引起的误差下限，然后切换到集中式方法以实现最终收敛，从而降低通信成本。该方法在所有数据异质性水平下几乎匹配理论通信复杂度下限，为现有方法提供了一种实用且理论基础扎实的替代方案。

ABSTRACT

A central question in federated learning (FL) is how to design optimization algorithms that minimize the communication cost of training a model over heterogeneous data distributed across many clients. A popular technique for reducing communication is the use of local steps, where clients take multiple optimization steps over local data before communicating with the server (e.g., FedAvg, SCAFFOLD). This contrasts with centralized methods, where clients take one optimization step per communication round (e.g., Minibatch SGD). A recent lower bound on the communication complexity of first-order methods shows that centralized methods are optimal over highly-heterogeneous data, whereas local methods are optimal over purely homogeneous data [Woodworth et al., 2020]. For intermediate heterogeneity levels, no algorithm is known to match the lower bound. In this paper, we propose a multistage optimization scheme that nearly matches the lower bound across all heterogeneity levels. The idea is to first run a local method up to a heterogeneity-induced error floor; next, we switch to a centralized method for the remaining steps. Our analysis may help explain empirically-successful stepsize decay methods in FL [Charles et al., 2020; Reddi et al., 2020]. We demonstrate the scheme's practical utility in image classification tasks.

研究动机与目标

• 为解决联邦学习中的通信成本挑战，特别是在尚无最优方法的中等数据异质性条件下。
• 弥合通信复杂度理论下限与实际联邦优化算法之间的差距。
• 设计一种结合本地与集中训练阶段的混合优化策略，以适应数据异质性。
• 为联邦学习中经验上成功的步长衰减策略提供理论依据。

提出的方法

• 该方法从本地优化阶段开始（例如，类似FedAvg的方法），客户端在与服务器通信前执行多次本地更新。
• 基于固有的数据分布偏移，将异质性引起的误差下限作为本地阶段的停止准则。
• 在达到误差下限后，算法切换到集中式优化方法（例如，小批量SGD）以完成最终收敛步骤。
• 当本地更新不再显著降低优化误差时，触发从本地训练到集中式训练的切换。
• 理论分析表明，这种两阶段方法在所有异质性水平下几乎匹配已知的通信复杂度下限。
• 该框架设计为与现有联邦学习算法兼容，可应用于标准训练流程。

实验结果

研究问题

• RQ1 结合本地与集中式方法的混合优化策略是否能在所有数据异质性水平下实现接近最优的通信复杂度？
• RQ2 在最小化通信成本的前提下，从本地训练切换到集中式训练的最佳时机是什么？
• RQ3 与现有算法（如FedAvg和SCAFFOLD）相比，该方法在收敛性和通信效率方面表现如何？
• RQ4 该多阶段方案的理论分析是否能解释联邦学习中步长衰减策略的经验成功？
• RQ5 该方法在图像分类等实际任务中，在现实数据异质性条件下是否保持性能？

主要发现

• 该多阶段优化方法在所有数据异质性水平下（从同质到高度异质）几乎匹配理论通信复杂度下限。
• 在中等异质性水平下，该方法在通信效率方面优于标准的本地方法（如FedAvg）和集中式方法（如小批量SGD）。
• 在误差下限点从本地训练切换到集中式训练，显著减少了达到收敛所需的通信轮数。
• 理论分析支持实践中观察到的步长衰减策略的有效性，为其成功提供了有原则的解释。
• 在图像分类任务上的实证评估证实了该方案的实际效用和通信成本的降低。