Asynchronous Federated Optimization

Introduction

Federated learning [Konevcnỳ et al.(2016)Konevcnỳ, McMahan, Yu, Richtárik, Suresh, and Bacon, McMahan et al.(2016)McMahan, Moore, Ramage, Hampson, et al.] enables training a global model on datasets partitioned across a massive number of resource-weak edge devices. Motivated by the modern phenomenon of distributed (often personal) data collected by edge devices at scale, federated learning can use the large amounts of training data from diverse users for better representation and generalization. Federated learning is also motivated by the desire for privacy preservation [Bonawitz et al.(2019)Bonawitz, Eichner, Grieskamp, Huba, Ingerman, Ivanov, Kiddon, Konecny, Mazzocchi, McMahan, et al., Bonawitz et al.(2017)Bonawitz, Ivanov, Kreuter, Marcedone, McMahan, Patel, Ramage, Segal, and Seth]. In some scenarios, on-device training without depositing data in the cloud may be legally required by regulations [Steve Anderson: HealthInsurance.org(1996), US Department of Education(2019), EU(2018)].

A federated learning system is often composed of servers and workers, with an architecture that is similar to parameter servers [Li et al.(2014a)Li, Andersen, Park, Smola, Ahmed, Josifovski, Long, Shekita, and Su, Li et al.(2014b)Li, Andersen, Smola, and Yu, Ho et al.(2013)Ho, Cipar, Cui, Lee, Kim, Gibbons, Gibson, Ganger, and Xing]. The workers (edge devices) train the models locally on private data. The servers aggregate the learned models from the workers and update the global model.

Federated learning has three key properties [Konevcnỳ et al.(2016)Konevcnỳ, McMahan, Yu, Richtárik, Suresh, and Bacon, McMahan et al.(2016)McMahan, Moore, Ramage, Hampson, et al.]: 1) Infrequent task activation. For the weak edge devices, learning tasks are executed only when the devices are idle, charging, and connected to unmetered networks [Bonawitz et al.(2019)Bonawitz, Eichner, Grieskamp, Huba, Ingerman, Ivanov, Kiddon, Konecny, Mazzocchi, McMahan, et al.]. 2) Infrequent communication. The connection between edge devices and the remote servers may frequently be unavailable, slow, or expensive (in terms of communication costs or battery power usage). 3) Non-IID training data. For federated learning, the data on different devices are disjoint, thus may represent non-identically distributed samples from the population.

Federated learning [McMahan et al.(2016)McMahan, Moore, Ramage, Hampson, et al., Bonawitz et al.(2019)Bonawitz, Eichner, Grieskamp, Huba, Ingerman, Ivanov, Kiddon, Konecny, Mazzocchi, McMahan, et al.] is most often implemented using the synchronous approach, which could be slow due to stragglers. When handling massive edge devices, there could be a large number of stragglers. As availability and completion time vary from device to device, due to limited computational capacity and battery time, the global synchronization is difficult, especially in the federated learning scenario.

Asynchronous training [Zinkevich et al.(2009)Zinkevich, Langford, and Smola, Lian et al.(2017)Lian, Zhang, Zhang, and Liu, Zheng et al.(2017)Zheng, Meng, Wang, Chen, Yu, Ma, and Liu] is widely used in traditional distributed stochastic gradient descent (SGD) for stragglers and heterogeneous latency [Zinkevich et al.(2009)Zinkevich, Langford, and Smola, Lian et al.(2017)Lian, Zhang, Zhang, and Liu, Zheng et al.(2017)Zheng, Meng, Wang, Chen, Yu, Ma, and Liu]. In this paper, we take the advantage of asynchronous training and combines it with federated optimization.

We propose a novel asynchronous algorithm for federated optimization. The key ideas are (i) to solve regularized local problems to guarantee convergence, and (ii) then use a weighted average to update the global model, where the mixing weight is set adaptively as a function of the staleness. Together, these techniques result in an effective asynchronous federated optimization procedure. The main contributions of our paper are listed as follows:

We propose a new asynchronous federated optimization algorithm and a prototype system design.

We prove the convergence of the proposed approach for a restricted family of non-convex problems.

We propose strategies for controlling the error caused by asynchrony. To this end, we introduce a mixing hyperparameter which adaptively controls the trade-off between the convergence rate and variance reduction according to the staleness.

We show empirically that the proposed algorithm converges quickly and often outperforms synchronous federated optimization in practical settings.

Problem formulation

Methodology

The training takes $T$ global epochs. In the $t^{\mbox{th}}$ epoch, the server receives a locally trained model $x_{new}$ from an arbitrary worker, and updates the global model by weighted averaging: $x_{t}=(1-\alpha)x_{t-1}+\alpha x_{new},$ where $\alpha\in(0,1)$ is the mixing hyperparameter. A system overview is illustrated in Figure 1.

The server and workers conduct updates asynchronously, i.e., the server immediately updates the global model whenever it receives a local model. The communication between the server and the workers is non-blocking. Thus, the server and workers can update the models at any time without synchronization, which is favorable when the devices have heterogeneous conditions.

