PaperlayerView exact revision on arXiv

Source-linked AI summary

FedPD: A Federated Learning Framework with Optimal Rates and Adaptivity to Non-IID Data

Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, Yang Liu

arXiv:2005.11418v3cs.LGstat.ML

TL;DR

CTA-based federated learning lacks clearly established theoretical properties under common assumptions that may not fit agent-generated data. This paper introduces FedPD from a primal-dual perspective and shows it achieves the best possible optimization and communication complexity for non-i.i.d. data while adapting communication to heterogeneity.

  • Problem

    Existing CTA federated-learning algorithms lack clearly established theoretical properties, while common bounded-gradient and i.i.d. assumptions may not fit data generated by individual agents.

    Source: §1 · pp. 1–2

  • Method

    The paper proposes FedPD, a CTA meta-algorithm built around a global-consensus reformulation and augmented-Lagrangian optimization.

    Source: §1 · p. 5, §3.1 · p. 7

  • Results

    The authors show BG and/or i.i.d. assumptions matter for FedAvg, while FedPD achieves the best possible optimization and communication complexity on non-i.i.d. data with adaptive communication.

    Source: §1 · p. 5

  • Takeaways & Limitations

    Under the δ-non-i.i.d. condition, communication savings and data heterogeneity exhibit a linear-logarithmic relationship.

    Source: §1 · p. 5

    The paper does not establish lower bounds on centralized gradient access or optimal graph dependency, instead focusing on potential communication savings from multiple local-processing steps.

    Source: §2 · p. 6

Abstract

from arXiv · show

Federated Learning (FL) has become a popular paradigm for learning from distributed data. To effectively utilize data at different devices without moving them to the cloud, algorithms such as the Federated Averaging (FedAvg) have adopted a "computation then aggregation" (CTA) model, in which multiple local updates are performed using local data, before sending the local models to the cloud for aggregation. However, these schemes typically require strong assumptions, such as the local data are identically independent distributed (i.i.d), or the size of the local gradients are bounded. In this paper, we first explicitly characterize the behavior of the FedAvg algorithm, and show that without strong and unrealistic assumptions on the problem structure, the algorithm can behave erratically for non-convex problems (e.g., diverge to infinity). Aiming at designing FL algorithms that are provably fast and require as few assumptions as possible, we propose a new algorithm design strategy from the primal-dual optimization perspective. Our strategy yields a family of algorithms that take the same CTA model as existing algorithms, but they can deal with the non-convex objective, achieve the best possible optimization and communication complexity while being able to deal with both the full batch and mini-batch local computation models. Most importantly, the proposed algorithms are {\it communication efficient}, in the sense that the communication pattern can be adaptive to the level of heterogeneity among the local data. To the best of our knowledge, this is the first algorithmic framework for FL that achieves all the above properties.

1 Introduction

The introduction frames CTA-based federated learning as a way to train a global model while limiting data sharing, but argues that its algorithms need closer analysis under non-i.i.d. data. It presents FedPD as a CTA framework designed for non-convex objectives, minimal assumptions, and communication adapted to heterogeneity.

Problem Formulation

  • Bounded-gradient and i.i.d. assumptions are restrictive: gradient bounds often fail, agent-generated data may not be i.i.d., and condition (5) essentially requires realizability.The paper defines δ-non-i.i.d. data by bounding differences between local and global gradients.Source: §1 · p. 2

Questions to address

  • The introduction poses four CTA questions: which local updates and aggregation improve performance, whether local steps reduce communication, and what performance is possible with minimal assumptions.The paper measures resources by local updates, communications, and sample accesses needed to reach an ϵ-stationary solution.Source: §1 · p. 3

Related Works

  • FedProx [12] avoids bounded gradients but retains an i.i.d. assumption, while VRL-SGD [16] achieves O(1/ϵ) communication without it; F-SVRG [17] and PR-SPIDER [18] do not follow CTA.These works also leave aggregation processing and communication-efficiency questions unaddressed; transmitting local gradients can leak private data.Source: §1 · p. 4

