PaperlayerView exact revision on arXiv

Source-linked AI summary

Federated Optimization in Heterogeneous Networks

Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, Virginia Smith

arXiv:1812.06127v5cs.LGstat.ML

TL;DR

Federated learning must handle systems and statistical heterogeneity, while FedAvg lacks flexible local work and realistic convergence guarantees. FedProx addresses these challenges with a proximal framework and improves convergence, including a 22% average absolute test-accuracy gain over FedAvg in highly heterogeneous settings.

  • Problem

    FedAvg does not support variable local work and lacks convergence guarantees for heterogeneous, non-identically distributed federated data.

  • Method

    FedProx generalizes FedAvg with a proximal term, permits variable local work, and provides convergence guarantees under a device dissimilarity assumption.

  • Results

    FedProx significantly improves convergence behavior across federated datasets, improving absolute testing accuracy relative to FedAvg by 22% on average in highly heterogeneous environments.

  • Takeaways & Limitations

    FedProx offers more stable and accurate training in realistic heterogeneous networks, including settings with stragglers.

  • Takeaways & Limitations

    FedProx’s convergence analysis requires a sufficient but not necessary condition on its proximal parameter, leaving other convergent settings unresolved.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Federated learning must address both systems heterogeneity and statistical heterogeneity, challenges that distinguish it from traditional distributed optimization. FedProx extends FedAvg with a proximal term to support heterogeneous device workloads and improve convergence stability and accuracy.

  • Motivation: Federated learning differs from traditional distributed optimization through high systems and statistical heterogeneity.These challenges arise in networks of remote devices training models.
  • Background: FedAvg became the de facto federated optimization method by combining local SGD updates, low participation, and server-side averaging.Each iteration uses E local SGD epochs on K devices before averaging their updates centrally.
  • Limitations: FedAvg does not let devices vary local work according to systems constraints and commonly drops devices that miss the E-epoch deadline.Its handling of heterogeneity therefore remains incomplete despite empirical success.
  • FedProx: FedProx adds a proximal term that accounts for partial information and heterogeneity, providing convergence guarantees while improving method stability and overall accuracy.The framework addresses systems and statistical heterogeneity theoretically and empirically.

2 BACKGROUND AND RELATED WORK

Federated optimization builds on distributed methods but must address both statistical and systems heterogeneity. Prior approaches leave gaps in realistic, nonconvex networks, motivating FedProx’s broader framework and convergence analysis.

  • Distributed optimization methods developed for large-scale machine learning are increasingly challenged by federated environments involving phones, sensors, and wearable devices.
  • FedAvg works empirically for nonconvex problems by averaging devices’ local SGD updates, but its local updating, partial participation, and heterogeneous data make analysis difficult.
  • Existing convergence analyses for statistical heterogeneity often assume full device participation and specific on-device solvers, limiting realism and solver flexibility.
  • Systems heterogeneity arises from differences in hardware, connectivity, and power, while dropping constrained devices can reduce effective participation and induce sampling bias.
  • FedProx extends FedAvg to heterogeneous federated environments while retaining similar privacy and computational benefits, analyzing convergence through statistical dissimilarity and systems constraints.Its analysis uses a statistical dissimilarity characterization between local functions.

3 FEDERATED OPTIMIZATION: METHODS