The detailed algorithm is shown in Algorithm 1. The model parameter $x_{\tau,h}^{i}$ is updated in the $h$th local iteration after receiving $x_{\tau}$, on the $i$th device. The data $z_{\tau,h}^{i}$ is randomly drawn in the $h$th local iteration after receiving $x_{\tau}$, on the $i$th device. $H_{\tau}^{i}$ is the number of local iterations after receiving $x_{\tau}$ on the $i$th device. $\gamma$ is the learning rate and $T$ is the total number of global epochs.

On the server side, the scheduler and the updater run asynchronously in parallel. The scheduler periodically triggers training tasks and controls the staleness ($t-\tau$ in the updater thread). The updater receives models from workers and updates the global model. Our architecture allows for multiple updater threads with read-write lock on the global model, which improves the throughput.

Intuitively, larger staleness results in greater error when updating the global model. For the local models with large staleness $(t-\tau)$, we can decrease $\alpha$ to mitigate the error caused by staleness. As shown in Algorithm 1, optionally, we use a function $s(t-\tau)$ to determine the value of $\alpha$. In general, $s(t-\tau)$ should be $1$ when $t=\tau$, and monotonically decrease when $(t-\tau)$ increases. There are many functions that satisfy such two properties, with different decreasing rate, e.g., $s_{a}(t-\tau)=\frac{1}{t-\tau+1}$. The options used in this paper can be found in Section 5.2.

Convergence analysis

First, we introduce some definitions and assumptions for our convergence analysis.

(Smoothness) A differentiable function $f$ is $L$-smooth if for $\forall x,y$, $f(y)-f(x)\leq\left\langle\nabla f(x),y-x\right\rangle+\frac{L}{2}\|y-x\|^{2},$ where $L>0$.

(Weak convexity) A differentiable function $f$ is $\mu$-weakly convex if the function $g$ with $g(x)=f(x)+\frac{\mu}{2}\|x\|^{2}$ is convex, where $\mu\geq 0$. $f$ is convex if $\mu=0$, and non-convex if $\mu>0$.

We have the following convergence guarantees. Detailed proofs can be found in the appendix.

Experiments

In this section, we empirically evaluate the proposed algorithm.

We conduct experiments on two benchmarks: CIFAR-10 [Krizhevsky and Hinton(2009)], and WikiText-2 [Merity et al.(2016)Merity, Xiong, Bradbury, and Socher]. The training set is partitioned onto $n=100$ devices. The mini-batch sizes are 50 and 20 respectively.

2 Evaluation setup

The baseline algorithm is FedAvg introduced by [McMahan et al.(2016)McMahan, Moore, Ramage, Hampson, et al.], which implements synchronous federated optimization. For FedAvg, in each epoch, $k=10$ devices are randomly selected to launch local updates. We also consider single-thread SGD as a baseline. For FedAsync, we simulate the asynchrony by randomly sampling the staleness $(t-\tau)$ from a uniform distribution.

We repeat each experiment 10 times and take the average. For CIFAR-10, we use the top-1 accuracy on the testing set as the evaluation metric. To compare asynchronous training and synchronous training, we consider “metrics vs. number of gradients”. The “number of gradients” is the number of gradients applied to the global model.

For convenience, we name Algorithm 1 as FedAsync. We also test the performance of FedAsync with adaptive mixing hyperparameters $\alpha_{t}=\alpha\times s(t-\tau)$, as outlined in Section 3. We employ the following three strategies for the weighting function $s(t-\tau)$ (parameterized by $a,b>0$): • Constant: $s(t-\tau)=1$. • Polynomial: $s_{a}(t-\tau)=(t-\tau+1)^{-a}$. • Hinge: s_{a,b}(t-\tau)=\begin{cases}1&\mbox{ift-\tau\leq b}\\ \frac{1}{a(t-\tau-b)+1}&\mbox{otherwise}\end{cases}.

For convenience, we refer to FedAsync with constant $\alpha$ as FedAsync+Const, FedAsync with polynomial adaptive $\alpha$ as FedAsync+Poly, and FedAsync with hinge adaptive $\alpha$ as FedAsync+Hinge.

3 Empirical results

We test FedAsync (asynchronous federated optimization in Algorithm 1) with different learning rates $\gamma$, regularization weights $\rho$, mixing hyperparameter $\alpha$, and staleness.

In Figure 2 and 3, we show how FedAsync converges when the number of gradients grows. We can see that when the overall staleness is small, FedAsync converges as fast as SGD, and faster than FedAvg. When the staleness is larger, FedAsync converges slower. In the worst case, FedAsync has similar convergence rate as FedAvg. When $\alpha$ is too large, the convergence can be unstable, especially for FedAsync+Const. The convergence is more robust when adaptive $\alpha$ is used.