Our Main Contributions

  • O(1/ϵ) is the best possible communication efficiency for CTA algorithms based on successive local gradient updates; extra local processing or general linear processing cannot improve it.The authors also show that bounded-gradient and/or i.i.d. assumptions are important for FedAvg to work as intended.Source: §1 · pp. 4–5
  • FedPD handles general non-convex problems, achieves best-possible optimization and communication complexity for non-i.i.d. data, and converges under only Assumptions A1–A2.The meta-algorithm follows the CTA protocol and can be implemented in several forms.Source: §1 · p. 5
  • The communication pattern adapts to data heterogeneity: under δ-non-i.i.d. data, communication savings and heterogeneity have a linear-logarithmic relationship.The authors describe this as the first FL algorithm to achieve all the stated properties.Source: §1 · p. 5

2 Addressing Open Questions

The section proves that local updates and linear model combination cannot guarantee fewer communication rounds in the worst case. It also shows that FedAvg can diverge without the paper’s additional assumptions, even on non-convex problems.

2 Addressing Open Questions

  • At least O(1/ϵ) communication rounds are required in the worst case for an algorithm in the described class to reach an ϵ-stationary solution, regardless of Q > 0.The class permits linear, possibly time-varying server and local update operators; the hard instance satisfies Assumptions 1–2.Source: §2 · p. 5
  • This is a communication lower bound, not a lower bound on centralized gradient access or optimal graph dependence; its hard instances have non-i.i.d. data with δ growing with T.For 0-non-i.i.d. data, O(1) communication rounds suffice, while the intermediate-heterogeneity case is addressed in Section 3.Source: §2 · p. 6
  • For any fixed η > 0 and Q > 1, FedAvg can diverge to infinity on a problem satisfying A1–A2 but not A3–A4.The paper notes that divergence can occur even with Q = 2 and a solution satisfying ∑ᵢ ∥fᵢ(ẋ)∥² = 0.Source: §2 · p. 6
  • With diminishing stepsizes, FedAvg converges under the BG assumption, but without BG the stepsize choice can be significantly restricted.The paper states this result for any diminishing stepsize under BG.Source: §2 · p. 6

3 The FedPD Framework

FedPD is a primal-dual framework that supports full or sampled local gradients under A1–A2 and makes communication adaptive to data heterogeneity. Its local-oracle design and convergence analysis specify communication and computation tradeoffs for non-convex problems.

3 The FedPD Framework

  • FedPD is a meta-algorithm that requires only A1–A2, supports full or sampled local gradients, and can adapt communication to the degree of non-i.i.d.-ness.For non-i.i.d. data, the proposed variants achieve the optimization and communication complexity bounds in Claim 2.1.Source: §3 · p. 6

3.1 The Proposed Algorithm

  • FedPD relaxes a global-consensus formulation into separable local primal-dual problems, using local oracles and a parameter p to control skipped aggregation and communication.The oracle placeholder allows different local-processing instantiations; examples include GD or SGD updates and a variance-reduction oracle.Source: §3.1 · p. 7
  • FedPD’s communication saving depends on ε/δ²: it can be significant when δ < ε, while the saving decreases as δ increases.The analysis further states that p is inversely proportional to δ² when δ² ∈ (O(ε), ∞), and p approaches 1 at a log-rate as δ² approaches 0.Source: §3.1 · p. 7, §3.2 · p. 9

