SCAFFOLD: Stochastic Controlled Averaging for Federated Learning

Introduction

Federated learning has emerged as an important paradigm in modern large-scale machine learning. Unlike in traditional centralized learning where models are trained using large datasets stored in a central server (Dean et al., 2012; Iandola et al., 2016; Goyal et al., 2017), in federated learning, the training data remains distributed over a large number of clients, which may be phones, network sensors, hospitals, or alternative local information sources (Konečnỳ et al., 2016b, a; McMahan et al., 2017; Mohri et al., 2019; Kairouz et al., 2019). A centralized model (referred to as server model) is then trained without ever transmitting client data over the network, thereby ensuring a basic level of privacy. In this work, we investigate stochastic optimization algorithms for federated learning.

The key challenges for federated optimization are 1) dealing with unreliable and relatively slow network connections between the server and the clients, 2) only a small subset of clients being available for training at a given time, and 3) large heterogeneity (non-iid-ness) in the data present on the different clients (Konečnỳ et al., 2016a). The most popular algorithm for this setting is FedAvg (McMahan et al., 2017). FedAvg tackles the communication bottleneck by performing multiple local updates on the available clients before communicating to the server. While it has shown success in certain applications, its performance on heterogeneous data is still an active area of research (Li et al., 2018; Yu et al., 2019; Li et al., 2019b; Haddadpour & Mahdavi, 2019; Khaled et al., 2020). We prove that indeed such heterogeneity has a large effect on FedAvg—it introduces a drift in the updates of each client resulting in slow and unstable convergence. Further, we show that this client-drift persists even if full batch gradients are used and all clients participate throughout the training.

As a solution, we propose a new Stochastic Controlled Averaging algorithm (SCAFFOLD) which tries to correct for this client-drift. Intuitively, SCAFFOLD estimates the update direction for the server model ($\bm{c}$) and the update direction for each client $\bm{c}_{i}$.We refer to these estimates as control variates and the resulting correction technique as stochastic controlled averaging. The difference $(\bm{c}-\bm{c}_{i})$ is then an estimate of the client-drift which is used to correct the local update. This strategy successfully overcomes heterogeneity and converges in significantly fewer rounds of communication. Alternatively, one can see heterogeneity as introducing ‘client-variance’ in the updates across the different clients and SCAFFOLD then performs ‘client-variance reduction’ (Schmidt et al., 2017; Johnson & Zhang, 2013; Defazio et al., 2014). We use this viewpoint to show that SCAFFOLD is relatively unaffected by client sampling.

Finally, while accommodating heterogeneity is important, it is equally important that a method can take advantage of similarities in the client data. We prove that SCAFFOLD indeed has such a property, requiring fewer rounds of communication when the clients are more similar.

We derive tighter convergence rates for FedAvg than previously known for convex and non-convex functions with client sampling and heterogeneous data.

We give matching lower bounds to prove that even with no client sampling and full batch gradients, FedAvg can be slower than SGD due to client-drift.

We propose a new Stochastic Controlled Averaging algorithm (SCAFFOLD) which corrects for this client-drift. We prove that SCAFFOLD is at least as fast as SGD and converges for arbitrarily heterogeneous data.

We show SCAFFOLD can additionally take advantage of similarity between the clients to further reduce the communication required, proving the advantage of taking local steps over large-batch SGD for the first time.

We prove that SCAFFOLD is relatively unaffected by the client sampling obtaining variance reduced rates, making it especially suitable for federated learning.

Finally, we confirm our theoretical results on simulated and real datasets (extended MNIST by Cohen et al. (2017)).

Related work. For identical clients, FedAvg coincides with parallel SGD analyzed by (Zinkevich et al., 2010) who proved asymptotic convergence. Stich (2018) and, more recently Stich & Karimireddy (2019); Patel & Dieuleveut (2019); Khaled et al. (2020), gave a sharper analysis of the same method, under the name of local SGD, also for identical functions. However, there still remains a gap between their upper bounds and the lower bound of Woodworth et al. (2018). The analysis of FedAvg for heterogeneous clients is more delicate due to the afore-mentioned client-drift, first empirically observed by Zhao et al. (2018). Several analyses bound this drift by assuming bounded gradients (Wang et al., 2019; Yu et al., 2019), or view it as additional noise (Khaled et al., 2020), or assume that the client optima are $\epsilon$-close (Li et al., 2018; Haddadpour & Mahdavi, 2019). In a concurrent work, (Liang et al., 2019) propose to use variance reduction to deal with client heterogeneity but still show rates slower than SGD and do not support client sampling. Our method SCAFFOLD can also be seen as an improved version of the distributed optimization algorithm DANE by (Shamir et al., 2014), where a fixed number of (stochastic) gradient steps are used in place of a proximal point update. A more in-depth discussion of related work is given in Appendix A. We summarize the complexities of different methods for heterogeneous clients in Table 2.

Setup

We formalize the problem as minimizing a sum of stochastic functions, with only access to stochastic samples:

The functions $f_{i}$ represents the loss function on client $i$. All our results can be easily extended to the weighted case.

We assume that $f$ is bounded from below by $f^{\star}$ and $f_{i}$ is $\beta$-smooth. Further, we assume $g_{i}(\bm{x}):=\nabla f_{i}(\bm{x};\zeta_{i})$ is an unbiased stochastic gradient of $f_{i}$ with variance bounded by $\sigma^{2}$. For some results, we assume $\mu\geq 0$ (strong) convexity. Note that $\sigma$ only bounds the variance within clients. We also define two non-standard terminology below.

$(G,B)$-BGD or bounded gradient dissimilarity: there exist constants $G\geq 0$ and $B\geq 1$ such that

If $\{f_{i}\}$ are convex, we can relax the assumption to

$\delta$-BHD or bounded Hessian dissimilarity:

Further, $f_{i}$ is $\delta$-weakly convex i.e. $\nabla^{2}f_{i}(\bm{x})\succeq-\delta I$.

The assumptions (A1) and (A2) are orthogonal—it is possible to have $G=0$ and $\delta=2\beta$, or $\delta=0$ but $G\gg 1$.

Convergence of FedAvg

In this section we review FedAvg and improve its convergence analysis by deriving tighter rates than known before. The scheme consists of two main parts: local updates to the model (1), and aggregating the client updates to update the server model (2). In each round, a subset of clients $\mathcal{S}\subseteq[N]$ are sampled uniformly. Each of these clients $i\in\mathcal{S}$ copies the current sever model $\bm{y}_{i}=\bm{x}$ and performs $K$ local updates of the form:

Here $\eta_{l}$ is the local step-size. Then the clients’ updates $(\bm{y}_{i}-\bm{x})$ are aggregated to form the new server model using a global step-size $\eta_{g}$ as:

We now state our novel convergence results for functions with bounded dissimilarity (proofs in Appendix D.2).

For $\beta$-smooth functions $\{f_{i}\}$ which satisfy ((A1)), the output of FedAvg has expected error smaller than $\epsilon$ for some values of $\eta_{l},\eta_{g},R$ satisfying

Strongly convex: $\eta_{g}\geq\sqrt{S}$, $\eta_{l}\leq\frac{1}{(1+B^{2})6\beta K\eta_{g}}$, and

General convex: $\eta_{g}\geq\sqrt{S}$, $\eta_{l}\leq\frac{1}{(1+B^{2})6\beta K\eta_{g}}$, and

Non-convex: $\eta_{g}\geq\sqrt{S}$, $\eta_{l}\leq\frac{1}{(1+B^{2})6\beta K\eta_{g}}$, and

where $D:=\lVert\bm{x}^{0}-\bm{x}^{\star}\rVert^{2}$ and $F:=f(\bm{x}^{0})-f^{\star}$.

It is illuminating to compare our rates with those of the simpler iid. case i.e. with $G=0$ and $B=1$. Our strongly-convex rates become $\frac{\sigma^{2}}{\mu SK\epsilon}+\frac{1}{\mu}$. In comparison, the best previously known rate for this case was by Stich & Karimireddy (2019) who show a rate of $\frac{\sigma^{2}}{\mu SK\epsilon}+\frac{S}{\mu}$. The main source of improvement in the rates came from the use of two separate step-sizes ($\eta_{l}$ and $\eta_{g}$). By having a larger global step-size $\eta_{g}$, we can use a smaller local step-size $\eta_{l}$ thereby reducing the client-drift while still ensuring progress. However, even our improved rates do not match the lower-bound for the identical case of $\frac{\sigma^{2}}{\mu SK\epsilon}+\frac{1}{K\mu}$ (Woodworth et al., 2018). We bridge this gap for quadratic functions in Section 6.

