Non-convex composite federated learning with heterogeneous data

TL;DR

This paper introduces a novel federated learning algorithm for non-convex, non-smooth problems that reduces communication costs and handles data heterogeneity effectively, with proven convergence guarantees and superior empirical performance.

Contribution

It proposes a decoupled, communication-efficient federated learning algorithm for non-convex composite problems with theoretical convergence analysis.

Findings

Outperforms state-of-the-art methods on synthetic datasets.

Achieves sublinear and linear convergence under general conditions.

Effectively handles client drift and data heterogeneity.

Abstract

We propose an innovative algorithm for non-convex composite federated learning that decouples the proximal operator evaluation and the communication between server and clients. Moreover, each client uses local updates to communicate less frequently with the server, sends only a single d-dimensional vector per communication round, and overcomes issues with client drift. In the analysis, challenges arise from the use of decoupling strategies and local updates in the algorithm, as well as from the non-convex and non-smooth nature of the problem. We establish sublinear and linear convergence to a bounded residual error under general non-convexity and the proximal Polyak-Lojasiewicz inequality, respectively. In the numerical experiments, we demonstrate the superiority of our algorithm over state-of-the-art methods on both synthetic and real datasets.