Federated Learning Under Second-Order Data Heterogeneity

Konstantin Mishchenko, Rui Li, Hongxiang Fan, Stylianos Venieris

Primary Area: optimization

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: federated learning, data heterogeneity, optimization, theory

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

TL;DR: We study the problem of nonconvex federated optimization under second-order data heterogeneity (bounded Hessian differences) and present an algorithm called SABER

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 nonconvex problems is $O\left(\delta\varepsilon^2\sqrt{M}\right)$. In addition, for problems satisfying $\mu$-Polyak-Lojasiewicz condition, the method converges linearly with communication complexity of $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 such as FedAvg on a few empirical problems.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 7609