Achieving Linear Speedup with Partial Worker Participation in Non-IID Federated Learning

Introduction

Federated Learning (FL) is a distributed machine learning paradigm that leverages a large number of workers to collaboratively learn a model with decentralized data under the coordination of a centralized server. Formally, the goal of FL is to solve an optimization problem, which can be decomposed as:

Non-independent-identically-distributed (non-i.i.d.) datasets across workers (data heterogeneity): In conventional distributed learning in data centers, the distribution for each worker’s local dataset can usually be assumed to be i.i.d., i.e., $D_{i}=D,\forall i\in\{1,...,m\}$. Unfortunately, this assumption rarely holds for FL since data are generated locally at the workers based on their circumstances, i.e., $D_{i}\neq D_{j}$, for $i\neq j$. It will be seen later that the non-i.i.d assumption imposes significant challenges in algorithm design for FL and their performance analysis.

Time-varying partial worker participation (systems non-stationarity): With the flexibility for workers’ participation in many scenarios (particularly in mobile edge computing), workers may randomly join or leave the FL system at will, thus rendering the active worker set stochastic and time-varying across communication rounds. Hence, it is often infeasible to wait for all workers’ responses as in traditional distributed learning, since inactive workers or stragglers will significantly slow down the whole training process. As a result, only a subset of the workers may be chosen by the server in each communication round, i.e., partial worker participation.

In recent years, the Federated Averaging method (FedAvg) and its variants (McMahan et al., 2016; Li et al., 2018; Hsu et al., 2019; Karimireddy et al., 2019; Wang et al., 2019a) have emerged as a prevailing approach for FL. Similar to the traditional distributed learning, FedAvg leverages local computation at each worker and employs a centralized parameter server to aggregate and update the model parameters. The unique feature of FedAvg is that each worker runs multiple local stochastic gradient descent (SGD) steps rather than just one step as in traditional distributed learning between two consecutive communication rounds. For i.i.d. datasets and the full worker participation setting, Stich (2018) and Yu et al. (2019b) proposed two variants of FedAvg that achieve a convergence rate of $\mathcal{O}(\frac{mK}{T}+\frac{1}{\sqrt{mKT}})$ with a bounded gradient assumption for both strongly convex and non-convex problems, where $m$ is the number of workers, $K$ is the local update steps, and $T$ is the total communication rounds. Wang & Joshi (2018) and Stich & Karimireddy (2019) further proposed improved FedAvg algorithms to achieve an $\mathcal{O}(\frac{m}{T}+\frac{1}{\sqrt{mKT}})$ convergence rate without bounded gradient assumption. Notably, for a sufficiently large $T$, the above rates become $\mathcal{O}(\frac{1}{\sqrt{mKT}})$This rate also matches the convergence rate order of parallel SGD in conventional distributed learning., which implies a linear speedup with respect to the number of workers.To attain $\epsilon$ accuracy for an algorithm, it needs to take $\mathcal{O}(\frac{1}{\epsilon^{2}})$ steps with a convergence rate $\mathcal{O}(\frac{1}{\sqrt{T}})$, while needing $\mathcal{O}(\frac{1}{m\epsilon^{2}})$ steps if the convergence rate is $\mathcal{O}(\frac{1}{\sqrt{mT}})$ (the hidden constant in Big-O is the same). In this sense, one achieves a linear speedup with respect to the number of workers. This linear speedup is highly desirable for an FL algorithm because the algorithm is able to effectively leverage the massive parallelism in a large FL system. However, with non-i.i.d. datasets and partial worker participation in FL, a fundamental open question arises: Can we still achieve the same linear speedup for convergence, i.e., $\mathcal{O}(\frac{1}{\sqrt{mKT}})$, with non-i.i.d. datasets and under either full or partial worker participation?

In this paper, we show the answer to the above question is affirmative. Specifically, we show that a generalized FedAvg with two-sided learning rates achieves linear convergence speedup with non-i.i.d. datasets and under full/partial worker participation. We highlight our contributions as follows:

For non-convex problems, we show that the convergence rate of the FedAvg algorithm on non-i.i.d. dataset are $\mathcal{O}(\frac{1}{\sqrt{mKT}}+\frac{1}{T})$ and $\mathcal{O}(\frac{\sqrt{K}}{\sqrt{nT}}+\frac{1}{T})$ for full and partial worker participation, respectively, where $n$ is the size of the partially participating worker set. This indicates that our proposed algorithm achieves a linear speedup for convergence rate for a sufficiently large $T$. When reduced to the i.i.d. case, our convergence rate is $\mathcal{O}(\frac{1}{TK}+\frac{1}{\sqrt{mKT}})$, which is also better than previous works. We summarize the convergence rate comparisons for both i.i.d. and non-i.i.d. cases in Table 1. It is worth noting that our proof does not require the bounded gradient assumption. We note that the SCAFFOLD algorithm (Karimireddy et al., 2019) also achieves the linear speedup but extra variance reduction operations are required, which lead to higher communication costs and implementation complexity. By contrast, we do not have such extra requirements in this paper.