We now compare FedAvg to two other algorithms FedProx by (Li et al., 2018) (aka EASGD by (Zhang et al., 2015)) and to SGD. Suppose that $G=0$ and $\sigma=0$ i.e. we use full batch gradients and all clients have very similar optima. In such a case, FedAvg has a complexity of $\frac{B^{2}}{\mu}$ which is identical to that of FedProx (Li et al., 2018). Thus, FedProx does not have any theoretical advantage.

Next, suppose that all clients participate (no sampling) with $S=N$ and there is no variance $\sigma=0$. Then, the above for strongly-convex case simplifies to ${\frac{G}{\mu\sqrt{\epsilon}}+\frac{B^{2}}{\mu}}$. In comparison, extending the proof of (Khaled et al., 2020) using our techniques gives a worse dependence on $G$ of ${\frac{G^{2}}{\mu KN\epsilon}+\frac{G}{\mu\sqrt{\epsilon}}}$. Similarly, for the non-convex case, our rates are tighter and have better dependence on $G$ than (Yu et al., 2019). However, simply running SGD in this setting would give a communication complexity of $\frac{\beta}{\mu}$ which is faster, and independent of similarity assumptions. In the next section we examine the necessity of such similarity assumptions.

2 Lower bounding the effect of heterogeneity

We now show that when the functions $\{f_{i}\}$ are distinct, the local updates of FedAvg on each client experiences drift thereby slowing down convergence. We show that the amount of this client drift, and hence the slowdown in the rate of convergence, is exactly determined by the gradient dissimilarity parameter $G$ in ((A1)).

We now examine the mechanism by which the client-drift arises (see Fig. 1). Let $\bm{x}^{\star}$ be the global optimum of $f(\bm{x})$ and $\bm{x}_{i}^{\star}$ be the optimum of each client’s loss function $f_{i}(\bm{x})$. In the case of heterogeneous data, it is quite likely that each of these $\bm{x}_{i}^{\star}$ is far away from the other, and from the global optimum $\bm{x}^{\star}$. Even if all the clients start from the same point $\bm{x}$, each of the $\bm{y}_{i}$ will move towards their client optimum $\bm{x}_{i}^{\star}$. This means that the average of the client updates (which is the server update) moves towards $\frac{1}{N}\sum_{i=1}^{N}\bm{x}_{i}^{\star}$. This difference between $\frac{1}{N}\sum_{i=1}^{N}\bm{x}_{i}^{\star}$ and the true optimum $\bm{x}^{\star}$ is exactly the cause of client-drift. To counter this drift, FedAvg is forced to use much smaller step-sizes which in turn hurts convergence. We can formalize this argument to prove a lower-bound (see Appendix D.4 for proof).

For any positive constants $G$ and $\mu$, there exist $\mu$-strongly convex functions satisfying (A1) for which FedAvg with $K\geq 2$, $\sigma=0$ and $N=S$ has an error

This implies that the $\frac{G}{\sqrt{\epsilon}}$ term is unavoidable even if there is no stochasticity. Further, because FedAvg uses $RKN$ stochastic gradients, we also have the statistical lower-bound of $\frac{\sigma^{2}}{\mu KN\epsilon}$. Together, these lower bounds prove that the rate derived in Theorem I is nearly optimal (up to dependence on $\mu$). In the next section, we introduce a new method SCAFFOLD to mitigate this client-drift.

SCAFFOLD algorithm

In this section we first describe SCAFFOLD and then discuss how it solves the problem of client-drift.

SCAFFOLD has three main steps: local updates to the client model (3), local updates to the client control variate (4), and aggregating the updates (5). We describe each in more detail.

Along with the server model $\bm{x}$, SCAFFOLD maintains a state for each client (client control variate $\bm{c}_{i}$) and for the server (server control variate $\bm{c}$). These are initialized to ensure that $\bm{c}=\frac{1}{N}\sum\bm{c}_{i}$ and can safely all be initialized to 0. In each round of communication, the server parameters $(\bm{x},\bm{c})$ are communicated to the participating clients $\mathcal{S}\subset[N]$. Each participating client $i\in\mathcal{S}$ initializes its local model with the server model $\bm{y}_{i}\leftarrow\bm{x}$. Then it makes a pass over its local data performing $K$ updates of the form:

Then, the local control variate $\bm{c}_{i}$ is also updated. For this, we provide two options:

Option I involves making an additional pass over the local data to compute the gradient at the server model $\bm{x}$. Option II instead re-uses the previously computed gradients to update the control variate. Option I can be more stable than II depending on the application, but II is cheaper to compute and usually suffices (all our experiments use Option II). The client updates are then aggregated and used to update the server parameters:

This finishes one round of communication. Note that the clients in SCAFFOLD are stateful and retain the value of $\bm{c}_{i}$ across multiple rounds. Further, if $\bm{c}_{i}$ is always set to 0, then SCAFFOLD becomes equivalent to FedAvg. The full details are summarized in Algorithm 1.

If communication cost was not a concern, the ideal update on client $i$ would be

Such an update essentially computes an unbiased gradient of $f$ and hence becomes equivalent to running FedAvg in the iid case (which has excellent performance). Unfortunately such an update requires communicating with all clients for every update step. SCAFFOLD instead uses control variates such that

Then, SCAFFOLD (3) mimics the ideal update (6) with

Thus, the local updates of SCAFFOLD remain synchronized and converge for arbitrarily heterogeneous clients.

Convergence of SCAFFOLD

We state the rate of SCAFFOLD without making any assumption on the similarity between the functions. See Appendix E for the full proof.

For any $\beta$-smooth functions $\{f_{i}\}$, the output of SCAFFOLD has expected error smaller than $\epsilon$ for $\eta_{g}=\sqrt{S}$ and some values of $\eta_{l},R$ satisfying

Strongly convex: $\eta_{l}\leq\min\left(\tfrac{1}{81\beta K\eta_{g}},\tfrac{S}{15\mu NK\eta_{g}}\right)$ and

General convex: $\eta_{l}\leq\tfrac{1}{81\beta K\eta_{g}}$ and

Non-convex: $\eta_{l}\leq\frac{1}{24K\eta_{g}\beta}\left(\frac{S}{N}\right)^{\frac{2}{3}}$ and

where $D:=\lVert\bm{x}^{0}-\bm{x}^{\star}\rVert^{2}$ and $F:=f(\bm{x}^{0})-f^{\star}$.

Let us first examine the rates without client sampling ($S=N$). For the strongly convex case, the number of rounds becomes $\frac{\sigma^{2}}{\mu NK\epsilon}+\frac{1}{\mu}$. This rate holds for arbitrarily heterogeneous clients unlike Theorem I and further matches that of SGD with $K$ times larger batch-size, proving that SCAFFOLD is at least as fast as SGD. These rates also match known lower-bounds for distributed optimization (Arjevani & Shamir, 2015) (up to acceleration) and are unimprovable in general. However in certain cases SCAFFOLD is provably faster than SGD. We show this fact in Section 6.

Now let $\sigma=0$. Then our rates in the strongly-convex case are $\frac{1}{\mu}+\frac{N}{S}$ and \big{(}\frac{N}{S}\big{)}^{\frac{2}{3}}\frac{1}{\epsilon} in the non-convex case. These exactly match the rates of SAGA (Defazio et al., 2014; Reddi et al., 2016c). In fact, when $\sigma=0$, $K=1$ and $S=1$, the update of SCAFFOLD with option I reduces to SAGA where in each round consists of sampling one client $f_{i}$. Thus SCAFFOLD can be seen as an extension of variance reduction techniques for federated learning, and one could similarly extend SARAH (Nguyen et al., 2017), SPIDER (Fang et al., 2018), etc. Note that standard SGD with client sampling is provably slower and converges at a sub-linear rate even with $\sigma=0$.

For simplicity, assume that $\sigma=0$ and consider the ideal update of (6) which uses the full gradient $\nabla f(\bm{y})$ every step. Clearly, this would converge at a linear rate even with $S=1$. FedAvg would instead use an update $\nabla f_{i}(\bm{y})$. The difference between the ideal update (6) and the FedAvg update (1) is $\lVert\nabla f_{i}(\bm{y})-\nabla f(\bm{y})\rVert$. We need a bound on the gradient-dissimilarity as in ((A1)) to bound this error. SCAFFOLD instead uses the update $\nabla f_{i}(\bm{y})-\bm{c}_{i}+\bm{c}$, and the difference from ideal update becomes

