Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Achieving Tighter Finite-Time Rates for Heterogeneous Federated Stochastic Approximation under Markovian Sampling

Feng Zhu, Aritra Mitra, Robert W. Heath

TL;DR

This work addresses federated stochastic approximation when data are time-correlated via Markov chains and local operators are heterogeneous. It introduces FedHSA, a drift-aware local update rule that uses a correction term based on the global operator to eliminate heterogeneity-induced bias without projection. Theoretical results show FedHSA achieves convergence to the correct root with centralized-rate performance and a linear speedup in sample complexity, surviving Markovian sampling through a mixing-time dependent analysis. Experiments in federated quadratic loss, TD learning with linear function approximation, and finite-sum quadratic loss corroborate the finite-time guarantees and the practical benefits of collaboration in heterogeneous FRL problems.

Abstract

Motivated by collaborative reinforcement learning (RL) and optimization with time-correlated data, we study a generic federated stochastic approximation problem involving $M$ agents, where each agent is characterized by an agent-specific (potentially nonlinear) local operator. The goal is for the agents to communicate intermittently via a server to find the root of the average of the agents' local operators. The generality of our setting stems from allowing for (i) Markovian data at each agent and (ii) heterogeneity in the roots of the agents' local operators. The limited recent work that has accounted for both these features in a federated setting fails to guarantee convergence to the desired point or to show any benefit of collaboration; furthermore, they rely on projection steps in their algorithms to guarantee bounded iterates. Our work overcomes each of these limitations. We develop a novel algorithm titled \texttt{FedHSA}, and prove that it guarantees convergence to the correct point, while enjoying an $M$-fold linear speedup in sample-complexity due to collaboration. To our knowledge, \emph{this is the first finite-time result of its kind}, and establishing it (without relying on a projection step) entails a fairly intricate argument that accounts for the interplay between complex temporal correlations due to Markovian sampling, multiple local steps to save communication, and the drift-effects induced by heterogeneous local operators. Our results have implications for a broad class of heterogeneous federated RL problems (e.g., policy evaluation and control) with function approximation, where the agents' Markov decision processes can differ in their probability transition kernels and reward functions.

Achieving Tighter Finite-Time Rates for Heterogeneous Federated Stochastic Approximation under Markovian Sampling

TL;DR

This work addresses federated stochastic approximation when data are time-correlated via Markov chains and local operators are heterogeneous. It introduces FedHSA, a drift-aware local update rule that uses a correction term based on the global operator to eliminate heterogeneity-induced bias without projection. Theoretical results show FedHSA achieves convergence to the correct root with centralized-rate performance and a linear speedup in sample complexity, surviving Markovian sampling through a mixing-time dependent analysis. Experiments in federated quadratic loss, TD learning with linear function approximation, and finite-sum quadratic loss corroborate the finite-time guarantees and the practical benefits of collaboration in heterogeneous FRL problems.

Abstract

Motivated by collaborative reinforcement learning (RL) and optimization with time-correlated data, we study a generic federated stochastic approximation problem involving agents, where each agent is characterized by an agent-specific (potentially nonlinear) local operator. The goal is for the agents to communicate intermittently via a server to find the root of the average of the agents' local operators. The generality of our setting stems from allowing for (i) Markovian data at each agent and (ii) heterogeneity in the roots of the agents' local operators. The limited recent work that has accounted for both these features in a federated setting fails to guarantee convergence to the desired point or to show any benefit of collaboration; furthermore, they rely on projection steps in their algorithms to guarantee bounded iterates. Our work overcomes each of these limitations. We develop a novel algorithm titled \texttt{FedHSA}, and prove that it guarantees convergence to the correct point, while enjoying an -fold linear speedup in sample-complexity due to collaboration. To our knowledge, \emph{this is the first finite-time result of its kind}, and establishing it (without relying on a projection step) entails a fairly intricate argument that accounts for the interplay between complex temporal correlations due to Markovian sampling, multiple local steps to save communication, and the drift-effects induced by heterogeneous local operators. Our results have implications for a broad class of heterogeneous federated RL problems (e.g., policy evaluation and control) with function approximation, where the agents' Markov decision processes can differ in their probability transition kernels and reward functions.

Paper Structure

This paper contains 17 sections, 13 theorems, 104 equations, 9 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Motivation and Related Work
  3. Setting and Motivation
  4. Setting
  5. Motivation for a new Federated SA Algorithm
  6. Proposed Algorithm: FedHSA
  7. Main Results and Discussion
  8. Challenges and Proof Sketch
  9. Experimental Results
  10. Federated Quadratic Loss Minimization Problem
  11. Federated TD Learning with Linear Function Approximation
  12. Federated Finite-Sum Minimization Problem with Quadratic Loss
  13. Conclusion
  14. Basic Results
  15. Analysis of FedHSA under I.I.D. Sampling
  16. ...and 2 more sections

Key Result

Proposition 1

Suppose $M >1$ and $H=2$. Consider the local SA update rule in Eq. eqn:localSA, and suppose the step-size $\eta$ is chosen such that the matrix $(I+2\eta \bar{A} + \eta^2 \bar{A}')$ is Schur-stable. Then, we have: $\lim_{t \to \infty} e_t = \eta v$, where $v = (1/M) (2 \bar{A} + \eta \bar{A}')^{-1}

Figures (9)

  • Figure 1: Illustration of the different statistical correlations in our heterogeneous federated SA setting.
  • Figure 2: Comparison between Local SA and FedHSA
  • Figure 3: Comparison between different numbers of agents for the FedHSA algorithm.
  • Figure 4: Comparison between Local SA and FedHSA with nonconvex objectives
  • Figure 5: Comparison between different numbers of agents for the FedHSA algorithm with nonconvex objectives.
  • ...and 4 more figures

Theorems & Definitions (24)

  • Proposition 1
  • Theorem 1
  • Theorem 2: Main Result
  • Corollary 1: Linear Speedup
  • proof
  • Corollary 2
  • proof
  • Lemma 1
  • proof
  • Lemma 2
  • ...and 14 more