In order to achieve a linear speedup, i.e., a convergence rate $\mathcal{O}(\frac{1}{\sqrt{mKT}})$, we show that the number of local updates $K$ can be as large as $T/m$, which improves the $T^{1/3}/m$ result previously shown in Yu et al. (2019a) and Karimireddy et al. (2019). As shown later in the communication complexity comparison in Table 1, a larger number of local steps implies relatively fewer communication rounds, thus less communication overhead. Interestingly, our results also indicate that the number of local updates $K$ does not hurt but rather help the convergence with a proper learning rates choice in full worker participation. This overcomes the limitation as suggested in Li et al. (2019b) that local SGD steps might slow down the convergence ($\mathcal{O}(\frac{K}{T})$ for strongly convex case). This result also reveals new insights on the relationship between the number of local steps and learning rate.

The rest of the paper is organized as follows. In Section 2, we review the literature to put our work in comparative perspectives. Section 3 presents the convergence analysis for our proposed algorithm. Section 4 discusses the implication of the convergence rate analysis. Section 5 presents numerical results and Section 6 concludes this paper. Due to space limitation, the details of all proofs and some experiments are provided in the supplementary material.

Related work

The federated averaging (FedAvg) algorithm was first proposed by McMahan et al. (2016) for FL as a heuristic to improve communication efficiency and data privacy. Since then, this work has sparked many follow-ups that focus on FL with i.i.d. datasets and full worker participation (also known as LocalSGD (Stich, 2018; Yu et al., 2019b; Wang & Joshi, 2018; Stich & Karimireddy, 2019; Lin et al., 2018; Khaled et al., 2019a; Zhou & Cong, 2017)). Under these two assumptions, most of the theoretical works can achieve a linear speedup for convergence, i.e., $\mathcal{O}(\frac{1}{\sqrt{mKT}})$ for a sufficiently large $T$, matching the rate of the parallel SGD. In addition, LocalSGD is empirically shown to be communication-efficient and enjoys better generalization performance (Lin et al., 2018). For a comprehensive introduction to FL, we refer readers to Li et al. (2019a) and Kairouz et al. (2019).

For non-i.i.d. datasets, many works (Sattler et al., 2019; Zhao et al., 2018; Li et al., 2018; Wang et al., 2019a; Karimireddy et al., 2019; Huang et al., 2018; Jeong et al., 2018) heuristically demonstrated the performance of FedAvg and its variants. On convergence rate with full worker participation, many works (Stich et al., 2018; Yu et al., 2019a; Wang & Joshi, 2018; Karimireddy et al., 2019; Reddi et al., 2020) can achieve linear speedup, but their convergence rate bounds could be improved as shown in this paper. On convergence rate with partial worker participation, Li et al. (2019b) showed that the original FedAvg can achieve $\mathcal{O}(K/T)$ for strongly convex functions, which suggests that local SGD steps slow down the convergence in the original FedAvg. Karimireddy et al. (2019) analyzed a generalized FedAvg with two-sided learning rates under strongly convex, convex and non-convex cases. However, as shown in Table 1, none of them indicates that linear speedup is achievable with non-i.i.d. datasets under partial worker participation. Note that the SCAFFOLD algorithm (Karimireddy et al., 2019) can achieve linear speedup but extra variance reduction operations are required, which lead to higher communication costs and implementation complexity. In this paper, we show that this linear speedup can be achieved without any extra requirements. For more detailed comparisons and other algorithmic variants in FL and decentralized settings, we refer readers to Kairouz et al. (2019).

Linear Speedup of the Generalized FedAvg with Two-Sided Learning Rates for Non-IID Datasets

In this paper, we consider a FedAvg algorithm with two-sided learning rates as shown in Algorithm 1, which is generalized from previous works (Karimireddy et al., 2019; Reddi et al., 2020). Here, workers perform multiple SGD steps using a worker optimizer to minimize the local loss on its own dataset, while the server aggregates and updates the global model using another gradient-based server optimizer based on the returned parameters. Specifically, between two consecutive communication rounds, each worker performs $K$ SGD steps with the worker’s local learning rate $\eta_{L}$. We assume an unbiased estimator in each step, which is denoted by $\mathbf{g}_{t,k}^{i}=\nabla F_{i}(\mathbf{x}_{t,k}^{i},\xi_{t,k}^{i})$, where $\xi_{t,k}^{i}$ is a random local data sample for $k$-th steps after $t$-th communication round at worker $i$. Then, each worker sends the accumulative parameter difference $\Delta_{t}^{i}$ to the server. On the server side, the server aggregates all available $\Delta_{t}^{i}$-values and updates the model parameters with a global learning rate $\eta$. The FedAvg algorithm with two-sided learning rates provides a natural way to decouple the learning of workers and server, thus utilizing different learning rate schedules for workers and the server. The original FedAvg can be viewed as a special case of this framework with server-side learning rate being one.

In what follows, we show that a linear speedup for convergence is achievable by the generalized FedAvg for non-convex functions on non-i.i.d. datasets. We first state our assumptions as follows.