Thus, the error is independent of how similar or dissimilar the functions $f_{i}$ are, and instead only depends on the quality of our approximation $\bm{c}_{i}\approx\nabla f_{i}(\bm{y})$. Since $f_{i}$ is smooth, we can expect that the gradient $\nabla f_{i}(\bm{y})$ does not change too fast and hence is easy to approximate. Appendix E translates this intuition into a formal proof.

Usefulness of local steps

In this section we investigate when and why taking local steps might be useful over simply computing a large-batch gradient in distributed optimization. We will show that when the functions across the clients share some similarity, local steps can take advantage of this and converge faster. For this we consider quadratic functions and express their similarity with the $\delta$ parameter introduced in ((A2)).

For any $\beta$-smooth quadratic functions $\{f_{i}\}$ with $\delta$ bounded Hessian dissimilarity ((A2)), the output of SCAFFOLD with $S=N$ (no sampling) has error smaller than $\epsilon$ for $\eta_{g}=1$ and some values of $\eta_{l},R$ satisfying

Strongly convex: $\eta_{l}\leq\tfrac{1}{15K\delta+8\beta}$ and

Weakly convex: $\eta_{l}\leq\tfrac{1}{15K\delta+8\beta}$ and

where we define $F:=(f(\bm{x}^{0})-f^{\star})$.

When $\sigma=0$ and $K$ is large, the complexity of SCAFFOLD becomes $\frac{\delta}{\mu}$. In contrast DANE, which being a proximal point method also uses large $K$, requires $(\frac{\delta}{\mu})^{2}$ rounds (Shamir et al., 2014) which is significantly slower, or needs an additional backtracking-line search to match the rates of SCAFFOLD (Yuan & Li, 2019). Further, Theorem IV is the first result to demonstrate improvement due to similairty for non-convex functions as far as we are aware.

Suppose that $\{f_{i}\}$ are identical. Recall that $\delta$ in ((A2)) measures the Hessian dissimilarity between functions and so $\delta=0$ for this case. Then Theorem IV shows that the complexity of SCAFFOLD is $\frac{\sigma^{2}}{\mu KN\epsilon}+\frac{1}{\mu K}$ which (up to acceleration) matches the i.i.d. lower bound of (Woodworth et al., 2018). In contrast, SGD with $K$ times larger batch-size would require $\frac{\sigma^{2}}{\mu KN\epsilon}+\frac{1}{\mu}$ (note the absence of $K$ in the second term). Thus, for identical functions, SCAFFOLD (and in fact even FedAvg) improves linearly with increasing number of local steps. In the other extreme, if the functions are arbitrarily different, we may have $\delta=2\beta$. In this case, the complexity of SCAFFOLD and large-batch SGD match the lower bound of Arjevani & Shamir (2015) for the heterogeneous case.

The above insights can be generalized to when the functions are only somewhat similar. If the Hessians are $\delta$-close and $\sigma=0$, then the complexity is $\frac{\beta+\delta K}{\mu K}$. This bound implies that the optimum number of local steps one should use is $K=\frac{\beta}{\delta}$. Picking a smaller $K$ increases the communication required whereas increasing it further would only waste computational resources. While this result is intuitive—if the functions are more ‘similar’, local steps are more useful—Theorem IV shows that it is the similarity of the Hessians which matters. This is surprising since the Hessians of $\{f_{i}\}$ may be identical even if their individual optima $\bm{x}_{i}^{\star}$ are arbitrarily far away from each other and the gradient-dissimilarity ((A1)) is unbounded.

Consider a simplified SCAFFOLD update with $\sigma=0$ and no sampling ($S=N$):

We would ideally want to perform the update $\bm{y}_{i}=\bm{y}_{i}-\eta\nabla f(\bm{y}_{i})$ using the full gradient $\nabla f(\bm{y}_{i})$. We reinterpret the correction term of SCAFFOLD $(\bm{c}-\bm{c}_{i})$ as performing the following first order correction to the local gradient $\nabla f_{i}(\bm{y}_{i})$ to make it closer to the full gradient $\nabla f(\bm{y}_{i})$:

Thus the SCAFFOLD update approximates the ideal update up to an error $\delta$. This intuition is proved formally for quadratic functions in Appendix F. Generalizing these results to other functions is a challenging open problem.

Experiments

We run experiments on both simulated and real datasets to confirm our theory. Our main findings are i) SCAFFOLD consistently outperforms SGD and FedAvg across all parameter regimes, and ii) the benefit (or harm) of local steps depends on both the algorithm and the similarity of the clients data.

Our simulated experiments uses $N=2$ quadratic functions based on our lower-bounds in Theorem II. We use full-batch gradients ($\sigma=0$) and no client sampling. Our real world experiments run logistic regression (convex) and 2 layer fully connected network (non-convex) on the EMNIST (Cohen et al., 2017). We divide this dataset among $N=100$ clients as follows: for $s\%$ similar data we allocate to each client $s\%$ i.i.d. data and the remaining $(100-s)\%$ by sorting according to label (cf. Hsu et al. (2019)).

We consider four algorithms: SGD, FedAvg SCAFFOLD and FedProx with SGD as the local solver (Li et al., 2018). On each client SGD uses the full local data to compute a single update, whereas the other algorithms take 5 steps per epoch (batch size is 0.2 of local data). We always use global step-size $\eta_{g}=1$ and tune the local step-size $\eta_{l}$ individually for each algorithm. SCAFFOLD uses option II (no extra gradient computations) and FedProx has fixed regularization $=1$ to keep comparison fair. Additional tuning of the regularization parameter may sometimes yield improved empirical performance.

2 Simulated results

The results are summarized in Fig. 3. Our simulated data has Hessian difference $\delta=1$ ((A2)) and $\beta=1$. We vary the gradient heterogeneity ((A1)) as $G\in$. For all valued of $G$, FedAvg gets slower as we increase the number of local steps. This is explained by the fact that client-drift increases as we increase the number of local steps, hindering progress. Further, as we increase $G$, FedAvg continues to slow down exactly as dictated by Thms. I and II. Note that when heterogeneity is small ($G=\beta=1$), FedAvg can be competitive with SGD.

SCAFFOLD is consistently faster than SGD, with $K=2$ being twice as fast and $K=10$ about 5 times faster. Further, its convergence is completely unaffected by $G$, confirming our theory in Thm. III. The former observation that we do not see linear improvement with $K$ is explained by Thm. IV since we have $\delta>0$. This sub linear improvement is still significantly faster than both SGD and FedAvg.

3 EMNIST results

We run extensive experiments on the EMNIST dataset to measure the interplay between the algorithm, number of epochs (local updates), number of participating clients, and the client similarity. Table 3 measures the benefit (or harm) of using more local steps, Table 4 studies the resilience to client sampling, and Table 5 reports preliminary results on neural networks. We are mainly concerned with minimizing the number of communication rounds. We observe that

SCAFFOLD is consistently the best. Across all range of values tried, we observe that SCAFFOLD outperforms SGD, FedAvg, and FedProx. The latter FedProx is always slower than the other local update methods, though in some cases it outperforms SGD. Note that it is possible to improve FedProx by carefully tuning the regularization parameter (Li et al., 2018). FedAvg is always slower than SCAFFOLD and faster than FedProx.

SCAFFOLD $>$ SGD $>$ FedAvg for heterogeneous clients. When similarity is $0\%$, FedAvg gets slower with increasing local steps. If we take more than 5 epochs, its performance is worse than SGD’s. SCAFFOLD initially worsens as we increase the number of epochs but then flattens. However, its performance is always better than that of SGD, confirming that it can handle heterogeneous data.

SCAFFOLD and FedAvg get faster with more similarity, but not SGD. As similarity of the clients increases, the performance of SGD remains relatively constant. On the other hand, SCAFFOLD and FedAvg get significantly faster as similarity increases. Further, local steps become much more useful, showing monotonic improvement with the increase in number of epochs. This is because with increasing the i.i.d.ness of the data, both the gradient and Hessian dissimilarity decrease.

SCAFFOLD is resilient to client sampling. As we decrease the fraction of clients sampled, SCAFFOLD ,and FedAvg only show a sub-linear slow-down. They are more resilient to sampling in the case of higher similarity.