This section formulates federated optimization over potentially non-identical, non-convex local objectives and contrasts FedAvg with FedProx. FedProx adds a proximal mechanism and permits variable local work to address statistical and systems heterogeneity while retaining FedAvg as a special case.

  • Federated Optimization: Federated learning minimizes a global objective assembled from local empirical risks over potentially differing data distributions and non-convex objectives.Local risks are defined over device-specific distributions D_k, with weights p_k that can be set from device sample counts n_k.
  • FedAvg: FedAvg selects K ≪ N devices, runs SGD for the same E local epochs and learning rate, and averages their resulting model updates.This fixed-work design can reduce communication, but many local epochs on heterogeneous objectives may hurt convergence or cause device dropout.
  • FedProx: FedProx tolerates partial work by allowing device- and iteration-specific γ-inexactness, so local computation can vary with systems constraints.Variable local iterations act as a proxy for γ-inexactness and support convergence analysis involving systems heterogeneity such as stragglers.
  • FedProx: FedProx adds a proximal term that keeps local updates closer to the global model, addressing statistical heterogeneity while safely incorporating variable local work.The modification avoids manually fixing the number of local epochs and supports more robust and stable convergence on heterogeneous datasets.
  • Relationship to FedAvg: FedAvg is a special case of FedProx when µ = 0, the local solver is SGD, and γ is constant.Because the modifications are lightweight, FedProx can be integrated into existing federated learning packages and systems.

4 FE DPR O X: CONVERGENCE ANALYSIS

This section develops FedProx’s convergence analysis under a bounded measure of local-function dissimilarity, accommodating statistical heterogeneity, inexact local updates, and partial device participation. It establishes sufficient objective decrease and convergence guarantees for non-convex and convex objectives, while exposing trade-offs among dissimilarity and algorithm parameters.

  • Local dissimilarity: FedProx introduces B-local dissimilarity to quantify heterogeneity among local objectives and uses bounded dissimilarity as a convergence assumption.The assumption generalizes IID data with bounded dissimilarity while allowing statistical heterogeneity; a bounded gradient-variance condition provides a simpler, more restrictive alternative.
  • Non-convex convergence: Under bounded dissimilarity, Theorem 4 establishes expected objective decrease for non-convex, smooth local functions with inexact updates and only K active devices per round.The analysis uses γ-inexactness, device sampling, and the condition ¯µ > 0, which is sufficient but not necessary for convergence.
  • Parameter conditions: Positive decrease requires γB < 1 and B√K < 1, quantifying the trade-off between device dissimilarity and the algorithm parameters γ and K.These conditions are sufficient for ρ > 0 in Theorem 4, not necessary conditions for FedProx convergence.
  • Convergence rate: Theorem 6 gives a convergence rate to approximate stationary solutions for general non-convex local objectives under the bounded-dissimilarity assumption.The result applies after T iterations when the assumptions of Theorem 4 hold at each iteration and the initial optimality gap is ∆.
  • Convex case and comparison: For convex objectives with exact local minimization, Corollary 7 characterizes convergence and indicates that increasingly high accuracy requires increasing µ as Bϵ grows.Under bounded variance, the required number of steps is O(L∆/ϵ^2), matching SGD’s asymptotic rate; local updating can nevertheless perform worse than distributed SGD on non-IID data.

5 EXPERIMENTS

Experiments across synthetic and real federated settings evaluate FedProx under systems and statistical heterogeneity. FedProx improves convergence by incorporating partial work and using a proximal term, with the latter especially helpful as statistical heterogeneity increases.

  • Experimental Setup: FedProx is evaluated on diverse tasks, models, real-world federated datasets, and synthetic data with controllable statistical heterogeneity and simulated systems heterogeneity.Systems heterogeneity is simulated by assigning different amounts of local work to devices.
  • Systems Heterogeneity: Systems heterogeneity worsens FedAvg convergence, while incorporating variable amounts of work from constrained devices with FedProx produces more stable convergence.The comparison uses FedProx with µ = 0 against dropping constrained devices under FedAvg.
  • Systems Heterogeneity: Allowing partial work improves convergence over FedAvg even when every device performs at most one local epoch or when statistical heterogeneity is absent.These additional settings are reported in the appendix and in experiments without statistical heterogeneity.
  • Statistical Heterogeneity: As statistical heterogeneity increases, FedProx with µ = 0, equivalent to FedAvg, converges worse, whereas µ > 0 is particularly useful in heterogeneous settings.The comparison uses four synthetic datasets without systems heterogeneity and fixes E to be 20.
  • Proximal Term: An appropriate proximal parameter µ increases stability for unstable methods and can force divergent methods to converge with or without systems heterogeneity.This effect is observed for 50% and 90% stragglers as well as 0% stragglers.
  • Heterogeneity Measurement: Increasing µ empirically reduces dissimilarity among local functions, while the measured gradient variance captures dataset heterogeneity as a proxy of performance.The measurement tracks E_k[∥∇F_k(w)−∇f(w)∥2], which is lower bounded by Bϵ.

