FedStaleWeight: Buffered Asynchronous Federated Learning with Fair Aggregation via Staleness Reweighting

In real-world scenarios, the task of training a model may be distributed among up to millions of agents, called edge devices, each of whom contribute their own limited (and perhaps private) set of data, a paradigm known as federated learning (FL). Our paper focuses on cross-device FL, in which there must be a large number of agents in order to extract meaningful results. The other type of FL, cross-silo, utilizes fewer agents because each one has more data and participates actively in the model training process, like a consortium of hospitals sharing information about a new disease [1].

Cross-device FL faces two major challenges when put into practice: scalability and privacy. In a massive distributed network, it’s impossible for every agent to be active at every time step, or for all agents to train their private models at the same speed. Naive aggregation techniques rely on concurrent reports from agents because, for instance, a simple average may be taken over reports in order to update the global model at regular intervals. By relaxing the concurrency constraint, we move away from synchronous FL, and the problem becomes how to best aggregate updates arriving at different points in time. On the topic of privacy, agents want their personal data to remain fully confidential; the global model must therefore evolve in a way that does not reveal individual contributions through gradient updates or the like [2]. Recent work in secure aggregation (SecAgg) and differential privacy provide frameworks for exploring the tradeoffs between privacy and utility [3].

Unlike its synchronous conterpart, Asynchronous federated learning (AFL) does not rely on clients to all finish and communicate their results in the same round and differs from its synchronous counterpart in the following ways:

• •

  Synchronicity: Clients are assumed to be heterogeneous and may return updates at different times or fail to return updates in time, necessitating mechanisms that do not wait for all participants before updating the global model through aggregation.

• •

  Staleness: As models may take longer to update, AFL mechanisms often are able to merge updates to models (i.e. model version $T$ is deployed to a client, the client takes a long time to fit, and it replies with its updated model while the global server has already aggregated updates and stepped to model version $T+1$). Emerging research in AFL hypothesizes that these stale updates are likely still useful and should be incorporated into learning.

• •

  Convergence challenges: AFL often experiences higher training instability based on the strategy chosen for aggregating stale updates, similar to issues encountered in off-policy reinforcement learning (RL).

We observe that the asynchronicity in AFL naturally leaves slow updating agents behind: since updates are no longer incorporated synchronously and are incorporated on demand, agents who are not able to update as quickly will contribute less to global model learning. However, naively up-weighting slow agent updates creates incentive issues: if our upweighting is too extreme, we incentivize fast updating agents to throttle and falsely report at slower speeds to increase their contribution to learning.

Thus, we explore the problem of fairly aggregating updates from clients in AFL to offset the unfairness introduced by compute heterogeneity in asynchronous federated learning. Several known synchronous approaches such as FedAvg and FedProx exist and more nuanced asynchronous approaches have been explored such as a fixed buffer length with averaging [2] or simply updating the central model in an asynchronous immediate manner [4]. To our knowledge, more nuanced weighting of stale updates or an adaptive approach has not yet been explored. Our contribution is the following:

1. 1.

   We reframe the problem of asynchronous federated learning aggregation as a welfare maximization problem and derive a gradient aggregation method which is strategy-proof to fast agents who may throttle their update rate to have their updates weighted higher while increasing the contribution of slower update producing agents to global model training.

2. 2.

   We show how our derived fair weighting can be computed via observing agent update staleness, a quantity naturally visible in buffered asynchronous federated learning, and present an algorithm for fair federated learning, FedStaleWeight.

3. 3.

   We provide a convergence guaruntee for FedStaleWeight in the smooth, non-convex setting, showing that the upweighting of stale model updates can still yield convergence.

4. 4.

   We empirically show in a series of realistic non-IID federated learning settings how current AFL aggregation techniques are biased towards clients with higher throughput and how FedStaleWeight maintains fairness and subsequently converges to higher global model accuracy faster. We release an open-source test bench for other researchers to explore further buffered AFL aggregation strategies.

The federated optimization problem with $m$ agents can be formulated as:

where $p_{i}$ is the importance of the $i$th agent (often assumed to be in proportion to the size of each client’s dataset, i.e. $p_{i}=n_{i}/n$ or equal for all agents $p_{i}=1/n$). $F_{i}$ represents agent $i$’s local objective function and is only available to agent $i$.