In Figure 4, we show how staleness affects the convergence of FedAsync, evaluated on CNN and CIFAR-10 dataset. Overall, larger staleness makes the convergence slower, but the influence is not catastrophic. Furthermore, the instability caused by large staleness can be mitigated by using adaptive $\alpha$. Using adaptive $\alpha$ always improves the performance, compared to using constant $\alpha$.

4 Discussion

In general, the convergence rate of FedAsync is between single-thread SGD and FedAvg. Larger $\alpha$ and smaller staleness make FedAsync closer to single-thread SGD. Smaller $\alpha$ and larger staleness makes FedAsync closer to FedAvg.

Empirically, we observe that FedAsync is generally insensitive to hyperparameters. When the staleness is large, we can tune $\alpha$ to improve the convergence. Without adaptive $\alpha$, smaller $\alpha$ is better for larger staleness. For adaptive $\alpha$, our best choice empirically was FedAsync+Hinge. FedAsync+Poly and FedAsync+Hinge have similar performance.

In summary, compared to FedAvg, FedAsync performs as good as, and in most cases better. When the staleness is small, FedAsync converges much faster than FedAvg. When the staleness is large, FedAsync still achieves similar performance as FedAvg.

Conclusion

We proposed a novel asynchronous federated optimization algorithm on non-IID training data. We proved the convergence for a restricted family of non-convex problems. Our empirical evaluation validated both fast convergence and staleness tolerance. An interesting future direction is the design of strategies to adaptively tune the mixing hyperparameters.

This work was funded in part by the following grants: NSF IIS 1909577, NSF CNS 1908888, NSF CCF 1934986 and a JP Morgan Chase Fellowship, along with computational resources donated by Intel, AWS, and Microsoft Azure.

References

Appendix A Proofs

Taking $\alpha=\frac{1}{\sqrt{H_{min}}}$, $\gamma=\frac{1}{\sqrt{T}}$, $T=H_{min}^{5}$, we have

Without loss of generality, we assume that in the $t^{\mbox{th}}$ epoch, the server receives the model $x_{new}$, with time stamp $\tau$. We assume that $x_{new}$ is the result of applying $H_{min}\leq H\leq H_{max}$ local updates to $x_{\tau}$ on the $i$th device. We also ignore $i$ in $x_{\tau,h}^{i}$ and $z_{\tau,h}^{i}$ for convenience.

Thus, using smoothness and strong convexity, conditional on $x_{\tau,h-1}$, for $\forall h\in[H]$ we have

Taking $\rho$ large enough such that $-(1+2\rho+\epsilon)V_{1}+\rho^{2}\|x_{\tau,h-1}-x_{\tau}\|^{2}-\frac{\rho}{2}\|x_{\tau,h-1}-x_{\tau}\|^{2}\geq 0,\forall x_{\tau,h-1},x_{\tau}$, and write $\nabla g_{x_{\tau}}(x_{\tau,h-1};z_{\tau,h})$ as $\nabla g_{x_{\tau}}(x_{\tau,h-1})$ for convenience, we have

where $a=\|x_{\tau,h-1}-x_{\tau}\|^{2}>0$, $b=-2V_{1}-\frac{1}{2}\|x_{\tau,h-1}-x_{\tau}\|^{2}$, $c=-(1+\epsilon)V_{1}$. Thus, we have $\gamma\left\langle\nabla G_{x_{\tau}}(x_{\tau,h-1}),\nabla g_{x_{\tau}}(x_{\tau,h-1})\right\rangle\leq\gamma\epsilon\|\nabla F(x_{\tau,h-1})\|^{2}$.

Using $\tau-(t-1)\leq K$, we have $\|x_{\tau}-x_{t-1}\|^{2}\leq\|(x_{\tau}-x_{\tau+1})+\ldots+(x_{t-1}-x_{t-1})\|^{2}\leq K\|x_{\tau}-x_{\tau+1}\|^{2}+\ldots+K\|x_{t-1}-x_{t-1}\|^{2}\leq\alpha^{2}\gamma^{2}K^{2}H_{max}^{2}{\mathcal{O}}(V_{2}).$

Also, we have $\|x_{\tau}-x_{t-1}\|\leq\|(x_{\tau}-x_{\tau+1})+\ldots+(x_{t-1}-x_{t-1})\|\leq\|x_{\tau}-x_{\tau+1}\|^{2}+\ldots+\|x_{t-1}-x_{t-1}\|^{2}\leq\alpha\gamma KH_{max}{\mathcal{O}}(\sqrt{V_{2}}).$

By rearranging the terms and telescoping, we have

where $H^{\prime}_{t}$ is the number of local iterations applied in the $t$th iteration.

By telescoping and taking total expectation, after $T$ global epochs, we have

Using $\delta=\frac{H_{max}}{H_{min}}$, and taking $\alpha=\frac{1}{\sqrt{H_{min}}}$, $\gamma=\frac{1}{\sqrt{T}}$, $T=H_{min}^{5}$, we have

Appendix B Experiment details

In Table 2, we show the detailed network structures of the CNN used in our experiments.