[Paper Review] Federated Optimization in Heterogeneous Networks

Li Tian, Anit Kumar Sahu+4 more|arXiv (Cornell University)|Dec 14, 2018

Privacy-Preserving Technologies in Data44 references451 citations

TL;DR

FedProx generalizes FedAvg to handle both system and statistical heterogeneity in federated learning by introducing a proximal term and allowing variable local updates, with convergence guarantees and empirically robust performance.

ABSTRACT

Federated Learning is a distributed learning paradigm with two key challenges that differentiate it from traditional distributed optimization: (1) significant variability in terms of the systems characteristics on each device in the network (systems heterogeneity), and (2) non-identically distributed data across the network (statistical heterogeneity). In this work, we introduce a framework, FedProx, to tackle heterogeneity in federated networks. FedProx can be viewed as a generalization and re-parametrization of FedAvg, the current state-of-the-art method for federated learning. While this re-parameterization makes only minor modifications to the method itself, these modifications have important ramifications both in theory and in practice. Theoretically, we provide convergence guarantees for our framework when learning over data from non-identical distributions (statistical heterogeneity), and while adhering to device-level systems constraints by allowing each participating device to perform a variable amount of work (systems heterogeneity). Practically, we demonstrate that FedProx allows for more robust convergence than FedAvg across a suite of realistic federated datasets. In particular, in highly heterogeneous settings, FedProx demonstrates significantly more stable and accurate convergence behavior relative to FedAvg---improving absolute test accuracy by 22% on average.

Motivation & Objective

• Motivate federated learning under both systems and statistical heterogeneity.
• Introduce FedProx as a generalization of FedAvg with a proximal term.
• Provide convergence guarantees under non-IID data and partial device participation.
• Demonstrate practical robustness and improved accuracy in heterogeneous settings.

Proposed method

• Present FedProx as a generalization of FedAvg that allows variable local work across devices.
• Add a proximal term to local subproblems: h_k(w; w^t) = F_k(w) + (μ/2)||w - w^t||^2.
• Introduce and utilize γ_k^t-inexactness to model inexact local solves with varying accuracy.
• Define B-local dissimilarity to quantify statistical heterogeneity across devices.
• Prove convergence guarantees under bounded dissimilarity and non-IID data.
• Show that FedProx reduces sensitivity to stragglers and improves stability.

Experimental results

Research questions

• RQ1 How does heterogeneity (both systems and data) affect convergence in federated optimization?
• RQ2 Can a proximal term and variable local work enable provable convergence under non-IID data and partial device participation?
• RQ3 What are the theoretical conditions ensuring sufficient descent in FedProx with heterogeneous devices?
• RQ4 How does FedProx compare to FedAvg in terms of robustness and accuracy under realistic heterogeneity?

Key findings

• FedProx improves stability and accuracy in highly heterogeneous settings, with absolute test accuracy increasing by 22% on average.
• The proximal term helps control the impact of variable local updates and aids convergence under heterogeneity.
• Convergence guarantees are provided for non-convex local objectives under bounded dissimilarity.
• The framework accommodates variable amounts of local work and still yields guaranteed decrease per round.
• FedProx generalizes FedAvg, which is recovered as a special case with μ = 0 and uniform γ.
• Empirical results on synthetic and real federated datasets validate robustness to stragglers and statistical heterogeneity.