Specifically, we study the buffered asynchronous federated learning setting, consisting of a singular central server with an update aggregation buffer and $m$ clients. The server only updates once $b$ updates arrive at the buffer, where $b$ is a tune-able parameter. Each client has their own private data and pulls the latest version of the global model, trains it on its private dataset and communicates the update back to the server, which is appended to the buffer to later be aggregated (Figure 1). Specifically, in the asynchronous setting, we introduce the challenge of staleness, where updates to update model at version $v$ may have been trained on models of version smaller than $v$.

In the buffered AFL setting, we denote staleness $\tau_{i}$ for an agent $i$’s update as the positive difference in version number between the current global model and the global model that the update was computed from. Formally, we denote a global model with version $n$ and with 0 client training steps as $w_{n,0}$. For a client $i$, we denote the global model with version $n^{\prime}$ before and after training on $q$ local steps as $w_{n^{\prime},0}$ and $w_{n^{\prime},q}$, respectively; the client’s update is denoted as $\Delta_{i}=w_{n^{\prime},q}-w_{n^{\prime},0}$. Given the current global model $w_{n,0}$ and update $\Delta_{i}$, the staleness is thus $\tau_{i}=n-n^{\prime}$.

Each FL client $i$ represents its true update reporting speed as a random variable $\mathbf{R_{i}}$ where we denote the mean reporting speed as $r_{i}=\mathbb{E}[\mathbf{R_{i}}]$. Each agent $i$ chooses to report a random variable $\tilde{\mathbf{R_{i}}}$ with mean $\tilde{r_{i}}=\mathbb{E}[\tilde{\mathbf{R_{i}}}]$ as their observed update reporting speed. We assume for simplicity that $\tilde{r_{i}}\in[0,r_{i}]$, following intuition of fixed hardware specs: in other words, clients cannot manipulate to be faster on average than their fastest (true) compute speed on average. Given reported update speeds, we naturally see that the expected proportion that each agent contributes to learning without re-weighting is:

A fair re-weighting $\alpha(\cdot)$ takes $p_{i}(\tilde{r_{i}})$ and returns a weighting close to equal for every participating learning agent. In other words, we seek an $\alpha(\cdot)$ that solves the following constrained welfare maximization problem:

Clients seek to maximize a utility representing their effective influence (and thus their accuracy), modelled as their re-weighting times the overall proportional influence of their reported mean update rate. We use this utility function to justify strategy-proofness in Section 4.2.1.

Recall that computing the update step $u^{*}$ for a client with loss $L(\mathbf{x})$ at point $\mathbf{x}$ can be framed as the solution of the following minimization game, where step size $\eta$ represents the uncertainty in using the approximation $L(\mathbf{x}+u^{*})\approx L(\mathbf{x})+u^{*,\top}\nabla L(\mathbf{x})$ for stochastic gradient descent:

Similarly, in federated learning, we wish to restrict our aggregated gradient to be no larger than the step sizes of any of the individual aggregated gradients for convergence reasons, necessitating that any solution $\alpha(\cdot)$ must have $\sum_{i=1}^{n}{\alpha(p_{i}(\tilde{r_{i}}))p(\tilde{r_{i}})}=1$ as noted above.

Naively, if $\alpha(\cdot)=1$ is chosen (all agents are weighted equally), we see that $p(\tilde{r_{i}})$ is strictly increasing with respect to $\tilde{r_{i}}$ and all agents will report $\tilde{r_{i}}=r_{i}$. However, we see that welfare is poor, since an agent’s influence $\alpha(p(\tilde{r_{i}}))p(\tilde{r_{i}})$ is exactly increasing with respect to its rate, meaning that slow agents are left behind as they cannot manipulate any higher than their true $r_{i}$.