The first two assumptions are standard in non-convex optimization (Ghadimi & Lan, 2013; Bottou et al., 2018). For Assumption 3, the bounded local variance is also a standard assumption. We use a universal bound $\sigma_{G}$ to quantify the heterogeneity of the non-i.i.d. datasets among different workers. In particular, $\sigma_{G}=0$ corresponds to i.i.d. datasets. This assumption is also used in other works for FL under non-i.i.d. datasets (Reddi et al., 2020; Yu et al., 2019b; Wang et al., 2019b) as well as in decentralized optimization (Kairouz et al., 2019). It is worth noting that we do not require a bounded gradient assumption, which is often assumed in FL optimization analysis.

In this subsection, we first analyze the convergence rate of the generalized FedAvg with two-sided learning rates under full worker participation, for which we have the following result:

Let constant local and global learning rates $\eta_{L}$ and $\eta$ be chosen as such that $\eta_{L}\leq\frac{1}{8LK}$ and $\eta\eta_{L}\leq\frac{1}{KL}$. Under Assumptions 1–3 and with full worker participation, the sequence of outputs $\{\mathbf{x}_{k}\}$ generated by Algorithm 1 satisfies:

where $\Phi\triangleq\frac{1}{c}[\frac{L\eta\eta_{L}}{2m}\sigma_{L}^{2}+\frac{5K\eta_{L}^{2}L^{2}}{2}(\sigma_{L}^{2}+6K\sigma_{G}^{2})]$, $c$ is a constant, $f_{0}\triangleq f(\mathbf{x}_{0})$, $f_{*}\triangleq f(\mathbf{x}_{*})$ and the expectation is over the local dataset samples among workers.

The convergence bound contains two parts: a vanishing term $\frac{f_{0}-f_{*}}{c\eta\eta_{L}KT}$ as $T$ increases and a constant term $\Phi$ whose size depends on the problem instance parameters and is independent of $T$. The vanishing term’s decay rate matches that of the typical SGD methods.

The first part of $\Phi$ (i.e., $\frac{L\eta\eta_{L}}{2m}\sigma_{L}^{2}$) is due to the local stochastic gradients at each worker, which shrinks at rate $\frac{1}{m}$ as $m$ increases. The cumulative variance of the $K$ local steps contributes to the second term in $\Phi$ (i.e., $\frac{5K\eta_{L}^{2}L^{2}}{2}(\sigma_{L}^{2}+6K\sigma_{G}^{2}))$, which is independent of $m$ and largely affected by the data heterogeneity. To make the second part small, an inverse relationship between the local learning rate and local steps should be satisfied, i.e., $\eta_{L}=\mathcal{O}(\frac{1}{K})$. Specifically, note that the global and local variances are quadratically and linearly amplified by $K$. This requires a sufficiently small $\eta_{L}$ to offset the variance between two successive communication rounds to make the second term in $\Phi$ small. This is consistent with the observation in strongly convex FL that a decaying learning rate is needed for FL to converge under non-i.i.d. datasets even if full gradients used in each worker (Li et al., 2019b). However, we note that our explicit inverse relationship between $\eta_{L}$ and $K$ in the above is new. Intuitively, the $K$ local steps with a sufficiently small $\eta_{L}$ can be viewed as one SGD step with a large learning rate.

With Theorem 1, we immediately have the following convergence rate for the generalized FedAvg algorithm with a proper choice of two-sided learning rates:

The generalized FedAvg algorithm with two-sided learning rates can achieve a linear speedup for non-i.i.d. datasets, i.e., a $\mathcal{O}(\frac{1}{\sqrt{mKT}})$ convergence rate as long as $T\geq mK$. Although many works have achieved this convergence rate asymptotically, we improve the maximum number of local steps $K$ to $T/m$, which is significantly better than the state-of-art bounds such as $T^{1/3}/m$ shown in (Karimireddy et al., 2019; Yu et al., 2019a; Kairouz et al., 2019). Note that a larger number of local steps implies relatively fewer communication rounds, thus less communication overhead. See also the communication complexity comparison in Table 1. For example, when $T=10^{6}$ and $m=100$ (as used in (Kairouz et al., 2019)), the local steps in our algorithm is $K\leq T/m=10^{4}$. However, $K\leq\frac{T^{1/3}}{m}=1$ means that no extra local steps can be taken to reduce communication costs.

When degenerated to the i.i.d. case ($\sigma_{G}=0$), the convergence rate becomes $\mathcal{O}(\frac{1}{TK}+\frac{1}{\sqrt{mKT}})$, which has a better first term in the bound compared with previous work as shown in Table 1.

2 Convergence analysis for partial worker participation

Partial worker participation in each communication round may be more practical than full worker participation due to many physical limitations of FL in practice (e.g., excessive delays because of too many devices to poll, malfunctioning devices, etc.). Partial worker participation can also accelerate the training by neglecting stragglers. We consider two sampling strategies proposed by Li et al. (2018) and Li et al. (2019b). Let $S_{t}$ be the participating worker index set at communication round $t$ with $|S_{t}|=n$, $\forall t$, for some $n\in(0,m]$. $S_{t}$ is randomly and independently selected either with replacement (Strategy 1) or without replacement (Strategy 2) sequentially according to the sampling probabilities $p_{i},\forall i\in[m]$. For each member in $S_{t}$, we pick a worker from the entire set $[m]$ uniformly at random with probability $p_{i}=\frac{1}{m},\forall i\in[m]$. That is, selection likelihood for anyone worker $i\in S_{t}$ is $p=\frac{n}{m}$. Then we have the following results:

