STL-SGD: Speeding Up Local SGD with Stagewise Communication Period

Authors

  • Shuheng Shen School of Computer Science and Technology, University of Science and Technology of China Ant Financial Services Group
  • Yifei Cheng School of Data Science, University of Science and Technology of China
  • Jingchang Liu Department of Computer Science and Engineering, Hong Kong University of Science and Technology
  • Linli Xu School of Computer Science and Technology, University of Science and Technology of China IFLYTEK Co., Ltd.

Keywords:

Distributed Machine Learning & Federated Learning, Optimization, Constraint Optimization

Abstract

Distributed parallel stochastic gradient descent algorithms are workhorses for large scale machine learning tasks. Among them, local stochastic gradient descent (Local SGD) has attracted significant attention due to its low communication complexity. Previous studies prove that the communication complexity of Local SGD with a fixed or an adaptive communication period is in the order of O (N3/2 T1/2) and O (N3/4 T3/4) when the data distributions on clients are identical (IID) or otherwise (Non-IID), where N is the number of clients and T is the number of iterations. In this paper, to accelerate the convergence by reducing the communication complexity, we propose STagewise Local SGD (STL-SGD), which increases the communication period gradually along with decreasing learning rate. We prove that STL-SGD can keep the same convergence rate and linear speedup as mini-batch SGD. In addition, as the benefit of increasing the communication period, when the objective is strongly convex or satisfies the Polyak-Lojasiewicz condition, the communication complexity of STL-SGD is O (N log T ) and O (N1/2 T1/2) for the IID case and the Non-IID case respectively, achieving significant improvements over Local SGD. Experiments on both convex and non-convex problems demonstrate the superior performance of STL-SGD.

Downloads

Published

2021-05-18

How to Cite

Shen, S., Cheng, Y., Liu, J., & Xu, L. (2021). STL-SGD: Speeding Up Local SGD with Stagewise Communication Period. Proceedings of the AAAI Conference on Artificial Intelligence, 35(11), 9576-9584. Retrieved from https://ojs.aaai.org/index.php/AAAI/article/view/17153

Issue

Section

AAAI Technical Track on Machine Learning IV