We observe that $\alpha(\cdot)$ should be monotonically decreasing with respect to its input, but not at a rate that decreases slower than $p(\cdot)$ increases. If the latter as true, faster agents would be incentivized to throttle their rates to improve their own weighted influence. The optimal $\alpha(\cdot)$ becomes $\alpha(\cdot)=\frac{1}{n(\cdot)}$, where $\alpha(p_{i}(\tilde{r_{i}}))=1/(np_{i}(\tilde{r_{i}}))$. Thus, with this choice, each agent’s weighted influence remains constant regardless of which reported $\tilde{\mathbf{R_{i}}}$ they choose, and thus there is no incentive to not report truthfully. We note that this weighting achieves weak truthfulness where agents are indifferent between reporting at any speed, since their utility stays the same regardless.

In buffered AFL, we assume the global aggregator has access only to the aggregation buffer and not the arrivals of individual agent updates, disallowing use of a weighting scheme that relies on estimating individual agent reporting schemes based on unique identifiers. We proceed to show how a fair weighting can be computed from staleness, an available local quantity, in asynchronous federated learning.

Given an initially empty buffer, we can compute the expected index of the first update from a given agent $i$ with mean rate $\tilde{r_{i}}$ by summing the expected number of updates from all agents in the expected time period for one update from agent $i$. For example, we note that an agent twice as fast as $i$ would have two updates expected in the same time as $i$’s first expected update. Recall that staleness corresponds to the difference between the version number of the current global model and the version number of the original global model that an agent’s update was reported from. We derive the following expression for expected staleness of agent $i$.

Reconciling this with our mechanism definitions above, we get the unweighted influence of agent $i$ given its expected staleness.

Recalling our derived weighting from the mechanism, we get a formula for our weighting as a function of expected staleness. We note that under a synchronous setting of zero expected staleness, we recover the exactly equal weighting used in synchronous FedAvg.

We present the following algorithm, FedStaleWeight, applying the above weighting to each aggregation. In practice, to compute the expected staleness for each agent $i$, $\mathbb{E}[\tau_{i}]$, we maintain a moving average of previous staleness for that agent. Our method computes a weighting to apply over the aggregation buffer, meaning it can be combined with methods like weighted secure aggregation [15] to protect client gradients.

We implement a generic testbench to empirically verify our algorithm above. The simulation environment consists of concurrently running clients with private data and and a centralized server and buffer, where an aggregation strategy can be specified.

1. 1.

   Clients run each on their own process thread concurrently, continuously pulling the latest available model from the global server and training it on their local private data before communicating the update to the server and repeating. Clients are initialized with both an individual slice of the overall dataset being evaluated (split into training and evaluation) and as a private runtime model representing a training delay distribution, from which the client process samples and sleeps for that duration to simulate local training delays.

2. 2.

   Server aggregates updates from the global buffer and updates the global model continuously. The server is initialized with a specified aggregation strategy and an aggregation buffer to pull updates from: once the buffer fills to a specific size, aggregation is performed and the global model is updated, after which clients will begin to pull that new model once they finish training. Clients broadcast their updates once training and their simulated delays complete, and their update is appended to the global buffer.

The testbench allows users to specify custom aggregation strategies by implementing a simple callback function, from which evaluation against other baselines can be computed, logged and compared. We also implement the FedAvg baseline in this testbench.

We test our algorithm empirically on two datasets, CIFAR10 and FashionMNIST. For both datasets, we model a system of 15 agents, with 10 “fast" agents with training delay randomly sampled from uniform distributions $D_{\text{fast}}=U(1,2)$ and 5 “slow" agents with training delays $D_{\text{slow}}=U(8,12)$. Data is distributed in a non-IID fashion and we reserve 20% of data across all labels for testing global model accuracy to measure overall generalization For the remainder, the fast agents are given data points with labels 4 through 9, IID-distributed across fast agents where no two agents have the same data point. The slow agents are given labels 0 through 3, similarly IID-distributed. The results are shown in Figure 2.

Specifically, we model a scenario where low-compute agents have exclusive access to a non-trivial subset of all data: in this setting, higher fairness equates to a model with higher global accuracy, demonstrated in our empirical results. We run FashionMNIST and CIFAR10 for 4,000 and 16,000 aggregations respectively. We evaluate FedAvg and our FedStaleWeight with a buffer size of $b=5$ and a local learning rate of $\eta_{\ell}=0.01$ with $Q=1$ local step. We clearly observe that our algorithm FedStaleWeight converges more quickly to a higher server test accuracy, indicating better generalization and less bias towards fast update producing agents.