On the Convergence of FedAvg on Non-IID Data

Abstract: Federated learning enables a large amount of edge computing devices to jointly learn a model without data sharing. As a leading algorithm in this setting, Federated Averaging (\texttt{FedAvg}) runs Stochastic Gradient Descent (SGD) in parallel on a small subset of the total devices and averages the sequences only once in a while. Despite its simplicity, it lacks theoretical guarantees under realistic settings. In this paper, we analyze the convergence of \texttt{FedAvg} on non-iid data and establish a convergence rate of O(1T) for strongly convex and smooth problems, where T is the number of SGDs. Importantly, our bound demonstrates a trade-off between communication-efficiency and convergence rate. As user devices may be disconnected from the server, we relax the assumption of full device participation to partial device participation and study different averaging schemes; low device participation rate can be achieved without severely slowing down the learning.  Our results indicate that heterogeneity of data slows down the convergence, which matches empirical observations. Furthermore, we provide a necessary condition for \texttt{FedAvg} on non-iid data: the learning rate η must decay, even if full-gradient is used; otherwise, the solution will be Ω(η) away from the optimal.

Reviews and author responses:

This paper analyzes the convergence of FedAvg, the most popular algorithm for federated learning. The highlight of the paper is removing the following two assumptions: (i) the data are iid across devices, and (ii) all the devices are active. For smooth and strongly convex problems, the paper proves an O(1/T) convergence rate to global optimum for learning rate decaying like 1/t with time. It is also shown that with constant learning rate eta, the solution found can be necessarily Omega(eta) away from the optimum (for a specific problem instance), thus justifying the decaying learning rate used in the positive result.

Federated learning has been an important and popular research area since it models a highly distributed and heterogeneous learning system in real world. Previous theoretical analysis of FedAvg was quite scarce and either made the iid data assumption or required averaging all the devices. This work is the first to prove a convergence guarantee without these two assumptions. In particular, it only requires averaging a (random) subset of devices each round, which is much more realistic than averaging all.

I don't quite have an intuition for why you need strong convexity. I hope the authors could explain this in words and maybe comment on what are the challenges of removing this assumption.


------
Thanks to the authors for their response.

We greatly appreciate the reviewer's effort. Thanks for your positive reviews.

We greatly appreciate the reviewer's effort. Here are our responses to your comments.

There are several benefits when we assume strong convexity.
First, it facilitates our analysis.
We have more ways to prove the convergence since there is no difference when we prove convergence for $\|w_t - w^*\|$ or $f(w_t) - f(w^*)$.
Second, strong convexity and smoothness mean fast convergence rate $O(1/T)$. 
Since FedAvg is a distributed variant of SGD, it couldn't achieve faster convergence than SGD.
We can remove the strong convexity assumption but the convergence rate is $O(1/\sqrt{T})$ (see Khaled et al. (2019)).
Third, in Theorems 1 and 2, we require the learning rate $\eta_t = 2/(\gamma + t) 1/\mu$, which makes the use of a positive $\mu$.

We greatly appreciate the reviewer's effort. Here are our responses to your comments.

1. Yes, transformed Scheme II is the scaling trick described at the end of Section 3.3. We will clarify its definition.

2. Since FedAvg is a distributed variant of SGD, it couldn't achieve faster convergence than SGD.
For strongly convex and smooth optimization problems, the best convergence of SGD is $\Theta(1/t)$.
In our case, if $\eta_t$ that is decaying slower than $\Theta(1/t)$, you might not obtain the convergence rate $O(1/t)$.
Based on Eq. (11), when $\eta_t = 1/(\mu\sqrt{t})$, we will obtain the convergence rate $O(1/\sqrt{t})$ by a similar argument.

3. We will correct these typos in its revision.

This paper presents convergence rates for straggler-aware averaged SGD for non-identically but independent distributed data. The paper is well-written and motivated with good discussions of the algorithm and the related works. The proof techniques involve bounding how much worse can the algorithm do because of non-identical distribution and introduction of stragglers into the standard analysis of SGD-like algorithms. The presented theory is useful, and also provides new insights such as a new sampling scheme and an inherent bias for the case of non-decaying step size. The empirical evaluation is adequate and well-presented. I think this paper is a strong contribution and should spark further discussions in the community. 

Federated learning is distinguished from the standard distributed learning in the following sense: 
1) training is distributed over a huge number (say N) of devices and communication between the central server and devices are slow.
2) The central server has no control of individual devices, and there are inactive devices that does not respond to the server; full participation of all devices is unrealistic.
3) The local data distribution at each device is different from each other; i.e., the data is non-iid.

Due to property 1), communication-efficient algorithms such as Federated Averaging (FedAvg) have been proposed and studied. FedAvg runs SGD in parallel on K (≤N) local devices using their local datasets, and updates the global parameter after E local iterations by aggregating the updates from the local devices.

Properties 2) and 3) makes analysis of FedAvg difficult, and previous results have proven convergence of FedAvg assuming that the data is iid and/or all devices are active. In contrast, this paper studies FedAvg on the non-iid data and inactive devices setting and shows that, with adequately chosen aggregation schemes and decaying learning rate, FedAvg on strongly convex and smooth functions converges with a rate of O(1/T). 

Overall, I enjoyed reading this paper and I would like to recommend acceptance. This is the first result showing convergence rate analysis of FedAvg under presence of properties 2) and 3), which is a nontrivial, important, and timely problem. The paper is well-written and reads smoothly, except for some minor typos. The convergence bounds provide insights of practical relevance, e.g., the optimal choice of E, the effect of K in convergence rate, etc. The authors also provide empirical results supporting their theoretical analysis.

Some questions I have in mind:
- What is "transformed Scheme II"? Is it the scaling trick described at the end of Section 3.3? The name appears in the experiment section before being defined.
- What happens if we choose \eta_t that is decaying but slower than O(1/t), say O(1/\sqrt t)? Can convergence be proved? If so, in what rate?

Minor typos:
- Footnote 3: know -> known
- Assumptions 1 & 2: f in $f(w)$ is math-bold
- Choice of sampling schemes: "If the system can choose to active..." -> activate
- mnist balanced and mnist unbalanced: the description after them suggests they should be switched
- Apdx D.1: widely -> wide, summary -> summarize

Meta-review: This manuscript analyzes the convergence of federated learning wit hstragellers, and provides convergence rates. The proof techniques involve bounding the effects of the non-identical distribution due to stragglers and related issues. The manuscript also includes a thorough empirical evaluation. Overall, the reviewers were quite positive about the manuscript, with a few details that should be improved. 

Decision: accept-talk