Federated Learning Under Second-Order Data Heterogeneity
Abstract: We consider the problem of Federated Learning over clients with heterogeneous data. We propose an algorithm called SABER that samples a subset of clients and tasks each client with its own local subproblem. SABER provably reduces client drift by incorporating an estimate of the global update direction and regularization into each client's subproblem. Under second-order data heterogeneity with parameter $\delta$, we prove that the method's communication complexity for non-convex problems is $\mathcal{O}\left(\delta\varepsilon^{-2}\sqrt{M}\right)$. In addition, for problems satisfying $\mu$-Polyak-\L{}ojasiewicz condition, the method converges linearly with communication complexity of $\mathcal{O}\left(\left(\frac{\delta}{\mu}\sqrt{M} + M\right)\log\frac{1}{\varepsilon}\right)$. To showcase the empirical performance of our method, we compare it to standard baselines including FedAvg, FedProx, and SCAFFOLD on image classification problems and demonstrate its superior performance in data-heterogeneous settings.
Submission Length: Regular submission (no more than 12 pages of main content)
Assigned Action Editor: ~Zachary_B._Charles1
Submission Number: 2494
Loading