[论文解读] FedPAGE: A Fast Local Stochastic Gradient Method for Communication-Efficient Federated Learning

Haoyu Zhao, Zhize Li+1 more|arXiv (Cornell University)|Aug 10, 2021

Stochastic Gradient Optimization Techniques参考文献 43被引用 8

一句话总结

FedPAGE 提出了一种通信高效的联邦学习算法，用最优的 PAGE 随机梯度方法替代 FedAvg 中的标准本地 SGD，实现了最先进的通信复杂度：在凸问题中为 $O(N^{3/4}/(S heta))$，在非凸问题中为 $O((\sqrt{N}+S)/(S heta^2))$，相较于 SCAFFOLD 分别提升了 $N^{1/4}$ 和 $N^{1/6}S^{1/3}$ 的因子，在典型的客户端采样设置下表现更优。

ABSTRACT

Federated Averaging (FedAvg, also known as Local-SGD) (McMahan et al., 2017) is a classical federated learning algorithm in which clients run multiple local SGD steps before communicating their update to an orchestrating server. We propose a new federated learning algorithm, FedPAGE, able to further reduce the communication complexity by utilizing the recent optimal PAGE method (Li et al., 2021) instead of plain SGD in FedAvg. We show that FedPAGE uses much fewer communication rounds than previous local methods for both federated convex and nonconvex optimization. Concretely, 1) in the convex setting, the number of communication rounds of FedPAGE is $O(\frac{N^{3/4}}{Sε})$, improving the best-known result $O(\frac{N}{Sε})$ of SCAFFOLD (Karimireddy et al.,2020) by a factor of $N^{1/4}$, where $N$ is the total number of clients (usually is very large in federated learning), $S$ is the sampled subset of clients in each communication round, and $ε$ is the target error; 2) in the nonconvex setting, the number of communication rounds of FedPAGE is $O(\frac{\sqrt{N}+S}{Sε^2})$, improving the best-known result $O(\frac{N^{2/3}}{S^{2/3}ε^2})$ of SCAFFOLD (Karimireddy et al.,2020) by a factor of $N^{1/6}S^{1/3}$, if the sampled clients $S\leq \sqrt{N}$. Note that in both settings, the communication cost for each round is the same for both FedPAGE and SCAFFOLD. As a result, FedPAGE achieves new state-of-the-art results in terms of communication complexity for both federated convex and nonconvex optimization.

研究动机与目标

• 通过改进标准 FedAvg 算法，降低联邦学习中的通信复杂度。
• 解决在异构、地理分布广泛的客户端之间去中心化、隐私保护训练中的通信瓶颈问题。
• 在凸与非凸优化设置下，实现更少通信轮次的更快收敛，同时保持模型精度。
• 利用最优的 PAGE 方法进行本地更新，实现优于先前本地方法（如 SCAFFOLD）的收敛速率。

提出的方法

• 用 PAGE（带误差校正的近端加速梯度）方法替代 FedAvg 中的本地 SGD 步骤，以提高收敛效率。
• 采用混合本地更新策略，结合方差减少的梯度估计与自适应步长，以降低噪声并加速收敛。
• 在每轮中采用客户端采样，采样子集大小为 $S$，其中 $S \leq \sqrt{N}$ 在实践中较为典型，以平衡通信成本与模型精度。
• 在光滑性与有界方差假设下，提出一种新颖的分析框架，推导出紧致的通信复杂度边界。
• 推导出全局与本地更新的最优步长 $\eta_g$ 与 $\eta_l$，以最小化通信轮次数。
• 应用李雅普诺夫风格分析，界定期望次优性与梯度范数，从而在凸与非凸设置下均提供收敛保证。

实验结果

研究问题

• RQ1 通过用更先进的随机梯度方法替代本地 SGD，能否显著降低 FedAvg 的通信复杂度？
• RQ2 在凸与非凸问题中，FedPAGE 与 SCAFFOLD 相比，在通信轮次方面表现如何？
• RQ3 在标准光滑性与有界方差假设下，FedPAGE 的理论通信复杂度是多少？
• RQ4 当每轮仅采样极少数客户端（$S \ll N$）时，FedPAGE 是否仍保持其优势？
• RQ5 FedPAGE 是否能在大规模联邦学习设置中，以更少的通信轮次实现更快收敛，同时保持模型精度？

主要发现

• 在凸设置下，FedPAGE 实现了 $O(N^{3/4}/(S\theta))$ 的通信复杂度，相较于 SCAFFOLD 的 $O(N/(S\theta))$ 提升了 $N^{1/4}$ 倍。
• 在非凸设置下，FedPAGE 将通信轮次减少至 $O((\sqrt{N}+S)/(S\theta^2))$，当 $S \leq \sqrt{N}$ 时，相较于 SCAFFOLD 的 $O(N^{2/3}/(S^{2/3}\theta^2))$ 提升了 $N^{1/6}S^{1/3}$ 倍。
• 该改进在与 SCAFFOLD 相同的每轮通信成本下实现，意味着 FedPAGE 在不增加每轮通信负载的前提下实现了更优性能。
• FedPAGE 即使在常见的实际联邦学习场景中常见的有界方差（BV）假设下，仍保持其理论优势。
• FedPAGE 在凸与非凸问题上的通信复杂度收敛速率均为最优，确立了新的最先进水平。
• 该方法对客户端采样具有鲁棒性，在每轮仅选择少量客户端 $S$ 的情况下仍表现良好，这在现实联邦学习中是典型场景。