Under Assumptions 1–3 with partial worker participation, the sequence of outputs $\{\mathbf{x}_{k}\}$ generated by Algorithm 1 with constant learning rates $\eta$ and $\eta_{L}$ satisfies:

where $f_{0}=f(\mathbf{x}_{0})$, $f_{*}=f(\mathbf{x}_{*})$, and the expectation is over the local dataset samples among workers.

For sampling Strategy 1, let $\eta$ and $\eta_{L}$ be chosen as such that $\eta_{L}\leq\frac{1}{8LK}$, $\eta\eta_{L}KL<\frac{n-1}{n}$ and $30K^{2}\eta_{L}^{2}L^{2}-\frac{L\eta\eta_{L}}{n}(90K^{3}L^{2}\eta_{L}^{2}+3K)<1$. It then holds that:

For sampling Strategy 2, let $\eta$ and $\eta_{L}$ be chosen as such that $\eta_{L}\leq\frac{1}{8LK}$, $\eta\eta_{L}KL\leq\frac{n(m-1)}{m(n-1)}$ and $10K^{2}\eta_{L}^{2}L^{2}-L\eta\eta_{L}\frac{m-n}{n(m-1)}(90K^{3}\eta_{L}^{2}L^{2}+3K)<1$. It then holds that:

From Theorem 2, we immediately have the following convergence rate for the generalized FedAvg algorithm with a proper choice of two-sided learning rates:

Let $\eta_{L}=\frac{1}{\sqrt{T}KL}$ and $\eta=\sqrt{Kn}$. The convergence rate of the generalized FedAvg algorithm under partial worker participation and both sampling strategies are:

The convergence rate bound for partial worker participation has the same structure but with a larger variance term. This implies that the partial worker participation through the uniform sampling does not result in fundamental changes in convergence (in order sense) except for an amplified variance due to fewer workers participating and random sampling. The intuition is that the uniform sampling (with/without replacement) for worker selection yields a good approximation of the entire worker distribution in expectation, which reduces the risk of distribution deviation due to the partial worker participation. As shown in Section 5, the distribution deviation due to fewer worker participation could render the training unstable, especially in highly non-i.i.d. cases.

The generalized FedAvg with partial worker participation under non-i.i.d. datasets can still achieve a linear speedup $\mathcal{O}(\frac{\sqrt{K}}{\sqrt{nT}})$ with proper learning rate settings as shown in Corollary 2. In addition, when degenerated to i.i.d. case ($\sigma_{G}=0$), the convergence rate becomes $\mathcal{O}(\frac{1}{TK}+\frac{1}{\sqrt{nKT}})$.

Here, we let $|S_{t}|=n$ only for ease of presentation and better readability. We note that this is not a restrictive condition. We can show that $|S_{t}|=n$ can be relaxed to $|S_{t}|\geq n,\forall t\in[T]$ and the same convergence rate still holds. In fact, our full proof in Appendix A.2 is for $|S_{t}|\geq n$.

Discussion

In light of above results, in what follows, we discuss several insights from the convergence analysis:

Convergence Rate: We show that the generalized FedAvg algorithm with two-sided learning rates can achieve a linear speedup, i.e., an $\mathcal{O}(\frac{1}{\sqrt{mKT}})$ convergence rate with a proper choice of hyper-parameters. Thus, it works well in large FL systems, where massive parallelism can be leveraged to accelerate training. The key challenge in convergence analysis stems from the different local loss functions (also called “model drift” in the literature) among workers due to the non-i.i.d. datasets and local steps. As shown above, we obtain a convergence bound for the generalized FedAvg method containing a vanishing term and a constant term (the constant term is similar to that of SGD). In contrast, the constant term in SGD is only due to the local variance. Note that, similar to SGD, the iterations do not diminish the constant term. The local variance $\sigma_{L}^{2}$ (randomness of stochastic gradients), global variability $\sigma_{G}^{2}$ (non-i.i.d. datasets), and the number of local steps $K$ (amplification factor) all contribute to the constant term, but the total global variability in $K$ local steps dominates the term. When the local learning rate $\eta_{L}$ is set to an inverse relationship with respect to the number of local steps $K$, the constant term is controllable. An intuitive explanation is that the $K$ small local steps can be approximately viewed as one large step in conventional SGD. So this speedup and the more allowed local steps can be largely attributed to the two-sided learning rates setting.

Number of Local Steps: Besides the result that the maximum number of local steps is improved to $K\leq T/m$, we also show that the local steps could help the convergence with the proper hyper-parameter choices, which supports previous numerical results (McMahan et al., 2016; Stich, 2018; Lin et al., 2018) and is verified in different models with different non-i.i.d. degree datasets in Section 5. However, there are other results showing the local steps slow down the convergence (Li et al., 2019b). We believe that whether local steps help or hurt the convergence in FL worths further investigations.

