[论文解读] STEM: A Stochastic Two-Sided Momentum Algorithm Achieving Near-Optimal Sample and Communication Complexities for Federated Learning

Prashant Khanduri, Pranay Sharma+5 more|arXiv (Cornell University)|Jun 19, 2021

Privacy-Preserving Technologies in Data参考文献 35被引用 20

一句话总结

该论文提出了一种用于联邦学习的随机双向动量算法 STEM，其在非凸设置下可实现近似最优的样本复杂度 $\tilde{\mathcal{O}}(\epsilon^{-3/2})$ 和通信复杂度 $\tilde{\mathcal{O}}(\epsilon^{-1})$，以寻找 $\epsilon$-平稳解。通过联合优化工作节点与服务器端基于动量的更新、自适应小批量大小和本地更新频率，STEM 建立了一条权衡曲线，维持了这些复杂度，优于 FedAvg 和先前的方法。

ABSTRACT

Federated Learning (FL) refers to the paradigm where multiple worker nodes (WNs) build a joint model by using local data. Despite extensive research, for a generic non-convex FL problem, it is not clear, how to choose the WNs' and the server's update directions, the minibatch sizes, and the local update frequency, so that the WNs use the minimum number of samples and communication rounds to achieve the desired solution. This work addresses the above question and considers a class of stochastic algorithms where the WNs perform a few local updates before communication. We show that when both the WN's and the server's directions are chosen based on a stochastic momentum estimator, the algorithm requires $ ilde{\mathcal{O}}(ε^{-3/2})$ samples and $ ilde{\mathcal{O}}(ε^{-1})$ communication rounds to compute an $ε$-stationary solution. To the best of our knowledge, this is the first FL algorithm that achieves such {\it near-optimal} sample and communication complexities simultaneously. Further, we show that there is a trade-off curve between local update frequencies and local minibatch sizes, on which the above sample and communication complexities can be maintained. Finally, we show that for the classical FedAvg (a.k.a. Local SGD, which is a momentum-less special case of the STEM), a similar trade-off curve exists, albeit with worse sample and communication complexities. Our insights on this trade-off provides guidelines for choosing the four important design elements for FL algorithms, the update frequency, directions, and minibatch sizes to achieve the best performance.

研究动机与目标

• 为解决在非凸联邦学习中联合优化工作节点与服务器端更新方向、小批量大小和本地更新频率这一开放问题。
• 在具有异构数据的随机去中心化设置下，同时最小化样本复杂度与通信复杂度。
• 建立本地更新频率与小批量大小之间的理论权衡曲线，以维持近似最优的收敛速率。
• 通过揭示 FedAvg（Local SGD）的内在权衡及其次优复杂度边界，为 FedAvg 提供新的洞见。
• 证明在标准假设下，工作节点与服务器端均采用动量更新可实现最优收敛。

提出的方法

• 提出 STEM，一种随机双向动量算法，其中工作节点和服务器均使用动量估计的梯度进行模型更新。
• 引入一个统一框架，联合优化本地更新频率 $I$、小批量大小 $b$ 和基于动量的更新方向。
• 采用通过动量实现方差减少的随机梯度估计器，以稳定训练并降低梯度噪声。
• 在非凸光滑性和有界梯度假设下，利用李雅普诺夫分析和类似强凸性的不等式推导收敛边界。
• 建立 $I$ 与 $b$ 之间的权衡曲线，使得 $\tilde{\mathcal{O}}(\epsilon^{-3/2})$ 的样本量和 $\tilde{\mathcal{O}}(\epsilon^{-1})$ 的通信轮次得以维持。
• 将 FedAvg 视为 STEM 的特例（无动量），并推导其复杂度边界以供比较。

实验结果

研究问题

• RQ1 联邦学习算法是否能在非凸设置下同时实现近似最优的样本复杂度与通信复杂度？
• RQ2 在具有随机梯度的联邦学习中，本地更新频率与小批量大小之间的最优权衡是什么？
• RQ3 在非凸联邦学习中，同时在工作节点与服务器端引入动量对收敛复杂度有何影响？
• RQ4 在与 STEM 相同的假设下，FedAvg（Local SGD）的理论样本复杂度与通信复杂度是多少？
• RQ5 能否开发一个统一框架，以联合优化更新方向、小批量大小和本地更新频率，从而实现最小资源消耗？

主要发现

• STEM 在非凸联邦学习中实现了 $\tilde{\mathcal{O}}(\epsilon^{-3/2})$ 的样本复杂度和 $\tilde{\mathcal{O}}(\epsilon^{-1})$ 的通信复杂度，以寻找 $\epsilon$-平稳解。
• 存在一条 $b$ 与 $I$ 之间的权衡曲线，当 $b$ 和 $I$ 沿此曲线选择时，STEM 能够维持近似最优的复杂度。
• 对于 FedAvg（STEM 的无动量变体），样本复杂度为 $\mathcal{O}(\epsilon^{-2})$，通信复杂度为 $\mathcal{O}(\epsilon^{-3/2})$，相比 STEM 显然是次优的。
• 当 $b = \mathcal{O}(1)$ 且 $I = (T/b^3K^3)^{1/4}$ 时，STEM 仍能维持 $\tilde{\mathcal{O}}(\epsilon^{-3/2})$ 的样本复杂度和 $\tilde{\mathcal{O}}(\epsilon^{-1})$ 的通信复杂度。
• 分析表明，工作节点与服务器端均使用动量是实现在随机联邦学习中最优收敛速率的关键因素。
• 在标准假设下，所推导的边界是紧致的，且 STEM 是首个在该设置下同时实现近似最优样本与通信复杂度的算法。