STEM: A Stochastic Two-Sided Momentum Algorithm Achieving Near-Optimal Sample and Communication Complexities for Federated Learning
STEM: A Stochastic Two-Sided Momentum Algorithm Achieving Near-Optimal Sample and Communication Complexities for Federated Learning
复制标题
DOI：
--
复制
发表时间：
2021-06
期刊：
影响因子：
--
通讯作者：
Prashant Khanduri;Pranay Sharma;Haibo Yang;Min-Fong Hong;Jia Liu;K. Rajawat;P. Varshney
中科院分区：
文献类型：
--
作者：
Prashant Khanduri;Pranay Sharma;Haibo Yang;Min-Fong Hong;Jia Liu;K. Rajawat;P. Varshney
联邦学习（FL）是指多个工作节点（WN）通过使用本地数据建立联合模型的范例。尽管广泛的研究，对于一般的非凸FL问题，它是不清楚的，如何选择WN的和服务器的更新方向，小批量的大小，和本地更新频率，使WN使用最少数量的样本和通信轮，以实现所需的解决方案。这项工作解决了上述问题，并考虑了一类随机算法，其中WN在通信之前执行一些本地更新。我们表明，当WN的和服务器的方向选择的基础上随机动量估计，该算法需要$\tilde{\mathcal{O}}（\displaystyle ^{-3/2}）$样本和$\tilde{\mathcal{O}}（\displaystyle ^{-1}）$通信轮计算$\tilde $-固定的解决方案。据我们所知，这是第一个FL算法，同时实现这样的{\it near-optimal}样本和通信复杂性。此外，我们发现，有一个权衡曲线之间的本地更新频率和本地minibatch大小，上面的样本和通信的复杂性可以保持。最后，我们表明，对于经典的FedAvg（a.k.a.局部SGD，这是STEM的动量较少的特殊情况），存在类似的权衡曲线，尽管具有更差的样本和通信复杂性。我们对这种权衡的见解为选择FL算法的四个重要设计元素提供了指导方针，即更新频率，方向和minibatch大小，以实现最佳性能。
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 $\tilde{\mathcal{O}}(\epsilon^{-3/2})$ samples and $\tilde{\mathcal{O}}(\epsilon^{-1})$ communication rounds to compute an $\epsilon$-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.