On Convergence of FedProx: Local Dissimilarity Invariant Bounds, Non-smoothness and Beyond

TL;DR

This paper develops a new convergence theory for FedProx in federated learning that is invariant to local dissimilarity and applicable to non-smooth problems, providing deeper insights and practical benefits.

Contribution

It introduces a local dissimilarity invariant convergence analysis for FedProx, extending understanding to non-smooth and non-convex federated optimization.

Findings

Convergence guarantees independent of local dissimilarity conditions

Applicable to non-smooth federated learning problems

Linear speedup with minibatch size and number of devices

Abstract

The FedProx algorithm is a simple yet powerful distributed proximal point optimization method widely used for federated learning (FL) over heterogeneous data. Despite its popularity and remarkable success witnessed in practice, the theoretical understanding of FedProx is largely underinvestigated: the appealing convergence behavior of FedProx is so far characterized under certain non-standard and unrealistic dissimilarity assumptions of local functions, and the results are limited to smooth optimization problems. In order to remedy these deficiencies, we develop a novel local dissimilarity invariant convergence theory for FedProx and its minibatch stochastic extension through the lens of algorithmic stability. As a result, we contribute to derive several new and deeper insights into FedProx for non-convex federated optimization including: 1) convergence guarantees independent on local dissimilarity type conditions; 2) convergence guarantees for non-smooth FL problems; and 3) linear speedup with respect to size of minibatch and number of sampled devices. Our theory for the first time reveals that local dissimilarity and smoothness are not must-have for FedProx to get favorable complexity bounds. Preliminary experimental results on a series of benchmark FL datasets are reported to demonstrate the benefit of minibatching for improving the sample efficiency of FedProx.