Number of Workers: We show that the convergence rate improves substantially as the the number of workers in each communication round increases. This is consistent with the results for i.i.d. cases in Stich (2018). For i.i.d. datasets, more workers means more data samples and thus less variance and better performance. For non-i.i.d. datasets, having more workers implies that the distribution of the sampled workers is a better approximation for the distribution of all workers. This is also empirically observed in Section 5. On the other hand, the sampling strategy plays an important role in non-i.i.d. case as well. Here, we adopt the uniform sampling (with/without replacement) to enlist workers to participate in FL. Intuitively, the distribution of the sampled workers’ collective datasets under uniform sampling yields a good approximation of the overall data distribution in expectation.

Note that, in this paper, we assume that every worker is available to participate once being enlisted. However, this may not always be feasible. In practice, the workers need to be in certain states in order to be able to participate in FL (e.g., in charging or idle states, etc. (Eichner et al., 2019)). Therefore, care must be taken in sampling and enlisting workers in practice. We believe that the joint design of sampling schemes and the generalized FedAvg algorithm will have a significant impact on the convergence, which needs further investigations.

Numerical Results

We perform extensive experiments to verify our theoretical results. We use three models: logistic regression (LR), a fully-connected neural network with two hidden layers (2NN) and a convolution neural network (CNN) with the non-i.i.d. version of MNIST (LeCun et al., 1998) and one ResNet model with CIFAR-10 (Krizhevsky et al., 2009). Due to space limitation, we relegate some experimental results in the supplementary material.

In this section, we elaborate the results under non-i.i.d. MNIST datasets for the 2NN. We distribute the MNIST dataset among $m=100$ workers randomly and evenly in a digit-based manner such that the local dataset for each worker contains only a certain class of digits. The number of digits in each worker’s dataset represents the non-i.i.d. degree. For $digits\_10$, each worker has training/testing samples with ten digits from to $9$, which is essentially an i.i.d. case. For $digits\_1$, each worker has samples only associated with one digit, which leads to highly non-i.i.d. datasets among workers. For partial worker participation, we set the number of workers $n=10$ in each communication round.

Impact of non-i.i.d. datasets: As shown in Figure 1(a), for the 2NN model with full worker participation, the top-row figures are for training loss versus communication round and the bottom-row are for test accuracy versus communication round. We can see that the generalized FedAvg algorithm converges under non-i.i.d. datasets with a proper learning rate choice in both cases. For five digits ($digits\_5$) in each worker’s dataset with full (partial) worker participation in Figure 1(a), the generalized FedAvg algorithm achieves a convergence speed comparable to that of the i.i.d. case ($digits\_10$). Another key observation is that non-i.i.d. datasets slow down the convergence under the same learning rate settings for both cases. The higher the non-i.i.d. degree, the slower the convergence speed. As the non-i.i.d. degree increases (from case $digits\_10$ to case $digits\_1$), it is obvious that the training loss is increasing and test accuracy is decreasing. This trend is more obvious from the zigzagging curves for partial worker participation. These two observations can also be verified for other models as shown in the supplementary material, which confirms our theoretical analysis.

Impact of worker number: As shown in Figure 1(b), we compare the training loss and test accuracy between full worker participation $n=100$ and partial worker participation $n=10$ with the same hyper-parameters. Compared with full worker participation, partial worker participation introduces another source of randomness, which leads to zigzagging convergence curves and slower convergence. This problem is more prominent for highly non-i.i.d. datasets. For full worker participation, it can neutralize the the system heterogeneity in each communication round. However, it might not be able to neutralize the gaps among different workers for partial worker participation. That is, the datasets’ distribution does not approximate the overall distribution well. Specifically, it is not unlikely that the digits in these datasets among all active workers are only a proper subset of the total 10 digits in the original MNIST dataset, especially with highly non-i.i.d. datasets. This trend is also obvious for complex models and complicated datasets as shown in the supplementary material. The sampling strategy here is random sampling with equal probability without replacement. In practice, however, the actual sampling of the workers in FL could be more complex, which requires further investigations.

Impact of local steps: One open question of FL is that whether the local steps help the convergence or not. In Figure 1(c), we show that the local steps could help the convergence for both full and partial worker participation. These results verify our theoretical analysis. However, Li et al. (2019b) showed that the local steps may hurt the convergence, which was demonstrated under unbalanced non-i.i.d. MNIST datasets. We believe that this may be due to the combined effect of unbalanced datasets and local steps rather than just the use of local steps only.

Comparison with SCAFFOLD: Lastly, we compare with the SCAFFOLD algorithm (Karimireddy et al., 2019) since it also achieves the same linear speedup effect under non-i.i.d. datasets. We compare communication rounds, total communication load, and estimated wall-clock time under the same settings to achieve certain test accuracy, and the results are reported in Table 2. The non-i.i.d. dataset is $digits\_2$ and the i.i.d. dataset is $digits\_10$. The learning rates are $\eta_{L}=0.1,\eta=1.0$, and number of local steps $K$ is $5$ epochs. We set the target accuracy $\epsilon=95\%$ for MNIST and $\epsilon=75\%$ for CIFAR-10. Note that the total training time contains two parts: i) the computation time for training the local model at each worker and ii) the communication time for information exchanges between the workers and the server. We assume the bandwidth $20$ MB/s for both uplink and downlink connections. For MNIST datasets, we can see that our algorithm is similar to or outperforms SCAFFOLD. This is because the numbers of communication rounds of both algorithms are relatively small for such simple tasks. For non-i.i.d. CIFAR-10, the SCAFFOLD algorithm takes slightly fewer number of communication rounds than our FedAvg algorithm to achieve $\epsilon=75\%$ thanks to its variance reduction. However, it takes more than 1.5 times of communication cost and wall-clock time compared to those of our FedAvg algorithm. Due to space limitation, we relegate the results of time proportions for computation and communication to Appendix B (see Figure 7).

