FedChain: Chained Algorithms for
Near-Optimal Communication Cost
in Federated Learning

In federated learning (FL) (McMahan et al., 2017; Kairouz et al., 2019; Li et al., 2020), distributed clients interact with a central server to jointly train a single model without directly sharing their data with the server. The training objective is to solve the following minimization problem:

where variable $x$ is the model parameter, $i$ indexes the clients (or devices), $N$ is the number of clients, and $F_{i}(x)$ is a loss that only depends on that client’s data. Typical FL deployments have two properties that make optimizing Eq. 1 challenging: ($i$) Data heterogeneity: We want to minimize the average of the expected losses $F_{i}(x)=\mathbb{E}_{z_{i}\sim\mathcal{D}_{i}}[f(x;z_{i})]$ for some loss function $f$ evaluated on client data $z_{i}$ drawn from a client-specific distribution $\mathcal{D}_{i}$. The convergence rate of the optimization depends on the heterogeneity of the local data distributions, $\{\mathcal{D}_{i}\}_{i=1}^{N}$, and this dependence is captured by a popular notion of client heterogeneity, $\zeta^{2}$, defined as follows:

This captures the maximum difference between a local gradient and the global gradient, and is a standard measure of heterogeneity used in the literature (Woodworth et al., 2020a; Gorbunov et al., 2020; Deng et al., 2020; Woodworth et al., 2020a; Gorbunov et al., 2020; Yuan et al., 2020; Deng et al., 2021; Deng & Mahdavi, 2021). ($ii$) Communication cost: In many FL deployments, communication is costly because clients have limited bandwidths. Due to these two challenges, most federated optimization algorithms alternate between local rounds of computation, where each client locally processes only their own data to save communication, and global rounds, where clients synchronize with the central server to resolve disagreements due to heterogeneity in the locally updated models.

Several federated algorithms navigate this trade-off between reducing communication and resolving data heterogeneity by modifying the amount and nature of local and global computation (McMahan et al., 2017; Li et al., 2018; Wang et al., 2019b; a; Li et al., 2019; Karimireddy et al., 2020b; a; Al-Shedivat et al., 2020; Reddi et al., 2020; Charles & Konečnỳ, 2020; Mitra et al., 2021; Woodworth et al., 2020a). These first-order federated optimization algorithms largely fall into one of two camps: ($i$) clients in local update methods send updated models after performing multiple steps of model updates, and ($ii$) clients in global update methods send gradients and do not perform any model updates locally. Examples of ($i$) include FedAvg (McMahan et al., 2017), SCAFFOLD (Karimireddy et al., 2020b), and FedProx (Li et al., 2018). Examples of ($ii$) include SGD and Accelerated SGD (ASG), where in each round $r$ the server collects from the clients gradients evaluated on local data at the current iterate $x^{(r)}$ and then performs a (possibly Nesterov-accelerated) model update.

As an illustrating example, consider the scenario when $F_{i}$’s are $\mu$-strongly convex and $\beta$-smooth such that the condition number is $\kappa={\beta}/{\mu}$ as summarized in Table 1, and also assume for simplicity that all clients participate in each communication round (i.e., full participation). Existing convergence analyses show that local update methods have a favorable dependence on the heterogeneity $\zeta$. For example, Woodworth et al. (2020a) show that FedAvg achieves an error bound of $\tilde{\mathcal{O}}((\kappa\zeta^{2}/\mu)R^{-2})$ after $R$ rounds of communication. Hence FedAvg achieves theoretically faster convergence for smaller heterogeneity levels $\zeta$ by using local updates. However, this favorable dependency on $\zeta$ comes at the cost of slow convergence in $R$. On the other hand, global update methods converge faster in $R$, with error rate decaying as $\tilde{\mathcal{O}}(\Delta\exp(-R/\sqrt{\kappa}))$ (for the case of ASG), where $\Delta$ is the initial function value gap to the (unique) optimal solution. This is an exponentially faster rate in $R$, but it does not take advantage of small heterogeneity even when client data is fully homogeneous.

Our main contribution is to design a novel family of algorithms that combines the strengths of local and global methods to achieve a faster convergence while maintaining the favorable dependence on the heterogeneity, thus achieving an error rate of $\tilde{\mathcal{O}}((\zeta^{2}/\mu)\exp(-R/\sqrt{\kappa}))$ in the strongly convex scenario. By adaptively switching between this novel algorithm and the existing best global method, we can achieve an error rate of $\tilde{\mathcal{O}}(\min\{\Delta,\zeta^{2}/\mu\}\exp(-R/\sqrt{\kappa}))$. We further show that this is near-optimal by providing a matching lower bound in Thm. 5.4. This lower bound tightens an existing one from (Woodworth et al., 2020a) by considering a smaller class of algorithms that includes those presented in this paper.