3.2 Convergence and Complexity Analysis

  • FedPD can match the same ε-accuracy as non-skipping communication using T(1−p) communication rounds, with p selected according to heterogeneity δ and target accuracy ε.For finite δ, the analysis selects ε1 = min{ε/(4C4), δ²} and T = 1/(2C2D0ε).Source: §3.2 · pp. 8–9
  • With T = O(1/ε) and ε1 = O(ε), local gradient accesses are O(1/ε × log(1/ε)) for GD and O(1/ε²) for SGD.These costs follow because the local augmented Lagrangian is strongly convex in xi; the theorem’s communication and computation results are described as best so far among FL algorithms for non-convex problems satisfying A1–A2.Source: §3.2 · p. 9

3.3 Connection with Other Algorithms

  • FedPD’s extra local dual variables help global convergence, whereas FedProx fails to converge to a global stationary solution and FedDANE violates the CTA protocol.FedPD is also distinguished from LAG and COLA: LAG transmits gradient differences, while FedPD’s δ-non-i.i.d. criterion measures similarity across local problems rather than imbalance in their Lipschitz gradients.Source: §3.3 · p. 9

4 Numerical Experiments

Synthetic experiments compare FedPD variants with FedAvg and FedProx on weakly and strongly non-i.i.d. data. FedPD variants remain effective under strong heterogeneity, while communication skipping can reduce rounds on weakly non-i.i.d. data.

Numerical Experiments

  • The experiments use non-convex penalized logistic regression on synthetic datasets with 400 samples per agent across 100 agents, generated as weakly or strongly non-i.i.d.The weakly non-i.i.d. dataset is generated i.i.d.; the strongly non-i.i.d. dataset uses a non-i.i.d. distribution.Source: §4 · p. 10
  • FedPD-SGD uses Q = 600, while FedPD-GD and FedPD-VR use Q = 8; all algorithms run for T = 600 iterations, with FedPD-VR using B = 1 and I = 20.FedPD-GD is tested with p = 0 and p = 0.5; FedAvg uses Q = 8 for GD and Q = 600 for SGD, and FedProx uses variance reduction for Q = 8 local iterations.Source: §4 · p. 10
  • On weakly non-i.i.d. data, FedProx and FedPD without communication skipping are comparable, while FedPD-GD skipping half the rounds reaches an error similar to FedAvg with fewer rounds.FedAvg with local GD does not converge to a stationary point with a constant stepsize when Q > 1.Source: §4 · p. 10
  • On strongly non-i.i.d. data, FedPD-VR and FedPD-GD without skipping converge to the global stationary point, whereas FedProx reaches error similar to FedAvg with local GD.Algorithms using stochastic solvers become less stable than in the weakly non-i.i.d. case.Source: §4 · p. 10

5 Conclusion … B.3 Properties.

This section presents FedPD as a CTA-based meta-algorithm with optimal communication and computation complexity that adapts communication to data heterogeneity. It also covers a communication-step lower bound for making the averaged iterate nonzero and properties of the component functions, including a nonpositivity property used in a proof.

5 Conclusion

  • FedPD is a CTA-based meta-algorithm whose algorithms achieve the best communication and computation complexity and adapt communication to data heterogeneity.The conclusion identifies these as desirable properties of the proposed framework.Source: §5 · p. 11

A Examples of Cost Functions Satisfy A5

  • The appendix shows that logistic regression satisfies A5 with a finite constant and that tanh’s δ is four times the logistic-regression δ.It also notes that the analysis covers cross-entropy losses with sigmoidal activations, including MLP, CNN, and RNN training problems.Source: A Examples of Cost Functions Satisfy A5 · pp. 14–15

B Proof of Claim 2.1

  • Without communication and aggregation, local first-order processing cannot make significant progress in reducing the original problem’s stationarity gap for certain constructed local functions.The claim allows an arbitrary number of local processing iterations; the model has a central aggregator and differs from the cited centralized and decentralized lower-bound settings.Source: B Proof of Claim 2.1 · p. 15

B.2 Main Constructions.

  • The proof constructs non-convex local functions from components that couple only neighboring coordinates.Each Θ(x, j) depends only on x[j−1] and x[j], and Figure 3 illustrates the construction across agents.Source: §B.2 · p. 16, Figure 3 · p. 17