Conclusions and future work

In this paper, we analyzed the convergence of a generlized FedAvg algorithm with two-sided learning rates on non-i.i.d. datasets for general non-convex optimization. We proved that the generalized FedAvg algorithm achieves a linear speedup for convergence under full and partial worker participation. We showed that the local steps in FL could help the convergence and we improve the maximum number of local steps to $T/m$. While our work sheds light on theoretical understanding of FL, it also opens the doors to many new interesting questions in FL, such as how to sample optimally in partial worker participation, and how to deal with active participant sets that are both time-varying and size-varying across communication rounds. We hope that the insights and proof techniques in this paper can pave the way for many new research directions in the aforementioned areas.

Acknowledgements

This work is supported in part by NSF grants CAREER CNS-1943226, CIF-2110252, ECCS-1818791, CCF-1934884, ONR grant ONR N00014-17-1-2417, and a Google Faculty Research Award.

References

Appendix A Appendix I: Proofs

In this section, we give the proofs in detail for full and partial worker participation in Section A.1 and Section A.2, respectively.

For convenience, we define $\bar{\Delta}_{t}\triangleq\frac{1}{m}\sum_{i=1}^{m}\Delta_{t}^{i}$. Under full device participation (i.e., $S_{t}=[m]$), it is clear that $\Delta_{t}=\frac{1}{m}\sum_{i=1}^{m}\Delta_{t}^{i}=\bar{\Delta}_{t}$.

Due to the smoothness in Assumption 1, taking expectation of $f(\mathbf{x}_{t+1})$ over the randomness at communication round $t$, we have:

Note that the term $A_{1}$ in (1) can be bounded as follows:

The term $A_{2}$ in (1) can be bounded as:

Substituting the inequalities in (2) of $A_{1}$ and (3) of $A_{2}$ into inequality (1), we have:

where $(a7)$ follows from $(\frac{\eta\eta_{L}}{2Km^{2}}-\frac{L\eta^{2}\eta_{L}^{2}}{2m^{2}})\geq 0$ if $\eta\eta_{L}\leq\frac{1}{KL}$, $(a8)$ holds because there exists a constant $c>0$ satisfying $(\frac{1}{2}-15K^{2}\eta_{L}^{2}L^{2})>c>0$ if $\eta_{L}<\frac{1}{\sqrt{30}KL}$.

Rearranging and summing from $t=0,\cdots,T-1$, we have:

where $\Phi=\frac{1}{c}[\frac{L\eta\eta_{L}}{2m}\sigma_{L}^{2}+\frac{5K\eta_{L}^{2}L^{2}}{2}(\sigma_{L}^{2}+6K\sigma_{G}^{2})]$. This completes the proof. ∎

A.2 Proof of Theorem 2

Due to the smoothness assumption in Assumption 1, taking expectation of $f(\mathbf{x}_{t+1})$ over the randomness at communication round t:

For strategy 1: We can bound $A^{{}^{\prime}}_{2}$ in (4) as follows.

Note $S_{t}$ is an index set (multiset) for independent sampling (equal probability) with replacement in which some elements may have the same value. Suppose $S_{t}=\{l_{1},\dots,l_{n}\}$.

By letting $\mathbf{t}_{i}=\sum_{j=0}^{K-1}\nabla F_{i}(\mathbf{x}_{t,j}^{i})$, we have:

where $(b3)$ is due to the independent sampling with replacement.

So we can bound $A^{{}^{\prime}}_{2}$ as follows.

Substituting the inequalities in ( 5) of $A_{1}^{{}^{\prime}}$ and ( 6) of $A_{2}^{{}^{\prime}}$ into inequality (4), we have:

where $(b6)$ follows from $\frac{(n-1)L\eta^{2}\eta_{L}^{2}}{2m^{2}n}-\frac{\eta\eta_{L}}{2Km^{2}}\leq 0$ if $\eta\eta_{L}KL\leq\frac{n-1}{n}$, $(b7)$is due to inequality (7) and $(b8)$ holds since there exists a constant $c>0$ such that $[\frac{1}{2}-15K^{2}\eta_{L}^{2}L^{2}-\frac{L\eta\eta_{L}}{2n}(90K^{3}L^{2}\eta_{L}^{2}+3K)]>c>0$ if $30K^{2}\eta_{L}^{2}L^{2}-\frac{L\eta\eta_{L}}{n}(90K^{3}L^{2}\eta_{L}^{2}+3K)<1$.

Note that the requirement of $|S_{t}|=n$ can be relaxed to $|S_{t}|\geq n$. With $p_{t}\geq n$ workers in $t$-th communication round, 8 is

