Best of Both Worlds in Federated LSA: Speedup When Possible, Personalization Always

The gist: PF-LSA introduces a minimalist algorithm that mixes each agent’s local stochastic update with the average update across agents, achieving best-of-both-worlds guarantees without prior knowledge of heterogeneity.

Personalized Federated Linear Stochastic Approximation (PF-LSA)

PF-LSA is a novel personalized federated method which mixes updates across all agents, interpolating between fully local learning and classical federated averaging with the relative importance determined by a single hyperparameter PF-LSA, a novel personalized federated method which mixes updates across all agents, interpolating between fully local learning and classical federated averaging (McMahan et al., 2017) with the relative importance of each part being determined by a single hyperparameter The algorithm is thus very minimalist, as it simply combines local updates with their global average, without imposing any explicit structure on agent similarity

Convergence Guarantees and Error Decomposition

The analysis of PF-LSA is based on a sharp decomposition of the error into consensus and disagreement components The consensus error decays rapidly, whereas the disagreement error decays more slowly but becomes negligible in low-heterogeneity regimes Theorem 4.4 makes the two-clock structure of PF-LSA explicit: a fast consensus clock and a slow personalization clock The resulting bound shows that when the local operators are sufficiently similar and the initial disagreement is negligible compared to the initial consensus error, PF-LSA can simultaneously achieve personalization and linear speedup in the number of agents

Regime Analysis: Low-Heterogeneity vs. High-Heterogeneity

The paper establishes distinct performance guarantees based on the level of heterogeneity, quantified by the parameter ζ In the low-heterogeneity regime, where ζ squared < a 2/4, PF-LSA can simultaneously achieve personalization and linear speedup in the number of agents Specifically, Corollary 4.6 shows that when the initial disagreement D0 is negligible compared to the initial consensus error, PF-LSA achieves a sample complexity of order Oe(1/(N ε)) Conversely, in the high-heterogeneity regime where ζ squared ≥ a 2/4, collaboration does not necessarily translate into a linear improvement of the overall convergence rate, but PF-LSA retains personalization and recovers the statistical behavior of single-agent learning

Application to Personalized Federated Temporal Difference (TD) Learning

