Source-linked AI summary
Lower Bounds and Optimal Algorithms for Personalized Federated Learning
Filip Hanzely, Slavomír Hanzely, Samuel Horváth, Peter Richtárik
TL;DR
The paper studies personalized federated learning because client data distributions can differ, while communication and local computation remain central complexity costs. It establishes lower bounds and develops accelerated and variance-reduced algorithms targeting them. The resulting methods are optimal in almost all regimes, with a local-Katyusha guarantee of optimal communication and, under a stated condition, optimal local-gradient complexity.
Problem
Personalized federated learning requires complexity guarantees for communication and local computation when clients hold heterogeneous data.
Method
The paper proves lower bounds for the personalized FL formulation and develops accelerated proximal-gradient and variance-reduced local-gradient methods.
Results
Local Katyusha achieves optimal communication complexity and also optimal total expected local-gradient complexity when √mλ = O(˜L).
Takeaways & Limitations
The work provides provably optimal personalized FL methods, including accelerated FedProx and accelerated variance-reduced FedAvg/Local SGD variants.
Takeaways & Limitations
Upper and lower bounds do not match for the local summand gradient oracle in the classical FL limit λ = ∞, where the known optimal method uses a stronger proximal oracle.
Abstract
from arXiv · showhide
In this work, we consider the optimization formulation of personalized federated learning recently introduced by Hanzely and Richtárik (2020) which was shown to give an alternative explanation to the workings of local {\tt SGD} methods. Our first contribution is establishing the first lower bounds for this formulation, for both the communication complexity and the local oracle complexity. Our second contribution is the design of several optimal methods matching these lower bounds in almost all regimes. These are the first provably optimal methods for personalized federated learning. Our optimal methods include an accelerated variant of {\tt FedProx}, and an accelerated variance-reduced version of {\tt FedAvg}/Local {\tt SGD}. We demonstrate the practical superiority of our methods through extensive numerical experiments.
1 Introduction
Federated learning trains models from data kept on geographically dispersed clients, making communication a major bottleneck. Personalized formulations address client heterogeneity by allowing local models to differ while penalizing their dissimilarity, and relate to local SGD/FedAvg.
- Federated learning setting: Federated learning fits models using data stored locally on numerous clients whose distributions may differ substantially.Clients communicate with a central server over potentially long physical distances, making communication more consequential than in-datacenter learning.
- Standard objective: Standard FL seeks the minimizer of the overall population loss, with each client loss depending only on its local data.
- Motivation for personalization: Personalized FL is motivated by cases where the population-loss minimizer is unsuitable for an individual client with a different data distribution.Next-word prediction on mobile keyboards is given as an example where personalized FL significantly outperformed non-personalized FL.
- Existing strategies: Personalization has been incorporated through multi-task learning, transfer learning, variational inference, and mixing local and global models.
- Mixing formulation: The paper studies a mixing objective that permits mutually different local models while penalizing their dissimilarity.This formulation is connected to distributed optimization and distributed transfer learning.
- Connection to local SGD: Stochastic gradient descent on the personalized formulation is essentially equivalent to FedAvg, and local gradient methods can outperform non-local methods on heterogeneous data.
2 Contributions
The paper establishes lower complexity bounds for personalized federated learning and develops algorithms that attain optimal communication or local-oracle complexity under broad conditions.
- Lower bounds: The work establishes lower bounds for both communication complexity and local computation in personalized FL.The lower-bound analysis covers communication and multiple local-oracle models, including proximal and gradient-based access.
- Algorithms: The proposed methods target optimal communication and local gradient complexity under different oracle and parameter regimes.The paper applies accelerated proximal-gradient approaches in two directions and considers inexact local solves using AGD or Katyusha.
- Limitations: An inexact APGD approach has extra logarithmic local-gradient factors, requires bounded iterates, and is suboptimal in communication when λ > L.These drawbacks motivate accelerating the L2SGD+ algorithm.
- Algorithms: AL2SGD+ is reported to achieve optimal communication complexity and local summand-gradient complexity in its applicable regime.The paper separately notes that the two bounds are not always achieved simultaneously.
- Optimality: The paper concludes that local algorithms are optimal for personalized FL with heterogeneous data, extending optimality beyond identical-data settings.This result is presented as support for using local methods in practical non-iid federated-learning scenarios.
- Optimality: Lower and upper communication complexities match regardless of the local oracle, while local-oracle calls match almost always.The stated exception is the summand-gradient oracle when λ > ˜L.
3 Lower complexity bounds
The paper formalizes communication and local-oracle complexity for personalized FL and establishes lower bounds under explicit oracle and algorithm assumptions. These bounds cover communication and three local-oracle types, with matching upper bounds reported up to logarithmic and constant factors in almost all cases.
- Oracle model: The analysis distinguishes communication-oracle queries from local-oracle calls, with local access provided through proximal, gradient, or summand-gradient information.A local oracle call is counted across all clients simultaneously.
- Optimality: Table 2 reports matching lower and upper complexities for communication regardless of the local oracle, while local-oracle bounds match almost always except for summand gradients when λ > ˜L.The table compares bounds up to logarithmic and constant factors.
- Oracle model: Algorithms are restricted to iterates lying in the span of previously observed oracle queries, though adversarial constructions could remove this assumption.The restriction is presented as a standard distributed-optimization assumption.
- Communication complexity: Theorem 3.1 constructs smooth, strongly convex instances yielding a lower bound on the communication rounds required to reach an ε-neighborhood of the optimum.The communication count is denoted C(k), the number of communication-oracle queries during the first k iterations.
- Local computation complexity: Lower bounds are also established for local proximal, gradient, and summand-gradient oracles; when λ = ∞, the results recover known bounds for the classical FL objective.For gradient oracles, identical starting points and a one-coordinate-per-local-call construction make parallelism ineffective.
4 Optimal algorithms
The paper develops accelerated and variance-reduced methods that match the personalized federated-learning lower bounds on communication and local computation in broad parameter regimes.
- Overview: The proposed algorithms are designed to match the lower complexity bounds on communication rounds and local steps.This section presents several methods for different oracle models and parameter regimes.
- Accelerated proximal gradient descent: APGD1 achieves optimal communication and local proximal complexity when λ ≤ L, while its alternative achieves optimal local gradient complexity across relative parameter regimes.The two variants apply acceleration to different decompositions of the regularized objective.
- Inexact accelerated proximal gradient descent: IAPGD replaces impractical exact local proximal solves with progressively more accurate local gradient-based solves using AGD or Katyusha.The methods support either full local gradients or finite-sum summand gradients.
- IAPGD with AGD: IAPGD+AGD attains optimal communication complexity and local gradient complexity up to logarithmic and constant factors.The local gradient complexity matches the lower bound in this sense.
- IAPGD with Katyusha: IAPGD+Katyusha has optimal communication complexity and achieves optimal total local-gradient complexity when √mλ = O(˜L).The Katyusha approach exploits finite-sum structure in local objectives, but requires bounded iterates.
- Accelerated L2SGD+: AL2SGD+ is an accelerated variance-reduced FedAvg variant whose communication complexity is optimal for all comparisons of ˜L and λ, with optimal local complexity when λ = O(˜L).It improves over previous methods by removing their parameter-regime restriction for communication optimality.
5 Experiments
Experiments test the proposed algorithms on homogeneous and heterogeneous logistic-regression data, and examine how λ affects APGD variants. Results support the theoretical convergence predictions and show practical improvements over L2SGD+.
- Homogeneous assignment: AL2SGD+ and IAPGD+Katyusha are significantly superior to L2SGD+ in communication rounds on the homogeneous data split.The experiment uses LIBSVM datasets with local finite-sum objectives and summand-gradient oracles.
- Heterogeneous assignment: Heterogeneous data assignment does not significantly influence convergence, which remains similar to the homogeneous case.The heterogeneous split sorts examples by labels before assigning them to clients.
- Experimental protocol: The experiments track relative suboptimality against communication rounds and local gradients computed over 10^3 communication rounds.For L2SGD+ and AL2SGD+, p = ρ = 1/m is selected, which is optimal up to a constant communication factor in the stated setup.
- Experimental protocol: Figure 1 compares communication and local computation complexity across IAPGD+Katyusha, AL2SGD+, and L2SGD+ on disjoint client datasets.The first row reports communication complexity and the second row reports local computation complexity.
C Proof of Theorem 3.1
The proof constructs quadratic personalized federated-learning instances whose optima contain coordinates that communication-limited algorithms cannot recover quickly. This yields a lower-bound argument based on sparse iterates versus dense optimal solutions.
- Hard-instance construction: The proof constructs local objectives whose aggregate personalized-FL objective is quadratic, with separate even- and odd-client constructions.The construction uses structured matrices and parameters selected to make the optimum analyzable.
- Hard-instance construction: The constructed functions depend on parameters a, b, and c, which are chosen so the optimal solution can be obtained explicitly.The smoothness of the objective can be controlled through the construction parameters.
- Characterizing the optimum: The proof identifies the optimum through first-order optimality conditions and derives recursive relations among its coordinates.These relations are established separately for even and odd numbers of clients.
- Lower-bound argument: The construction chooses parameters so iterates follow a geometric coordinate pattern, with w_i = γ^(i−1)w_1.An eigenvector construction makes the iterates multiples of a common vector.
- Lower-bound argument: After K communication rounds, the oracle permits at most K + 1 nonzero coordinates in x^K, whereas the optimum x⋆ has many nonzero coordinates.The resulting distance between the iterate and optimum supplies the communication lower bound.
D Proofs for Section 4
The proofs establish convergence guarantees for the paper's algorithms by combining properties of strongly convex local subproblems with bounds on approximate proximal solves. The resulting communication complexity is stated for reaching an ε-approximate solution.
- Theorem D.2: The proof expresses the error of the approximate proximal-operator evaluation at iteration k and applies a proposition from prior work.The resulting bound is obtained under condition (22).
- Proof strategy: Algorithm 1's convergence analysis uses the strong convexity of the objective and controls the local subproblem inaccuracies across iterations.The proof takes expectations after bounding the relevant error terms.
- Theorem D.2: Theorem D.2 bounds the expected communications required by Algorithm 1 to reach an ε-approximate solution.The theorem assumes each local subproblem is solved to suboptimality ε_k by a possibly randomized iterative method satisfying the stated condition.
D.1.2 Proof of Theorem 4.2
Theorem 4.2's proof bounds the iterations and local gradient computations needed to reach ε-suboptimality. It uses bounded level-set geometry and accelerated local solves for strongly convex subproblems.
- Bounded-region argument: The proof restricts the iterates to a bounded level-set-derived set S, whose finite diameter D controls the distance terms.The sets S′ and S are defined from the objective level set and its extrapolated combinations.
- Local solver: Each local subproblem is n(µ + λ)-strongly convex, allowing local AGD runs with c1 + c2k iterations.The proof uses this local solver schedule to establish the overall complexity result.
- Complexity bound: Reaching ε-suboptimality requires k = O iterations, followed by a bound on the total number of local gradient computations.The local-gradient complexity is stated immediately after the iteration bound.
D.1.3 Proof of Theorem 4.3
The proof applies Katyusha to smooth, strongly convex local objectives and uses its constant expected local stochastic-gradient cost to derive the total local-gradient complexity.
- The local objective is (˜L + λ)-smooth and (µ + λ)-strongly convex.
- Katyusha requires O(1) local stochastic gradient evaluations on average.
- The proof combines Katyusha’s local evaluation cost with the preceding complexity bound to obtain total local gradient complexity.
D.2 Towards the Proof of Theorem 4.4
This section develops the proof using variance-reduced stochastic gradients, a Lyapunov-function inequality, and expected accounting of local-gradient and communication costs.
- Algorithm 4 uses a variance-reduced stochastic gradient estimator under smoothness and strong-convexity assumptions.
- Proposition D.4 introduces a Lyapunov function and establishes the inequality used to derive Algorithm 4’s iteration complexity.
- AL2SGD+ has an explicitly derived communication complexity and local stochastic-gradient complexity.
- AL2SGD+ is identified as a special case of L-Katyusha, allowing direct application of the cited theorem after matching smoothness quantities.
- A single AL2SGD+ iteration uses 0, 1, m, or m + 1 local stochastic-gradient evaluations depending on the ξ and ξ′ cases.
- The expected communication cost per iteration is O(ρ + p(1 −p)), combining full-gradient and remaining communication events.
D.2.1 Proof of Theorem 4.4
The section specializes AL2SGD+ complexity expressions and situates the result among lower-bound literature and a concurrent accelerated FedProx variant, FedSplit.
- Proof of Theorem 4.4: For ρ = p(1 −p) and p = λ/(λ+˜L), the proof specializes AL2SGD+’s total communication complexity.
- Proof of Theorem 4.4: The section also states AL2SGD+’s local stochastic-gradient complexity for p = λ/(λ+˜L) and ρ = 1.
- Related lower bounds: Prior lower-bound work covers sequential optimization and the potential gains from parallelism, including applications in federated learning.
- FedSplit comparison: FedSplit is a concurrent accelerated FedProx variant designed for problem (1), whereas Algorithm 1 addresses problem (2).
- FedSplit comparison: The cited comparison notes that FedSplit does not argue optimality, while this work does, and describes FedSplit as a local algorithm with the correct fixed point for problem (1).