That is, the same convergence rate can be guaranteed if at least $n$ workers in each communication round (no need to be exactly $n$).

Rearranging and summing from $t=0,\cdots,T-1$, we have the convergence for partial device participation with sampling strategy 1 as follows:

where \Phi=\frac{1}{c}\big{[}\frac{L\eta\eta_{L}}{2n}\sigma_{L}^{2}+\frac{3KL\eta\eta_{L}}{2n}\sigma_{G}^{2}+(\frac{5K\eta_{L}^{2}L^{2}}{2}+\frac{15K^{2}\eta\eta_{L}^{3}L^{3}}{2n})(\sigma_{L}^{2}+6K\sigma_{G}^{2})\big{]} and $c$ is a constant.

For strategy 2: Under the strategy of independent sampling with equal probability without replacement. We bound $A^{{}^{\prime}}_{2}$ as follows.

Substituting the inequalities in (5) of $A_{1}^{{}^{\prime}}$ and (9) of $A_{2}^{{}^{\prime}}$ into inequality (4), we have:

Then we bound $A_{3}^{{}^{\prime}}$ as follows.

By letting $\mathbf{t}_{i}=\sum_{j=0}^{K-1}\nabla F_{i}(\mathbf{x}_{t,j}^{i})$, we have:

where $(b14)$ follows from the fact that $\|\sum_{i\in[m]}\mathbf{t}_{i}\|^{2}=\sum_{i\in[m]}m\|\mathbf{t}_{i}\|^{2}-\frac{1}{2}\sum_{i\neq j}\|\mathbf{t}_{i}-\mathbf{t}_{j}\|^{2}$, and $(b15)$ is due to the fact that $(\frac{\eta\eta_{L}}{2Km^{2}}-\frac{L\eta^{2}\eta_{L}^{2}(n-1)}{2mn(m-1)})\geq 0$ if $\eta\eta_{L}KL\leq\frac{n(m-1)}{m(n-1)}$.

where $(b16)$ holds because there exists a constant $c>0$ satisfying $(\frac{1}{2}-5K^{2}\eta_{L}^{2}L^{2}-L\eta\eta_{L}\frac{m-n}{2n(m-1)}(90K^{3}\eta_{L}^{2}L^{2}+3K))>c>0$ if $10K^{2}\eta_{L}^{2}L^{2}-L\eta\eta_{L}\frac{m-n}{n(m-1)}(90K^{3}\eta_{L}^{2}L^{2}+3K)<1$.

Note that the requirement of $|S_{t}|=n$ can be relaxed to $|S_{t}|\geq n$. With $p_{t}\geq n$ workers in $t$-th communication round, 10 is

That is, the same convergence rate can be guaranteed if at least $n$ workers in each communication round (no need to be exactly $n$).

Rearranging and summing from $t=0,\cdots,T-1$, we have the convergence for partial device participation with sampling strategy 2 as follows:

where \Phi=\frac{1}{c}\big{[}\frac{L\eta\eta_{L}}{2n}\sigma_{L}^{2}+3KL\eta\eta_{L}\frac{m-n}{2n(m-1)}\sigma_{G}^{2}+(\frac{5K\eta_{L}^{2}L^{2}}{2}+15K^{2}\eta\eta_{L}^{3}L^{3}\frac{m-n}{2n(m-1)})(\sigma_{L}^{2}+6K\sigma_{G}^{2})\big{]} and $c$ is a constant. This completes the proof. ∎

For strategies 1 and 2, the estimator $\Delta_{t}$ is unbiased, i.e.,

Proof of Lemma 1. Let $S_{t}=\{t_{1},\cdots,t_{n}\}$ with size $n$. Both for sampling strategies 1 and 2, each sampling distribution is identical. Then we have:

A.3 Auxiliary Lemmas

For any step-size satisfying $\eta_{L}\leq\frac{1}{8LK}$, we can have the following results:

In order for this paper to be self-contained, we restate the proof of Lemma 4 in (Reddi et al., 2020) here.

For any worker $i\in[m]$ and $k\in[K]$, we have:

Appendix B Appendix II: Experiments

We provide the full detail of the experiments. We uses non-i.i.d. versions for MNIST and CIFAR-10, which are described as follows:

We study image classification of handwritten digits 0-9 in MNIST and modify the MNIST dataset to a non-i.i.d. version.

To impose statistical heterogeneity, we split the data based on the digits ($p$) they contain in their dataset. We distribute the data to $m=100$ workers such that each worker contains only a certain class of digits with the same number of training/test samples. For example, for $p=1$, each worker only has training/testing samples with one digit, which causes heterogeneity among different workers. For $p=10$, each worker has samples with 10 digits, which is essentially i.i.d. case. In this way, we can use the digits in worker’s local dataset to represent the non-i.i.d. degree qualitatively. In each communication round, 100 workers run $K$ epochs locally in parallel and then the server samples $n$ workers for aggregation and update. We make a grid-search experiments for the hyper-parameters as shown in Table 3.

We run three models: multinomial logistic regression, fully-connected network with two hidden layers (2NN) (two 200 neurons hidden layers with ReLU followed by an output layer), convolutional neural network (CNN), as shown in Table 4. The results are shown in Figures 2, 3 and 4.

