STEM: A Stochastic Two-Sided Momentum Algorithm Achieving Near-Optimal Sample and Communication Complexities for Federated Learning

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 $\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.

References in corpus (18)

• Communication-Efficient Learning of Deep Networks from Decentralized Data
• Federated Learning with Non-IID Data
• Federated Optimization: Distributed Machine Learning for On-Device Intelligence
• On the Convergence of Adam and Beyond
• Federated Optimization in Heterogeneous Networks
• Don't Use Large Mini-Batches, Use Local SGD
• Local SGD Converges Fast and Communicates Little
• Cooperative SGD: A unified Framework for the Design and Analysis of Communication-Efficient SGD Algorithms
• On the Linear Speedup Analysis of Communication Efficient Momentum SGD for Distributed Non-Convex Optimization
• Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning
• Variance Reduced Local SGD with Lower Communication Complexity
• SlowMo: Improving Communication-Efficient Distributed SGD with Slow Momentum
• Is Local SGD Better than Minibatch SGD?
• Minibatch vs Local SGD for Heterogeneous Distributed Learning
• On the Computation and Communication Complexity of Parallel SGD with Dynamic Batch Sizes for Stochastic Non-Convex Optimization
• Parallel Restarted SPIDER -- Communication Efficient Distributed Nonconvex Optimization with Optimal Computation Complexity
• Faster Non-Convex Federated Learning via Global and Local Momentum
• Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learning

Cited by in corpus (1)

• Federated Learning Survey: A Multi-Level Taxonomy of Aggregation Techniques, Experimental Insights, and Future Frontiers