We propose FedChain, a unifying chaining framework for federated optimization that enjoys the benefits of both local update methods and global update methods. For a given total number of rounds $R$, FedChain first uses a local-update method for a constant fraction of rounds and then switches to a global-update method for the remaining rounds. The first phase exploits client homogeneity when possible, providing fast convergence when heterogeneity is small. The second phase uses the unbiased stochastic gradients of global update methods to achieve a faster convergence rate in the heterogeneous setting. The second phase inherits the benefits of the first phase through the iterate output by the local-update method. An instantiation of FedChain consists of a combination of a specific local-update method and global-update method (e.g., FedAvg $\to$ SGD).

Contributions. We propose the FedChain framework and analyze various instantiations under strongly convex, general convex, and nonconvex objectives that satisfy the PL condition.¹¹1$F$ satisfies the $\mu$-PL condition if $2\mu(F(x)-F(x^{*}))\leq\|\nabla F(x)\|^{2}$. Achievable rates are summarized in Tables 1, LABEL:, 2, LABEL: and 4. In the strongly convex setting, these rates nearly match the algorithm-independent lower bounds we introduce in Theorem 5.4, which shows the near-optimality of FedChain. For strongly convex functions, chaining is optimal up to a factor that decays exponentially in the condition number $\kappa$. For convex functions, it is optimal in the high-heterogeneity regime, when $\zeta>\beta DR^{1/2}$. For nonconvex PL functions, it is optimal for constant condition number $\kappa$. In all three settings, FedChain instantiations improve upon previously known worst-case rates in certain regimes of $\zeta$. We further demonstrate the empirical gains of our chaining framework in the convex case (logistic regression on MNIST) and in the nonconvex case (ConvNet classification on EMNIST (Cohen et al., 2017) and ResNet-18 classification on CIFAR-100 (Krizhevsky, 2009)).

Federated optimization proceeds in rounds. We consider the partial participation setting, where at the beginning of each round, $S\in{\mathbb{Z}}_{+}$ out of total $N$ clients are sampled uniformly at random without replacement. Between each global communication round, the sampled clients each access either (i) their own stochastic gradient oracle $K$ times, update their local models, and return the updated model, or (ii) their own stochastic function value oracle $K$ times and return the average value to the server. The server aggregates the received information and performs a model update. For each baseline optimization algorithm, we analyze the sub-optimality error after $R$ rounds of communication between the clients and the server. Suboptimality is measured in terms of the function value $\mathbb{E}F(\hat{x})-F(x^{*})$, where $\hat{x}$ is the solution estimate after $R$ rounds and $x^{*}=\operatorname*{arg\,min}_{x}F(x)$ is a (possibly non-unique) optimum of $F$. We let the estimate for the initial suboptimality gap be $\Delta$ (Assumption B.9), and the initial distance to a (not necessarily unique) optimum be $D$ (Assumption B.10). If applicable, $\epsilon$ is the target expected function value suboptimality.

We study three settings: strongly convex $F_{i}$’s, convex $F_{i}$’s and $\mu$-PL $F$; for formal definitions, refer to App. B. Throughout this paper, we assume that the $F_{i}$’s are $\beta$-smooth (Assumption B.4). If $F_{i}$’s are $\mu$-strongly convex (Assumption B.1) or $F$ is $\mu$-PL (Assumption B.3), we denote $\kappa={\beta}/{\mu}$ as the condition number. $\mathcal{D}_{i}$ is the data distribution of client $i$. We define the heterogeneity of the problem as $\zeta^{2}:=\max_{i\in[N]}\sup_{x}\|\nabla F(x)-\nabla F_{i}(x)\|^{2}$ in Assumption B.5. We assume unless otherwise specified that $\zeta^{2}>0$, i.e., that the problem is heterogeneous. We assume that the client gradient variance is upper bounded by $\sigma^{2}$ (Assumption B.6). We also define the analogous quantities for function value oracle queries: $\zeta^{2}_{F}$ Assumption B.8, $\sigma^{2}_{F}$ Assumption B.7. We use the notation $\tilde{\mathcal{O}},\tilde{\Omega}$ to hide polylogarithmic factors, and $\mathcal{O},\Omega$ if we are only hiding constant factors.