B.3 Properties.

  • The constructed h_q gradients are Lipschitz with constant at most 27π; g and each g_i have constant 27π, while f and each f_i share a constant U.These properties are stated for the component and average functions used in the construction.Source: §B.3 · pp. 18, 21
  • In the Lemma 3 proof, the relevant partial-derivative expressions are nonpositive for all i and j.The proof also notes that one case is nonpositive because Ψ′ and Φ are nonnegative.Source: §B.3 · p. 23
  • From zero initialization, any positive-coefficient linear aggregation requires at least T communication steps to make the averaged iterate nonzero, regardless of local computation count.The aggregation coefficients may vary over time.Source: §B.3 · p. 23
  • Each aggregation advances the averaged iterate’s nonzero support by one coordinate, whereas multiple local updates alone do not advance it.At the first stage, only one local variable can gain a second nonzero entry; after aggregation, the next local variable can gain a third.Source: §B.3 · p. 24

B.4 Main Result for Claim 2.1.

Claim 2.1 gives a lower bound on communication stages for algorithms with linear operators, regardless of their number of local updates. The result also extends to sampled local gradients because they preserve the batch-gradient dependency structure.

B.4 Main Result for Claim 2.1.

  • For any algorithm following rules (10) with linear V^t and W^t, some problem satisfying Assumptions 1–2 requires at least the stated number of communication stages to reach the stated error, regardless of local updates.Source: §B.4 · p. 24
  • Sampled local gradients preserve the batch-gradient dependency structure and average-iterate nonzero pattern, so the batch-case estimates also apply.The passage states that the local function remains the same as in the batch case, preserving other estimates as well.Source: §B.4 · p. 25

C Proof of Claim 2.2

The proof constructs a two-agent FedAvg instance satisfying Assumptions 1 and 2 that diverges for Q > 1; when one agent’s samples are identical, local SGD produces the same update as GD.

C Proof of Claim 2.2

  • The example uses N = 2 and f(x) = 0 for all x, and satisfies both Assumptions 1 and 2.The proof defines y = [x1, x2]^T and a diagonal matrix D to express Q rounds of FedAvg.Source: C Proof of Claim 2.2 · pp. 25–26
  • For any Q > 1, FedAvg diverges from every nonzero initial point in the constructed example, whose iteration matrix has an eigenvalue greater than 1.The proof identifies the matrix eigenvalues as 0 and a value greater than 1.Source: C Proof of Claim 2.2 · p. 26
  • When one agent’s samples have the same loss function, local SGD gives the same update as local GD.The passage gives the example of agent 1 having two samples, each with loss function x^2.Source: C Proof of Claim 2.2 · p. 26

D Results showing the role of GB for FedAvg with diminishing stepsizes · D.1 Proof of Claim D.1

This section isolates A3’s role in FedAvg: under the stated stepsize conditions, average gradients converge to zero with A3, while without A3 the method can diverge.

D Results showing the role of GB for FedAvg with diminishing stepsizes

D.1 Proof of Claim D.1

  • The Claim D.1 proof fixes the outer stepsize in (0,1/L), bounds nonzero-index local stepsizes by 1/(2QL), and applies Lemma 5.Source: §D.1 · p. 28
  • A telescoping sum over outer iterations and a diminishing choice for nonzero-index stepsizes make the proof’s right-hand side converge to zero.Source: §D.1 · pp. 28–29

E Proof of Claim D.2 · F Proofs for Results in Section 3

The proof constructs a two-client FedAvg problem satisfying Assumptions 1 and 2, then shows that a specified step-size schedule makes the algorithm diverge.

E Proof of Claim D.2

F.1 Proof of Theorem 1 … F.1.4 Proof of Lemma 9

The FedPD Theorem 1 proof for Oracle I is organized around lemmas bounding local augmented-Lagrangian descent, error propagation, local-to-global augmented-Lagrangian differences, and the gap to the original loss.