SCAFFOLD outperforms FedAvg on non-convex experiments. We see that SCAFFOLD is better than FedAvg in terms of final test accuracy reached, though interestingly FedAvg seems better than SGD even when similarity is 0. However, much more extensive experiments (beyond current scope) are needed before drawing conclusions.

Conclusion

Our work studied the impact of heterogeneity on the performance of optimization methods for federated learning. Our careful theoretical analysis showed that FedAvg can be severely hampered by gradient dissimilarity, and can be even slower than SGD. We then proposed a new stochastic algorithm (SCAFFOLD) which overcomes gradient dissimilarity using control variates. We demonstrated the effectiveness of SCAFFOLD via strong convergence guarantees and empirical evaluations. Further, we showed that while SCAFFOLD is always at least as fast as SGD, it can be much faster depending on the Hessian dissimilarity in our data. Thus, different algorithms can take advantage of (and are limited by) different notions of dissimilarity. We believe that characterizing and isolating various dissimilarities present in real world data can lead to further new algorithms and significant impact on distributed, federated, and decentralized learning.

We thank Filip Hanzely and Jakub Konečnỳ for discussions regarding variance reduction techniques and Blake Woodworth, Virginia Smith and Kumar Kshitij Patel for suggestions which improved the writing.

References

Appendix

As stated earlier, federated learning involves learning a centralized model from distributed client data. This centralized model benefits from all client data and can often result in a beneficial performance e.g. in including next word prediction (Hard et al., 2018; Yang et al., 2018), emoji prediction (Ramaswamy et al., 2019), decoder models (Chen et al., 2019b), vocabulary estimation (Chen et al., 2019a), low latency vehicle-to-vehicle communication (Samarakoon et al., 2018), and predictive models in health (Brisimi et al., 2018). Nevertheless, federated learning raises several types of issues and has been the topic of multiple research efforts studying the issues of generalization and fairness (Mohri et al., 2019; Li et al., 2019a), the design of more efficient communication strategies (Konečnỳ et al., 2016b, a; Suresh et al., 2017; Stich et al., 2018; Karimireddy et al., 2019; Basu et al., 2019), the study of lower bounds (Woodworth et al., 2018), differential privacy guarantees (Agarwal et al., 2018), security (Bonawitz et al., 2017), etc. We refer to Kairouz et al. (2019) for an in-depth survey of this area.

For identical clients, FedAvg coincides with parallel SGD analyzed by (Zinkevich et al., 2010) who proved asymptotic convergence. Stich (2018) and, more recently Stich & Karimireddy (2019); Patel & Dieuleveut (2019); Khaled et al. (2020), gave a sharper analysis of the same method, under the name of local SGD, also for identical functions. However, there still remains a gap between their upper bounds and the lower bound of Woodworth et al. (2018). The analysis of FedAvg for heterogeneous clients is more delicate due to the afore-mentioned client-drift, first empirically observed by Zhao et al. (2018). Several analyses bound this drift by assuming bounded gradients (Wang et al., 2019; Yu et al., 2019), or view it as additional noise (Khaled et al., 2020), or assume that the client optima are $\epsilon$-close (Li et al., 2018; Haddadpour & Mahdavi, 2019). In a concurrent work, (Liang et al., 2019) propose to use variance reduction to deal with client heterogeneity but still show rates slower than SGD. We summarize the communication complexities of different methods for heterogeneous clients in Table 2.

The use of control variates is a classical technique to reduce variance in Monte Carlo sampling methods (cf. (Glasserman, 2013)). In optimization, they were used for finite-sum minimization by SVRG (Johnson & Zhang, 2013; Zhang et al., 2013a) and then in SAGA (Defazio et al., 2014) to simplify the linearly convergent method SAG (Schmidt et al., 2017). Numerous variations and extensions of the technique are studied in (Hanzely & Richtárik, 2019). Starting from (Reddi et al., 2016a), control variates have also frequently been used to reduce variance in finite-sum non-convex settings (Reddi et al., 2016c; Nguyen et al., 2018; Fang et al., 2018; Tran-Dinh et al., 2019). Further, they are used to obtain linearly converging decentralized algorithms under the guise of ‘gradient-tracking’ in (Shi et al., 2015; Nedich et al., 2016) and for gradient compression as ‘compressed-differences’ in (Mishchenko et al., 2019). Our method can be viewed as seeking to remove the ‘client-variance’ in the gradients across the clients, though there still remains additional stochasticity as in (Kulunchakov & Mairal, 2019), which is important in deep learning (Defazio & Bottou, 2019).

The problem of client-drift we described is a common phenomenon in distributed optimization. In fact, classic techniques such as ADMM mitigate this drift, though they are not applicable in federated learning. For well structured convex problems, CoCoA uses the dual variable as the control variates, enabling flexible distributed methods (Smith et al., 2018). DANE by (Shamir et al., 2014) obtain a closely related primal only algorithm, which was later accelerated by Reddi et al. (2016b) and recently extended to federated learning (Li et al., 2020). SCAFFOLD can be viewed as an improved version of DANE where a fixed number of (stochastic) gradient steps are used in place of a proximal point update. In a similar spirit, distributed variance reduction techniques have been proposed for the finite-sum case (Lee et al., 2015; Konečnỳ et al., 2016a; Cen et al., 2019). However, these methods are restricted to finite-sums and are not applicable to the stochastic setting studied here.

Appendix B Technicalities

We examine some additional definitions and introduce some technical lemmas.

We make precise a few definitions and explain some of their implications.

$f_{i}$ is $\mu$-convex for $\mu\geq 0$ and satisfies:

Here, we allow that $\mu=0$ (we refer to this case as the general convex case as opposed to strongly convex). It is also possible to generalize all proofs here to the weaker notion of PL-strong convexity (Karimi et al., 2016).

$g_{i}(\bm{x}):=\nabla f_{i}(\bm{x};\zeta_{i})$ is unbiased stochastic gradient of $f_{i}$ with bounded variance

Note that ((A4)) only bounds the variance within the same client, but not the variance across the clients.

$\{f_{i}\}$ are $\beta$-smooth and satisfy:

The assumption ((A5)) also implies the following quadratic upper bound on $f_{i}$

If additionally the function $\{f_{i}\}$ are convex and $\bm{x}^{\star}$ is an optimum of $f$, ((A5)) implies (via Nesterov (2018), Theorem 2.1.5)

Further, if $f_{i}$ is twice-differentiable, ((A5)) implies that $\lVert\nabla^{2}f_{i}(\bm{x})\rVert\leq\beta$ for any $\bm{x}$.

B.2 Some technical lemmas

Now we cover some technical lemmas which are useful for computations later on. The two lemmas below are useful to unroll recursions and derive convergence rates. The first one is a slightly improved (and simplified) version of (Stich, 2019, Theorem 2). It is straightforward to remove the additional logarithmic terms if we use a varying step-size (Kulunchakov & Mairal, 2019, Lemma 13).

For every non-negative sequence $\{d_{r-1}\}_{r\geq 1}$ and any parameters $\mu>0$, $\eta_{\max}\in(0,1/\mu]$, $c\geq 0$, $R\geq\frac{1}{2\eta_{\max}\mu}$, there exists a constant step-size $\eta\leq\eta_{\max}$ and weights $w_{r}:=(1-\mu\eta)^{1-r}$ such that for $W_{R}:=\sum_{r=1}^{R+1}w_{r}$,

By substituting the value of $w_{r}$, we observe that we end up with a telescoping sum and estimate

When $R\geq\frac{1}{2\mu\eta}$, $(1-\mu\eta)^{R}\leq\exp(-\mu\eta R)\leq\frac{2}{3}$. For such an $R$, we can lower bound $\eta W_{R}$ using

This proves that for all $R\geq\frac{1}{2\mu\eta}$,

The lemma now follows by carefully tuning $\eta$. Consider the following two cases depending on the magnitude of $R$ and $\eta_{\max}$:

Suppose $\frac{1}{2\mu R}\leq\eta_{\max}\leq\frac{\log(\max(1,\mu^{2}Rd_{0}/c))}{\mu R}$. Then we can choose $\eta=\eta_{\max}$,

Instead if $\eta_{\max}>\frac{\log(\max(1,\mu^{2}Rd_{0}/c))}{\mu R}$, we pick $\eta=\frac{\log(\max(1,\mu^{2}Rd_{0}/c))}{\mu R}$ to claim that

