Local Stochastic Approximation: A Unified View of Federated Learning and Distributed Multi-Task Reinforcement Learning Algorithms

In this paper, we study local stochastic approximation (SA), a distributed variant of the classic SA originally introduced by Robbins and Monro [1] for solving the root-finding problems under corrupted measurements of an operator (or function). We consider the setting where there are a group of agents communicating indirectly through a centralized coordinator. The goal of the agents is to find the root of the operator, which is composed of the local operators at the agents. For solving this problem, each agent iteratively runs a number of local SA steps based on its own data, whose iterates are then averaged at the centralized coordinator. This algorithm, presented in detail in Section 2, is relatively simple and efficient for solving problems requiring a large amount of data that are distributed to different agents.

We are motivated by the broad applications of local SA in solving problems in federated learning [2, 3] and multi-task reinforcement learning [4, 5, 6]. These two areas share a common framework, where multiple agents (clients or workers) collaboratively solve a machine learning/reinforcement learning problem under the coordination of a centralized server [2, 3, 4, 5, 6]. In stead of sharing the data collected at the local devices/environments to the server, the agents run local updates of their models/policy based on their data, whose results are then aggregated at the server with the goal to find the global learning objective. In these contexts, local SA can be used to formulate the popular algorithms studied in these two areas, as shown in Section 4.

Our focus, in this paper, is on the theoretical aspects of the finite-time performance of the local stochastic approximation. Our goal is to characterize the convergence rate of this method when the data at the agents are heterogeneous and dependent. In particular, we consider the case when the data at each agent is generated from a Markov process as often considered in the context of multi-task reinforcement learning. Our setting generalizes the existing works in the literature, where the local data is assumed i.i.d. Under fairly standard assumptions, our main contribution is to show that the convergence rates of local SA is within a logarithmic factor of the comparable bounds for independent data. As illustrated in Section 4, the results in this paper provide a unified view for the finite-time bounds of federated learning and multi-task reinforcement learning algorithms under different settings.

We consider distributed learning framework, where there are a group of $N$ agents communicating indirectly through a centralized coordinator. Associated with each agent $i$ is a local operator $F_{i}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$. The goal of the agents is to find the solution $\theta^{*}$ satisfying

where each $F_{i}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ is given as

Here ${\cal X}_{i}$ is a statistical sample space with probability distribution $\pi_{i}$ at agent $i$. We assume that each agent $i$ has access to operator $F_{i}$ only through its samples $\{F_{i}(\cdot,X_{i}^{k})\}$, where $\{X_{i}^{k}\}$ is a sequence of samples of the random variable $X_{i}$. We are interested in the case where each sequence $\{X_{i}^{k}\}$ is generated from an ergodic Markov process, whose stationary distribution is $\pi_{i}$. Moreover, the sequences $\{X_{i}^{k}\}$ are independent across $i$. For solving problem (1), our focus is to study the local stochastic approximation formally stated in Algorithm 1. In our algorithm, we implicitly assume that there is an oracle that returns to agent $i$ the value $F_{i}(\theta;X_{i})$ for a given $\theta$ and $X_{i}$.

Algorithm 1 is relatively simple to implement, which can be explained as follows. In this algorithm, each agent $i$ maintains a copy $\theta_{i}$ of the solution $\theta^{*}$ and the centralized coordinator maintains ${\bar{\theta}}$ to estimate the average of $\theta_{i}$. At any iteration $k\geq 0$, each agent $i$ first received ${\bar{\theta}}^{k}$ from the centralized coordinator and initalizes its iterate $\theta_{i}^{k,0}={\bar{\theta}}^{k}$. Here $\theta_{i}^{k,t}$ denotes the iterate at iteration $k$ and local time $t\in[0,\ldots,H-1]$. Agent $i$ then runs a number $H$ of local stochastic approximation steps using the time-varying step sizes $\alpha_{k}$ and based on its local data $\{X_{i}^{k+t}\}$. After $H$ local steps, the agents then send their new local updates $\theta_{i}^{k,H}$ to the centralized coordinator to update ${\bar{\theta}}^{k+1}$ by taking the average of these local values.

In this section, we study the finite-time performance of Algorithm 1 when each operator $F_{i}$ is strongly monotone. We provide an upper bound for the convergence rate of the mean square error $\mathbb{E}\left[\|{\bar{\theta}}^{k}-\theta^{*}\|^{2}\right]$ to zero for both constant and time-varying step sizes. In particular, under constant step size, $\alpha_{k}=\alpha$ for some constant $\alpha$ and $k\gtrsim\log(1/\alpha)$, this convergence occurs at a rate

where $\tau(\alpha)$ denotes the mixing time of the underlying Markov chain. On the other hand, under time-varying step sizes $\alpha_{k}\sim 1/(k+1)$ this rate happens at

for some positive integer ${\cal K}^{*}$ depending on the problem parameters $C,L,B,\mu$, which we will define shortly. First, our rates scale linearly with the number of local steps $H$. As expected, when $H$ goes to $\infty$, each agent only consider its local SA without communicating with the centralized coordinator. In this case, one would expect that each agent only finds the root of its own operator. Second, one can view the constant $B$ as the variance of the noise. Finally, our rates match the ones of local SGD under i.i.d noise [27], except for a $\log$ factor reflecting the Markovian randomness.

We start our analysis by introducing the following technical assumptions and notation used in this section. Given a constant $\alpha>0$, we denote by $\tau_{i}(\alpha)$ the mixing time associated with the Markov chain $\{X_{i}^{k}\}$, i.e., the following condition holds

where $\|\cdot\|_{TV}$ is the total variance distance and $\mathbb{P}_{i}^{k}(X_{i}^{0},X_{i})$ is the probability that $X_{i}^{k}=X_{i}$ when we start from $X_{i}^{0}$. The mixing time represents the time $X_{i}^{k}$ getting close to the stationary distribution $\pi_{i}$. In addition, we denote by $\tau(\alpha)=\max_{i}\tau_{i}(\alpha)$.

For our results derived in this section, we consider the following assumptions. For simplicity, we assume that these assumptions hold to the rest of this paper.

Assumption 1 implies that the Markov chain $\{X_{i}^{k}\}$ has geometric mixing time, which depends on the second largest eigenvalue of the transition probability matrix of $\{X_{i}^{k}\}$ [31]. This assumption holds in various applications, e.g, in incremental optimization [32] and in reinforcement learning modeled by Markov decision processes with a finite number of states and actions [33]. In addition, Assumption 1 is used in the existing literature to study the finite-time performance of SA and its distributed variants under Markov randomness; see [13, 16, 20, 19, 34, 35] and the references therein.

Assumption 2 is often used in local SGD methods in federated learning [27]. This assumption also holds in the context of reinforcement learning considered in Section 4. Assumption 3 implies that $F_{i}$ are strongly monotone (or strongly convex in the context of SGD).

The following result is a consequence of Assumptions 1 and 2, which is shown in [20, Lemma $3.2$]. This lemma basically states that each Markov chain $X_{i}$ has a geometric mixing time, which translates to the operator $F_{i}$ due to the Lipschitz condition.

Finally, since ${\cal X}_{i}$ is finite for all $i\in[N]$, Assumption 2 also gives the following result.