We start with a toy example (Fig. 1) to illustrate FedChain. Consider two strongly convex client objectives: $F_{1}(x)=({1}/{2})(x-1)^{2}$ and $F_{2}(x)=(x+1)^{2}$. The global objective $F(x)=({F_{1}(x)+F_{2}(x)})/{2}$ is their average. Fig. 1 (top) plots the objectives, and Fig. 1 (bottom) displays their gradients. Far from the optimum, due to strong convexity, all client gradients point towards the optimal solution; specifically, this occurs when  $x\in(-\infty,-1]\cup[1,\infty)$. In this regime, clients can use local steps (e.g., FedAvg) to reduce communication without sacrificing the consistency of per-client local updates. On the other hand, close to the optimum, i.e., when $x\in(-1,1)$, some client gradients may point away from the optimum. This suggests that clients should not take local steps to avoid driving the global estimate away from the optimum. Therefore, when we are close to the optimum, we use an algorithm without local steps (e.g. SGD), which is less affected by client gradient disagreement.

This intuition seems to carry over to the nonconvex setting: Charles et al. (2021) show that over the course of FedAvg execution on neural network StackOverflow next word prediction, client gradients become more orthogonal.

FedChain: To exploit the strengths of both local and global update methods, we propose the federated chaining (FedChain) framework in Alg. 1. There are three steps: (1) Run a local-update method $\mathcal{A}_{\text{local}}$, like FedAvg. (2) Choose the better point between the output of $\mathcal{A}_{\text{local}}$ (which we denote $\hat{x}_{1/2}$), and the initial point $\hat{x}_{0}$ to initialize (3) a global-update method $\mathcal{A}_{\text{global}}$ like SGD. Note that when heterogeneity is large, $\mathcal{A}_{\text{local}}$ can actually output an iterate with higher suboptimality gap than $\hat{x}_{0}$. Hence, selecting the point (between $\hat{x}_{0}$ and $\hat{x}_{1/2}$) with a smaller $F$ allows us to adapt to the problem’s heterogeneity, and initialize $\mathcal{A}_{\text{global}}$ appropriately to achieve good convergence rates. To compare $\hat{x}_{1/2}$ and the initial point $\hat{x}_{0}$, we approximate $F(\hat{x}_{0})$ and $F(\hat{x}_{1/2})$ by averaging; we compute $\frac{1}{S{K}}\sum_{i\in\mathcal{S},k\in[K]}f(x;\hat{z}_{i,k})$ for $\hat{x}_{1/2},\hat{x}_{0}$, where $K$ is also the number of local steps per client per round in $\mathcal{A}_{\text{local}}$, and $\mathcal{S}$ is a sample of $S$ clients. In practice, one might consider adaptively selecting how many rounds of $\mathcal{A}_{\text{local}}$ to run, which can potentially improve the convergence by a constant factor. Our experiments in § J.1 show significant improvement when using only 1 round of $\mathcal{A}_{\text{local}}$ with a large enough $K$ for convex losses.

For the formal statement, see Thm. F.2. We show and discuss the near-optimality of FedAvg $\to$ ASG under strongly convex and PL conditions (and under full participation) in Section 5, where we introduce matching lower bounds. Under strong-convexity shown in Table 1, FedAvg $\to$ ASG, when compared to ASG, converts the $\Delta\exp(-{R}/{\sqrt{\kappa}})$ term into a $\min\{\Delta,{\zeta^{2}}/{\mu}\}\exp(-{R}/{\sqrt{\kappa}})$ term, improving over ASG when heterogeneity moderately small: ${\zeta^{2}}/{\mu}<\Delta$. It also significantly improves over FedAvg, as $\min\{\Delta,{\zeta^{2}}/{\mu}\}\exp(-{R}/{\sqrt{\kappa}})$ is exponentially faster than $\kappa({\zeta^{2}}/{\mu})R^{-2}$. Under the PL condition (which is not necessarily convex) the story is similar (Table 4), except we use FedAvg $\to$ SGD as our representative algorithm, as Nesterov acceleration is not known to improve under non-convex settings.

In the general convex case (Table 2), let $\beta=D=1$ for the purpose of comparisons. Then FedAvg $\to$ ASG’s convergence rate is $\min\{1/R^{2},\zeta^{1/2}/R\}+\sqrt{1-S/N}(\zeta/\sqrt{SR})+\sqrt[4]{1-S/N}\sqrt{\zeta}/\sqrt[4]{SR}$. If $\zeta<\frac{1}{R^{2}}$, then $\zeta^{1/2}/R<1/R^{2}$ and if $\zeta<\sqrt{S/R^{7}}$, then $\sqrt[4]{1-S/N}\sqrt{\zeta}/\sqrt[4]{SR}<1/R^{2}$, so the FedAvg $\to$ ASG convergence rate is better than the convergence rate of ASG under the regime $\zeta<\min\{1/R^{2},\sqrt{S/R^{7}}\}$. The rate of Karimireddy et al. (2020b) for FedAvg (which does not require $S=N$) is strictly worse than ASG, and so has a worse convergence rate than FedAvg $\to$ ASG if $\zeta<\min\{1/R^{2},\sqrt{S/R^{7}}\}$. Altogether, if $\zeta<\min\{1/R^{2},\sqrt{S/R^{7}}\}$, FedAvg $\to$ ASG achieves the best known worst-case rate. Finally, in the $S=N$ case, FedAvg $\to$ ASG does not have a regime in $\zeta$ where it improves over both ASG and FedAvg (the analysis of Woodworth et al. (2020a), which requires $S=N$) at the same time. It is unclear if this is due to looseness in the analysis.

