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

The paper studies how computation-then-aggregation federated learning behaves under non-i.i.d. data and weak assumptions. It analyzes FedAvg and proposes the primal-dual FedPD framework, which retains CTA while targeting optimal complexities and communication adapted to data heterogeneity. The authors conclude that FedPD achieves these properties under only A1–A2 and supports non-convex objectives and varied local computation.

  • Problem

    CTA federated learning algorithms commonly rely on bounded gradients or i.i.d. data, while their behavior and attainable resource complexities under weaker assumptions remain insufficiently understood.

  • Method

    The paper analyzes FedAvg and develops FedPD, a primal-dual meta-algorithm with interchangeable local-processing oracles that retains the CTA protocol.

  • Results

    FedPD is reported to handle general non-convex problems, achieve best optimization and communication complexity under non-i.i.d. data, and converge under only A1–A2.

  • Takeaways & Limitations

    FedPD’s communication pattern can adapt to the degree of local-data heterogeneity while supporting full-batch or sampled local gradients.

  • Takeaways & Limitations

    The lower-bound analysis targets communication savings with local processing and does not establish lower bounds for centralized gradient access or optimal graph dependence.

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 paper examines theoretical gaps in computation-then-aggregation federated learning, especially under non-i.i.d. data and weak assumptions. It introduces FedPD, a primal-dual framework intended to improve robustness, complexity, and communication adaptivity.

  • Motivation: Federated learning uses a computation-then-aggregation protocol in which agents perform local updates before sending models to a server.The global model is trained from distributed data without requiring agents to reveal too much local information.
  • Limitations of existing analyses: FedAvg analyses often rely on restrictive assumptions, including bounded gradients and identically distributed local data.The paper notes that bounded gradients may fail for common least-squares objectives, while i.i.d. assumptions are unsuitable for data generated by individual agents.
  • Research questions: The paper asks how local updates, aggregation, communication frequency, and assumptions affect CTA algorithms’ stability and resource complexity.The resources considered are local updates, transmissions to the server, and accessed local samples before reaching an ϵ-stationary solution.
  • Theoretical findings: The authors report that successive local gradient-descent steps can cause divergence, while more general linear aggregation cannot always reduce communication below O(1/ϵ) rounds under A1–A2.They also state that bounded-gradient and i.i.d. assumptions are important for FedAvg to work as intended in non-convex settings.
  • FedPD framework: FedPD retains the CTA model while targeting general non-convex objectives, best optimization and communication complexity under non-i.i.d. data, and convergence under only A1–A2.Its framework supports different local-processing oracles and both full or sampled local gradients.
  • Communication adaptivity: FedPD adapts its communication pattern to local-data heterogeneity, with communication savings and δ-non-i.i.d.-ness exhibiting a linear-logarithmic relationship.The paper presents this as a framework-level property and claims it achieves the stated combination of properties.

2 Addressing Open Questions

The paper examines fundamental limits and failure modes of CTA-based federated learning, showing that local computation and linear aggregation cannot generally guarantee communication efficiency or convergence without stronger assumptions.

  • Q2–Q3: At least O(1/ϵ) communication rounds may be necessary for any linear CTA algorithm to reach an ϵ-stationary solution, regardless of the number of local updates.This lower bound holds for a class of non-convex problems satisfying the stated baseline assumptions.
  • Scope of analysis: The analysis targets communication savings from multiple local-processing steps, not lower bounds on centralized gradient access or optimal graph dependence.The stated scope focuses on communication rounds in the CTA setting.
  • Q2–Q3: The lower-bound construction requires communication to make progress when local objectives are sufficiently non-i.i.d., whereas 0-non-i.i.d. data can require only O(1) rounds.The intermediate regime, where δ lies between 0 and infinity, motivates analyzing adaptive communication savings.
  • Q1 and Q4: FedAvg can diverge to infinity on a problem satisfying A1–A2 when the BG and i.i.d. assumptions fail, even with Q > 1 and constant stepsize.The claim applies for any fixed η > 0 and Q > 1.
  • Q1 and Q4: For non-convex problems, FedAvg can diverge even when Q = 2 and a solution has zero summed local-gradient norm; diminishing stepsizes do not remove all assumption sensitivity.Under BG, FedAvg converges for any diminishing stepsize, while without BG the stepsize choice can be significantly restricted.

