Escaping saddle points in heterogeneous federated learning via distributed SGD with communication compression

• Main Authors: CHEN, Sijin, LI, Zhize, CHI, Yuejie
• Format: text
• Language: English
• Published: Institutional Knowledge at Singapore Management University 2024
• Subjects: Federated learning; Machine learning; Communication compression; Artificial Intelligence and Robotics
• Online Access: https://ink.library.smu.edu.sg/sis_research/9493; https://ink.library.smu.edu.sg/context/sis_research/article/10493/viewcontent/1_s2.0_S2214140524001233_pvoa_cc_by_nc.pdf
• Institution: Singapore Management University

We consider the problem of finding second-order stationary points in the optimization of heterogeneous federated learning (FL). Previous works in FL mostly focus on first-order convergence guarantees, which do not rule out the scenario of unstable saddle points. Meanwhile, it is a key bottleneck of FL to achieve communication efficiency without compensating the learning accuracy, especially when local data are highly heterogeneous across different clients. Given this, we propose a novel algorithm PowerEF-SGD that only communicates compressed information via a novel error-feedback scheme. To our knowledge, PowerEF-SGD is the first distributed and compressed SGD algorithm that provably escapes saddle points in heterogeneous FL without any data homogeneity assumptions. In particular, PowerEF-SGD improves to second-order stationary points after visiting first-order (possibly saddle) points, using additional gradient queries and communication rounds only of almost the same order required by first-order convergence, and the convergence rate shows a linear-speedup pattern in terms of the number of workers. Our theory improves/recovers previous results, while extending to much more tolerant settings on the local data. Numerical experiments are provided to complement the theory.