For the formal statement, see Thm. F.4. SSNM (Zhou et al., 2019) is the Nesterov accelerated version of SAGA (Defazio et al., 2014).

The main contribution of using variance reduced methods in $\mathcal{A}_{\text{global}}$ is the removal of the sampling error (the terms depending on the sampling heterogeneity error $(1-S/N)(\zeta^{2}/S)$) from the convergence rates, in exchange for requiring a round complexity of at least $N/S$. To illustrate this, observe that in the strongly convex case, FedAvg $\to$ SGD has convergence rate $\min\{\Delta,\zeta^{2}/\mu\}\exp(-R/\kappa)+(1-S/N)(\zeta^{2}/SR)$ and FedAvg $\to$ SAGA (FedAvg $\to$ SGD’s variance-reduced counterpart) has convergence rate $\min\{\Delta,\zeta^{2}/\mu\}\exp(-\min\{1/\kappa,S/N\}R)$, dropping the $(1-S/N)(\zeta^{2}/\mu SR)$ sampling heterogeneity error term in exchange for harming the rate of linear convergence from $1/\kappa$ to $\min\{1/\kappa,S/N\}$.

This same tradeoff occurs in finite sum optimization (the problem $\min_{x}(1/n)\sum_{i=1}^{n}\psi_{i}(x)$, where $\psi_{i}$’s (typically) represent losses on data points and the main concern is computation cost), which is what variance reduction is designed for. In finite sum optimization, variance reduction methods such as SVRG (Johnson & Zhang, 2013) and SAGA (Defazio et al., 2014; Reddi et al., 2016) achieve linear rates of convergence (given strong convexity of $\psi_{i}$’s) in exchange for requiring at least $N/S$ updates. Because we can treat FL as an instance of finite-sum optimization (by viewing Eq. 1 as a finite sum of objectives and $\zeta^{2}$ as the variance between $F_{i}$’s), these results from variance-reduced finite sum optimization can be extended to federated learning. This is the idea behind SCAFFOLD (Karimireddy et al., 2020b).

It is not always the case that variance reduction in $\mathcal{A}_{\text{global}}$ achieves better rates. In the strongly convex case if $(1-S/N)(\zeta^{2}/\mu S\epsilon)>N/S$, then variance reduction gets gain, otherwise not.

A proof of the lower bound is in App. G, and the corollary follows immediately from the fact that $\mu$-strong convexity implies $\mu$-PL. This result tightens the lower bound in Woodworth et al. (2020a), which proves a similar lower bound but in terms of a much larger class of functions with heterogeneity bounded in $\zeta_{*}=(1/N)\sum_{i=1}^{N}\|\nabla F_{i}(x^{*})\|^{2}$. To the best of our knowledge, all achievable rates that can take advantage of heterogeneity require $(\zeta,c)$-heterogeneity (which is a smaller class than $\zeta_{*}$-heterogeneity), which are incomparable with the existing lower bound from (Woodworth et al., 2020a) that requires $\zeta_{*}$-heterogeneity. By introducing the class of distributed distance-conserving algorithms (which includes most of the algorithms we are interested in), Thm. 5.4 allows us, for the first time, to identify the optimality of the achievable rates as shown in Tables 1, 2, and 4.

Note that this lower bound allows full participation and should be compared to the achievable rates with $S=N$; comparisons with variance reduced methods like FedAvg $\to$ SAGA and FedAvg $\to$ SSNM are unnecessary. Our lower bound proves that FedAvg $\to$ ASG is optimal up to condition number factors shrinking exponentially in $\sqrt{\kappa}$ among algorithms satisfying Definition 5.2 and Definition 5.3. Under the PL-condition, the situation is similar, except FedAvg $\to$ SGD loses $\sqrt{\kappa}$ in the exponential versus the lower bound. On the other hand, there remains a substantial gap between FedAvg $\to$ ASG and the lower bound in the general convex case. Meaningful progress in closing the gap in the general convex case is an important future direction.

We evaluate the utility of our framework on strongly convex and nonconvex settings. We compare the communication round complexities of four baselines: FedAvg, SGD, ASG, and SCAFFOLD, the stepsize decaying variants of these algorithms (which are prefixed by M- in plots) and the various instantiations of FedChain (Algo. 1).