3 The FedPD Framework

FedPD reformulates federated learning through a primal-dual consensus framework, allowing flexible local optimization and communication skipping. Its variants support full or sampled gradients, adapt communication to heterogeneity, and attain strong communication and computation guarantees under weak assumptions.

  • Framework: FedPD relaxes global consensus constraints so the augmented Lagrangian separates across local agent variables and dual variables.This separability enables local processing between communication rounds while retaining primal-dual coordination.
  • Communication adaptivity: The communication-skipping parameter p controls the fraction of rounds omitted, while ϵ/δ^2 determines the potential communication saving under data heterogeneity.The saving becomes substantial when δ is smaller than ϵ and decreases as δ increases.
  • Framework: FedPD uses local processing oracles that can implement GD, SGD, or variance-reduced solvers while preserving the framework’s update pattern.Oracle I supports GD and SGD; Oracle II applies variance reduction to reduce sample complexity.
  • Complexity: FedPD reaches ϵ-stationarity with O(1/ϵ) communication rounds without skipping, while skipping reduces the rounds to T(1−p) at the same accuracy.The no-skipping case sets T proportional to 1/ϵ; the skipped-communication case preserves the same ϵ-accuracy.
  • Complexity: The required local gradient accesses are O(1/ϵ × log(1/ϵ)) for GD and O(1/ϵ^2) for SGD.These bounds combine T = O(1/ϵ) stages with the local optimization effort needed to reach ϵ accuracy.
  • Comparison: The paper characterizes FedPD’s communication and computation guarantees as the best so far for non-convex FL satisfying A1–A2.FedPD also introduces dual variables that record local-global model gaps, unlike FedProx’s local proximal objective.

4 Numerical Experiments

Experiments compare FedPD variants with FedAvg and FedProx on penalized logistic regression under weakly and strongly non-i.i.d. data. FedPD remains effective under strong heterogeneity and can reduce communication through skipped rounds.

  • Setup: The experiments evaluate FedPD-SGD, FedPD-GD, and FedPD-VR against FedAvg and FedProx on synthetic penalized logistic-regression data.The data include weakly non-i.i.d. and strongly non-i.i.d. settings, with different local solver configurations.
  • Weak heterogeneity: With weakly non-i.i.d. data, FedProx and FedPD-GD without skipping have comparable convergence, while FedAvg with local GD fails to converge for Q > 1 and constant stepsize.FedPD-GD skipping half the communication rounds achieves a similar error to FedAvg using fewer communication rounds.
  • Strong heterogeneity: Under strongly non-i.i.d. data, stochastic-solver methods become less stable, while FedPD-VR and FedPD-GD without skipping still converge to the global stationary point.FedProx reaches an error similar to FedAvg with local GD in this setting.

5 Conclusion

The paper studies the CTA protocol theoretically and develops FedPD as a primal-dual framework for non-convex federated learning. Its conclusions emphasize optimal complexity and communication adaptivity to data heterogeneity.

  • Conclusion: The paper’s main conclusion is that FedPD combines strong communication and computation complexity with communication patterns adaptive to data heterogeneity.This conclusion covers the framework’s theoretical properties under the studied FL setting.
  • Lower-bound construction: The lower-bound argument extends to sampled gradients because sampled local gradients preserve the dependency structure of their averaged versions.Thus the progression argument is not limited to full local-gradient access.
  • Lower-bound construction: The construction uses local functions whose dependencies couple neighboring coordinates through Θ(·), producing the staged progression used in the lower-bound argument.Figure 3 represents coordinate dependencies with shaded blocks and dashed boxes containing coupled variables.
  • Function properties: The component functions in the construction are designed to remain smooth, with Lipschitz first-order derivatives having constant 27π.The stated Lipschitz property applies to both the average function and each component function.
  • Lower-bound construction: The constructed non-convex instances show that at least T communication steps are needed before a linear combination of local variables can become nonzero.Each aggregation advances the nonzero pattern by one coordinate, whereas multiple local updates do not advance it.