The results are instantiated for Personalized Federated TD learning, where each agent aims to estimate its local value function V pi c(s) as a linear combination of features The analysis yields two sample-complexity guarantees corresponding to the low- and high-heterogeneity regimes determined by the level of heterogeneity ζTD Corollary 5.2 provides a bound for the low-heterogeneity regime, showing that PF-LSA satisfies 1/N P c=1 Eθ T c - θ⋆ c squared ≤ ε provided T ≳ 1/(1−γ)ν max n (4/ (1−γ), 32(1+γ) 2(1-γ)ν, 288σ bar 2 ε,TD, (1-γ)νNϵ 2) × max (log 4e 0 2ϵ, 2N log (D0 + ζ TD(1−γ) 2ν 2Ne 0 2)/ϵo The high-heterogeneity regime is covered by Corollary C.4, which provides a similar sample complexity guarantee under the condition ζ 2 TD ≥ (1−γ) 2ν 2/16

Empirical Validation and Conclusion

Numerical experiments demonstrate that PF-LSA remains stable and converges towards the personalized solutions across all considered levels of heterogeneity PF-LSA successfully accelerates the learning process, consistently reaching the low-error regime in significantly fewer iterations compared to Single agent LSA The experiments highlight that PF-LSA combines the advantages of both approaches: it uses global information through the averaged update, which accelerates learning, while retaining a local update component, which guarantees convergence to each agent’s own solution Overall, Figure 2 provides empirical evidence for the best-of-both-worlds behavior predicted by our theory: PF-LSA benefits increasingly from additional agents when collaboration is useful, while converging to each agent’s personalized solution across all heterogeneity levels

Technical Lemmas

Lemma D.3 states that for any c ∈ [N], and x ∈ R d, and η ∈ (0, η∞), it holds that ⟨x, A c bar x⟩ ≥ ax squared Lemma D.2 proves the relationship between the norms of the averaged and individual vectors Lemma B.3 shows that for any c ∈ [N], and x ∈ R d, and η ∈ (0, η∞), it holds that ⟨x, A c bar x⟩ ≥ ax squared Lemma D.1 provides the relationship between successive iterates Lemma B.4 proves the bound for the low-heterogeneity regime, showing that when ζ squared < a 2/4, PF-LSA satisfies 4 1 − ηβa squared t D0 + 4βζ 2a squared e 0 squared + 6ηaβ + 5N σ bar 2 ε Lemma B.3 provides the bound for the high-heterogeneity regime, showing that when ζ squared ≥ a 2/4, PF-LSA satisfies (1 − ηa)t e 0 squared + 16ζ 2a squared 1 − ηβa squared t D0 + βe 0 squared + ζ 2a 3Nϵ a / (168ζ 2σ bar ε) Lemma B.4 proves the bound for the high-heterogeneity regime, showing that when ζ squared ≥ a 2/4, PF-LSA satisfies 1N P c=1 Ee T c - θ⋆ c squared ≤ (1 − ηa)t e 0 squared + 16ζ 2a 2 1 − ηβa 2 t D0 + βe 0 squared + ζ 2a 3Nϵ a / (168ζ 2σ bar ε) Corollary C.4 provides the sample and communication complexity for the high-heterogeneity regime, showing that PF-LSA satisfies 1N P c=1 Ee T c - θ⋆ c squared ≤ ε provided T ≳ 2(1−γ)ν max (4/ (1−γ), 32(1+γ) 2(1-γ)ν, 1344ζ 2 TDσ bar ε,TD) × max (log 3e 0 2ϵ, 2N log 192ζ 2 TD D0 + 1N e 0 2(1−γ) 2ν 2ϵ) Corollary C.

1. A personalized federated linear stochastic approximation (PF-LSA) algorithm can be implemented to train distinct models for each agent by mixing their local stochastic updates with a global average update, achieving full personalization even when incorporating global information from other agents. This allows the system to solve distinct linear fixed-point equations, each corresponding to an agent-specific learning problem.

2. The PF-LSA framework enables agents to achieve a linear speedup in the number of agents when their learning problems are sufficiently similar, specifically when the heterogeneity level is low, as shown by the result: "when the local operators are sufficiently similar and the discrepancy between the agents’ optimal solutions is negligible compared with the initial consensus error, PF-LSA can simultaneously achieve personalization and a linear speedup in the number of agents."

3. The system can be deployed in reinforcement learning or temporal difference learning settings by instantiating PF-LSA for Personalized Federated TD learning, enabling it to estimate its local value function for each agent by solving their respective projected Bellman equations, thereby preserving personalization even under environmental heterogeneity.

4. The algorithm provides a robust convergence guarantee across varying levels of heterogeneity; in the high-heterogeneity regime, PF-LSA retains personalization while recovering the statistical complexity of independent learning, ensuring that even when local problems are substantially different, the best one can reasonably expect is to recover the statistical complexity of independent learning.

We study personalized federated linear stochastic approximation (LSA), a framework which notably encompass personalized temporal difference learning. In this setting, heterogeneous agents collaborate to solve distinct linear fixed-point equations, each corresponding to an agent-specific learning problem. A central open question in personalized learning is whether a single method can adapt to an unknown level of heterogeneity by converging to each agent's personalized solution in all regimes while achieving a linear speedup in the number of agents when their learning problems are sufficiently similar. We answer this question affirmatively by introducing PF-LSA, a minimalist algorithm that mixes each agent's local stochastic update with the average update across agents, at no additional computational cost relative to standard federated methods. We prove that PF-LSA, achieves best-of-both-worlds guarantees without any prior knowledge on the level of heterogeneity. Our analysis is based on a sharp decomposition of the error into consensus and disagreement components. The consensus error decays rapidly, whereas the disagreement error decays more slowly but becomes negligible in low-heterogeneity regimes.