The next lemma is an extension of (Stich & Karimireddy, 2019, Lemma 13), (Kulunchakov & Mairal, 2019, Lemma 13) and is useful to derive convergence rates for general convex functions ($\mu=0$) and non-convex functions.

For every non-negative sequence $\{d_{r-1}\}_{r\geq 1}$ and any parameters $\eta_{\max}\geq 0$, $c\geq 0$, $R\geq 0$, there exists a constant step-size $\eta\leq\eta_{\max}$ and weights $w_{r}=1$ such that,

Similar to the strongly convex case (Lemma 1), we distinguish the following cases:

When $R+1\leq\frac{d_{0}}{c_{1}\eta_{\max}^{2}}$, and $R+1\leq\frac{d_{0}}{c_{2}\eta_{\max}^{3}}$ we pick $\eta=\eta_{\max}$ to claim

In the other case, we have $\eta_{\max}^{2}\geq\frac{d_{0}}{c_{1}(R+1)}$ or $\eta_{\max}^{3}\geq\frac{d_{0}}{c_{2}(R+1)}$. We choose \eta=\min\bigg{\{}\sqrt{\frac{d_{0}}{c_{1}(R+1)}},\sqrt{\frac{d_{0}}{c_{2}(R+1)}}\bigg{\}} to prove

$\lVert\bm{v}_{i}+\bm{v}_{j}\rVert^{2}\leq(1+a)\lVert\bm{v}_{i}\rVert^{2}+(1+\tfrac{1}{a})\lVert\bm{v}_{j}\rVert^{2}$ for any $a>0$, and

$\lVert\sum_{i=1}^{\tau}\bm{v}_{i}\rVert^{2}\leq\tau\sum_{i=1}^{\tau}\lVert\bm{v}_{i}\rVert^{2}$.

The proof of the first statement for any $a>0$ follows from the identity:

For the second inequality, we use the convexity of $\bm{x}\rightarrow\lVert\bm{x}\rVert^{2}$ and Jensen’s inequality

Next we state an elementary lemma about expectations of norms of random vectors.

Expanding the above expression using relaxed triangle inequality (Lemma 3) proves the first claim:

For the second statement, $\xi_{i}$ is not deterministic and depends on $\Xi_{i-1},\dots,\Xi_{1}$. Hence we have to resort to the cruder relaxed triangle inequality to claim

and then use the tighter expansion of the second term:

The cross terms in the above expression have zero mean since $\{\Xi_{i}-\xi_{i}\}$ form a martingale difference sequence. ∎

Appendix C Properties of convex functions

We now study two lemmas which hold for any smooth and strongly-convex functions. The first is a generalization of the standard strong convexity inequality ((A3)), but can handle gradients computed at slightly perturbed points.

The following holds for any $\beta$-smooth and $\mu$-strongly convex function $h$, and any $\bm{x},\bm{y},\bm{z}$ in the domain of $h$:

Given any $\bm{x}$, $\bm{y}$, and $\bm{z}$, we get the following two inequalities using smoothness and strong convexity of $h$:

Further, applying the relaxed triangle inequality gives

Combining all the inequalities together we have

The lemma follows since $\beta\geq\mu$. ∎

Here, we see that a gradient step is a contractive operator.

For any $\beta$-smooth and $\mu$-strongly convex function $h$, points $\bm{x},\bm{y}$ in the domain of $h$, and step-size $\eta\leq\frac{1}{\beta}$, the following is true

Recall our bound on the step-size $\eta\leq\frac{1}{\beta}$ which implies that $(\eta^{2}\beta-2\eta)\leq-\eta$. Finally, apply the $\mu$-strong convexity of $h$ to get

Appendix D Convergence of FedAvg

We outline the FedAvg method in Algorithm 2. In round $r$ we sample $\mathcal{S}^{r}\subseteq[N]$ clients with $\lvert\mathcal{S}^{r}\rvert=S$ and then perform the following updates:

Starting from the shared global parameters $\bm{y}^{r}_{i,0}=\bm{x}^{r-1}$, we update the local parameters for $k\in[K]$

Compute the new global parameters using only updates from the clients $i\in\mathcal{S}^{r}$ and a global step-size $\eta_{g}$:

Finally, for some weights $\{w_{r}\}$, we output

Recall our bound on the gradient dissimilarity:

If $\{f_{i}\}$ are convex, we can relax the assumption to

We defined two variants of the bounds on the heterogeneity depending of whether the functions are convex or not. Suppose that the functions $f$ is indeed convex as in ((A3)) and $\beta$-smooth as in ((A5)), then it is straightforward to see that (13) implies (14). Thus for convex functions, ((A1)) is mildly weaker. Suppose that the functions $\{f_{1},\dots,f_{N}\}$ are convex and $\beta$-smooth. Then (14) is satisfied with $B^{2}=2$ since

Thus, $(G,B)$-BGD (14) is equivalent to the heterogeneity assumption of (Mishchenko et al., 2019) with $G^{2}=\sigma_{f}^{2}$. Instead, if we have the stronger assumption ((A1)) but the functions are possibly non-convex, then $G=\epsilon$ corresponds to the local dissimilarity defined in (Li et al., 2018). Note that assuming $G$ is negligible is quite strong and corresponds to the strong-growth condition (Vaswani et al., 2019).

D.2 Rates of convergence (Theorem I)

We first restate Theorem I with some additional details and then see its proof.

Suppose that the functions $\{f_{i}\}$ satisfies assumptions (A4), (A5), and (A1). Then, in each of the following cases, there exist weights $\{w_{r}\}$ and local step-sizes $\eta_{l}$ such that for any $\eta_{g}\geq 1$ the output of FedAvg (12) $\bar{\bm{x}}^{R}$ satisfies

Strongly convex: $f_{i}$ satisfies ((A3)) for $\mu>0$, $\eta_{l}\leq\frac{1}{8(1+B^{2})\beta K\eta_{g}}$, $R\geq\frac{8(1+B^{2})\beta}{\mu}$ then

General convex: $f_{i}$ satisfies ((A3)) for $\mu=0$, $\eta_{l}\leq\frac{1}{(1+B^{2})8\beta K\eta_{g}}$, $R\geq 1$ then

Non-convex: $f_{i}$ satisfies ((A1)) and $\eta_{l}\leq\frac{1}{(1+B^{2})8\beta K\eta_{g}}$, then

where $M^{2}:=\sigma^{2}(1+\tfrac{S}{\eta_{g}^{2}})+K(1-\tfrac{S}{N})G^{2}$, $D:=\lVert\bm{x}^{0}-\bm{x}^{\star}\rVert$, and $F:=f(\bm{x}^{0})-f^{\star}$.

D.3 Proof of convergence

We will only prove the rate of convergence for convex functions here. The corresponding rates for non-convex functions are easy to derive following the techniques in the rest of the paper.

where $\mathcal{E}_{r}$ is the drift caused by the local updates on the clients defined to be

We can directly apply Lemma 5 with $h=f_{i}$, $\bm{x}=\bm{y}_{i,k-1}$, $\bm{y}=\bm{x}^{\star}$, and $\bm{z}=\bm{x}$ to the first term $\mathcal{A}_{1}$

For the second term $\mathcal{A}_{2}$, we repeatedly apply the relaxed triangle inequality (Lemma 4)