B.4 Main Result for Claim 2.1.

Claim 2.1 establishes a lower bound showing that generic CTA algorithms with linear server and client operators may require many communication stages to reach the stated error. The construction applies to both full-batch and sampled local gradients.

  • Main result: The claim gives a problem satisfying Assumptions 1–2 that requires at least the stated number of stages, regardless of the number of local updates.These stages correspond equivalently to aggregation and communication rounds.
  • Main result: The lower-bound argument applies to algorithms whose V^t(·) and W^t(·) are linear operators.
  • Sampled gradients: For sampled local gradients, the progression of the averaged model’s non-zero pattern is exactly the same as with batch gradients.The local sampled gradients have the same dependency on x as their averaged version.
  • Sampled gradients: Other estimates, including the bound on f(0) − inf f(x), remain unchanged because the sampled-gradient local function matches the preceding local function.

C Proof of Claim 2.2

The proof constructs a two-agent problem satisfying the stated assumptions and expresses FedAvg over multiple local updates as a matrix iteration. Its eigenvalues show divergence whenever more than one local update is performed.

  • Construction: The construction uses N = 2 and an objective f(x) = 0 for all x, while satisfying Assumptions 1 and 2.
  • Matrix iteration: Q rounds of FedAvg are represented through the vector y = [x1, x2]T and diagonal matrix D = [1 −η, 0; 0, 1 + η].
  • Divergence: For any Q > 1, the iteration diverges from every non-zero initial point because its relevant eigenvalue exceeds one.The eigenvalues are described as 0 and ((1+η)^Q+(1−η)^Q)/2 > 1.
  • Sampled gradients: When each agent’s samples are identical, using local SGD produces the same update as local GD.

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

The results isolate A3 as the condition enabling FedAvg convergence with diminishing local stepsizes. Under A1–A3, the average gradient converges to zero; without A3, FedAvg can diverge for any Q > 1.

  • Convergence with A3: With A1–A3, FedAvg converges to zero average gradient when the initial stepsize is in (0, 1/L) and later local stepsizes are bounded and diminish to zero.The later stepsizes satisfy ηr,q ≤ min{1/[2(Q−1)L], η/Q} for q ≠ 0.
  • Failure without A3: Without A3, FedAvg with local-GD updates can diverge for any Q > 1, despite the other assumptions from Claim D.1 remaining in force.
  • Proof strategy: The proof uses Lemma 5, Jensen’s inequality, A1, and A3 to derive a telescoping bound on the average gradient.
  • Proof strategy: Choosing the non-initial local stepsizes as a sequence tending to zero makes the right-hand side of the bound converge to zero.

E Proof of Claim D.2

Claim D.2 proves divergence under diminishing stepsizes when A3 is absent. A matrix representation of the two-agent construction yields an eigenvalue strictly larger than one.

  • Construction: The proof again uses a two-agent problem satisfying Assumptions 1 and 2 with f(x) = 0 for all x.
  • Matrix iteration: With y = [x1, x2]T and Dr = [1 −ηr, 0; 0, 1 + ηr], Q FedAvg rounds are written as a product of round-dependent matrix updates.
  • Divergence: Choosing ηr = 1/√r except ηkQ+1 = 1/2 produces an eigenvalue λ2 strictly larger than one for Q > 1.
  • Divergence: Because λ2 > 1, the resulting FedAvg iteration diverges.

F.1 Proof of Theorem 1

The proof establishes FedPD’s convergence through descent of local augmented Lagrangians, control of communication-skipping errors, and comparison with the global objective.

  • Descent analysis: FedPD’s proof begins by establishing sufficient descent of the local augmented Lagrangian under Assumption A1 and an inexact local-solution condition.This is stated as Lemma 6 for Algorithm 4 with Oracle I.
  • Communication skipping: A separate error-propagation lemma bounds discrepancies between local models and local copies of the global model when communication is skipped.The bound is derived under Assumptions A1 and A5.
  • Virtual sequence: The analysis introduces a virtual sequence whose global copy is synchronized whenever communication and aggregation occur.This sequence connects intermittent aggregation to the augmented-Lagrangian analysis.
  • Objective comparison: The proof bounds the difference between local and global augmented Lagrangians, then bounds the original loss using the global augmented Lagrangian.The latter comparison requires Assumptions A1 and A2 and an ϵ1-accurate local solution.