6 CONCLUSION · A COMPLETE PROOFS · A.1 Proof of Theorem 4

FedProx addresses systems and statistical heterogeneity through variable local work and a stabilizing proximal term, with convergence guarantees under device dissimilarity and straggler considerations. The proof of Theorem 4 combines strong convexity, bounded dissimilarity, smoothness, local Lipschitz continuity, and random device sampling to establish objective decrease bounds.

  • 6 CONCLUSION: FedProx tackles systems and statistical heterogeneity by allowing variable local work and using a proximal term to stabilize optimization.The framework provides convergence guarantees under a device dissimilarity assumption while accounting for stragglers.
  • A.1 Proof of Theorem 4: The proof uses γ-inexactness of each local solver to define the local-update error.This establishes the starting relation for the proof of Theorem 4.
  • A.1 Proof of Theorem 4: Strong convexity of h_k bounds the exact and approximate local solutions relative to the current global iterate.The argument invokes h_k’s μ̄-strong convexity in both bounds.
  • A.1 Proof of Theorem 4: The bounded dissimilarity assumption controls the deviation terms arising in the local-update and aggregate-update bounds.The proof explicitly attributes the relevant inequalities to bounded dissimilarity.
  • A.1 Proof of Theorem 4: A sufficiently large penalty parameter μ yields objective decrease proportional to ∥∇f(w_t)∥^2 for the full-device update.The algorithm instead approximates this update using K randomly selected devices.
  • A.1 Proof of Theorem 4: Local Lipschitz continuity of f bounds the error introduced when approximating the full-device update.The proof uses the local Lipschitz constant L_0 for this step.
  • A.1 Proof of Theorem 4: Taking expectation over randomly chosen updating devices requires bounding the corresponding random-update term.The expectation is specifically over the random choice of devices in round t.
  • A.1 Proof of Theorem 4: Combining the preceding inequalities and substituting the random-device and dissimilarity bounds completes the main proof bound.The final combination uses equations (8), (10), (9), and (14), together with α = 1.

A.2 Proof for Bounded Variance

This section establishes bounded variance as an equivalent alternative to the dissimilarity condition and restates FedProx’s non-convex convergence result under that assumption. The theorem provides an expected decrease in the global objective under specified iterate and algorithm-parameter conditions.

  • Bounded variance equivalence: Corollary 10 establishes an equivalence between the bounded variance and dissimilarity assumptions under Assumption 1.The bounded variance case is introduced as the basis for reformulating the convergence analysis.
  • Non-convex convergence: Theorem 11 restates non-convex FedProx convergence using the bounded variance assumption.It builds on the assertions of Theorem 4 while replacing the dissimilarity condition with bounded variance.
  • Non-convex convergence: Under the theorem’s iterate and Algorithm 2 parameter conditions, the expected global objective decreases at iteration t.The result applies when the iterate satisfies the stated gradient-norm condition and µ, K, and γ are appropriately chosen.
  • Proof strategy: The proof derives Theorem 11 from Theorem 4 through Corollary 10’s relationship between the two assumptions.The selected device set St is defined as the K devices chosen at iteration t.

A.3 Proof of Corollary 7