The last step used Assumption $(G,B)$-BGD assumption (14) that $\frac{1}{N}\sum_{i=1}^{N}\lVert\nabla f_{i}(\bm{x})\rVert^{2}\leq G^{2}+2\beta B^{2}(f(\bm{x})-f^{\star})$. The extra $(1-\frac{S}{N}$ improvement we get is due to sampling the functions $\{f_{i}\}$ without replacement. Plugging back the bounds on $\mathcal{A}_{1}$ and $\mathcal{A}_{2}$,

(bounded drift) Suppose our functions satisfies assumptions ((A1)) and ((A3))–((A5)). Then the updates of FedAvg for any step-size satisfying $\eta_{l}\leq\frac{1}{(1+B^{2})8\beta K\eta_{g}}$ have bounded drift:

If $K=1$, the lemma trivially holds since $\bm{y}_{i,0}=\bm{x}$ for all $i\in[N]$ and $\mathcal{E}_{r}=0$. Assume $K\geq 2$ here on. Recall that the local update made on client $i$ is $\bm{y}_{i,k}=\bm{y}_{i,k-1}-\eta_{l}g_{i}(\bm{y}_{i,k-1})$. Then,

Adding the statements of Lemmas 7 and 8, we get

D.4 Lower bound for FedAvg (Theorem II)

We first formalize the class of algorithms we look at before proving out lower bound.

We assume that FedAvg is run with $\eta_{g}=1$, $K>1$, and arbitrary possibly adaptive positive step-sizes $\{\eta_{1},\dots,\eta_{R}\}$ are used with $\eta_{r}\leq\frac{1}{\mu}$ and fixed within a round for all clients. Further, the server update is a convex combination of the client updates with non-adaptive weights.

Below we state a more formal version of Theorem II.

For any positive constants $G$, $\mu$, there exist $\mu$-strongly convex functions satisfying (A1) for which that the output of FedAvg satisfying (A6) has the error for any $r\geq 1$:

Consider the following simple one-dimensional functions for any given $\mu$ and $G$:

Let us start FedAvg from $x^{0}>0$. A single local update for $f_{1}$ and $f_{2}$ in round $r\geq 1$ is respectively

Then, straightforward computations show that the update at the end of round $r$ is of the following form for some averaging weight $\alpha\in$

Since $\alpha$ was picked obliviously, we can assume that $\alpha\leq 0.5$. If indeed $\alpha>0.5$, we can swap the definitions of $f_{1}$ and $f_{2}$ and the sign of $x^{0}$. With this, we can simplify as

Observe that in the above expression, the right hand side is increasing with $\eta_{r}$—this represents the effect of the client drift and increases the error as the step-size increases. The left hand side decreases with $\eta_{r}$—this is the usual convergence observed due to taking gradient steps. The rest of the proof is to show that even with a careful balancing of the two terms, the effect of $G$ cannot be removed. Lemma 9 performs exactly such a computation to prove that for any $r\geq 1$,

We finish the proof by noting that $f(x^{r})=\frac{\mu}{2}(x^{r})^{2}$. ∎

Suppose that for all $r\geq 1$, $\eta_{r}\leq\frac{1}{\mu}$ and the following is true:

Then, there exists a constant $c>0$ such that for any sequence of step-sizes $\{\eta^{r}\}$:

Define $\gamma_{r}=\mu\eta_{r}R(K-1)$. Such a $\gamma_{r}$ exists and is positive since $K\geq 2$. Then, $\gamma_{r}$ satisfies

The second inequality follows because $\eta_{r}\leq\frac{1}{\mu}$ implies that $(1-(1-2\mu\eta_{r})^{\tau})$ is always positive. If $\gamma_{r}\geq R/8$, then we have a constant $c_{1}\in(0,1/32)$ which satisfies

On the other hand, if $\gamma_{r}<R/8$, we have a tighter inequality

The last step used Bernoulli’s inequality and the fact that $K-1\leq K/2$ for $K\geq 2$. Observe that in the above expression, the right hand side is increasing with $\gamma_{r}$—this represents the effect of the client drift and increases the error as the step-size increases. The left hand side decreases with $\gamma_{r}$—this is the usual convergence observed due to taking gradient steps. The rest of the proof is to show that even with a careful balancing of the two terms, the effect of $G$ cannot be removed.

Suppose that all rounds after $r_{0}\geq 0$ have a small step-size i.e. $\gamma_{r}\leq R/8$ for all $r>r_{0}$ and hence satisfies (16). Then we will prove via induction that

For $r=r_{0}$, (17) is trivially satisfied. Now for $r>r_{0}$,

The first step is because of (16) and the last step uses the induction hypothesis. The second step considers two cases for $\gamma_{r}$: either $\gamma_{r}\leq\frac{1}{8}$ and $(1-\tfrac{1}{2R})\geq(1-\tfrac{1}{2R})$, or $\gamma_{r}^{2}\geq\frac{1}{64}$. Finally note that $c^{r}\geq\frac{1}{2}$ using Bernoulli’s inequality. We have hence proved

Now suppose $\gamma_{r_{0}}>R/8$. Then (15) implies that $x^{R}\geq\frac{cG}{\mu R}$ for some constant $c>0$. If instead no such $r_{0}\geq 1$ exists, then we can set $r_{0}=0$. Now finally observe that the previous proof did not make any assumption on $R$, and in fact the inequality stated above holds for all $r\geq 1$. ∎

Appendix E Convergence of SCAFFOLD

We first restate the convergence theorem more formally, then prove the result for the convex case, and then for non-convex case. Throughout the proof, we will focus on the harder option II. The proofs for SCAFFOLD with option I are nearly identical and so we skip them.

Suppose that the functions $\{f_{i}\}$ satisfies assumptions (A4) and (A5). Then, in each of the following cases, there exist weights $\{w_{r}\}$ and local step-sizes $\eta_{l}$ such that for any $\eta_{g}\geq 1$ the output (22) of SCAFFOLD satisfies:

Strongly convex: $f_{i}$ satisfies ((A3)) for $\mu>0$, $\eta_{l}\leq\min\left(\tfrac{1}{81\beta K\eta_{g}},\tfrac{S}{15\mu NK\eta_{g}}\right)$, $R\geq\max(\frac{162\beta}{\mu},\frac{30N}{S})$ then

Thus, the asymptotic rates of SCAFFOLD for general convex functions only incurs an additive term of the order of $O(\sqrt{\frac{N}{S}}\frac{1}{R})$. For strongly convex functions, we only see the affects in the logarithmic terms.

We will rewrite SCAFFOLD using notation which is convenient for the proofs: $\{\bm{y}_{i}\}$ represent the client models, $\bm{x}$ is the aggregate server model, and $\bm{c}_{i}$ and $\bm{c}$ are the client and server control variates. For an equivalent description which is easier to implement, we refer to Algorithm 1. The server maintains a global control variate $\bm{c}$ as before and each client maintains its own control variate $\bm{c}_{i}$. In round $r$, a subset of clients $\mathcal{S}^{r}$ of size $S$ are sampled uniformly from $\{1,\dots,N\}$. Suppose that every client performs the following updates

Starting from the shared global parameters $\bm{y}^{0}_{i,r}=\bm{x}^{r-1}$, we update the local parameters for $k\in[K]$

Update the control iterates using (option II):

We update the local control variates only for clients $i\in\mathcal{S}^{r}$

Compute the new global parameters and global control variate using only updates from the clients $i\in\mathcal{S}^{r}$:

Finally, for some weights $\{w_{r}\}$, we output

Note that the clients are agnostic to the sampling and their updates are identical to when all clients are participating. Also note that the control variate choice (19) corresponds to (option II) of Algorithm 1. Further, the updates of the clients $i\notin\mathcal{S}^{r}$ is forgotten and is defined only to make the proofs easier. While actually implementing the method, only clients $i\in\mathcal{S}^{r}$ participate and the rest remain inactive (see Algorithm 1).

We will first bound the variance of SCAFFOLD update in Lemma 12, then see how sampling of clients effects our control variates in Lemma 13, and finally bound the amount of client-drift in Lemma 14. We will then use these three lemmas to prove the progress in a single round in Lemma 15. Combining this progress with Lemmas 1 and 2 gives us the desired rates.

Before proceeding with the proof of our lemmas, we need some additional definitions of the various errors we track. As before, we define the effective step-size to be

We define client-drift to be how much the clients move from their starting point:

Because we are sampling the clients, not all the client control-variates get updated every round. This leads to some ‘lag’ which we call control-lag:

We study how the variance of the server update can be bounded.

The inequality before the last used the smoothness of $\{f_{i}\}$. The last inequality which separates the mean and the variance is an application of Lemma 4: the variance of $(\frac{1}{KS}\sum_{k,i\in\mathcal{S}}g_{i}(\bm{y}_{i,k-1}))$ is bounded by $\sigma^{2}/KS$. Similarly, $\bm{c}_{j}$ as defined in (19) for any $j\in[N]$ has variance smaller than $\sigma^{2}/K$ and hence the variance of $(\frac{1}{S}\sum_{i\in\mathcal{S}}\bm{c}_{i})$ is smaller than $\sigma^{2}/KS$.

We have previously related the variance of the server update to the control lag. We now examine how the control-lag grows each round.

Recall that after round $r$, the control update rule (19) implies that $\bm{c}^{r}_{i}$ is set as per

Plugging the above expression in the definition of $\mathcal{C}_{r}$ we get

The final step applied Jensen’s inequality twice. We can then further simplify using the relaxed triangle inequality as

We will now bound the final source of error which is the client-drift.

Suppose our step-sizes satisfy $\eta_{l}\leq\frac{1}{81\beta K\eta_{g}}$ and $f_{i}$ satisfies assumptions (A3)–(A5). Then, for any global $\eta_{g}\geq 1$ we can bound the drift as

First, observe that if $K=1$, $\mathcal{E}_{r}=0$ since $\bm{y}_{i,0}=\bm{x}$ for all $i\in[N]$ and that $\mathcal{C}_{r-1}$ and the right hand side are both positive. Thus the lemma is trivially true if $K=1$. For $K>1$, we build a recursive bound of the drift.Starting from the definition of the update (18) and then applying the relaxed triangle inequality, we can expand

The final step follows from the relaxed triangle inequality (Lemma 3). Applying the contractive mapping Lemma 6 for $\eta_{l}\leq 1/\beta$ shows

Once again using our relaxed triangle inequality to expand the other term $\mathcal{T}_{3}$, we get

The last step used the smoothness of $f_{i}$. Combining the bounds on $\mathcal{T}_{2}$ and $\mathcal{T}_{3}$ in the original inequality and using $a=\frac{1}{K-1}$ gives

Recall that with the choice of $\bm{c}_{i}$ in (19), the variance of $c_{i}$ is less than $\frac{\sigma^{2}}{K}$. Separating its mean and variance gives

Unrolling the recursion (25), we get the following for any $k\in\{1,\dots,K\}$

The inequality $(K-1)((1+\tfrac{1}{K-1})^{K}-1)\leq 3K$ can be verified for $K=2,3$ manually. For $K\geq 4$,

Again averaging over $k$ and multiplying by $3\beta$ yields

E.2 Convergence of SCAFFOLD for non-convex functions (Theorem III)

Now that we have a bound on all errors, we can describe our progress.

Starting from our server update equation,

We can then apply Lemma 12 to bound the second moment of the server update as

The term $\mathcal{T}_{4}$ can be bounded by using perturbed strong-convexity (Lemma 5) with $h=f_{i}$, $\bm{x}=\bm{y}_{i,k-1}$, $\bm{y}=\bm{x}^{\star}$, and $\bm{z}=\bm{x}$ to get

Plugging $\mathcal{T}_{4}$ back, we can further simplify the expression to get

Now recall that Lemma 14 bounds the client-drift:

We now analyze the most general case of SCAFFOLD with option II on functions which are potentially non-convex. Just as in the non-convex proof, we will first bound the variance of the server update in Lemma 16, the change in control lag in Lemma 17 and finally we bound the client-drift in Lemma 18. Combining these three together gives us the progress made in one round in Lemma 19. The final rate is derived from the progress made using Lemma 2.

Recall that in round $r$, we update the control variate as (19)

We introduce the following notation to keep track of the ‘lag’ in the update of the control variate: define a sequence of parameters $\{\bm{\alpha}_{i,k-1}^{r-1}\}$ such that for any $i\in[N]$ and $k\in[K]$ we have $\bm{\alpha}_{i,k-1}^{0}:=\bm{x}^{0}$ and for $r\geq 1$,

By the update rule for control variates (19) and the definition of $\{\bm{\alpha}_{i,k-1}^{r-1}\}$ above, the following property always holds:

We can then define the following $\Xi_{r}$ to be the error in control variate for round $r$:

Also recall the closely related definition of client drift caused by local updates:

Let us analyze how the control variates effect the variance of the aggregate server update.

Recall that that the server update satisfies

Taking norm on both sides and separating mean and variance, we proceed as

We now analyze the ‘lag’ in the control variates due to us sampling only a small subset of clients each round. Because we cannot rely on convexity anymore but only on the Lipschitzness of the gradients, the control-lag increases faster in the non-convex case.

The proof proceeds similar to that of Lemma 13 except that we cannot rely on convexity. Recall that after round $r$, the definition of $\bm{\alpha}_{i,k-1}^{r}$ (26) implies that

Plugging the above expression in the definition of $\Xi_{r}$ we get

We can expand the second term $\mathcal{T}_{6}$ with the relaxed triangle inequality to claim

We will expand the first term $\mathcal{T}_{5}$ to claim for a constant $b\geq 0$ to be chosen later

where we used Young’s inequality which holds for any $b\geq 0$. Combining the bounds for $\mathcal{T}_{5}$ and $\mathcal{T}_{6}$,

We will next bound the client drift $\mathcal{E}_{r}$. For this, convexity is not crucial and we will recover a very similar result to Lemma 14 only use the Lipschitzness of the gradient.

Suppose our step-sizes satisfy $\eta_{l}\leq\frac{1}{24\beta K\eta_{g}}$ and $f_{i}$ satisfies assumptions (A4)–(A5). Then, for any global $\eta_{g}\geq 1$ we can bound the drift as

First, observe that if $K=1$, $\mathcal{E}_{r}=0$ since $\bm{y}_{i,0}=\bm{x}$ for all $i\in[N]$ and that $\Xi_{r-1}$ and the right hand side are both positive. Thus the Lemma is trivially true if $K=1$ and we will henceforth assume $K\geq 2$. Starting from the update rule (18) for $i\in[N]$ and $k\in[K]$

The inequalities above follow from repeated application of the relaxed triangle inequalities and the $\beta$-Lipschitzness of $f_{i}$. Averaging the above over $i$, the definition of $\bm{c}=\frac{1}{N}\sum_{i}\bm{c}_{i}$ and $\Xi_{r-1}$ (27) gives

Appendix F Usefulness of local steps (Theorem IV)

Given that we can bound all sources of error, we can finally prove the progress made in each round.

Starting from the smoothness of $f$ and taking conditional expectation gives

The third inequality follows from the observation that $-ab=\frac{1}{2}((b-a)^{2}-a^{2})-\frac{1}{2}b^{2}\leq\frac{1}{2}((b-a)^{2}-a^{2})$ for any $a,b\in$, and the last from the $\beta$-Lipschitzness of $f_{i}$. Now we use Lemma 17 to bound $\Xi_{r}$ as

Let us state our rates of convergence for SCAFFOLD which interpolates between identical and completely heterogeneous clients. In this section, we always set $\eta_{g}=1$ and assume all clients participate $(S=N)$.

Suppose that the functions $\{f_{i}\}$ are quadratic and satisfy assumptions (A4), (A5) and additionally (A2). Then, for global step-size $\eta_{g}=1$ in each of the following cases, there exist probabilities $\{p_{k}^{r}\}$ and local step-size $\eta_{l}$ such that the output (29) of SCAFFOLD when run with no client sampling ($S=N$) using update (28) satisfies:

Strongly convex: $f_{i}$ satisfies ((A3)) for $\mu>0$, $\eta_{l}\leq\min(\frac{1}{10\beta},\frac{1}{22\delta K},\frac{1}{10\mu K})$, $R\geq\max(\frac{20\beta}{\mu},\frac{44\delta K+20\mu K}{\mu},20K)$ then

General convex: $f$ satisfies $\nabla^{2}f\succeq-\delta I$, $\eta_{l}\leq\min(\frac{1}{10\beta},\frac{1}{22\delta K})$, and $R\geq 1$, then

Note that if $\delta=0$, we match (up to acceleration) the lower bound in (Woodworth et al., 2018). While certainly $\delta=0$ when the functions are identical as studied in (Woodworth et al., 2018), our upper-bound is significantly stronger since it is possible that $\delta=0$ even for highly heterogeneous functions. For example, objective perturbation (Chaudhuri et al., 2011; Kifer et al., 2012) is an optimal mechanism to achieve differential privacy for smooth convex objectives (Bassily et al., 2014). Intuitively, objective perturbation relies on masking each client’s gradients by adding a large random linear term to the objective function. In such a case, we would have high gradient dissimilarity but no Hessian dissimilarity.

Our non-convex convergence rates are the first of their kind as far as we are aware—no previous work shows how one can take advantage of similarity for non-convex functions. However, we should note that non-convex quadratics do not have a global lower-bound on the function value $f^{\star}$. We will instead assume that $f^{\star}$ almost surely lower-bounds the value of $f(\bm{x}^{R})$, implicitly assuming that the iterates remain bounded.

In the rest of this section, we will focus on proving Theorem VIII. We will show how to bound variance in Lemma 23, bound the amount of drift in Lemma 22, and show progress made in one step in Lemma 24. In all of these we do not use convexity, but strongly rely on the functions being quadratics. Then we combine these to derive the progress made by the server in one round—for this we need weak-convexity to argue that averaging the parameters does not hurt convergence too much. As before, it is straight-forward to derive rates of convergence from the one-round progress using Lemmas 1 and 2.

F.1 Additional notation and assumptions

For any matrix $M$ and vector $\bm{v}$, let $\lVert\bm{v}\rVert_{M}^{2}:=\bm{v}^{\top}M\bm{v}$. Since all functions in this section are quadratics, we can assume w.l.o.g they are of the following form:

for some $\{\bm{x}_{i}^{\star}\}$ and $\bm{x}^{\star}$, $A:=\frac{1}{N}\sum_{i=1}^{N}A_{i}$. We also assume that $A$ is a symmetric matrix though this requirement is easily relaxed. Note that this implies $f(\bm{x}^{\star})=0$ and that $\nabla f_{i}(\bm{x})=A(\bm{x}-\bm{x}^{\star}_{i})$. If $\{f_{i}\}$ are additionally convex, we have that $\bm{x}_{i}^{\star}$ is the optimum of $f_{i}$ and $\bm{x}^{\star}$ the optimum of $f$. However, this is not necessarily true in general.

We will also focus on a simplified version of SCAFFOLD where in each round $r$, client $i$ performs the following update starting from $\bm{y}_{i,0}^{r}\leftarrow\bm{x}^{r-1}$:

where the second part is specialized to quadratics and the expectation is conditioned over everything before current step $k$ of round $r$. At the end of each round, as before, $\bm{x}^{r}=\frac{1}{N}\sum_{i=1}^{N}\bm{y}^{r}_{i,K}$. The final output of the algorithm is chosen using probabilities $\{p^{r}_{k}\}$ as

Note that we are now possibly outputting iterates computed within a single round and that $\bm{x}^{r}=\bm{x}^{r}_{K}$. Beyond this, the update above differs from our usual SCAFFOLD in two key aspects: a) it uses gradients computed at $\bm{x}^{r-1}$ as control variates instead of those at either $\bm{x}^{r-2}$ (as in option I) or $\bm{y}_{i,k}^{r-1}$ (as in option II), and b) it uses full batch gradients to compute its control variates instead of stochastic gradients. The first issue is easy to fix and our proof extends to using both option I or option II using techniques in Section E. The second issue is more technical—using stochastic gradients for control variates couples the randomness across the clients in making the local-updates biased. While it may be possible to get around this (cf. (Lei & Jordan, 2017; Nguyen et al., 2017; Tran-Dinh et al., 2019)), we will not attempt to do so in this work. Note that if $K$ local update steps typically represents running multiple epochs on each client. Hence one additional epoch to compute the control variate $\nabla f_{i}(\bm{x})$ does not significantly add to the cost.

