---
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 (`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 `FedAvg` on non-iid data and establish a convergence rate of $\mathcal{O}(\frac{1}{T})$ 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 indicates that heterogeneity of data slows down the convergence, which matches empirical observations. Furthermore, we provide a necessary condition for `FedAvg` on non-iid data: the learning rate $\eta$ must decay, even if full-gradient is used; otherwise, the solution will be $\Omega (\eta)$ away from the optimal.'
author:
- |
    Xiang Li[^1]\
    School of Mathematical Sciences\
    Peking University\
    Beijing, 100871, China\
    `smslixiang@pku.edu.cn`\
    Kaixuan Huang[$^*$]{}\
    School of Mathematical Sciences\
    Peking University\
    Beijing, 100871, China\
    `hackyhuang@pku.edu.cn`\
    Wenhao Yang[$^*$]{}\
    Center for Data Science\
    Peking University\
    Beijing, 100871, China\
    `yangwenhaosms@pku.edu.cn`\
    Shusen Wang\
    Department of Computer Science\
    Stevens Institute of Technology\
    Hoboken, NJ 07030, USA\
    `shusen.wang@stevens.edu`\
    Zhihua Zhang\
    School of Mathematical Sciences\
    Peking University\
    Beijing, 100871, China\
    `zhzhang@math.pku.edu.cn`
bibliography:
- 'refer.bib'
title: 'On the Convergence of FedAvg on Non-IID Data'
---

Introduction {#sec:intro}
============

Federated Learning (FL), also known as federated optimization, allows multiple parties to collaboratively train a model without data sharing [@konevcny2015federated; @shokri2015privacy; @mcmahan2016communication; @konevcny2017stochastic; @sahu2018convergence; @zhuo2019federated]. Similar to the centralized parallel optimization [@jakovetic2013distributed; @li2014scaling; @li2014communication; @shamir2014communication; @zhang2015disco; @meng2016mllib; @reddi2016aide; @richtarik2016distributed; @smith2016cocoa; @zheng2016general; @wang2018giant], FL let the user devices (aka worker nodes) perform most of the computation and a central parameter server update the model parameters using the descending directions returned by the user devices. Nevertheless, FL has three unique characters that distinguish it from the standard parallel optimization [@li2019federated].

First, the training data are [massively distributed]{} over an incredibly large number of devices, and the connection between the central server and a device is slow. A direct consequence is the slow communication, which motivated communication-efficient FL algorithms [@mcmahan2016communication; @smith2017federated; @sahu2018convergence; @sattler2019robust]. Federated averaging (`FedAvg`) is the first and perhaps the most widely used FL algorithm. It runs $E$ steps of SGD in parallel on a small sampled subset of devices and then averages the resulting model updates via a central server once in a while.[^2] In comparison with SGD and its variants, `FedAvg` performs more local computation and less communication.

Second, unlike the traditional distributed learning systems, the FL system does not have control over users’ devices. For example, when a mobile phone is turned off or WiFi access is unavailable, the central server will lose connection to this device. When this happens during training, such a non-responding/inactive device, which is called a straggler, appears tremendously slower than the other devices. Unfortunately, since it has no control over the devices, the system can do nothing but waiting or ignoring the stragglers. Waiting for all the devices’ response is obviously infeasible; it is thus impractical to require all the devices be active.

Third, the training data are non-iid[^3], that is, a device’s local data cannot be regarded as samples drawn from the overall distribution. The data available locally fail to represent the overall distribution. This does not only bring challenges to algorithm design but also make theoretical analysis much harder. While `FedAvg` actually works when the data are non-iid [@mcmahan2016communication], `FedAvg` on non-iid data lacks theoretical guarantee even in convex optimization setting.

There have been much efforts developing convergence guarantees for FL algorithm based on the assumptions that (1) the data are iid and (2) all the devices are active. @khaled2019first [@yu2018parallel; @wang2019adaptive] made the latter assumption, while @zhou2017convergence [@stich2018local; @wang2018cooperative; @woodworth2018graph] made both assumptions. The two assumptions violates the second and third characters of FL. Previous algorithm `Fedprox` [@sahu2018convergence] doesn’t require the two mentioned assumptions and incorporates `FedAvg` as a special case when the added proximal term vanishes. However, their theory fails to cover `FedAvg`.

#### Notation.

Let $N$ be the total number of user devices and $K$ ($\leq N$) be the maximal number of devices that participate in every round’s communication. Let $T$ be the total number of every device’s SGDs, $E$ be the number of local iterations performed in a device between two communications, and thus $\frac{T}{E}$ is the number of communications.

#### Contributions.

For strongly convex and smooth problems, we establish a convergence guarantee for `FedAvg` without making the two impractical assumptions: (1) the data are iid, and (2) all the devices are active. To the best of our knowledge, this work is the first to show the convergence rate of `FedAvg` without making the two assumptions.

We show in Theorem \[thm:full\], \[thm:w\_replace\], and \[thm:wo\_replace\] that `FedAvg` has $\OM (\frac{1}{T})$ convergence rate. In particular, Theorem \[thm:wo\_replace\] shows that to attain a fixed precision $\epsilon$, the number of communications is $$\frac{T}{E}
\: = \:
\OM \left[ \frac{1}{\epsilon} \Bigg( \left(1+\frac{1}{K}\right) EG^2 + \frac{\myave ^2\sigma_k^2 + \Gamma + G^2}{E} + G^2 \Bigg)\right].$$ Here, $G$, $\Gamma$, $p_k$, and $\sigma_k$ are problem-related constants defined in Section \[sec:convergence:notation\]. The most interesting insight is that $E$ is a knob controlling the convergence rate: neither setting $E$ over-small ($E=1$ makes `FedAvg` equivalent to SGD) nor setting $E$ over-large is good for the convergence.

This work also makes algorithmic contributions. We summarize the existing sampling[^4] and averaging schemes for `FedAvg` (which do not have convergence bounds before this work) and propose a new scheme (see Table \[tab:conv\]). We point out that a suitable sampling and averaging scheme is crucial for the convergence of `FedAvg`. To the best of our knowledge, we are the first to theoretically demonstrate that `FedAvg` with certain schemes (see Table \[tab:conv\]) can achieve $\OM(\frac{1}{T})$ convergence rate in non-iid federated setting. We show that heterogeneity of training data and partial device participation slow down the convergence. We empirically verify our results through numerical experiments.

Our theoretical analysis requires the decay of learning rate (which is known to hinder the convergence rate.) Unfortunately, we show in Theorem \[thm:failure\] that the decay of learning rate is necessary for `FedAvg` with $E>1$, even if full gradient descent is used.[^5] If the learning rate is fixed to $\eta$ throughout, `FedAvg` would converge to a solution at least $\Omega (\eta (E-1))$ away from the optimal. To establish Theorem \[thm:failure\], we construct a specific $\ell_2$-norm regularized linear regression model which satisfies our strong convexity and smoothness assumptions.

             Paper                             Sampling                                               Averaging                                Convergence rate
  --------------------------- ------------------------------------------ ------------------------------------------------------------------- --------------------
   @mcmahan2016communication        $\SM_t \sim \mathcal{U}(N,K) $        $ \sum_{k\notin \SM_t} p_k \w_t + \sum_{k \in \SM_t} p_k \w_t^k $           -
     @sahu2018convergence      $\SM_t \sim \mathcal{W}(N,K,\mathbf{p})$                 $\frac{1}{K}\sum_{k\in \SM_t} \w_t^k$                 $\OM(\frac{1}{T})$
             Ours                   $\SM_t \sim \mathcal{U}(N,K) $                   $\sum_{k \in \SM_t} p_k \frac{N}{K} \w_t^k$              $\OM(\frac{1}{T})$

  : Sampling and averaging schemes for `FedAvg`. $\SM_t \sim \mathcal{U}(N,K)$ means $\SM_t$ is a size-$K$ subset uniformly sampled **without replacement** from $[N]$. $\SM_t \sim \mathcal{W}(N,K,\mathbf{p})$ means $\SM_t$ contains $K$ elements that are iid sampled **with replacement** from $[N]$ with probabilities $\{p_k\}$. In the latter scheme, $\SM_t$ is not a set. []{data-label="tab:conv"}

#### Paper organization.

In Section \[sec:alg\], we elaborate on `FedAvg`. In Section \[sec:convergence\], we present our main convergence bounds for `FedAvg`. In Section \[sec:example\], we construct a special example to show the necessity of learning rate decay. In Section \[sec:related\], we discuss and compare with prior work. In Section \[sec:experiment\], we conduct empirical study to verify our theories. All the proofs are left to the appendix.

Federated Averaging (`FedAvg`) {#sec:alg}
==============================

#### Problem formulation.

In this work, we consider the following distributed optimization model: $$\label{eq:loss}
\min_{\w} \;
\Big\{ F (\w )  \, \triangleq \,
\sum_{k=1}^N \, p_k F_k(\w) \Big\},$$ where $N$ is the number of devices, and $p_k$ is the weight of the $k$-th device such that $p_k  \ge  0$ and $\sum_{k=1}^N p_k = 1$. Suppose the $k$-th device holds the $n_k$ training data: $x_{k, 1}, x_{k, 2}, \cdots , x_{k, n_k}$. The local objective $F_k (\cdot)$ is defined by $$\label{eq:local-f}
F_k (\w) \triangleq \frac{1}{n_k} \sum_{j=1}^{n_k}\ell(\w; x_{k, j}),$$ where $\ell (\cdot ; \cdot)$ is a user-specified loss function.

#### Algorithm description.

Here, we describe one around (say the $t$-th) of the *standard* `FedAvg` algorithm. First, the central server **broadcasts** the latest model, $\w_t$, to all the devices. Second, every device (say the $k$-th) lets $\w_t^k = \w_t$ and then performs $E$ ($\geq 1$) **local updates**: $$\w_{t+i+1}^k \: \longleftarrow \: \w_{t+i}^k - \eta_{t+i} \nabla F_k(\w_{t+i}^k, \xi_{t+i}^k) , i=0,1,\cdots,E-1$$ where $\eta_{t+i}$ is the learning rate (a.k.a. step size) and $\xi_{t+i}^k$ is a sample uniformly chosen from the local data. Last, the server **aggregates** the local models, $\w_{t+E}^1 , \cdots, \w_{t+E}^N$, to produce the new global model, $\w_{t+E}$. Because of the non-iid and partial device participation issues, the aggregation step can vary.

#### IID versus non-iid.

Suppose the data in the $k$-th device are i.i.d. sampled from the distribution $\DM_k$. Then the overall distribution is a mixture of all local data distributions: $\DM = \sum_{k=1}^N p_k \DM_k$. The prior work [@zhang2015deep; @zhou2017convergence; @stich2018local; @wang2018cooperative; @woodworth2018graph] assumes the data are iid generated by or partitioned among the $N$ devices, that is, $\DM_k = \DM$ for all $k \in [N]$. However, real-world applications do not typically satisfy the iid assumption. One of our theoretical contributions is avoiding making the iid assumption.

#### Full device participation.

The prior work [@coppola2015iterative; @zhou2017convergence; @stich2018local; @yu2018parallel; @wang2018cooperative; @wang2019adaptive] requires the [full device participation]{} in the aggregation step of `FedAvg`. In this case, the aggregation step performs $$\w_{t+E} \: \longleftarrow \: \sum_{k=1}^N p_k \, \w_{t+E}^k .$$ Unfortunately, the full device participation requirement suffers from serious “straggler’s effect” (which means everyone waits for the slowest) in real-world applications. For example, if there are thousands of users’ devices in the FL system, there are always a small portion of devices offline. Full device participation means the central server must wait for these “stragglers”, which is obviously unrealistic.

#### Partial device participation.

This strategy is much more realistic because it does not require all the devices’ output. We can set a threshold $K$ ($1\leq K < N$) and let the central server collect the outputs of the first $K$ responded devices. After collecting $K$ outputs, the server stops waiting for the rest; the $K+1$-th to $N$-th devices are regarded stragglers in this iteration. Let $\SM_t$ ($| \SM_t | = K$) be the set of the indices of the first $K$ responded devices in the $t$-th iteration. The aggregation step performs $$\w_{t+E} \: \longleftarrow \: \frac{N}{K} \sum_{k\in \SM_t } p_k  \, \w_{t+E}^k .$$ It can be proved that $\frac{N}{K}  \sum_{k\in \SM_t } p_k $ equals one in expectation.

#### Communication cost.

The `FedAvg` requires two rounds communications— one broadcast and one aggregation— per $E$ iterations. If $T$ iterations are performed totally, then the number of communications is $\lfloor \frac{2T}{E} \rfloor$. During the broadcast, the central server sends $\w_t$ to all the devices. During the aggregation, all or part of the $N$ devices sends its output, say $\w_{t+E}^k$, to the server.

Convergence Analysis of `FedAvg` in Non-iid Setting {#sec:convergence}
===================================================

In this section, we show that `FedAvg` converges to the global optimum at a rate of $\OM(1/T)$ for strongly convex and smooth functions and non-iid data. The main observation is that when the learning rate is sufficiently small, the effect of $E$ steps of local updates is similar to one step update with a larger learning rate. This coupled with appropriate sampling and averaging schemes would make each global update behave like an SGD update. Partial device participation ($K < N$) only makes the averaged sequence $\{\w_t\}$ have a larger variance, which, however, can be controlled by learning rates. These imply the convergence property of `FedAvg` should not differ too much from SGD. Next, we will first give the convergence result with full device participation (i.e., $K=N$) and then extend this result to partial device participation (i.e., $K<N$).

Notation and Assumptions {#sec:convergence:notation}
------------------------

We make the following assumptions on the functions $F_1, \cdots , F_N$. Assumption \[asm:smooth\] and \[asm:strong\_cvx\] are standard; typical examples are the $\ell_2$-norm regularized linear regression, logistic regression, and softmax classifier.

\[asm:smooth\] $F_1, \cdots, F_N$ are all $L$-smooth: for all $\v$ and $\w$, $F_k(\v)  \leq F_k(\w) + (\v - \w)^T \nabla F_k(\w) + \frac{L}{2} \| \v - \w\|_2^2$.

\[asm:strong\_cvx\] $F_1, \cdots, F_N$ are all $\mu$-strongly convex: for all $\v$ and $\w$, $F_k(\v)  \geq F_k(\w) + (\v - \w)^T \nabla F_k(\w) + \frac{\mu }{2} \| \v - \w\|_2^2$.

Assumptions \[asm:sgd\_var\] and \[asm:sgd\_norm\] have been made by the works [@zhang2013communication; @stich2018local; @stich2018sparsified; @yu2018parallel].

\[asm:sgd\_var\] Let $\xi_t^k$ be sampled from the $k$-th device’s local data uniformly at random. The variance of stochastic gradients in each device is bounded: $\EB \left\| \nabla F_k(\w_t^k,\xi_t^k) - \nabla F_k(\w_t^k) \right\|^2 \le \sigma_k^2$ for $k=1,\cdots,N$.

\[asm:sgd\_norm\] The expected squared norm of stochastic gradients is uniformly bounded, i.e., $\EB \left\| \nabla F_k(\w_t^k,\xi_t^k) \right\|^2  \le G^2$ for all $k=1,\cdots,N$ and $t=1,\cdots, T-1$

#### Quantifying the degree of non-iid (heterogeneity).

Let $F^*$ and $F_k^*$ be the minimum values of $F$ and $F_k$, respectively. We use the term $\Gamma = F^* - \myave F_k^*$ for quantifying the degree of non-iid. If the data are iid, then $\Gamma$ obviously goes to zero as the number of samples grows. If the data are non-iid, then $\Gamma$ is nonzero, and its magnitude reflects the heterogeneity of the data distribution.

Convergence Result: Full Device Participation {#subsec:convres_full}
---------------------------------------------

Here we analyze the case that all the devices participate in the aggregation step; see Section \[sec:alg\] for the algorithm description. Let the `FedAvg` algorithm terminate after $T$ iterations and return $\w_T$ as the solution. We always require $T$ is evenly divisible by $E$ so that `FedAvg` can output $\w_T$ as expected.

\[thm:full\] Let Assumptions \[asm:smooth\] to \[asm:sgd\_norm\] hold and $L, \mu, \sigma_k, G$ be defined therein. Choose $\kappa = \frac{L}{\mu}$, $\gamma = \max\{8\kappa, E\}$ and the learning rate $\eta_t = \frac{2}{\mu (\gamma+t)}$. Then $\texttt{FedAvg}$ with [full device participation]{} satisfies $$\label{eq:bound_K=N}
    \EB \left[ F(\w_T)\right] - F^* 
    \: \leq \: 
    \frac{\kappa}{\gamma +T-1} \left( \frac{2B}{\mu} + \frac{\mu \gamma}{2} \EB \|\w_1 - \w^*\|^2 \right),$$ where $$\label{eq:bound_B}
    B = \sum_{k=1}^N p_k^2 \sigma_k^2 + 6L \Gamma + 8  (E-1)^2G^2.$$

Convergence Result: Partial Device Participation {#subsec:conv_partial}
------------------------------------------------

As discussed in Section \[sec:alg\], partial device participation has more practical interest than full device participation. Let the set $\SM_t$ ($\subset [N]$) index the active devices in the $t$-th iteration. To establish the convergence bound, we need to make assumptions on $\SM_t$.

Assumption \[asm:w\_replace\] assumes the $K$ indices are selected from the distribution $p_k$ independently and with replacement. The aggregation step is simply averaging. This is first proposed in [@sahu2018convergence], but they did not provide theoretical analysis.

\[Scheme I\] \[asm:w\_replace\] Assume $\SM_t$ contains a subset of $K$ indices randomly selected [with replacement]{} according to the sampling probabilities $p_1, \cdots , p_N$. The aggregation step of `FedAvg` performs $\w_{t} \longleftarrow \frac{1}{K} \sum_{k\in \SM_t } \w_{t}^k $.

\[thm:w\_replace\] Let Assumptions \[asm:smooth\] to \[asm:sgd\_norm\] hold and $L, \mu, \sigma_k, G$ be defined therein. Let $\kappa, \gamma$, $\eta_t$, and $B$ be defined in Theorem \[thm:full\]. Let Assumption \[asm:w\_replace\] hold and define $C =  \frac{4}{K} E^2G^2$. Then $$\label{eq:bound_K<N}
    \EB \left[ F(\w_T)\right] - F^* 
    \: \leq \: 
    \frac{\kappa}{\gamma +T - 1} \left( \frac{2(B+C)}{\mu} + \frac{\mu \gamma}{2} \EB \|\w_1 - \w^*\|^2 \right),$$

Alternatively, we can select $K$ indices from $[N]$ uniformly at random without replacement. As a consequence, we need a different aggregation strategy. Assumption \[asm:wo\_replace\] assumes the $K$ indices are selected uniformly without replacement and the aggregation step is the same as in Section \[sec:alg\]. However, to guarantee convergence, we require an additional assumption of balanced data.

\[asm:wo\_replace\] Assume $\SM_t$ contains a subset of $K$ indices uniformly sampled from $[N]$ [without replacement]{}. Assume the data is balanced in the sense that $p_1 = \cdots = p_N = \frac{1}{N}$. The aggregation step of `FedAvg` performs $\w_{t} \longleftarrow \frac{N}{K} \sum_{k\in \SM_t } p_k \, \w_{t}^k $.

\[thm:wo\_replace\] Replace Assumption \[asm:w\_replace\] by Assumption \[asm:wo\_replace\] and $C$ by $C = \frac{N-K}{N-1} \frac{4}{K} E^2G^2$. Then the same bound in Theorem \[thm:w\_replace\] holds.

Scheme II requires $p_1 = \cdots = p_N = \frac{1}{N}$ which obviously violates the unbalance nature of FL. Fortunately, this can be addressed by the following transformation. Let $\widetilde{F}_k(\w)= p_kNF_k(\w)$ be a scaled local objective $F_k$. Then the global objective becomes a simple average of all scaled local objectives: $$F(\w) 
\: = \: \sum_{k=1}^N p_k {F}_k(\w)
\: = \: \frac{1}{N} \sum_{k=1}^N \widetilde{F}_k(\w) .$$ Theorem \[thm:wo\_replace\] still holds if $L, \mu, \sigma_k, G$ are replaced by $\widetilde{L} \triangleq \nu L$, $\widetilde{\mu} \triangleq \varsigma \mu$, $\widetilde{\sigma}_k = \sqrt{\nu} \sigma$, and $\widetilde{G} = \sqrt{\nu} G$, respectively. Here, $\nu =  N \cdot \max_{k}p_k$ and $ \varsigma = N \cdot \min_{k} p_k $.

Discussions
-----------

#### Choice of $E$.

Since $\| \w_0 - \w^* \|^2 \le \frac{4}{\mu^2} G^2$ for $\mu$-strongly convex $F$, the dominating term in is $$\label{eq:bound_dom}
\mathcal{O}\left(\frac{\sum _{k=1}^N p_k^2 \sigma_k^2 + L \Gamma + \left(1+\frac{1}{K}\right)E^2G^2 + \gamma G^2}{\mu T}\right).$$ Let $T_\epsilon$ denote the number of required steps for `FedAvg` to achieve an $\epsilon$ accuracy. It follows from that the number of required communication rounds is roughly[^6] $$\label{eq:communication_round}
\frac{T_\epsilon}{E} 
\: \propto \:   \left(1+\frac{1}{K}\right)EG^2 + \frac{\myave ^2\sigma_k^2 + L \Gamma + \kappa G^2}{E} + G^2 .$$ Thus, $\frac{T_\epsilon}{E}$ is a function of $E$ that first decreases and then increases, which implies that over-small or over-large $E$ may lead to high communication cost and that the optimal $E$ exists.

@stich2018local showed that if the data are iid, then $E$ can be set to $\OM (\sqrt{T} )$. However, this setting does not work if the data are non-iid. Theorem \[thm:full\] implies that $E$ must not exceed $\Omega (\sqrt{T})$; otherwise, convergence is not guaranteed. Here we give an intuitive explanation. If $E$ is set big, then $\w_t^k$ can converge to the minimizer of $F_k$, and thus `FedAvg` becomes the one-shot average [@zhang2013communication] of the local solutions. If the data are non-iid, the one-shot averaging does not work because weighted average of the minimizers of $F_1, \cdots , F_N$ can be very different from the minimizer of $F$.

#### Choice of $K$.

@stich2018local showed that if the data are iid, the convergence rate improves substantially as $K$ increases. However, under the non-iid setting, the convergence rate has a weak dependence on $K$, as we show in Theorems \[thm:w\_replace\] and \[thm:wo\_replace\]. This implies `FedAvg` is unable to achieve linear speedup. We have empirically observed this phenomenon (see Section \[sec:experiment\]). Thus, in practice, the participation ratio $\frac{K}{N}$ can be set small to alleviate the straggler’s effect without affecting the convergence rate.

#### Choice of sampling schemes.

We considered two sampling and averaging schemes in Theorems \[thm:w\_replace\] and \[thm:wo\_replace\]. Scheme I selects $K$ devices according to the probabilities $p_1, \cdots, p_N$ with replacement. The non-uniform sampling results in faster convergence than uniform sampling, especially when $p_1, \cdots, p_N$ are highly non-uniform. If the system can choose to activate any of the $N$ devices at any time, then Scheme I should be used.

However, oftentimes the system has no control over the sampling; instead, the server simply uses the first $K$ returned results for the update. In this case, we can assume the $K$ devices are uniformly sampled from all the $N$ devices and use Theorem \[thm:wo\_replace\] to guarantee the convergence. If $p_1, \cdots, p_N$ are highly non-uniform, then $\nu =  N \cdot \max_{k}p_k$ is big and $ \varsigma = N \cdot \min_{k} p_k $ is small, which makes the convergence of `FedAvg` slow. This point of view is empirically verified in our experiments.

Necessity of Learning Rate Decay {#sec:example}
================================

In this section, we point out that diminishing learning rates are crucial for the convergence of `FedAvg` in the non-iid setting. Specifically, we establish the following theorem by constructing a ridge regression model (which is strongly convex and smooth).

\[thm:failure\] We artificially construct a strongly convex and smooth distributed optimization problem. With full batch size, $E>1$, and any [*fixed*]{} step size, `FedAvg` will converge to sub-optimal points. Specifically, let $\tilde{\w}^*$ be the solution produced by `FedAvg` with a small enough and constant $\eta$, and $\w^*$ the optimal solution. Then we have $$\|\tilde{\w}^* - \w^* \|_2 = \Omega (  (E-1)\eta ) \cdot \| \w^* \|_2.$$ where we hide some problem dependent constants.

Theorem \[thm:failure\] and its proof provide several implications. First, the decay of learning rate is necessary of `FedAvg`. On the one hand, Theorem \[thm:full\] shows with $E>1$ and a decaying learning rate, `FedAvg` converges to the optimum. On the other hand, Theorem \[thm:failure\] shows that with $E>1$ and any fixed learning rate, `FedAvg` does not converges to the optimum.

Second, `FedAvg` behaves very differently from gradient descent. Note that `FedAvg` with $E=1$ and full batch size is exactly the `Full Gradient Descent`; with a proper and fixed learning rate, its global convergence to the optimum is guaranteed [@nesterov2013introductory]. However, Theorem \[thm:failure\] shows that `FedAvg` with $E>1$ and full batch size cannot possibly converge to the optimum. This conclusion doesn’t contradict with Theorem 1 in [@khaled2019first], which, when translated into our case, asserts that $\tilde{\w}^*$ will locate in the neighborhood of $\w^*$ with a constant learning rate.

Third, Theorem \[thm:failure\] shows the requirement of learning rate decay is not an artifact of our analysis; instead, it is inherently required by `FedAvg`. An explanation is that constant learning rates, combined with $E$ steps of possibly-biased local updates, form a sub-optimal update scheme, but a diminishing learning rate can gradually eliminate such bias.

The efficiency of `FedAvg` principally results from the fact that it performs several update steps on a local model before communicating with other workers, which saves communication. Diminishing step sizes often hinders fast convergence, which may counteract the benefit of performing multiple local updates. Theorem \[thm:failure\] motivates more efficient alternatives to `FedAvg`.

Related Work {#sec:related}
============

Federated learning (FL) was first proposed by [@mcmahan2016communication] for collaboratively learning a model without collecting users’ data. The research work on FL is focused on the communication-efficiency [@konevcny2016federated; @mcmahan2016communication; @sahu2018convergence; @smith2017federated] and data privacy [@bagdasaryan2018backdoor; @bonawitz2017practical; @geyer2017differentially; @hitaj2017deep; @melisexploiting]. This work is focused on the communication-efficiency issue.

`FedAvg`, a synchronous distributed optimization algorithm, was proposed by [@mcmahan2016communication] as an effective heuristic. @sattler2019robust [@zhao2018federated] studied the non-iid setting, however, they do not have convergence rate. A contemporaneous and independent work [@xie2019asynchronous] analyzed asynchronous `FedAvg`; while they did not require iid data, their bound do not guarantee convergence to saddle point or local minimum. @sahu2018convergence proposed a federated optimization framework called `FedProx` to deal with statistical heterogeneity and provided the convergence guarantees in non-iid setting. `FedProx` adds a proximal term to each local objective. When these proximal terms vanish, `FedProx` is reduced to `FedAvg`. However, their convergence theory requires the proximal terms always exist and hence fails to cover `FedAvg`.

When data are iid distributed and all devices are active, `FedAvg` is referred to as `LocalSGD`. Due to the two assumptions, theoretical analysis of `LocalSGD` is easier than `FedAvg`. @stich2018local demonstrated `LocalSGD` provably achieves the same linear speedup with strictly less communication for strongly-convex stochastic optimization. @coppola2015iterative [@zhou2017convergence; @wang2018cooperative] studied `LocalSGD` in the non-convex setting and established convergence results. @yu2018parallel [@wang2019adaptive] recently analyzed `LocalSGD` for non-convex functions in heterogeneous settings. In particular, @yu2018parallel demonstrated `LocalSGD` also achieves $\mathcal{O}(1/\sqrt{NT})$ convergence (i.e., linear speedup) for non-convex optimization. @lin2018don empirically shows variants of `LocalSGD` increase training efficiency and improve the generalization performance of large batch sizes while reducing communication. For `LocalGD` on non-iid data (as opposed to `LocalSGD`), the best result is by the contemporaneous work (but slightly later than our first version) [@khaled2019first]. @khaled2019first used fixed learning rate $\eta$ and showed $\OM (\frac{1}{T})$ convergence to a point $\OM (\eta^2 E^2)$ away from the optimal. In fact, the suboptimality is due to their fixed learning rate. As we show in Theorem \[thm:failure\], using a fixed learning rate $\eta$ throughout, the solution by `LocalGD` is at least $\Omega ((E-1) \eta)$ away from the optimal.

If the data are iid, distributed optimization can be efficiently solved by the second-order algorithms [@mahajan2018efficient; @reddi2016aide; @shamir2014communication; @wang2018giant; @zhang2015disco] and the one-shot methods [@lee2017communication; @lin2017distributed; @wang2019sharper; @zhang2013communication; @zhang2015divide]. The primal-dual algorithms [@hong2018gradient; @smith2016cocoa; @smith2017federated] are more generally applicable and more relevant to FL.

Numerical Experiments {#sec:experiment}
=====================

#### Models and datasets

We examine our theoretical results on a logistic regression with weight decay $\lambda=1e-4$. This is a stochastic convex optimization problem. We distribute MNIST dataset [@lecun1998gradient] among $N=100$ workers in a non-iid fashion such that each device contains samples of only two digits. We further obtain two datasets: `mnist balanced` and `mnist unbalanced`. The former is balanced such that the number of samples in each device is the same, while the latter is highly unbalanced with the number of samples among devices following a power law. To manipulate heterogeneity more precisly, we synthesize unbalanced datasets following the setup in @sahu2018convergence and denote it as `synthetic(\alpha, \beta)` where $\alpha$ controls how much local models differ from each other and $\beta$ controls how much the local data at each device differs from that of other devices. We obtain two datasets: `synthetic(0,0)` and `synthetic(1,1)`. Details can be found in Appendix \[sec:appen:exp\].

\

#### Experiment settings

For all experiments, we initialize all runnings with $\w_0 = 0$. In each round, all selected devices run $E$ steps of SGD in parallel. We decay the learning rate at the end of each round by the following scheme $\eta_t = \frac{\eta_0}{1+t}$, where $\eta_0$ is chosen from the set $\{1, 0.1, 0.01\}$. We evaluate the averaged model after each global synchronization on the corresponding global objective. For fair comparison, we control all randomness in experiments so that the set of activated devices is the same across all different algorithms on one configuration.

#### Impact of $E$

We expect that $T_\epsilon/E$, the required communication round to achieve curtain accuracy, is a hyperbolic finction of $E$ as equ (\[eq:communication\_round\]) indicates. Intuitively, a small $E$ means a heavy communication burden, while a large $E$ means a low convergence rate. One needs to trade off between communication efficiency and fast convergence. We empirically observe this phenomenon on unbalanced datasets in Figure \[fig:test\_figure\_a\]. The reason why the phenomenon does not appear in `mnist balanced` dataset requires future investigations.

#### Impact of $K$

Our theory suggests that a larger $K$ may slightly accelerate convergence since $T_\epsilon/E $ contains a term $\mathcal{O}\left( \frac{EG^2}{K}\right)$. Figure \[fig:test\_figure\_b\] shows that $K$ has limited influence on the convergence of `FedAvg` in `synthetic(0,0)` dataset. It reveals that the curve of a large enough $K$ is slightly better. We observe similar phenomenon among the other three datasets and attach additional results in Appendix \[sec:appen:exp\]. This justifies that when the variance resulting sampling is not too large (i.e., $B \gg C$), one can use a small number of devices without severely harming the training process, which also removes the need to sample as many devices as possible in convex federated optimization.

#### Effect of sampling and averaging schemes.

We compare four schemes among four federated datasets. Since the original scheme involves a history term and may be conservative, we carefully set the initial learning rate for it. Figure \[fig:test\_figure\_c\] indicates that when data are balanced, Schemes I and II achieve nearly the same performance, both better than the original scheme. Figure \[fig:test\_figure\_d\] shows that when the data are unbalanced, i.e., $p_k$’s are uneven, Scheme I performs the best. Scheme II suffers from some instability in this case. This is not contradictory with our theory since we don’t guarantee the convergence of Scheme II when data is unbalanced. As expected, transformed Scheme II performs stably at the price of a lower convergence rate. Compared to Scheme I, the original scheme converges at a slower speed even if its learning rate is fine tuned. All the results show the crucial position of appropriate sampling and averaging schemes for `FedAvg`.

Conclusion
==========

Federated learning becomes increasingly popular in machine learning and optimization communities. In this paper we have studied the convergence of `FedAvg`, a heuristic algorithm suitable for federated setting. We have investigated the influence of sampling and averaging schemes. We have provided theoretical guarantees for two schemes and empirically demonstrated their performances. Our work sheds light on theoretical understanding of `FedAvg` and provides insights for algorithm design in realistic applications. Though our analyses are constrained in convex problems, we hope our insights and proof techniques can inspire future work.

Acknowledgements {#acknowledgements .unnumbered}
================

Li, Yang and Zhang have been supported by the National Natural Science Foundation of China (No. 11771002 and 61572017), Beijing Natural Science Foundation (Z190001), the Key Project of MOST of China (No. 2018AAA0101000), and Beijing Academy of Artificial Intelligence (BAAI).

[^1]: Equal contribution.

[^2]: In original paper [@mcmahan2016communication], $E$ epochs of SGD are performed in parallel. For theoretical analyses, we denote by $E$ the times of updates rather than epochs.

[^3]: Throughout this paper, “non-iid” means data are not identically distributed. More precisely, the data distributions in the $k$-th and $l$-th devices, denote $D_k$ and $D_l$, can be different.

[^4]: Throughout this paper, “sampling” refers to how the server chooses $K$ user devices and use their outputs for updating the model parameters. “Sampling” does not mean how a device randomly selects training samples.

[^5]: It is well know that the full gradient descent (which is equivalent to `FedAvg` with $E=1$ and full batch) do not require the decay of learning rate.

[^6]: Here we use $\gamma = \OM(\kappa + E)$.