F.1 Proof of Theorem 1

  • For Oracle I, the proof uses condition (67); replacing it with the SGD condition does not significantly change the proof or results.The authors therefore use condition (67) throughout the Theorem 1 proof.Source: §F.1 · p. 30
  • The paper outlines its main lemmas and says their proofs are in Sections F.1.1–F.1.4; the first lemma establishes sufficient descent of the local AL function.Source: §F.1 · p. 30

F.1.1 Proof of Lemma 6

  • Under A1, FedPD with Algorithm 4 (Oracle I) has a bounded change in the local augmented Lagrangian when the local problem satisfies condition (67).The proof bounds three differences corresponding to the algorithm’s update steps.Source: §F.1 · p. 30, §F.1.1 · p. 31

F.1.2 Proof of Lemma 7

  • Under A1 and A5, Lemma 7 bounds the propagation of differences among local models and local copies of the global models when communication is skipped.The displayed bounds include the current errors and terms involving δ, √ϵ1, and the communication probability p.Source: §F.1 · p. 31, §F.1.2 · p. 34

F.1.3 Proof of Lemma 8

  • Lemma 8 bounds the difference between the local and global augmented Lagrangians for FedPD with Oracle I when condition (67) holds.Its proof uses Jensen’s inequality and bounds a term involving △x^r+1_0.Source: §F.1 · p. 31, §F.1.3 · p. 34

F.1.4 Proof of Lemma 9

  • Under A1 and A2, Lemma 9 bounds the difference between the original loss and the augmented Lagrangian when the local problem is solved to ϵ1 accuracy.The proof obtains the result by averaging over N agents.Source: §F.1 · p. 31, §F.1.4 · p. 35

F.1.5 Proof of Theorem 1 … G.2 Algorithm Convergence and Complexity

The paper proves Theorem 1 under a stepsize restriction and notes that its results extend to SGD local solvers in expectation. It then introduces a variance-reduced FedPD oracle and gives its convergence, communication, and sample-complexity guarantees.

F.1.5 Proof of Theorem 1

  • Theorem 1’s proof requires C2 through C8 to be positive, imposing a stepsize upper bound; with SGD local solvers, the results hold in expectation.The SGD local problem must satisfy condition (68); the proof is omitted because only the local-solver error term is used.Source: §F.1.5 · p. 39

F.1.6 Constants used in the proofs

  • The constants C2 through C8 used in the proof are listed, with all terms positive under the stated stepsize restriction.The restriction is the same one established at the end of the Theorem 1 proof.Source: §F.1.6 · p. 39, §F.1.5 · p. 39

G FedPD with Variance Reduction

G.1 Algorithm description

  • Oracle II periodically computes a full gradient, then performs Q local steps using B samples each, with Q chosen as any positive integer.The full gradient is needed only every I iterations.Source: §G.1 · p. 40
  • Oracle II applies variance reduction to the entire primal-dual iteration, not just the local subproblem, and sufficiently large I yields an O(M) sample-complexity reduction.Source: §G.1 · p. 40

G.2 Algorithm Convergence and Complexity

  • Under A1–A2, Theorem 3 gives Oracle II’s convergence result; its communication complexity is O(1/ϵ), independent of Q.The theorem also specifies p = 0 and a stepsize choice, with C9 > 0.Source: §G.2 · pp. 40–41
  • Oracle II’s sample complexity is O(M + MT/I + 2TQBN); choosing I = M/QN keeps the stepsize unchanged to maintain the same convergence speed.The passage describes T/I + 1 full-gradient evaluations, each using M samples, and TQ mini-batch evaluations, each using 2B samples per node.Source: §G.2 · p. 41

G.3 Proof of Theorem 3 … H Additional Numerical Results

