Achieving Linear Speedup for Composite Federated Learning

TL;DR

This paper introduces FedNMap, a novel federated learning method that achieves linear speedup for nonconvex composite problems by effectively handling nonsmooth regularizers and data heterogeneity.

Contribution

FedNMap is the first method to establish linear speedup in nonconvex composite federated learning, combining normal map updates with local correction strategies.

Findings

Achieves linear speedup with respect to number of clients and local updates.

Handles nonsmooth regularizers in federated learning.

Works under standard assumptions including weak convexity and bounded variance.

Abstract

This paper proposes FedNMap, a normal map-based method for composite federated learning, where the objective consists of a smooth loss and a possibly nonsmooth regularizer. FedNMap leverages a normal map-based update scheme to handle the nonsmooth term and incorporates a local correction strategy to mitigate the impact of data heterogeneity across clients. Under standard assumptions, including smooth local losses, weak convexity of the regularizer, and bounded stochastic gradient variance, FedNMap achieves linear speedup with respect to both the number of clients $n$ and the number of local updates $Q$ for nonconvex losses, both with and without the Polyak-{\L}ojasiewicz (PL) condition. To our knowledge, this is the first result establishing linear speedup for nonconvex composite federated learning.