Federated Optimization for Heterogeneous Networks

Federated learning involves training machine learning models in massively distributed networks. While Federated Averaging (FedAvg) is the leading optimization method for training non-convex models in this setting, its behavior is not well understood in realistic federated settings when learning across statistically heterogeneous devices, i.e., where each device collects data in a non-identical fashion. In this work, we introduce a framework to tackle statistical heterogeneity, FedProx, which encompasses FedAvg as a special case. We provide convergence guarantees for FedProx through a device dissimilarity assumption that allows us to characterize heterogeneity in the network. Finally, we perform a detailed empirical evaluation across a suite of federated datasets, validating our theoretical analysis and demonstrating the improved robustness and stability of the generalized FedProx framework relative to FedAvg for learning in heterogeneous networks.