Theorem 3 is proved for FedPD with Algorithm 4 (Oracle II) by bounding the local augmented-Lagrangian descent, stochastic-gradient error, and gap to the original loss, then combining these bounds across outer iterations.

G.3 Proof of Theorem 3

  • The proof assumes p = 0, meaning no delayed communication, and uses r for the outer-stage index and q for the local-update index.Source: §G.3 · p. 41

G.3.1 Proof of Lemma 10

  • Under Assumption 1, Lemma 10 bounds the change in the local augmented Lagrangian for FedPD with Algorithm 4 (Oracle II).Source: §G.3 · p. 41

G.3.2 Proof of Lemma 11

  • Under Assumption 1 and sampling according to (13), Lemma 11 bounds the expected squared difference in stochastic gradient estimates for Algorithm 4 (Oracle II).Source: §G.3 · p. 41

G.3.3 Proof of Lemma 12

  • Under A1 and A2, Lemma 12 bounds the difference between the original loss and the augmented Lagrangian.Source: §G.3 · p. 42

G.3.4 Proof of Theorem 3

  • Theorem 3’s proof sums expected descent across outer iterations, averages over agents, and applies the stochastic-gradient and loss-gap bounds.It then combines the resulting bounds to complete the theorem proof.Source: §G.3.4 · pp. 47–49
  • The Theorem 3 proof also requires a step-size restriction on η, which the text says is guaranteed.The supplied passage truncates the displayed restriction, so its exact bound cannot be stated.Source: §G.3.4 · p. 49

H.1 Penalized Logistic Regression · H.2 Handwritten Character Classification

The experiments evaluate FedPD on penalized logistic regression with weakly and strongly non-i.i.d. data and on naturally heterogeneous FEMNIST data. They report convergence comparisons with FedAvg and FedProx, including communication-saving tradeoffs.

H.1 Penalized Logistic Regression

  • The penalized regression experiment compares weakly non-i.i.d. random agent data with strongly non-i.i.d. data generated from agent-specific models, using 400 samples on each of 100 agents.For the stronger heterogeneity case, feature vectors follow a standard Normal distribution, local-model components are uniform on [−10, 10], and labels include uniform noise.Source: §H.1 · p. 49
  • With weakly non-i.i.d. data, FedPD-VR converges fastest by sample complexity, while FedPD-GD can match FedAvg’s error with fewer communication rounds by skipping half of them.FedProx is second fastest by sample complexity; FedProx and FedPD with R = 1 are comparable by communication rounds. FedAvg with local GD and a constant stepsize does not converge to a stationary point when Q > 1.Source: Figure 4 · p. 50, §H.1 · p. 50
  • With strongly non-i.i.d. data, stochastic solvers become less stable, while FedPD-VR and FedPD-GD with R = 1 still converge to a global stationary point.FedProx achieves an error similar to FedAvg with local GD.Source: §H.1 · p. 50, Figure 5 · p. 51

H.2 Handwritten Character Classification

  • FEMNIST contains handwritten digits and letters from different writers; the experiment uses 29,214 samples from 100 writers distributed across 90 agents and a four-layer CNN with 62 output classes.The data are naturally non-i.i.d.; the first ten agents receive data from two writers each and the rest from one writer each, making the distribution unbalanced.Source: §H.2 · p. 51
  • On FEMNIST, FedAvg is slower than FedPD and FedProx, while FedProx performs similarly to FedPD when R = 1.The comparisons cover training and test results in Figures 7 and 8.Source: §H.2 · p. 52, Figure 7 · p. 53, Figure 8 · p. 53
  • Reducing FedPD’s communication frequency lowers final accuracy and raises final loss, but the accuracy drop is not significant and performance is better at the same communication-round count.FedAvg also becomes unstable and performs significantly worse when tested with a larger stepsize of 0.01.Source: §H.2 · p. 52, Figure 7 · p. 53, Figure 8 · p. 53

Related papers

Loading 2005.11418v3…