Gradient and Variable Tracking with Multiple Local SGD for Decentralized Non-Convex Learning (2302.01537v1)

Abstract: Stochastic distributed optimization methods that solve an optimization problem over a multi-agent network have played an important role in a variety of large-scale signal processing and machine leaning applications. Among the existing methods, the gradient tracking (GT) method is found robust against the variance between agents' local data distribution, in contrast to the distributed stochastic gradient descent (SGD) methods which have a slowed convergence speed when the agents have heterogeneous data distributions. However, the GT method can be communication expensive due to the need of a large number of iterations for convergence. In this paper, we intend to reduce the communication cost of the GT method by integrating it with the local SGD technique. Specifically, we propose a new local stochastic GT (LSGT) algorithm where, within each communication round, the agents perform multiple SGD updates locally. Theoretically, we build the convergence conditions of the LSGT algorithm and show that it can have an improved convergence rate of $\mathcal{O}(1/\sqrt{ET})$, where $E$ is the number of local SGD updates and $T$ is the number of communication rounds. We further extend the LSGT algorithm to solve a more complex learning problem which has linearly coupled variables inside the objective function. Experiment results demonstrate that the proposed algorithms have significantly improved convergence speed even under heterogeneous data distribution.