Finally, we define the following sequence of positive numbers for notation convenience:

F.2 Lemmas tracking errors

We see how averaging can reduce variance. A similar argument was used in the special case of one-shot averaging in (Zhang et al., 2013b).

Suppose $\{f_{i}\}$ are quadratic functions and assumption (A4) is satisfied. Then let $\bm{x}_{k}^{r}$ and $\bm{y}_{i,k}^{r}$ be vectors in step $k$ and round $r$ generated using (28)—(29). Then,

The third equality was because $\{\bm{y}_{i,k}\}$ are independent of each other conditioned on everything before round $r$. ∎

We next see the effect of averaging on function values.

Suppose that $f$ is $\delta$ general-convex, then we have:

Since $f$ is $\delta$-general convex, it follows that the function $f(\bm{z})+\delta(1+\tfrac{1}{K})^{K-k}\lVert\bm{z}-\bm{x}\rVert^{2}_{2}$ is convex in $\bm{z}$ for any $k\in[K]$. The lemma now follows directly from using convexity and the definition of $\bm{x}^{r}_{k}=\frac{1}{N}\bm{y}_{i,k}^{r}$. ∎

We see how the client drift of SCAFFOLD depends on $\delta$.

For the update (28), assuming ((A2)) and that $\{f_{i}\}$ are quadratics, the following holds for any $\eta\leq\frac{1}{21\delta K}$