The proof analyzes the convex case with exact subproblem solutions and derives a gradient-dependent decrease under a batch-size condition. Using conditional expectations and telescoping, it bounds the iterations needed to obtain a solution with squared gradient norm below ϵ.

  • Convex case: The proof additionally considers the regime 1 << B ≤ 0.5K.This assumption is used to rewrite the decrease bound.
  • Iteration bound: Taking expectations after conditioning on the previous iterate and telescoping bounds the iterations needed to generate one solution with squared gradient norm below ϵ.The stated iteration complexity is O( LB2∆.

B CONNECTIONS TO OTHER SINGLE-MACHINE AND DISTRIBUTED METHODS · C SIMULATION DETAILS AND ADDITIONAL EXPERIMENTS

FedProx’s proximal objective connects to related optimization and meta-learning methods, while its bounded dissimilarity assumption builds on prior analyses of gradient diversity and heterogeneity. These connections differ in motivation and applicability to federated settings.

  • B CONNECTIONS TO OTHER SINGLE-MACHINE AND DISTRIBUTED METHODS: The paper frames the proximal term and bounded dissimilarity assumption as previously studied aspects whose motivations differ from those in the proposed work.The discussion is provided to clarify FedProx’s relation to prior optimization literature.
  • B CONNECTIONS TO OTHER SINGLE-MACHINE AND DISTRIBUTED METHODS: FedProx’s proximal term resembles EASGD, but EASGD uses a more complex moving average and is limited to SGD as its local solver.Both methods use the term to help prevent large parameter deviations on each device or machine.
  • B CONNECTIONS TO OTHER SINGLE-MACHINE AND DISTRIBUTED METHODS: DANE and AIDE also use proximal local objectives, augment them with gradient correction, and assume full device participation each communication round.The full-participation assumption is impractical in federated settings.
  • B CONNECTIONS TO OTHER SINGLE-MACHINE AND DISTRIBUTED METHODS: FedProx is connected to meta-learning, where similar proximal terms have been studied to improve few-shot learning performance.The cited meta-learning studies investigate proximal terms in the context of few-shot learning tasks.
  • B CONNECTIONS TO OTHER SINGLE-MACHINE AND DISTRIBUTED METHODS: The bounded dissimilarity assumption in FedProx has appeared in related forms in prior optimization analyses.Examples include Schmidt & Roux (2013), Yin et al. (2018), and Vaswani et al. (2019).
  • B CONNECTIONS TO OTHER SINGLE-MACHINE AND DISTRIBUTED METHODS: Yin et al. use bounded similarity to assert gradient diversity and quantify mini-batch SGD mean-square-error scaling benefits for IID data.Their setting concerns IID data rather than the federated setting described for FedProx.

C.1 Datasets and Models

The experiments use diverse non-synthetic federated datasets from prior work and LEAF, alongside synthetic data designed to isolate heterogeneity’s effect on convergence. Tasks span heterogeneous image classification benchmarks and synthetic non-IID settings with device-level sample imbalance.

  • Dataset suite: The study combines curated non-synthetic federated datasets with synthetic data to directly test how heterogeneity affects convergence.The non-synthetic datasets include prior federated-learning datasets and datasets proposed in LEAF.
  • Synthetic datasets: Synthetic experiments use (α, β)=(0,0), (0.5,0.5), and (1,1) to generate three non-identically distributed datasets across 30 devices.The number of samples per device follows a power law.
  • MNIST: MNIST models classify handwritten digits 0-9 with multinomial logistic regression across 1000 devices, each containing samples from only 2 digits.Inputs are flattened 784-dimensional (28 × 28) images, and per-device sample counts follow a power law.
  • FEMNIST: FEMNIST uses multinomial logistic regression on heterogeneous partitions of a 62-class EMNIST dataset, with only 5 classes assigned to each of 200 devices.The study subsamples 10 lower case characters (‘a’-‘j’); inputs are flattened 784-dimensional (28 × 28) images.

C.2 Implementation Details · C.3 Additional Experiments and Full Results

The implementation uses SGD-based FedProx with uniform device sampling and data-weighted update averaging for fair comparison with FedAvg. Experiments use fixed device participation, dataset-specific learning rates, an 80/20 local train-test split, commodity hardware, and TensorFlow 1.10.1.

  • C.2 Implementation Details: FedProx uses SGD locally, samples devices uniformly, and averages updates proportionally to each device’s local data points for comparison with FedAvg.The sampling scheme is not supported by the analysis, but similar relative FedProx-versus-FedAvg behavior is observed with or without it.
  • C.2 Implementation Details: The federated setup with 1 server and N devices is simulated on a commodity machine.The supplied hardware description specifies 2 Intel Xeon E5-2650 v4 CPUs and 8 NVidia 1080Ti GPUs.
  • C.2 Implementation Details: 10 devices are selected per round across all experiments, with local data split into 80% training and 20% testing sets.Learning rates are selected by grid search based on FedAvg and are not decayed across rounds.
  • C.2 Implementation Details: 0.01 is the learning rate for all synthetic data experiments, while MNIST, FEMNIST, Shakespeare, and Sent140 use 0.03, 0.003, 0.8, and 0.3, respectively.The supplied passage reports these dataset-specific learning rates after stating that the learning rate is not decayed through all rounds.
  • C.2 Implementation Details: All code is implemented in TensorFlow Version 1.10.1.Full implementation details are provided at github.com/litian96/FedProx.

C.3.1 Effects of Systems Heterogeneity on IID Data

On perfect IID synthetic data, allowing devices to perform partial work has little effect on convergence because FedAvg is robust to device failure.

  • Effects of Systems Heterogeneity on IID Data: The experiment evaluates partial work on perfect IID synthetic data.The dataset is identified as Synthetic IID.
  • Effects of Systems Heterogeneity on IID Data: FedAvg remains robust to device failure with IID data, so incorporating partial solutions from stragglers has little effect on convergence.This conclusion concerns the IID setting described in Figure 5.

C.3.2 Complete Results · C.3.3 Adaptively setting µ · C.3.4 Comparing Two Device Sampling Schemes

The complete results show FedProx’s accuracy gains and consistent dissimilarity behavior across heterogeneous settings, while additional experiments evaluate partial solutions, adaptive µ, and device-sampling schemes. Adaptive µ performs competitively, and sampling proportional to local data slightly outperforms uniform sampling with weighted averaging.

  • C.3.2 Complete Results: The dissimilarity characterization is consistent with practical performance, including training loss, across five datasets.The measurement covers five datasets, including four real datasets, while excluding systems heterogeneity by considering no device dropouts.
  • C.3.2 Complete Results: With E = 1, experiments evaluate how allowing partial solutions affects training loss and testing accuracy under systems heterogeneity.At E = 1, local updates deviate less from the global model, reducing the likelihood that statistical heterogeneity harms convergence.
  • C.3.2 Complete Results: When E = 1, tolerating partial solutions with FedProx (µ = 0) performs better than dropping stragglers with FedAvg.The comparison concerns loss under varied systems-heterogeneity settings.
  • C.3.3 Adaptively setting µ: An adaptive µ heuristic initializes µ at 1 for Synthetic-IID and 0 for the other three non-IID synthetic datasets.The initialization is adversarial to the methods; µ changes by 0.1 based on loss increases or five consecutive rounds of decreases.
  • C.3.3 Adaptively setting µ: The simple adaptive-µ heuristic achieves competitive performance across four synthetic datasets.The full results show the heuristic increasing µ after loss increases and decreasing it after five consecutive rounds of loss decreases.
  • C.3.4 Comparing Two Device Sampling Schemes: Sampling devices proportional to local data points and simply averaging models performs slightly better than uniform sampling with data-size-proportional model weights.Both schemes use E = 20 uniform local work per device, and the comparison reports loss, accuracy, and dissimilarity.
  • C.3.4 Comparing Two Device Sampling Schemes: Under either device-sampling scheme, µ = 1 demonstrates more stable performance than settings with lower µ.The experiment compares the schemes using FedProx on synthetic data.
Loading 1812.06127v5…