F.1.5 Proof of Theorem 1

Theorem 1 is extended with an alternative variance-reduced local oracle that lowers sample complexity while preserving convergence guarantees under the stated assumptions.

  • Guarantees: The proof of Theorem 1 requires a sufficiently small stepsize and uses inexact local solutions; with SGD, the resulting guarantees hold in expectation.The proof’s local-solver dependence is summarized through the error term.
  • Contribution: An alternative FedPD oracle is introduced specifically to provide lower sample complexity.The section explicitly characterizes this as the purpose of the alternative oracle.
  • Oracle II: Oracle II periodically computes a full gradient, then performs Q local mini-batch updates, with Q allowed to be any positive integer.The full gradient is computed every I iterations, while each local step uses B samples.
  • Oracle II: The variance-reduction technique is applied to the entire primal-dual iteration rather than only to the local augmented-Lagrangian subproblem.Consequently, full-gradient evaluation is needed only periodically.
  • Complexity: With p = 0, the communication complexity is T = O(1/ϵ), independent of Q.This is the communication-complexity statement in Remark 5.
  • Complexity: The total sample complexity combines full-gradient and mini-batch evaluations as O(M + MT/I + 2TQBN).The stated balancing choice is I = M/(QN).

G.3 Proof of Theorem 3

Theorem 3 analyzes FedPD with Oracle II by combining augmented-Lagrangian descent, stochastic-gradient variance control, and a bound relating the original loss to the augmented Lagrangian.

  • Proof setup: The proof analyzes descent across outer iterations for Oracle II under no delayed communication, setting p = 0.The outer index is r and the local update index is q.
  • Descent analysis: Lemma 10 bounds the change in the local augmented Lagrangian under Assumption A1.This supplies the descent component of the theorem proof.
  • Variance control: Lemma 11 bounds the expected squared norm of the stochastic-gradient estimation error for randomly sampled data.This controls the variance introduced by Oracle II’s mini-batch updates.
  • Objective comparison: Lemma 12 bounds the difference between the original loss and the augmented Lagrangian under Assumptions A1 and A2.The proof then averages this relation over agents.
  • Theorem completion: The theorem combines these bounds through repeated expectation, summation over outer iterations, and averaging across agents.The proof concludes after establishing the required positive constants and stepsize conditions.

H.1 Penalized Logistic Regression

The experiments evaluate FedPD on synthetic penalized logistic regression with differing heterogeneity and on naturally non-i.i.d. handwritten-character data, using FedAvg and FedProx as baselines.

  • Synthetic regression: The synthetic experiment uses non-convex penalized logistic regression with weakly and strongly non-i.i.d. data-generation settings.The strongly non-i.i.d. setting generates local models independently over [−10, 10].
  • Weak heterogeneity: With weak heterogeneity, FedProx and FedPD without communication skipping are comparable, while FedAvg with constant stepsize fails to converge for Q > 1.FedPD can skip half the communication and retain a similar error to FedAvg with fewer communication rounds.
  • Strong heterogeneity: Under strong heterogeneity, stochastic-solver methods become less stable, while FedPD-VR and FedPD-GD with R = 1 still converge to the global stationary point.FedProx achieves an error similar to FedAvg with local gradient descent.
  • Handwritten-character classification: On FEMNIST, the evaluation uses data from writers, preserving non-i.i.d. and unbalanced distributions across 90 agents.The experimental subset contains 29,214 samples from 100 writers.
  • Handwritten-character classification: For neural-network training, FedPD uses Oracle II with periodic full-gradient computation and local updates between communication rounds.The reported settings include I = 20 and Q = 2 in the main-text experiment.
Loading 2005.11418v3…