Note that if $K=1$, then the first inequality directly proves the lemma. For the second inequality, we assumed $K\geq 2$ and then applied our relaxed triangle inequality. By assumption (A2), we have the following for $\eta\delta\leq 1$

Using the bound on the step-size $\eta\leq\frac{1}{21\delta K}$ gives

Simple computations now give the Lemma statement for all $K\geq 1$. ∎

We will see how to bound the variance of the output.

Consider the update (28) for quadratic $\{f_{i}\}$ with $\eta\leq\max(\frac{1}{2\delta K},\frac{1}{\beta})$. Then, if further ((A2)), ((A5)) and ((A4)) are satisfied, we have

Further if $\{f_{i}\}$ are strongly convex satisfying ((A3)), we have

We can rewrite the update step (28) as below:

Similarly, the expected iterate satisfies the same equation without the $\zeta_{i,k}$

This implies that the difference satisfies

We can relate this to the function value as follows:

The last inequality used smoothness of $f$ and the one before that relied on the independence of $\zeta_{i,k}$. Now, if $f_{i}$ is general convex we have for $\eta\leq\frac{1}{2\delta K}$ that $I-\eta A_{i}\preceq(1+\frac{1}{2K})I$ and hence

This proves our second statement of the lemma. For strongly convex functions, we have for $\eta\leq\frac{1}{\beta}$,

F.3 Lemmas showing progress

Now we focus only on a single client and monitor their progress.

Suppose ((A2)), ((A5)) and ((A4)) hold, and $\{f_{i}\}$ are quadratics. Then, the following holds for the update (28) with $\eta\leq\min(\frac{1}{10\beta},\frac{1}{22\delta K},\frac{1}{\mu K})$ with $\mu=0$ is $f$ is non-convex or general-convex

Recall that $\xi_{i,k}^{r}\geq 0$ is defined to be

The second to last inequality used that $\lVert\,\cdot\,\rVert^{2}_{A}\leq\beta\lVert\,\cdot\,\rVert^{2}_{2}$ by ((A5)) and that $\lVert(A-A_{i})(\,\cdot\,)\rVert^{2}_{2}\leq\delta^{2}\lVert\,\cdot\,\rVert^{2}_{2}$ by ((A2)). The final inequality used that $\eta\leq\max(\frac{1}{10\beta},\frac{1}{22\delta K})$. Now, multiplying Lemma 22 by $\delta(1+\frac{1}{K})^{K-k}\leq\tfrac{20\delta}{7}$ we have

Adding this to our previous equation gives the following recursive bound:

The bound on our step-size $\eta\leq\min(\frac{1}{10\beta},\frac{1}{22\delta K})$ implies that $\tfrac{3\eta}{2}-2\eta^{2}\beta-20\delta K\eta^{2}\geq\tfrac{\eta}{3}$ and recall that $\delta\leq 2\beta$. This proves first statement of the lemma for non-strongly convex functions ($\mu=0$). If additionally $f$ is strongly-convex with $\mu>0$, we have

This can be used to tighten the inequality as follows

If $\eta\leq\frac{1}{\mu K}$, then $(1-\tfrac{1}{5K})\leq(1-\tfrac{\mu\eta}{6})$ and we have the strongly-convex version of the first statement.

Now we combine the progress made by each client in one step to calculate the server progress.

Suppose ((A2)), ((A5)) and ((A4)) hold, and $\{f_{i}\}$ are quadratics. Then, the following holds for the update (28) with $\eta\leq\min(\frac{1}{10\beta},\frac{1}{21\delta K},\frac{1}{10\mu K})$ and weights $w_{k}:=(1-\tfrac{\mu\eta}{6})^{1-k}$:

Set $\mu=0$ if $\{f_{i}\}$s are not strongly-convex (is only general-convex).

Let us do the non-convex (and general convex) case first. By summing over Lemma 24 we have

proving the second part of the Lemma for weights $w_{k}=1$. The proof of strongly convex follows a very similar argument. Unrolling Lemma 24 using weights $w_{k}:=(1-\tfrac{\mu\eta}{6})^{1-k}$ gives

As in the general-convex case, we can use Lemmas 21, 20 and 23 to prove that

The proof of Theorem VIII follows by appropriately unrolling Lemma 25. For general-convex functions, we can simply use Lemma 2 with the probabilities set as $p_{k}^{r}=\frac{1}{KR}$. For strongly-convex functions, we use $p_{k}^{r}\propto(1-\frac{\mu\eta}{6})^{1-rk}$ and follow the computations in Lemma 1.