B.2 CIFAR-10

Unless stated otherwise, we use the following default parameter setting: the server learning rate and client learning rate are set to $\eta=1.0$ and $\eta_{L}=0.1$, respectively. The local epochs is set to $K=10$. The total number of clients is set to 100, and the clients partition number is set to $n=10$. We use the same strategy to distribute the data over clients as suggested in McMahan et al. (2016). For the i.i.d. setting, we evenly partition all the training data among all clients, i.e., each client observes 500 data; for the non-i.i.d. setting, we first sort the training data by label, then divide all the training data into 200 shards of size 250, and randomly assign two shards to each client. For the CIFAR-10 dataset, we train our classifier with the ResNet model. The results are shown in Figure 5 and Figure 6.

B.3 Discussion

Impact of non-i.i.d. datasets: Figure 2 shows the results of training loss (top) and test accuracy (bottom) for three models under different non-i.i.d. datasets with full and partial worker participation on MNIST. We can see that the FedAvg algorithm converges under non-i.i.d. datasets with a proper learning rate choice in these cases. We believe that the major challenge in FL is the non-i.i.d. datasets. For these datasets with a lower degree of non-i.i.d., the FedAvg algorithm can achieve a good result compared with the i.i.d. case. For example, when the local dataset in each worker has five digits ($p=5$) with full (partial) worker participation, the FedAvg algorithm achieves a convergence speed comparable with that of the i.i.d. case ($p=10$). This result can be observed in Figure 2 for all three models. As the degree of non-i.i.d. datasets increases, its negative impact on the convergence is becoming more obvious. The higher the degree of non-i.i.d., the slower the convergence speed. As the non-i.i.d. degree increases (from case $p=10$ to case $p=1$), it is obvious that the training loss is increasing and test accuracy is decreasing. For these with high degree of non-i.i.d., the convergence curves oscillate and are highly unstable. This trend is more obvious for complex models such for CNN in Figure 2(c).

Impact of worker number: For full worker participation, the server can have an accurate estimation of the system heterogeneity after receiving the updates for all workers and neutralize this heterogeneity in each communication round. However, partial worker participation introduces another source of randomness, which leads to zigzagging convergence curves and slower convergence. In each communication round, the server can only receive a subset of workers based on the sampling strategy. So the server could only have a coarse estimation of the system heterogeneity and might not be able to neutralize the heterogeneity among different workers for partial worker participation. This problem is more prominent for highly non-i.i.d. datasets. It is not unlikely that the digits in these datasets among all active workers are only a proper subset of the total $10$ digits in the original MNIST dataset, especially with highly non-i.i.d. datasets. For example, for $p=1$ with $10$ workers in each communication round, it is highly likely that the datasets formed by these ten workers only includes certain small number of digits (say, $4$ or $5$) rather than total $10$ digits. But for $p=5$, it is the opposite, that is, the digits in these datasets among these $10$ workers are highly likely to be $10$. So in each communication round, the server can mitigate system heterogeneity since it covers the training samples with all $10$ digits. This trend is more obvious for complex models and datasets given the dramatic drop of test accuracy in the result of CIFAR-10 in Figure 5.

The sample strategy here is random sampling with equal probability without replacement. In practice, the workers need to be in certain states in order to be able to participate in FL (e.g., in charging or idle states, etc.(Eichner et al., 2019)). Therefore, care must be taken in sampling and enlisting workers in practice. We believe that the joint design of sampling schemes, number of workers and the FedAvg algorithm will have a significant impact on the convergence, which needs further investigations.

Impact of local steps: Figure 3 and Figure 4 shows the results of training loss (top) and test accuracy (bottom) for three models under different local steps with full and partial worker participation respectively. Figure 6 shows the impact of local steps in CIFAR-10. One open question of FL is that whether the local steps help the convergence or not. Li et al. (2019b) showed a convergence rate $\mathcal{O}(\frac{K}{T})$, i.e., the local steps may hurt the convergence for full and partial worker participation. In this two figures, we can see that local steps could help the convergence for both full and partial worker participation. However, it only has a slight effect on the convergence compared to the effects of non-i.i.d. datasets and number of workers.

Comparison with SCAFFOLD: We compare SCAFFOLD (Karimireddy et al., 2019) with the generalized FedAVg algorithm in this paper in terms of communication rounds, total communication overloads and estimated wall-clock time to achieve certain test accuracy in Table 2. We run the experiments using the same GPU (NVIDIA V100) to ensure the same conditions. Here, we give a specific comparison for these two algorithms under exact condition. Note that we divide the total training time to two parts: the computation time when the worker trains the local model and the communication time when information exchanges between the worker and server. We only compare the computation time and communication time with a fixed bandwidth $20MB/s$ for both uploading and downloading connections. As shown in Figure 7, to achieve $\epsilon=75\%$, SCAFFOLD performs less communication round due to the variance reduction techniques. That is, it spends less time on computation. However, it needs to communicates as twice as the FedAvg since the control variate to perform variance reduction in each worker needs to update in each round. In this way, the communication time would be largely prolonged.